4 ms·
The Python-only version can be sped up using vectorized numpy operations: import time from PIL import Image import numpy as np def PyMandelbrot(size
by pbourke 5y ago
The Python-only version can be sped up using vectorized numpy operations:
import time
from PIL import Image
import numpy as np
def PyMandelbrot(size):
t1 = time.time()
xPts = np.arange(-1.5, 0.5, 2.0 / size).reshape((1, size))
yPts = (np.arange(-1, 1, 2.0 / size) * 1j).reshape((size, 1))
cs = xPts + yPts
z = np.zeros_like(cs)
red = np.zeros_like(cs, dtype=np.uint8)
mask = np.full_like(cs, True, dtype=np.bool_)
max_idx = 0
while max_idx < 199:
np.putmask(z, mask, np.multiply(z, z) + cs)
np.putmask(red, mask, red + 1)
mask &= np.abs(z) < 2
max_idx += 1
dt = time.time() - t1
print(f"dt={dt:.2f}")
zeros = np.zeros_like(red)
pixels = np.dstack((red, zeros, zeros))
Image.fromarray(pixels).show()
if __name__ == '__main__':
PyMandelbrot(1_280)
On my machine (Core i9, Win 10, CPython 3.7):
> python .\mandelbrot_orig.py
dt=32.01
> python .\mandelbrot.py
dt=5.86
I'm sure my version can be improved upon.
- kzrdude 5y agoThanks for sharing, that's a big improvement! The details are easier to see, for me, in grayscale, that is using: Image.fromarray(255 - red).show() I'm learning as I go here, using numpy functions I've never used before, but we can use more in-place modifying operations like this (np.multiply, np.add with where= and out=) and it decreases runtime of the benchmark by 25%. Then, after that, this problem is embarassingly parallel and thanks to numpy we can parallelize your code very well, using all cores (improves runtime for me from 3.5 to 0.92 seconds) Using Python3.8.: import concurrent.futures import time from PIL import Image import numpy as np 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 np.multiply(z, z, where=mask, out=z) np.add(z, cs, where=mask, out=z) np.add(red, 1, where=mask, out=red) mask &= np.abs(z) < 2 return red def py_mandelbrot(size): t1 = time.time() yPts, xPts = np.ogrid[-1j:1j:size * 1j, -1.5:0.5:size * 1j] cs = xPts + yPts # over-split because chunks are not equal in time taken nsplit = 32 nthreads = 8 css = np.array_split(cs, nsplit) with concurrent.futures.ThreadPoolExecutor(max_workers=nthreads) as executor: reds_out = list(executor.map(compute_mandelbrot, css)) red = np.concatenate(reds_out) dt = time.time() - t1 print(f"dt={dt:.2f}") Image.fromarray(255 - red).show() if __name__ == '__main__': py_mandelbrot(1280) Python makes it easy to get computed chunks in the order they are submitted, using executor.map. I still tried out what it would look like if I concatenated the chunks in the order they were received. That's when I noticed that emptier tiles were computed faster. So then, we can split the problem into more fine-grained chunks and submit more of them while keeping the number of threads the same, for a small speedup since some chunks are faster to compute.
- pbourke 5y agoAwesome! That brings it down to 0.57s on my machine - a 56x improvement from the original. Thanks for bringing it further.
- BiteCode_dev 5y agoThose those who wonder how it can make it faster since it's using threads and Python has the GIL, numpy releases the gil while doing computation, so thread can actually use multiple cores if you speed most of your time in numpy world.
- kzrdude 5y agoFor this code we can also easily try the process executor instead of the thread executor - then we avoid the GIL. In this case, it has an approximate 10% performance penalty, probably since we have to serialize data back and forth, can't share memory.
- BiteCode_dev 5y agoYou can since 3.8 :) https://docs.python.org/3/library/multiprocessing.shared_memory.html https://docs.python.org/3/library/multiprocessing.shared_mem... If you use a data structure that supports the buffer protocol (such as a numpy array), you won't even pay a serialization cost. Python has made a LOT of progress, concurrency wise. Between asyncio, process pools, threads releasing the GIL and shared memory, you can get much perfs with the stdlib alone. Then of course you have Dask if you need more later.
- kzrdude 5y agoTo 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.
- BiteCode_dev 5y agoYeah if you see a for loop on a numpy array in a code base, it's game over.