Mandelbrot/Julia Set computation accelerated using NumPy and OpenCL C, and rendered to video using OpenCV

julia_set.py

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()