18

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

Ciro Santilli OurBigBook.com
  • 347,512
  • 102
  • 1,199
  • 985
Prokop Hapala
  • 2,424
  • 2
  • 30
  • 59

5 Answers5

11

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

enter image description here

Community
  • 1
  • 1
Hooked
  • 84,485
  • 43
  • 192
  • 261
  • thank you a lot! It is quite slow, so it would be better if there would be such function directly hardcoded numpy (I mean something accelerated in the same way as other array operations in numpy - without iterating over array by python loop ). But the important is that it works. – Prokop Hapala Jun 13 '13 at 17:58
  • @ProkopHapala Actually most of the work _is_ done with numpy, except for the call to `hls_to_rgb` which you could probably vectorize. You can make it _much much_ faster by changing the number of points `N`, the speed should be proportional to N^2. – Hooked Jun 13 '13 at 19:09
11

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()
nadapez
  • 2,603
  • 2
  • 20
  • 26
  • 1
    When you do c.swapaxes(0,2) you’re altering the orientation of the original image. I needed an additional swapaxes(0,1) afterwards for my image to render correctly. However, I used meshgrid instead of ogrid, as I don’t know exactly how the latter works; perhaps this made a difference? – Sascha Baer May 09 '18 at 14:38
  • 2
    @Adarain Or we change line 15 from `c = c.swapaxes(0,2)` to `c = c.transpose(1,2,0)` – max0r Apr 16 '20 at 07:15
8

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)

enter image description here

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)

enter image description here

Hooked
  • 84,485
  • 43
  • 192
  • 261
  • 1
    This looks nice, however it's just for plotting functions where I know analytinc prescription. It is not my case. I need something to plot complex data with dicrete sampling which I read from text file and store in 2D narray. I don't have explicit functionoal prescription for this data which could be sampled in any point. – Prokop Hapala Jun 12 '13 at 13:10
3

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
sk29910
  • 2,326
  • 1
  • 18
  • 23
w4m
  • 301
  • 1
  • 10
0

You can also use PIL.image to convert it

import PIL.Image

def colorize(z):
  z = Zxx
  n,m = z.shape

  A = (np.angle(z) + np.pi) / (2*np.pi)
  A = (A + 0.5) % 1.0 * 255
  B = 1.0 - 1.0/(1.0+abs(z)**0.3)
  B = abs(z)/ z.max() * 255
  H = np.ones_like(B)
  image = PIL.Image.fromarray(np.stack((A, B, np.full_like(A, 255)), axis=-1).astype(np.uint8), "HSV") # HSV has range 0..255 for all channels
  image = image.convert(mode="RGB")

  return np.array(image)
Kingsly2407
  • 476
  • 4
  • 8