mathplotlib imshow complex 2D array

问题

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 "saturation" and phase as "Hue" ( anyway Hue is nothing else than phase in RBG color space). http://en.wikipedia.org/wiki/HSL_and_HSV

But as far as I know imshow does accept only scalar values which are then mapped using some colorscale. There is nothing like ploting real RGB pictures?

I thing it would be easy just implement a version which accepts 2D array of tuples (vectors) of 3 floating point numbers or ndarray of floats of shape [:,:,3]. I guess this would be generally usefful feature. It would be also usefull for plotting real RGB colord images, such as textures outputted from OpenCL

回答1:

this does almost the same of @Hooked code but very much faster.

import numpy as np
from numpy import pi
import pylab as plt
from colorsys import hls_to_rgb

def colorize(z):
    r = np.abs(z)
    arg = np.angle(z) 

    h = (arg + pi)  / (2 * pi) + 0.5
    l = 1.0 - 1.0/(1.0 + r**0.3)
    s = 0.8

    c = np.vectorize(hls_to_rgb) (h,l,s) # --> tuple
    c = np.array(c)  # -->  array of (3,n,m) shape, but need (n,m,3)
    c = c.swapaxes(0,2) 
    return c

N=1000
x,y = np.ogrid[-5:5:N*1j, -5:5:N*1j]
z = x + 1j*y

w = 1/(z+1j)**2 + 1/(z-2)**2
img = colorize(w)
plt.imshow(img)
plt.show()

回答2:

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

回答3:

The library mpmath uses matplotlib to produce beautiful images of the complex plane. On the complex plane you usually care about the poles, so the argument of the function gives the color (hence poles will make a spiral). Regions of extremely large or small values are controlled by the saturation. From the docs:

By default, the complex argument (phase) is shown as color (hue) and the magnitude is show as brightness. You can also supply a custom color function (color). This function should take a complex number as input and return an RGB 3-tuple containing floats in the range 0.0-1.0.

Example:

import mpmath
mpmath.cplot(mpmath.gamma, points=100000)

Another example showing the zeta function, the trivial zeros and the critical strip:

import mpmath
mpmath.cplot(mpmath.zeta, [-45,5],[-25,25], points=100000)

回答4:

You can use matplotlib.colors.hsv_to_rgb instead of colorsys.hls_to_rgb. The matplotlib function is about 10 times faster! See the results below:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import hsv_to_rgb
import time

def Complex2HSV(z, rmin, rmax, hue_start=90):
    # get amplidude of z and limit to [rmin, rmax]
    amp = np.abs(z)
    amp = np.where(amp < rmin, rmin, amp)
    amp = np.where(amp > rmax, rmax, amp)
    ph = np.angle(z, deg=1) + hue_start
    # HSV are values in range [0,1]
    h = (ph % 360) / 360
    s = 0.85 * np.ones_like(h)
    v = (amp -rmin) / (rmax - rmin)
    return hsv_to_rgb(np.dstack((h,s,v)))

here is the method of picked answer by @nadapez:

from colorsys import hls_to_rgb
def colorize(z):
    r = np.abs(z)
    arg = np.angle(z) 

    h = (arg + np.pi)  / (2 * np.pi) + 0.5
    l = 1.0 - 1.0/(1.0 + r**0.3)
    s = 0.8

    c = np.vectorize(hls_to_rgb) (h,l,s) # --> tuple
    c = np.array(c)  # -->  array of (3,n,m) shape, but need (n,m,3)
    c = c.swapaxes(0,2) 
    return c

Testing the results from the two method with 1024*1024 2darray:

N=1024
x, y = np.ogrid[-4:4:N*1j, -4:4:N*1j]
z = x + 1j*y

t0 = time.time()
img = Complex2HSV(z, 0, 4)
t1 = time.time()
print "Complex2HSV method: "+ str (t1 - t0) +" s"

t0 = time.time()
img = colorize(z)
t1 = time.time()
print "colorize method: "+ str (t1 - t0) +" s"

This result on my old laptop:

Complex2HSV method: 0.250999927521 s
colorize method: 2.03200006485 s

来源：https://stackoverflow.com/questions/17044052/mathplotlib-imshow-complex-2d-array