Increase performance creating Mandelbrot set in python I created a program in python that generates an image of the mandelbrot set. The only problem I have is that the program is quite slow, it takes about a quarter of an hour to generate the following image of 2000 by 3000 pixels: 



I first created a matrix of complex numbers using numpy according to amount of pixels. I also created an array for the image generation.

import numpy as np
from PIL import Image

z = 0

real_axis = np.linspace(-2,1,num=3000)
imaginary_axis = np.linspace(1,-1,num=2000)

complex_grid = [[complex(np.float64(a),np.float64(b)) for a in real_axis] for b in imaginary_axis]

pixel_grid = np.zeros((2000,3000,3),dtype=np.uint8)


Then I check whether each complex number is in the mandelbrot set or not and give it an RGB colour code accordingly.

for complex_list in complex_grid:
    for complex_number in complex_list:
        for iteration in range(255):
            z = z**2 + complex_number
            if (z.real**2+z.imag**2)**0.5 &gt; 2:
                pixel_grid[complex_grid.index(complex_list),complex_list.index(complex_number)]=[iteration,iteration,iteration]
                break
            else:
                continue
        z = 0


Finally I generate the image using the PIL library.

mandelbrot = Image.fromarray(pixel_grid)
mandelbrot.save("mandelbrot.png")


I am using jupyter notebook and python 3. Hopefully some of you can help me improve the performance of this program or other aspects of it.
I'm going to reuse some parts of the answer I recently posted here on Code Review.

Losing your Loops


  (Most) loops are damn slow in Python. Especially multiple nested loops.
  
  NumPy can help to vectorize your code, i.e. in this case that more
  of the looping is done in the C backend instead of in the Python
  interpreter. I  would highly recommend to have a listen to the talk
  Losing your Loops: Fast Numerical Computing with NumPy by Jake
  VanderPlas.


All those loops used to generate the complex grid followed by the nested loops used to iterate over the grid and the image are slow when left to the Python interpreter. Fortunately, NumPy can take quite a lot of this burden off of you.

For example

real_axis = np.linspace(-2, 1, num=3000)
imaginary_axis = np.linspace(1, -1, num=2000)
complex_grid = [[complex(np.float64(a),np.float64(b)) for a in real_axis] for b in imaginary_axis]


could become

n_rows, n_cols = 2000, 3000
complex_grid_np = np.zeros((n_rows, n_cols), dtype=np.complex)
real, imag = np.meshgrid(real_axis, imaginary_axis)
complex_grid_np.real = real
complex_grid_np.imag = imag


No loops, just plain simple NumPy.

Same goes for

for complex_list in complex_grid:
    for complex_number in complex_list:
        for iteration in range(255):
            z = z**2 + complex_number
            if (z.real**2+z.imag**2)**0.5 &gt; 2:
                pixel_grid[complex_grid.index(complex_list),complex_list.index(complex_number)]=[iteration,iteration,iteration]
                break
            else:
                continue
        z = 0


can be transformed to

z_grid_np = np.zeros_like(complex_grid_np)
elements_todo = np.ones((n_rows, n_cols), dtype=bool)
for iteration in range(255):
    z_grid_np[elements_todo] = \
        z_grid_np[elements_todo]**2 + complex_grid_np[elements_todo]
    mask = np.logical_and(np.absolute(z_grid_np) &gt; 2, elements_todo)
    pixel_grid_np[mask, :] = (iteration, iteration, iteration)
    elements_todo = np.logical_and(elements_todo, np.logical_not(mask))


which is just a single loop instead of three nested ones. Here, a little more trickery was needed to treat the break case the same way as you did. elements_todo is used to only compute updates on the z value if it has not been marked as done. There might also be a better solution without this.

I added the following lines

complex_grid_close = np.allclose(np.array(complex_grid), complex_grid_np)
pixel_grid_close = np.allclose(pixel_grid, pixel_grid_np)
print("Results were similar: {}".format(all((complex_grid_close, pixel_grid_close))))


to validate my results against your reference implementation.

The vectorized code is about 9-10x faster on my machine for several n_rows/n_cols combinations I tested. E.g. for n_rows, n_cols = 1000, 1500:

Looped generation took 61.989842s
Vectorized generation took 6.656926s
Results were similar: True


Lose a dimension

An aspect I somehow slightly ignored while looking at your code was that you're essentially creating a grayscale image since all of your color channel values are the same. Accounting for this, you can easily reduce the size of data the program has handle to from, in your case, 3000x2000x3 to 3000x2000. This will likely help your program to be more cache efficient, although I'm not an expert in this field.



Edit/Appendix: Further timings

Including the "no square root" optimization as suggested by trichoplax in his answer and Peter Cordes in the comments like so

mask = np.logical_and((z_grid_np.real**2+z_grid_np.imag**2) &gt; 4, elements_todo)


will give you about another second and a half for n_rows, n_cols = 1000, 1500, i.e. about 12x the speed of the original solution

10 loops, best of 5: 4.98 s per loop
10 loops, best of 5: 4.28 s per loop (in grayscale, 14x)


A quick implementation of Reinderien's hint towards the symmetry of the Mandelbrot set will again add a factor of about two to that.

10 loops, best of 5: 2.54 s per loop (~24x)
10 loops, best of 5: 2.07 s per loop (in grayscale, ~30x)


However, my quick hacking approach did not lead to an output that was completely within the tolerance of np.allclose compared to the original one. Funnily, it seems to be off by one at a single pixel, but visually still the same. Since this post is already quite long, I will leave the reimplementation as an exercise to the reader.

Depending on your needs, you might also go down with the floating point precision from 64bit to 32bit on addition to all the previously mentioned optimizations.

10 loops, best of 5: 1.49 s per loop (~41x)
10 loops, best of 5: 1.14 s per loop (in grayscale, ~54x)