mathplotlib imshow complex 2D array

Question

Is there any good way how to plot 2D array of complex numbers as image in mathplotlib ?

It makes very much sense to map magnitude of complex number as \"brightness\" or

Answer 1

Adapting the plotting code from mpmath you can plot a numpy array even if you don't known the original function with numpy and matplotlib. If you do know the function, see my original answer using mpmath.cplot.

from colorsys import hls_to_rgb

def colorize(z):
    n,m = z.shape
    c = np.zeros((n,m,3))
    c[np.isinf(z)] = (1.0, 1.0, 1.0)
    c[np.isnan(z)] = (0.5, 0.5, 0.5)

    idx = ~(np.isinf(z) + np.isnan(z))
    A = (np.angle(z[idx]) + np.pi) / (2*np.pi)
    A = (A + 0.5) % 1.0
    B = 1.0 - 1.0/(1.0+abs(z[idx])**0.3)
    c[idx] = [hls_to_rgb(a, b, 0.8) for a,b in zip(A,B)]
    return c

From here, you can plot an arbitrary complex numpy array:

N = 1000
A = np.zeros((N,N),dtype='complex')
axis_x = np.linspace(-5,5,N)
axis_y = np.linspace(-5,5,N)
X,Y = np.meshgrid(axis_x,axis_y)
Z = X + Y*1j

A = 1/(Z+1j)**2 + 1/(Z-2)**2

# Plot the array "A" using colorize
import pylab as plt
plt.imshow(colorize(A), interpolation='none',extent=(-5,5,-5,5))
plt.show()