In the first post, I gave Claude Code my old black-hole research code and one prompt, and Claude organized it again.
I still wanted to improve the code, since it could take a lot of time to run for high quality images.
The biggest change: Apple Silicon GPU ray tracing in Python, with MLX and a custom Metal kernel.
The result: the black hole ray tracer runs 20 to 60 times faster than the numba code on 8 CPU cores. A 2048 × 2048 image, with more than 4 million light rays, takes 6 to 11 seconds on a MacBook.
Why run a black hole ray tracer on a GPU?
The code calculates the shadow of a black hole and the gravitational-lensing image of the sky behind it. It sends one light ray for each pixel. Each ray is calculated independently, so a GPU can trace thousands of rays at the same time.
But each ray takes a different time. The integrator (RKF45) uses an adaptive step: rays near the black hole take many small steps, and rays far away end quickly. With NumPy-style array operations, all rays wait for the slowest ray at each step.
So I did not use arrays. Each tracer got a custom Metal kernel, compiled at run time by MLX (mx.fast.metal_kernel). One GPU thread traces one ray and stops when its ray ends.
Note: Numba was not an option, since numba GPU support is for NVIDIA (CUDA) only.
A challenge to overcome: float32 on Apple Silicon
Apple GPUs do not support float64. The original physics code was written for float64, so the port needed some changes:
- a minimum step size and limits on steps and retries, so that no ray runs forever;
- analytic derivatives where a finite difference in float32 gives only noise;
- a CPU fallback: a ray that fails on the GPU is traced again on the CPU, in float64.
The CPU stays the reference. The GPU is the fast path.
How fast is the Apple Silicon GPU?
| Case | CPU (numba) | GPU (MLX) | Speedup |
|---|---|---|---|
| Schwarzschild with disk | 1.82 s | 0.079 s | 23× |
| Schwarzschild, no disk | 1.90 s | 0.063 s | 30× |
| Kerr, a = 0.9 | 0.81 s | 0.019 s | 43× |
64 × 64 pixels, median of 3 warm runs on Apple Silicon.
At large sizes the gain is what matters: a 2048 × 2048 Kerr image takes 6 to 11 seconds from start to end. That makes it easy to try new backgrounds, spins and observer positions.
Is the GPU result still correct?
Fast is not useful if it is wrong. Tests compare the GPU with the CPU: 0 of 4096 pixels changed class for Kerr and for Schwarzschild with a disk, and the median difference in the end coordinates is approximately 5e-4. A second test suite compares the shadows with the theory: the Schwarzschild edge agrees within half a pixel, and the traced Kerr edge agrees with the analytic curve.
Deep dive and resources
For the full technical details see the deep dive: Black Hole Shadows on the GPU: Apple Silicon Ray Tracing.
The code is open source on GitHub. The physics is in the paper on arXiv.
