Main menu

Computing the Mandelbrot Set using NumPy Array Operations

If you use NumPy for numerical computations, which I recommend if you program in Python, one of the key things to do is to process entire arrays in one go as much as possible (this is also true for MATLAB). Using these so-called vectorized operations makes sure that the operation is run using compiled code instead of interpreted Python code, since the array and matrix operations behind NumPy (and MATLAB) are implemented in compiled languages such as C.

As an example, I’ll create a Python program that produces the classical Mandelbrot fractal. This fractal is computed by defining, for each \(c\in\mathbb{C}\), a polynomial \(f_c(z)=z^2+c\). This polynomial is then iterated starting with \(z=0\), that is, \(f_c^n(0)\) is computed, and any point \(c\) for which \(f_c^n(0)\) does not “escape to infinity” for \(n\to\infty\) is part of the Mandelbrot set (there are some example values of \(c\) in the article Mandelbrot Set).

In order to compute this polynomial for all pixels of an image at the same time, we create matrices \({\bf Z}\) and \({\bf C}\), and then do the computation

\[{\bf Z}^{(t+1)}={\bf Z}^{(t)}*\,{\bf Z}^{(t)}+\,{\bf C},\]

where \(*\) means element-wise matrix multiplication. The matrix \({\bf Z}^{(0)}\) must be initialized to all zeros, so that’s easy.

The Matrix C

Each element of the matrix \({\bf C}\) must be initialized to the coordinate of the pixel in the complex plane that it represents, i.e., the pixel \((x,y)\) must get the value \(x+yi\). This makes sure that the correct function \(f_c\) is computed at each point. A straightforward way to do that is to create a row of \(X\) coordinates and a column of \(Y\) coordinates. For example, for the points with integer coordinates in the range \([-2,2]\), this would be

\[\begin{pmatrix}
-2 & -1 & 0 & 1 & 2
\end{pmatrix}\]

and

\[\begin{pmatrix}
2 \\
1 \\
0 \\
-1 \\
-2
\end{pmatrix}.\]

By replicating both to a full \(5\times5\) matrix and adding them, with the second one multiplied by \(i\), we get

In practice, a basic plot of the Mandelbrot fractal can be made on a grid that extends from \(-2\) to \(1\) in the \(X\) direction, and \(-1\) to \(1\) in the \(Y\) direction. The magnitude of the elements of a \(480\times320\) matrix with this range is shown in Figure 1.

This code literally translates the vector and matrix operations from the example above. There are alternatives such as mgrid that could be used, but I think that the construction using tile explicitly is clearer in this case.

Iterating the Polynomial

We now have the matrices \({\bf Z}\) and \({\bf C}\) ready to start iterating. This is where it gets a bit tricky. The basic code is completely trivial:

Z = np.zeros((n, m), dtype=complex)for i inrange(100):
Z = Z * Z + C

However, there’s a problem. The values that “escape to infinity” grow so quickly that they overflow the maximum float value in no time, which first results in Inf values and then quickly in NaN values. To avoid this, I’ll add a mask M of pixels that are potentially in the Mandelbrot set, but from which we will remove pixels if we discover that they are not. Mathematically, it can be proven that pixels values with a magnitude larger than 2 will escape and cannot be part of the set. Hence, the code can be adapted as follows.

There are two important constructs here. The first one is the notation Z[M]. This is so-called boolean indexing. It selects the elements of the matrix Z for which M contains True. This enables you to only continue iterating on those pixels that have not escaped yet. The second line updates M itself. After each iteration, we determine all pixels that have escaped, using the expression np.abs(Z) > 2. We then use boolean indexing again to remove these pixels from M, using the expression M[np.abs(Z) > 2] = False.

Hence, using array operations allows computing the Mandelbrot set in only a few lines of code, and with much better performance than by iterating over all pixels with Python for loops. Figure 2 shows the result of running the above code (so, for 100 iterations) and then plotting M.

Figure 2. Approximation of the Mandelbrot set after 100 iterations.

This black-and-white image shows the Mandelbrot set in black. For more information on the difference between this and the usual colorful images that you possibly expected here, see again the article Mandelbrot Set.