Python Mandelbrot Set Batch Rendering

mandelbrot.py

# Copyright 2018 John E. Farrier

# References:

# https://gist.github.com/jfpuget/60e07a82dece69b011bb
# https://linas.org/art-gallery/escape/smooth.html
# https://matplotlib.org/devdocs/api/_as_gen/matplotlib.colors.LightSource.html
# https://matplotlib.org/examples/color/colormaps_reference.html
# https://matplotlib.org/examples/pylab_examples/shading_example.html
# https://matplotlib.org/examples/showcase/mandelbrot.html
# https://stackoverflow.com/questions/60208/replacements-for-switch-statement-in-python
# https://www.ibm.com/developerworks/community/blogs/jfp/entry/How_To_Compute_Mandelbrodt_Set_Quickly?lang=en
# https://www.ibm.com/developerworks/community/blogs/jfp/entry/My_Christmas_Gift

import numpy

import time

import matplotlib
from matplotlib import colors
import matplotlib.pyplot

from enum import Enum, auto

class MandelbrotType(Enum):
    Classic = auto()
    Detailed1 = auto()
    Detailed2 = auto()
    Detailed3 = auto()
    Detailed4 = auto()
    Detailed5 = auto()
    Detailed6 = auto()
    Detailed7 = auto()
    Detailed8 = auto()

# Returns xMin, xMax, yMin, yMax, width resolution, height resolution, and the maximum number of iterations.
def MandelbrotTypeValues(x):
   return {
        MandelbrotType.Classic: (-2.5, 1.0, -1.5, 1.5, 7000, 7000, 100),
        MandelbrotType.Detailed1: (-0.74897, -0.74892, 0.065053, 0.065103, 7000, 7000, 256),
        MandelbrotType.Detailed2: (-0.748977, -0.748927, 0.065053, 0.065103, 4000, 4000, 256),
        MandelbrotType.Detailed3: (-0.748977, -0.748927, 0.065043, 0.065093, 4000, 4000, 256),
        MandelbrotType.Detailed4: (-0.748977, -0.748927, 0.065043, 0.065093, 4000, 4000, 256),
        MandelbrotType.Detailed5: (-0.74897, -0.74892, 0.065053, 0.065103, 7000, 7000, 512),
        MandelbrotType.Detailed6: (-0.748977, -0.748927, 0.065053, 0.065103, 4000, 4000, 1024),
        MandelbrotType.Detailed7: (-0.748977, -0.748927, 0.065043, 0.065093, 4000, 4000, 512),
        MandelbrotType.Detailed8: (-0.748977, -0.748927, 0.065043, 0.065093, 4000, 4000, 1024),
    }[x]

def MandelbrotSet(xmin, xmax, ymin, ymax, width, height, maxiter, horizon=2.0):
    X = numpy.linspace(xmin, xmax, width, dtype=numpy.float64)
    Y = numpy.linspace(ymin, ymax, height, dtype=numpy.float64)
    C = X + Y[:, None]*1j
    N = numpy.zeros(C.shape, dtype=int)
    Z = numpy.zeros(C.shape, numpy.complex64)
    for n in range(maxiter):
        I = numpy.less(abs(Z), horizon)
        N[I] = n
        Z[I] = Z[I]**2 + C[I]
    N[N == maxiter-1] = 0
    return Z, N

if __name__ == '__main__':
    # Build all the images when we run.
    for mandelbrotImageType in MandelbrotType:
        timerStart = time.process_time()
        print("Computing %s..." % (mandelbrotImageType,))
       
        mandelbrotHorizon = 2.0 ** 40
        mandelbrotXMin, mandelbrotXMax, mandelbrotYMin, mandelbrotYMax, mandelbrotWidth, mandelbrotHeight, mandelbrotMaxIterations = MandelbrotTypeValues(mandelbrotImageType)
        
        imageDpi = 50
        imageWidth = 256
        imageHeight = 256*mandelbrotHeight/mandelbrotWidth
        
        mandelbrotLogHorizon = numpy.log(numpy.log(mandelbrotHorizon))/numpy.log(2)
        Z, N = MandelbrotSet(mandelbrotXMin, mandelbrotXMax, mandelbrotYMin, mandelbrotYMax, mandelbrotWidth, mandelbrotHeight, mandelbrotMaxIterations, mandelbrotHorizon)

        with numpy.errstate(invalid='ignore'):
            M = numpy.nan_to_num(N + 1 - numpy.log(numpy.log(abs(Z)))/numpy.log(2) + mandelbrotLogHorizon)

        fig = matplotlib.pyplot.figure(figsize=(imageWidth, imageHeight), dpi=imageDpi)
        axes = fig.add_axes([0.0, 0.0, 1.0, 1.0], frameon=False, aspect=1)

        # Shading
        
        # Where is the light source located?
        light = colors.LightSource(azdeg=220, altdeg=5)

        # How should shading be applied?  (jet, nipy_spectral, gist_ncar, terrain, ocean, viridis, cubehelix) (overlay, hsv)
        M = light.shade(M, cmap=matplotlib.pyplot.cm.nipy_spectral, vert_exag=1.5, norm=colors.PowerNorm(0.2), blend_mode='overlay')

        matplotlib.pyplot.imshow(M, extent=[mandelbrotXMin, mandelbrotXMax, mandelbrotYMin, mandelbrotYMax])
        axes.set_xticks([])
        axes.set_yticks([])

        text = ("Digital in Blue")
        axes.text(mandelbrotXMin + .025, mandelbrotYMin + .025, text, color="white", fontsize=10, alpha=0.25)

        print("    Saving Mandelbrot for %s..." % (mandelbrotImageType,))
        matplotlib.pyplot.savefig("dib_%s.png" % (mandelbrotImageType,))
        timerElapsed = time.process_time() - timerStart
        print("    Done [%.2f]" % (timerElapsed,))
        # matplotlib.pyplot.show()