N-dimensional Mandelbrot sets in 100 lines of code

fractals.py

"""
This file shows how to generate n-dimensional mandelbrot sets with 
the clifford library. It uses the mathematics from the paper
Generating Fractals with Geometric Algebra by Rich Wareham and 
Joan Lasenby https://doi.org/10.1007/s00006-010-0265-1
No effort has been put into performance optimisation
"""

import numpy as np
import clifford


def generate_generalised_mandelbrot(bounds, nstep=100, nmax=1000):
    """
    Evaluate the generalised mandelbrot set within the bounds
    the dimension of the set is given by len(bounds).
    nstep controls how to quantise the space within the bounds
    nmax is the number of steps to take to check if a point is 
    in the set
    """

    # Generate a clifford algebra of the desired dimension
    ndims = len(bounds)
    layout, blades = clifford.Cl(ndims)

    # Pick the prefered directions
    er = blades['e1']
    ei = blades['e2']

    def comp_prod(a, b):
        """
        The equivalent complex product
        """
        return a*er*b

    def complex_func(x, c):
        """
        The complex function
        """
        return comp_prod(x,x) + c

    def generalised_mandelbrot(c, nmax):
        """
        Generates the generalised mandelbrot set 
        """
        x = 0*c
        for k in range(nmax):
            x = complex_func(x, c)
            if (x|x)[0] > 4:
                return k
        return nmax

    def ind_list_to_c(ind_list, step_size_list):
        """
        Generate the point associated with an ind_list
        """
        c = np.zeros(2**ndims)
        for i,v in enumerate(ind_list):
            c[1 + i] = (cmin[i] + step_size_list[i]*v)
        return layout.MultiVector(value=c)

    # Now evaluate on the domain
    print('Generating a ', ndims, 'dimensional Mandelbrot set')
    cmin = np.amin(bounds, axis=1)
    output = np.zeros([nstep]*ndims, dtype=int)
    step_size_list = [(b[1] - b[0])/nstep for b in bounds]
    n = 0
    for ind_list in np.ndindex(output.shape):
        c = ind_list_to_c(ind_list, step_size_list)
        output[ind_list] = generalised_mandelbrot(c, nmax)
        n += 1
        if n%10000 == 0:
            print(n, nstep**ndims)
    return output


def test_2d_image():
    """ Generate and plot a 2d Mandelbrot set """
    import matplotlib.pyplot as plt
    # from matplotlib import cm
    output = generate_generalised_mandelbrot([[-2, 2], [-2, 2]], nstep=200, nmax=100)
    plt.imshow(-output, cmap='Greens')
    plt.show()


def test_3d_voxels():
    """ Generate and plot a 3d Mandelbrot set """
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D 
    voxels = generate_generalised_mandelbrot([[-2, 2], [-2, 2], [-2, 2]], nstep=50, nmax=20)
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    ax.voxels(voxels==np.max(voxels), edgecolor='k')
    plt.show()


if __name__ == '__main__':
    test_2d_image()
    test_3d_voxels()