Numpy and the Mandelbrot Set
I’ve now rendered my first Mandelbrot set, which was unexpectedly easy.
The concept of how to render the Mandelbrot set really only crystallized in my mind when I saw the following image:
This is in fact the lead illustration on Wikipedia’s Complex Number page (redistributed here under the Creative Commons Attribution ShareAlike 3.0 License). To my mind it provides a clearer demonstration of how the Mandelbrot set is formed than anything on Wikipedia’s Mandelbrot Set page.
Points in the Mandelbrot set are just complex numbers with the real and imaginary parts mapped to the x and y axis respectively (known as the complex plane or Argand diagram). A simple repeated calculation then determines whether a given point is part of the Mandelbrot set:
That is, starting with z at the origin, square z and add a complex constant c. The process of squaring the result and then adding the constant c is repeated until it is apparent that the the point z is either orbiting close to the origin (therefore indicating that the constant c is a member of the Mandelbrot set) or diverging to infinity (c is not a part of the Mandelbrot set). Simply plot on the complex plane which values of c are included in this set and the classic Mandelbrot set fractal should appear before your eyes as if by magic.
The procedures for such operations as multiplying and adding complex numbers are quite straight forward. You can get the gist of the kind of things that may happen when multiplying by a complex number by examining the image above (The point rotates on the complex plane). However, it is even easier when programming in a language such as Python which has built-in support for complex numbers.
Consider the following interactive Python session
>>> a = (2 + 3j) >>> b = (1 - 8j) >>> a * b (26-13j) >>> z = a**2 + a >>> z = z**2 + a >>> z (-214-87j)
The first four lines assign two complex numbers a and b and then display their product. The next four lines show exactly the kind of calculation one would need to do to find the Mandelbrot set. Easy.
Another useful tool is Python’s built in abs function. For real numbers this returns the absolute value of that number. For a complex number however, abs returns the modulus of the number. That is, it’s distance from the origin on the complex plane.
Abs is extremely useful in calculating the Mandelbrot set, as any orbit which strays more than distance 2 from the origin is known to be not part of the Mandelbrot set. Further testing of the point c is unnecessary and the next point can then be tested.
My first Python script for rendering the Mandelbrot set simply looped through each pixel in the image one at a time. Each pixel is assigned a complex number from a 4 x 4 grid. That is, a number in the square bounded by (-2 – 2j) and (2 + 2j). For each pixel I then applied the “square and add a constant” dynamic for a certain number of iterations to determine if it is a member of the Mandelbrot set.
This technique works, but is quite slow. Here is the code in its entirety:
#! /usr/bin/python import png import numpy as np import datetime def mandelbrot_set(size, exit_limit): # initialize array img_array = np.zeros([size, size], int) for y in range(size): for x in range(size): c = complex(x / float(size) * 4 - 2, y / float(size) * 4 - 2) z = c for i in range(exit_limit): z = (z**2) + c img_array[y, x] += 1 if abs(z) > 2: # z is escaping to infinity, so point is not in set break else: # if loop is exausted, point is inside the set img_array[y, x] = 0 return img_array if __name__ == "__main__": start = datetime.datetime.now() # size of final image: size = 1600 # after this many iterations, point is assumed to be within the # mandelbrot set: exit_limit = 64 imgWriter = png.Writer(size, size, greyscale=True, alpha=False, bitdepth=16) img_array = mandelbrot_set(size, exit_limit) # normalise array (16 bit range) img_array = np.array(img_array/float(img_array.max())*(2**16), int) # save to final render to png file f = open("mandelbrot.png", "wb") imgWriter.write(f, img_array) f.close() print datetime.datetime.now()-start print "Done."
If you are wondering where the different shades of color come from, It it simply that the lighter shades take more iterations for z to escape to a distance of 2 from the origin. The black pixels in the middle of the shape indicate that z does not escape.
Now how could I speed this program up? Looping through each pixel is inherently slow. I decided to try again, this time using a Numpy array of complex numbers (yes, complex numbers are a valid data-type in Numpy!). Now I no longer need to loop through the pixel of my image, only through the iterations of the “square and add a constant” dynamic.
This reduces the rendering time to about a third of that needed for the pixel-looping solution above! The only drawback seems to be that this method can use a lot of memory while its running. This is a problem if you want to create a very large image. For example, attempting to render a 3200 x 3200 image results in an error on my computer just due to the vast size of the complex number array that is required. That problem could easily be solved by rendering the image in several smaller chunks that are more RAM friendly. The image array is okay at that size (it’s just integers after all), so there would be no need to be manually stitching the pieces together afterwards!
Here is the code for my Numpy solution:
import numpy as np import png from datetime import datetime start = datetime.now() size = 800 iterations = 63 # turn range(0, size**2) into an array of complex values encompassing the render area complex_val = lambda p: complex(p%size/float(size)*4-2, p/size/float(size)*4-2) c = np.array([complex_val(p) for p in range(size**2)])\ .reshape(size, size) # This array is used to mask out values that have already escaped: operation_mask = np.ones((size, size), np.bool) # This will become the final image: iteration_array = np.zeros((size, size), np.int32) # our starting "z" values. Note I don't start at '0' because 0**2 + c is clearly c. z = np.copy(c) for x in range(iterations): print "iteration %d" % x z[operation_mask] = (z[operation_mask]**2) + c[operation_mask] iteration_array[operation_mask] += 1 # remove values form the "operation mask" if their calculation has diverged operation_mask[operation_mask] = np.absolute(z[operation_mask]) < 2 # finally, paint points in the mandelbrot set black. iteration_array[operation_mask] = 0 imgwriter = png.Writer(size, size, greyscale=True, alpha=False, bitdepth=8) f = open('m.png', 'wb') imgwriter.write(f, iteration_array) f.close() print datetime.now() - start
Where to from here? There a few things I could try next:
- Smooth coloration – get rid of the color banding
- Interactive zooming and exploration, using PyGame perhaps
- Render the Buddhabrot attractor, because that looks cool
- More speed optimizations.