To bring "another language into it", i.e using numba for a JIT for python, this is another massive speedup for me, from 0.92 seconds to 0.26 seconds. The new code (needs some pieces from the old):
import numba
from numba import complex128, boolean
@numba.vectorize([complex128(complex128, complex128)])
def square_add(z, c):
"compute z² + c"
return z * z + c
@numba.vectorize([boolean(complex128)])
def not_escaped(z):
"compute |z|² < 4"
return ((z.conjugate() * z).real < 4.)
def compute_mandelbrot(cs):
"""
cs: complex grid subset
"""
z = np.zeros_like(cs)
red = np.zeros_like(cs, dtype=np.uint8)
mask = np.ones_like(cs, dtype=bool)
for _ in range(199):
# z² + c -> z
square_add(z, cs, where=mask, out=z)
np.add(red, 1, where=mask, out=red)
# mask &= |z| < 2
not_escaped(z, where=mask, out=mask)
return red
We can compare calling into numba with calling into luajit. In this code we see how much we get for free when using numpy ufuncs! All we did with numba is to define functions for the composite operations - one for z² + c and one for |z|² < 4 - which means we avoid storing intermediate arrays and do more in fewer array traversals.