Skip to content

Instantly share code, notes, and snippets.

@bergkvist
Last active July 12, 2023 08:43
Show Gist options
  • Save bergkvist/494846d6ef2bd97e93bf552f2341dc7b to your computer and use it in GitHub Desktop.
Save bergkvist/494846d6ef2bd97e93bf552f2341dc7b to your computer and use it in GitHub Desktop.
Mandelbrot/Julia Set computation accelerated using NumPy and OpenCL C, and rendered to video using OpenCV
import numpy as np
import pyopencl as cl
import matplotlib.pyplot as plt
import cv2
W, H = (1500, 1500)
(cxmin, cxmax), (cymin, cymax) = ((-1.5, 1.5), (-1.5, 1.5))
duration, fps = (10, 25)
frame_count = int(fps * duration)
cx, cy = np.meshgrid(np.linspace(cxmin, cxmax, W, dtype=np.float32),
np.linspace(cymin, cymax, H, dtype=np.float32))
out = np.empty_like(cx, dtype=np.uint8)
ctx = cl.create_some_context()
cx_buffer = cl.Buffer(ctx, cl.mem_flags.READ_ONLY | cl.mem_flags.COPY_HOST_PTR, hostbuf=cx)
cy_buffer = cl.Buffer(ctx, cl.mem_flags.READ_ONLY | cl.mem_flags.COPY_HOST_PTR, hostbuf=cy)
out_buffer = cl.Buffer(ctx, cl.mem_flags.WRITE_ONLY, size=out.nbytes)
program = cl.Program(ctx, """
__kernel void mandelbrot(__global float* cx, __global float* cy, __global char *out, __global float *q) {
int i = get_global_id(0);
float x = cx[i];
float y = cy[i];
float tmp = 0;
int iterations;
for (iterations = 0; iterations < 50; ++iterations) {
tmp = x*x - y*y + 0.0;
y = 2*x*y + q[0];
x = tmp;
if (x*x + y*y > 4.0) break;
}
out[i] = iterations;
}
""").build()
mandelbrot = program.mandelbrot
qs = np.linspace(0.0, 1.0, frame_count, dtype=np.float32)
q_buffer = cl.Buffer(ctx, cl.mem_flags.READ_ONLY | cl.mem_flags.COPY_HOST_PTR, hostbuf=qs[[3*len(qs)//4]])
with cl.CommandQueue(ctx) as queue:
mandelbrot(queue, (W*H,), None, cx_buffer, cy_buffer, out_buffer, q_buffer)
cl.enqueue_copy(queue, out, out_buffer)
plt.imshow(out, cmap='viridis', extent=(cxmin, cxmax, cymin, cymax))
video = cv2.VideoWriter('out.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, (W, H), False)
for i in range(fps * duration):
q_buffer = cl.Buffer(ctx, cl.mem_flags.READ_ONLY | cl.mem_flags.COPY_HOST_PTR, hostbuf=qs[[i]])
with cl.CommandQueue(ctx) as queue:
mandelbrot(queue, (W*H,), None, cx_buffer, cy_buffer, out_buffer, q_buffer)
cl.enqueue_copy(queue, out, out_buffer)
video.write(5*out)
video.release()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment