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.