home

Generating the Mandelbrot set using Python

Generating the mandelbrot set in Python is surprisingly easy, here is my code:

import time

from PIL import Image
import numpy as np
from matplotlib import cm


def mandelbrot(c: complex, z: int = 0, max_iter: int = 100) -> int:
    """
    Returns the number of iterations needed for z to reach 2 or higher
    If 2 isn't reached in `max_iter` iterations returns 0
    """
    n = 0
    for _ in range(max_iter):
        if abs(z) >= 2:
            return n
        n += 1
        z = z ** 2 + c
    return 0


if __name__ == "__main__":
    # No checks are performed to ensure that the width and height matches the
    # x_range and y_range. Ensure manually to avoid non-quadratic pixels.
    width, height = 1000, 1000
    x_range, y_range = (-2.25, 0.75), (-1.5, 1.5)
    depth = 250
    # see https://matplotlib.org/stable/tutorials/colors/colormaps.html
    # for other colourmaps
    colourmap = cm.magma
    cm_repetition = depth  # colourmap will repeat after x iterations

    start = time.time()

    x_step = abs(x_range[1] - x_range[0]) / width
    x_coords = np.arange(*x_range, x_step)

    y_step = abs(y_range[1] - y_range[0]) / height
    y_coords = np.arange(*y_range, y_step)

    get_depth = lambda x, y: mandelbrot(complex(x, y), max_iter=depth)
    arr = np.array([[get_depth(x, y) for x in x_coords] for y in y_coords])

    img = Image.fromarray(np.uint8(colourmap((arr % cm_repetition) / cm_repetition) * 255))

    print("Rendering finished, took {:.2f} seconds".format(time.time() - start))

    img.show()
    img.save(r"./mandelbrot.png")

Now this code isn’t very fast (though faster than i initially expected). Just running the code as is will result in the following output (note both images have been losslessly compressed to save on bandwidth):

Render of Mandelbrot set

If you have a little patience and change some of the parameters around you can generate beautiful images like this one (the center of the swirl remains black because depite a depth of 5000 iterations):

Render of Zoomed in Mandelbrot showing an infinite swirl

All in all i’m very satisfied with this little script, allthough i might come back and make it faster by adding some multithreading in order to be able to recreate some of the really deep zooms.