Mandelbrot rotate-zoom generator

mandelbrotzoom.py

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
import multiprocessing

num, den, dpi, fps, dur = 16, 9, 120, 30, 10
print("%dx%d@%d, %d frames" % (num * dpi, den * dpi, fps, fps * dur))

fig = plt.figure(figsize=(num, den))
ax = plt.axes([0, 0, 1, 1], frameon=False)
ax.set_xticks([])
ax.set_yticks([])
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
plt.autoscale(tight=True)

real, imag = np.meshgrid(np.linspace(-2 * num / den, 0, num * dpi),
                         np.linspace(-1, 1, den * dpi))
rotzoom = (np.cos(np.pi / 256) + 1j * np.sin(np.pi / 256)) * (1 - 1 / 128)


def compute_frame(i):
    grid = (real + 1j * imag) * rotzoom**i
    z = np.zeros_like(grid)
    p = np.zeros_like(grid, dtype=int)
    for _ in range(255):
        z2 = z**2
        p += (z2 < 4).astype(int)
        z = z2 + grid
    return p


pool = multiprocessing.Pool(8)
ims = [[plt.imshow(p, animated=True, cmap='jet')]
       for p in pool.map(compute_frame, range(fps * dur))]
pool.close()
pool.join()

ani = animation.ArtistAnimation(fig,
                                ims,
                                interval=1000 / fps,
                                blit=True,
                                repeat_delay=1000)
writer = animation.FFMpegWriter(fps=fps,
                                codec='libx264',
                                extra_args=[
                                    '-pix_fmt', 'yuv420p', '-profile:v',
                                    'high', '-tune', 'animation', '-preset',
                                    'veryslow', '-crf', '10'
                                ])
ani.save('mandelbrotzoom.mp4', dpi=dpi, writer=writer)