使用matplotlib.contour处理复杂数据

Question

我正在尝试使用matplotlib从一个复杂的数组中显示等高线图。这个数组是一个2x2的复杂矩阵，是通过类似C语言的方法生成的：

 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))

其中的和必须加起来。如何才能使用等高线图需要的meshgrid形式来实现呢？再次感谢。

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()

实数

明确的实数

虚数