import numpy as np
import matplotlib.pyplot as plt

d, h = 800, 600  # pixel density (= image width) and image height
n, r = 200, 500  # number of iterations and escape radius (r > 2)

x = np.linspace(0, 2, num=d+1)
y = np.linspace(0, 2 * h / d, num=h+1)

A, B = np.meshgrid(x - 1, y - h / d)
C = 2.0 * (A + B * 1j) - 0.5

def iteration(C):
    S, T = np.zeros(C.shape), np.zeros(C.shape)
    Z, dZ = np.zeros_like(C), np.zeros_like(C)

    def iterate(C, S, T, Z, dZ):
        S, T = S + np.exp(- abs(Z)), T + 1
        Z, dZ = Z * Z + C, 2 * Z * dZ + 1
        return S, T, Z, dZ

    for i in range(n):
        M = abs(Z) < r
        S[M], T[M], Z[M], dZ[M] = iterate(C[M], S[M], T[M], Z[M], dZ[M])

    return S, T, Z, dZ

S, T, Z, dZ = iteration(C)
D = np.zeros(C.shape)

fig = plt.figure(figsize=(12.8, 9.6))
fig.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95)

ax1 = fig.add_subplot(2, 2, 1)
ax1.imshow(S ** 0.1, cmap=plt.cm.twilight_shifted, origin="lower")

N = abs(Z) >= r  # normalized iteration count
T[N] = T[N] - np.log2(np.log(abs(Z[N])) / np.log(r))

ax2 = fig.add_subplot(2, 2, 2)
ax2.imshow(T ** 0.1, cmap=plt.cm.twilight_shifted, origin="lower")

N = abs(Z) > 2  # exterior distance estimation
D[N] = np.log(abs(Z[N])) * abs(Z[N]) / abs(dZ[N])

ax3 = fig.add_subplot(2, 2, 3)
ax3.imshow(D ** 0.1, cmap=plt.cm.twilight_shifted, origin="lower")

N, thickness = D > 0, 0.01  # boundary detection
D[N] = np.maximum(1 - D[N] / thickness, 0)

ax4 = fig.add_subplot(2, 2, 4)
ax4.imshow(D ** 2.0, cmap=plt.cm.binary, origin="lower")

fig.savefig("Mandelbrot_numpy_set_1.png", dpi=200)