How to create a phase plot for a 2D array of complex numbers with matplotlib?
Asked Answered
S

6

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

Sorrento answered 11/6, 2013 at 12:25 Comment(1)
See this answer for another question: https://mcmap.net/q/21398/-is-there-any-way-to-use-bivariate-colormaps-in-matplotlibAutotomy
C
12

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()
Cordova answered 6/1, 2014 at 20:29 Comment(3)
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?Mylohyoid
@Adarain Or we change line 15 from c = c.swapaxes(0,2) to c = c.transpose(1,2,0)Ignoble
Further vectorizationGoldie
E
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

Endor answered 12/6, 2013 at 14:34 Comment(2)
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.Sorrento
@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.Endor
E
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

Endor answered 11/6, 2013 at 15:27 Comment(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.Sorrento
P
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
Pasia answered 18/3, 2016 at 11:9 Comment(0)
B
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)
Bondholder answered 2/5, 2021 at 17:54 Comment(0)
G
0

Vectorization of nadapez's answer, x2-3 speedup, below.

The best bet is numba or Cython, but if one's stuck with vectorization, this works better. Note I also divide by max(abs(z)), don't recall why, but feel free to remove. I also changed swapaxes to make shape (*s, 3), as expected by plt.imshow.

Function

import numpy as np

ONE_THIRD = 1. / 3.
TWO_THIRD = 2. / 3.
ONE_SIXTH = 1. / 6.

def _v2(m1, m2, m2_m_m1, hue):
    hue = np.mod(hue, 1)
    # masks
    idxs0 = np.where((hue < ONE_SIXTH))
    idxs1 = np.where((ONE_SIXTH <= hue) * (hue < 0.5))
    idxs2 = np.where((0.5 <= hue) * (hue < TWO_THIRD))
    idxs3 = np.where((hue >= TWO_THIRD))

    # masked quantities
    hue_0 = hue[idxs0]
    hue_2 = hue[idxs2]

    m1_0 = m1[idxs0]
    m1_2 = m1[idxs2]
    m1_3 = m1[idxs3]

    m2_1 = m2[idxs1]

    m2_m_m1_0 = m2_m_m1[idxs0]
    m2_m_m1_2 = m2_m_m1[idxs2]

    # compute out
    out = np.zeros_like(m1)
    out[idxs0] = m1_0 + m2_m_m1_0*hue_0*6.
    out[idxs1] = m2_1
    out[idxs2] = m1_2 + m2_m_m1_2*(TWO_THIRD - hue_2)*6.
    out[idxs3] = m1_3
    return out


def hls_to_rgb_vec(h, l, s):
    # compute masks
    l_leq_05 = (l <= 0.5)
    l_gt_05  = np.logical_not(l_leq_05)
    ll_05 = l[l_leq_05]
    lg_05 = l[l_gt_05]

    # compute m2, m1
    m2 = np.zeros_like(l)
    m2[l_leq_05] = ll_05 * (1. + s)
    m2[l_gt_05] = lg_05*(1. - s) + s

    m1 = 2.*l - m2

    m2_m_m1 = m2 - m1

    # compute output
    out = np.zeros(l.shape + (3,), dtype=l.dtype)
    out[..., 0] = _v2(m1, m2, m2_m_m1, h+ONE_THIRD)
    out[..., 1] = _v2(m1, m2, m2_m_m1, h)
    out[..., 2] = _v2(m1, m2, m2_m_m1, h-ONE_THIRD)
    return out

Test + bench

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

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

    c = np.vectorize(hls_to_rgb)(h, l, s)
    c = np.array(c)
    c = c.transpose(1, 2, 0)
    return c

for N in (5, 20, 50, 100, 200, 500, 1000):
    print(f'\nN={N}')
    z = np.random.randn(N, N + 1) + 1j*np.random.randn(N, N + 1)

    # assert equality
    o0 = colorize(z)
    o1 = colorize_v2(z)
    assert np.allclose(o0, o1, atol=0)

    # bench
    %timeit colorize(z)
    %timeit colorize_v2(z)

Only N=5 is slower on my CPU, by x2.

Note: as far as hls_to_rgb_vec is concerned independent of colorize_v2, s != 0 is assumed (since for colorize_v2 it's always 0.8), and non-scalar s isn't tested.

Goldie answered 20/10, 2023 at 13:15 Comment(0)

© 2022 - 2025 — McMap. All rights reserved.