使用 matplotlib.contour 处理复杂数据

Question

我正在尝试使用复杂数组中的 matplotlib 显示等高线图。该数组是一个 2x2 复数矩阵，由 (C like) 方法生成：

 for i in max_y:
   for j in max_x:
      pos_x = pos_x + step
      z = complex(pos_x,pos_y)
        c_arr[i][j] = complex_function(z)
   pos_y = pos_y + step

我想用轮廓图绘制这个 c_arr（实部），但到目前为止，我唯一能从轮廓图中得到的是

TypeError: Input z must be a 2D array.

c_arr.real 是一个二维数组，不管我用x、y、pos_x还是pos_y做一个网格，结果总是一样的。 matplotlib 的文档告诉我如何使用它，但没有告诉我使用它所需的数据类型，所以我觉得自己被蒙在鼓里了。

编辑：感谢您的回答。我现在的问题是我必须以这种形式从函数中获取复数值：

 def f(z):
      return np.sum(np.arange(n)*np.sqrt(z-1)**np.arange(n))

总和必须相加。如何使用轮廓需要的网格形式来实现这一点？再次感谢。

Answer 1

matplotlib.pyplot.contour() 允许使用复值输入数组。它隐式地从数组中提取实际值：

#!/usr/bin/env python
import numpy as np
from matplotlib import pyplot as plt

# generate data
x = np.r_[0:100:30j]
y = np.r_[0:1:20j]
X, Y = np.meshgrid(x, y)
Z = X*np.exp(1j*Y) # some arbitrary complex data

# plot it
def plotit(z, title):
    plt.figure()
    cs = plt.contour(X,Y,z) # contour() accepts complex values
    plt.clabel(cs, inline=1, fontsize=10) # add labels to contours
    plt.title(title)
    plt.savefig(title+'.png')

plotit(Z, 'real')
plotit(Z.real, 'explicit real')
plotit(Z.imag, 'imagenary')

plt.show()

真实

显式实数

图像