Z^2的实部和虚部的等值线图

产生一个复数$Z$, $Z$的实部和虚部分别从$-100$变化到$100$，在一幅图的两幅子图中分别画$Z^2$的实部和虚部的等值线图.

import numpy as np
import matplotlib.pyplot as plt

def polt_Z(a):
    x_vals = np.linspace(-100,100,200)
    y_vals = np.linspace(-100,100,200)
    x,y=np.meshgrid(x_vals,y_vals)

    z = x + y*1j
    z_sq = z**2     
    Z = [z_sq.real,z_sq.imag]
    
    cs = plt.contour(x,y,Z[a], 15, colors='black')
    plt.contourf(x,y,Z[a], 15)
    plt.clabel(cs, inline=1, fontsize=10)
    plt.colorbar()

plt.figure()
for i in range(1,2+1):
    plt.subplot(1,2,i)
    polt_Z(i-1)

左图为实部，右图为虚部