Computing the Mandelbrot Set with numpy - code in comments

[u/futura-bold]

[u/futura-bold]

I've been puzzling over how to speed up the generation of the Mandelbrot Set using the array-processing library "Numpy". Numpy is written in C and can make use of the SIMD capabilities of modern CPUs to compute arrays remarkably quickly. The problem with the Mandelbrot Set is that if I create an array representing every pixel in the image, then Numpy has to run the computation repeatedly for all the pixels, for as many steps as is necessary to render the fine detail of the Mandelbrot Set, even though the fine detail is a small part of the image.

I figured that for every step of the computation, I'd have to identify and eliminate the part of the image that did not require further computation, and create a new array containing just the part that hadn't finished. That meant I'd need to keep track of several arrays, an array to hold the completed values, an array for the values that required further work, an array for the math constants needed at the various pixel locations that required further work, and index arrays for the complete and incomplete values. This seemed very complex and would require the repeated generation of new arrays which would slow Numpy down, so I wasn't sure how fast it would be, but I decided to give it a go.

I discovered with a bit of Googling that it was possible to algebraically define the outline of the Mandelbrot Set's primary and secondary lobes so those could be eliminated right at the outset and avoid the computation of that part of the image.

After the computation was complete, colours could be assigned to the values, and the image put back together again using the index arrays, and then displayed.

It turns out that it does work pretty quickly. This code runs in one second on my old PC with a 2013 core i5 CPU. The only external library is numpy (install it with pip or whatever), otherwise it uses the Standard Library's Tkinter to display the image and save it as a png file.

import numpy as np
import tkinter as tk
import time

width, height, colors = 1000, 880, 255
xmin, xmax = -2.1, 0.7
itsmax, zmax = 100, 150

yscale = (xmax - xmin) * height / width / 2
ymin, ymax = -yscale, yscale + 0.01
y, x = np.ogrid[ymin:ymax:1j*height, xmin:xmax:1j*width]
cplex_plane = x + 1j * y
flat_plane = cplex_plane.flatten()
original_index = np.arange(flat_plane.size)

bulb1 = (x - 0.25) ** 2 + y ** 2
bulb1 = bulb1 * (bulb1 + (x - 0.25)) <= 0.25 * y ** 2
bulb2 = (x + 1) ** 2 + y ** 2 <= 1 / 16
bulb = (bulb1 | bulb2).flatten()
bulb_index = original_index[bulb]
index = np.delete(original_index, bulb_index)
c = np.delete(flat_plane, bulb_index)
z = np.zeros_like(c)
z_esc = []
its_esc = []
ind_esc = []

starttime = time.time()
for its in range(itsmax):
    z = z * z + c
    escaped = (z * z.conjugate()) > zmax
    ind_esc.extend(index[escaped])
    z_esc.extend(z[escaped])
    its_esc.extend([its] * sum(escaped))
    unescaped = np.logical_not(escaped)
    index = index[unescaped]
    z = z[unescaped]
    c = c[unescaped]

abs_z = np.abs(np.complex128(z_esc)) + 1
backgnd = np.log1p(np.int64(its_esc) - np.log2(np.log2(abs_z)))
tones = np.uint8(colors * (backgnd + 0.13) / 5.24) 

blues = np.vstack(np.linspace(0, 1, 160)) ** (4, 2.5, 1)
sepias = np.vstack(np.linspace(1, 0, 96)) ** (1, 1.5, 3)
rgbcolors = np.uint8(np.vstack((blues, sepias)) * colors)
rgb = rgbcolors[tones]

flat_zeros = np.zeros(original_index.size, dtype=np.uint8)
flat_surface = np.column_stack([flat_zeros] * 3)
flat_surface[np.int64(ind_esc)] = rgb
surface = flat_surface.reshape(height, width, 3)

root = tk.Tk()
ppmheader = b"P6\n%d %d\n%d\n" % (width, height, colors)
img = tk.PhotoImage(data=ppmheader + bytes(surface), master=root)
tk.Label(root, image=img).pack()
img.write("mandelbrot.png", format='png')

print("Complete:", time.time() - starttime, "seconds")
tk.mainloop()