3 ms·
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 co
by kzrdude 5y ago
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.
- pbourke 5y agoI wonder if it would help to decorate compute_mandelbrot with @numba.jit(parallel=True) and change your call to range to a numba.prange. Ironically I’m going through this process at work this week to improve the performance of scipy.stats.ks_2samp
- kzrdude 5y agoI had this update lying around now too. Same signature function, this one actually was another 200% improvement: # guvectorize also supports target='parallel' for automatic parallelization @numba.guvectorize([(complex128[:], uint8[:])], "(n)->(n)") def compute_mandelbrot_gu(cs, out): "compute number of iterations of z² + c before it escapes" for i in range(cs.shape[0]): c = cs[i] z = 0j count = 0 while count < 200: z = z * z + c if (z * z.conjugate()).real > 4.: break count += 1 out[i] = count Whether to use numba's parallel for guvectorize or mine, it doesn't matter so much, they perform about the same. Parallelizing on a higher level is generally favorable to me. One thought here: the number of 100%+ improvements made here just underscores how slow python is to start with :). And yes, everyone in comments here complaining about numpy are sort of right, since the ultimate end here is just a flat loop without numpy, but as a numpy function.