Use matplotlib.contour with complex data

I'm trying to show a contour plot using matplotlib from a complex array. The array is a 2x2 complex matrix, generated by the (C like) method:

 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

I would like to plot this c_arr (real part) using contourplot, but so far the only thing that I can get from contour is

TypeError: Input z must be a 2D array.

The c_arr.real is a 2D array, and doesn't matter if I make a grid with x, y, or pos_x, or pos_y, the result is always the same. The docs from matplotlib tells me how to use it, but not the datatypes necessary to use it, so I feel left in the dark.

EDIT: Thanks for the answer. My problem now is that I have to get the complex values from 开发者_开发问答a function in this form:

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

where the sum must be added up. How can this be accomplished using the meshgrid form that contour needs? Thanks again.

matplotlib.pyplot.contour() allows complex-valued input arrays. It extracts real values from the array implicitly:

#!/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()

real

explicit real

imagenary