mandelbrot.py

# -*- coding: utf-8 -*-
import numpy as np
from sys import argv
from matplotlib import pyplot as plt
from matplotlib import colors

def mandelbrot(c, maxiter, h, log_log_h):
    output = np.empty(c.shape, np.float64)
    z = np.empty(c.shape, np.complex64)
    az = np.zeros(c.shape, np.complex64)
    notdone = np.ones(c.shape, np.bool)

    for n in range(maxiter):
        notdone[notdone] = np.less_equal(z.real*z.real + z.imag*z.imag, h)[notdone]
        z[notdone] = z[notdone]**2 + c[notdone]
        az[notdone] = np.abs(z[notdone])
        output[notdone] = n - np.log(np.log(az[notdone])).real / np.log(2) + log_log_h

    return output

def mandelbrot_set(xmin, xmax, ymin, ymax, width, height, maxiter):
    r1 = np.linspace(xmin, xmax, width, dtype=np.float64)
    r2 = np.linspace(ymin, ymax, height, dtype=np.float64)
    c = r1 + r2[:, None]*1j

    return mandelbrot(c, maxiter, 2.0**40, np.log(np.log(2.0**40))/np.log(2))

def mandelbrot_image(xmin, xmax, ymin, ymax, width, height, maxiter):
    dpi = 72
    img_width = width * dpi
    img_height = height * dpi

    fig, ax = plt.subplots(figsize=(width, height), dpi=dpi)
    ticks = np.arange(0, img_width, 3 * dpi)
    x_ticks = xmin + (xmax-xmin)*ticks / img_width
    y_ticks = ymin + (ymax-ymin)*ticks / img_width
    plt.xticks(ticks, x_ticks)
    plt.yticks(ticks, y_ticks)

    ax.imshow(
        mandelbrot_set(xmin, xmax, ymin, ymax, img_width, img_height, maxiter),
        cmap='magma',
        origin='lower',
        norm=colors.PowerNorm(0.3)
    )

    fig.savefig("mandelbrot.png")

# Run with e.g.
# $ python mandelbrot.py -2.0 0.5 -1.25 1.25 10 10 80
mandelbrot_image(float(argv[1]), float(argv[2]), float(argv[3]), float(argv[4]), int(argv[5]), int(argv[6]), int(argv[7]))