Hacker News Reader

Comment by pbourke | Hacker News Reader

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.

View on HN

Thanks 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.

Those 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.

For 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.

You can since 3.8 :)

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.

Awesome! That brings it down to 0.57s on my machine - a 56x improvement from the original.

Thanks for bringing it further.

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.

I 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

I 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.

Yeah if you see a for loop on a numpy array in a code base, it's game over.

Reply on news.ycombinator.com