Fractal Generation#
According to wikipedia,
A fractal is a geometric shape containing detailed structure at arbitrarily small scales…
Julia Sets We’ll focus on Julia Sets, a type of fractal made using complex numbers. One Julia set is the set of all complex numbers \(z\) for a given \(c = a + bi\), such that the sequence \(z_{n+1} = z_n^2 + c\) remains bounded.
We define a fuction \(F(z, c)\) that updates the complex number \(z\) iteratively:
Initialization of the complex number variable \(z\).
Iteratively update the value of \(z\) based upon the function \(z_{n+1} = z_n^2 + c\).
Return the number of iterations necessary to determine whether \(z\) is bounded or unbounded (return a nan/0).
Often, we set a threshold to prevent infinite iteration, which can be one or both of
we surpass a value of \(z\) (in the examples below, iteration stops when \(|z|>4\), and/or
we surpass a predefined number of iterations.
Based upon either method, \(z\) can be defined as bounded or unbounded (iteration trends towards infinity).
Visualisation for Julia Sets
To visualize the Julia Set Fractal, the initial \(z\) value can be defined as the location of a pixel in a 2-dimensional image: the real portion of the initial complex number the x pixel index, and the imaginary value the y pixel index (or vice-versa).
For each pixel, we can initialize \(z\) based upon its index and plug the values into \(z_{n+1} = z_n^2 + c\) to determine whether the result is bounded or unbounded :
(x,y) |
x1 |
x2 |
… |
width |
|---|---|---|---|---|
y1 |
F(x1 + y1*i) |
F(x2 + y1*i) |
… |
F(width + y1*i) |
y2 |
F(x1 + y2*i) |
F(x2 + y2*i) |
… |
F(width + y2*i) |
… |
… |
… |
… |
… |
height |
F(x1 + height*i) |
F(x2 + height*i) |
… |
F(width + height*i) |
Finally, one could color the image based upon resultant values:
Black for unbounded
Use colormap to display the normalized value of the resulting number of iterations (e.g. 0-255 for 8-bit color depth)
The final fractan can change based on
the constant \(c\),
the maximum number of iterations allowed
the resolution of the grid
the extents of the grid (e.g. \(x,y \in [-2, 2]\) versus \(x,y \in [-0.1, 0.1]\))
An interactive plotter can be found here.
CPU and Numpy Implementation#
import numpy as np
from time import process_time
RESOLUTION = 768
EPSILON = 1e-6
MESH_MIN, MESH_MAX = -2, 2
MESH_SIZE = MESH_MAX - MESH_MIN
MESH_RE, MESH_IM = np.meshgrid(
np.arange(MESH_MIN, MESH_MAX, (MESH_SIZE + EPSILON) / (RESOLUTION - 1)),
np.arange(MESH_MIN, MESH_MAX, (MESH_SIZE + EPSILON) / (RESOLUTION - 1))[::-1]
)
def compute_cpu_julia_grid(np_zmesh_re, np_zmesh_im, constant_real, constant_imag):
"""
Compute the Julia Set fractal using the given complex constant and a pre-computed mesh.
:param np_zmesh_re: Numpy array of real parts of the complex mesh
:param np_zmesh_im: Numpy array of imaginary parts of the complex mesh
:param constant_real: The real part of the complex constant 'c'
:param constant_imag: The imaginary part of the complex constant 'c'
:return: Numpy array representing the Julia Set fractal
"""
nr, nc = np_zmesh_re.shape
max_escape_iter = 1000
fractal_image = np.zeros((nr, nc))
for r in range(nr):
for c in range(nc):
a = np_zmesh_re[r, c]
b = np_zmesh_im[r, c]
temp_real, temp_imag = 0, 0
for iteration in range(1, max_escape_iter):
if a * a + b * b > 4.0:
break # its going to diverge
else:
temp_real = a * a - b * b + constant_real
temp_imag = 2 * a * b + constant_imag
a = temp_real
b = temp_imag
fractal_image[r, c] = np.log2(float(iteration)) / np.log2(max_escape_iter)
return fractal_image
t0 = process_time()
a, b = -0.8, 0.156
julia_fractal = compute_cpu_julia_grid(
MESH_RE, MESH_IM,
a, b
)
print(f"CPU Runtime: {process_time() - t0:.2f}s")
import matplotlib.pyplot as plt
def plot_julia_grid(julia_fractal, c=None, colormap='magma'):
fig, ax = plt.subplots(figsize=(5, 5))
ax.imshow(julia_fractal, cmap=colormap)
# text in the top left corner
if c is not None:
ax.text(0.05, 0.95, f"c = {c[0]:.2f} + i{c[1]:.2f}", transform=ax.transAxes, ha='left', va='top', color='white', bbox=dict(boxstyle='round', facecolor='black', alpha=0.9), fontsize=10)
ax.axis('off')
# remove white space around the image
fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
return fig
fig = plot_julia_grid(julia_fractal, c=(a,b), colormap='magma')
GPU and Cupy Implementation#
Although we could do this using Jax/cupy’s pythonic interface, we’ll write an actual GPU kernel to do the computation (specific code for the GPU). This is a bit more technical, but it’s also more flexible and allows us to do things with the GPU (like accessing shared memory, etc).
See more details on the Cupy docs.
import cupy as cp
# Create a GPU/CuPy grid of complex values (`cp.*` instead of `np.*`)
GPU_MESH_RE, GPU_MESH_IM = cp.meshgrid(
cp.arange(MESH_MIN, MESH_MAX, (MESH_SIZE + EPSILON) / (RESOLUTION - 1), dtype=cp.float32),
cp.arange(MESH_MIN, MESH_MAX, (MESH_SIZE + EPSILON) / (RESOLUTION - 1), dtype=cp.float32)[::-1]
)
# The following code will be complied into a CUDA kernel and uploaded to the GPU
compute_gpu_julia_grid = cp.ElementwiseKernel(
'float32 complex_grid_re, float32 complex_grid_im, float32 constant_real, float32 constant_imag',
'float32 out',
'''
int time;
float temp_real = 0.0;
float temp_imag = 0.0;
float zn_real = complex_grid_re;
float zn_imag = complex_grid_im;
for(time = 1; time < 1000; time++)
{
if(zn_real * zn_real + zn_imag * zn_imag > 4.0)
break;
else
{
temp_real = zn_real * zn_real - zn_imag * zn_imag + constant_real;
temp_imag = 2 * zn_real * zn_imag + constant_imag;
zn_real = temp_real;
zn_imag = temp_imag;
}
}
out = log2f((float)time) / log2f(1000.0f);
''',
'gpu_znplusc'
)
Writing a kernel can be better sometimes as we won’t have to incur any costs of having the cde get compiled during execution, and we won’t have to deal with transferring data/memory during computations (only at the end of the computations).
t0 = process_time()
out = compute_gpu_julia_grid(GPU_MESH_RE, GPU_MESH_IM, a, b)
print(f"GPU Runtime: {process_time() - t0:.2f}s")
c_vals = [
(-0.8, 0.156),
(-0.2, -0.4),
(-0.4, 0.6),
( 0.4, -0.1),
(0, 0.7885),
(0.335, 0.335)
]
for i, (re, im) in enumerate(c_vals):
out = compute_gpu_julia_grid(GPU_MESH_RE, GPU_MESH_IM, re, im)
fig = plot_julia_grid(out.get(), c=(re, im), colormap='magma')
Plot runtimes#
from tqdm.auto import tqdm
from time import process_time
def collect_cpu_runtimes(resolution_sizes, a=-0.8, b=0.156):
runtimes = []
for res in tqdm(resolution_sizes, desc='CPU'):
mesh_re, mesh_im = np.meshgrid(
np.arange(MESH_MIN, MESH_MAX, (MESH_SIZE + EPSILON) / (res - 1)),
np.arange(MESH_MIN, MESH_MAX, (MESH_SIZE + EPSILON) / (res - 1))[::-1]
)
t0 = process_time()
compute_cpu_julia_grid(mesh_re, mesh_im, a, b)
runtimes.append(process_time() - t0)
return runtimes
def collect_gpu_runtimes(resolution_sizes, a=-0.8, b=0.156):
runtimes = []
for res in tqdm(resolution_sizes, desc='GPU'):
mesh_re, mesh_im = cp.meshgrid(
cp.arange(MESH_MIN, MESH_MAX, (MESH_SIZE + EPSILON) / (res - 1), dtype=cp.float32),
cp.arange(MESH_MIN, MESH_MAX, (MESH_SIZE + EPSILON) / (res - 1), dtype=cp.float32)[::-1]
)
t0 = process_time()
compute_gpu_julia_grid(mesh_re, mesh_im, a, b)
runtimes.append(process_time() - t0)
return runtimes
resolution_sizes = [32, 64, 128, 256, 512]
cpu_runtimes = collect_cpu_runtimes(resolution_sizes)
gpu_runtimes = collect_gpu_runtimes(resolution_sizes)
import matplotlib.pyplot as plt
def plot_runtimes(resolution_sizes, cpu_runtimes, gpu_runtimes):
fig, ax = plt.subplots(figsize=(5, 5))
ax.plot(resolution_sizes, cpu_runtimes, label='CPU')
ax.plot(resolution_sizes, gpu_runtimes, label='GPU')
ax.set_xlabel('Resolution')
ax.set_ylabel('Runtime (s)')
ax.set_yscale('log')
ax.set_xscale('log')
ax.legend(fontsize=14)
return fig
fig = plot_runtimes(resolution_sizes, cpu_runtimes, gpu_runtimes)
fig.savefig("runtimes_fractal.png")
Animation#
We can make a small animation by allowing \(c\) to change as:
for \(\theta \in [0, 2\pi]\) and given some selected constant \(|c|\) (for example \(|c| = 0.7885\)).
(you could also try animating by zooming in/out on the fractal by changing the extents of the grid)
import imageio
def make_julia_gif(c_mag, outname='julia.gif'):
theta_vals = np.linspace(0, 2*np.pi, 100)
with imageio.get_writer(outname, mode='I', loop=0) as writer:
for theta in tqdm(theta_vals, desc="Generating GIF"):
a = c_mag * np.cos(theta)
b = c_mag * np.sin(theta)
out = compute_gpu_julia_grid(GPU_MESH_RE, GPU_MESH_IM, a, b)
fig = plot_julia_grid(out.get(), colormap='magma')
fig.savefig('temp.png', dpi=100)
writer.append_data(imageio.imread_v2('temp.png'))
plt.close(fig)
print(f"Saved GIF to {outname}")
make_julia_gif(0.7885, outname='julia_7885.gif')
make_julia_gif(0.335, outname='julia_335.gif')
Questions#
What is the GPU doing that makes it faster than the CPU?
Why is the GPU implementation getting slower as the resolution increases?
What will happen if the grid-size is increased?
What are some differences between the CPU and GPU implementations?
Can you write this in pythonic Jax/cupy code? Is it faster?
Can you write a kernel that uses shared memory to speed up the computation?
I’ll release the answers after the workshop. (Please remind me if i don’t lol)
Share your plots/code/cool GIFs down below (in the website version of this page)! 😄