Используйте 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 (действительная часть) с использованием contourplot, но пока единственное, что я могу получитьиз контура -

TypeError: Input z must be a 2D array.

C_arr.real - это 2D-массив, и не имеет значения, если я сделаю сетку с 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()

real

явное действительное

образный