Matplotlib julia animation

julia.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

resolution = 400
exponent = 2
iterations = 400


def julia(exponent, constant, size, iterations):
    (x_min, x_max) = (-2, 2)
    (y_min, y_max) = (-2, 2)
    X, Y = np.meshgrid(np.linspace(x_min, x_max, size),
                       np.linspace(y_min, y_max, size))
    Z = X+1j*Y
    diverged = np.zeros((size, size), dtype=np.int32)
    for i in range(iterations):
        Z[np.where(diverged == 0)] = np.power(Z[np.where(diverged == 0)],
                                              exponent)+constant
        diverged[np.where((abs(Z) > 2) & (diverged == 0))] = i+1
    return diverged


def update(*args):
    im.set_data(julia(exponent, np.exp(1j*args[0]), resolution, iterations))
    return im,


if __name__ == '__main__':
    writer = animation.writers['ffmpeg']
    writer = writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)
    a = julia(exponent, 1, resolution, iterations)
    fig = plt.figure()
    plt.axis('off')
    im = plt.imshow(a, animated=True)
    ani = animation.FuncAnimation(fig, update, frames=np.linspace(0, 2*np.pi,
                                                                  100),
                                  interval=50)
    ani.save('julia.mp4', writer=writer)