# Make sure @jupyter-widgets/jupyterlab-manager
# and jupyter-matplotlib
# are installed and enabled in the extension manager.
%matplotlib widget
import matplotlib.pyplot as plt
import numba
# Extra performance libraries
import numpy as np
# Plotting
from ipywidgets import interact
maxiterations = 50
# 300 x 400 matrix of complex numbers from [-1.5j, 1.5j] x [-2, 2]
c = np.sum(np.broadcast_arrays(*np.ogrid[-1.5j:1.5j:300j, -2:2:400j]), axis=0)
def fractal_numpy(c, maxiterations):
f = np.zeros_like(c, dtype=np.int32)
z = c.copy()
for i in range(1, maxiterations + 1):
z = z**2 + c # Compute z
diverge = abs(z**2) > 2**2 # Divergence criteria
z[diverge] = 2 # Keep number size small
f[~diverge] = i # Fill in non-diverged iteration number
return f
@numba.njit((numba.complex128[:, :], numba.int32))
def fractal_numba(c, maxiterations):
fractal = np.zeros_like(c, dtype=np.int32)
for yi in range(c.shape[0]):
for xi in range(c.shape[1]):
z = cxy = c[yi, xi]
for i in range(1, maxiterations + 1):
z = z**2 + cxy
if abs(z) > 2:
break
fractal[yi, xi] = i
return fractal
Change the numpy calculation to the numba one. Do you see a difference?
fig, ax = plt.subplots()
mesh = ax.imshow(c.real, vmin=0, vmax=1)
ax.set_xlabel("Re(x)")
ax.set_ylabel("Im(y)")
@interact(centerx=(-2.0, 2.0, 0.01), centery=(-2.0, 2.0, 0.1), scale=(-5.0, 2, 0.01))
def interactive_fractal(centerx=0.38, centery=-0.6, scale=0.25):
maxiterations = 50
scale = 10**scale
c = np.sum(
np.broadcast_arrays(
*np.ogrid[
(centery - scale) * 1j : (centery + scale) * 1j : 400j,
(centerx - scale) : (centerx + scale) : 400j,
]
),
axis=0,
)
f = fractal_numpy(c, maxiterations)
mesh.set_data(f / 50)
mesh.set_extent(
(centerx - scale, centerx + scale, centery - scale, centery + scale)
)