将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

我想使用contourplot绘制这个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()

实

显式真实

imagenary