Python is a great language for writing experimental code quickly. With PyGame, one can add graphics to visualize the experiments effortlessly. I always enjoy using PyGame for rendering simple physics simulation in realtime. However, as the simulation scales, Python can hardly keep up. In the past, I would have used hours to rewrite the heavy loops in C. Now, I can use Numba to “jit” the heavy loops.
I have a simple wave physics simulation that uses Hooke’s Law in a mass-spring system. A string of particles are connected together by mass-less springs. When plunked, the string would oscillate just like a guitar string. For readers who are interested in the full source code, it is available in GitHub. The simulation writes the oscillation data to a file. With another small script, the data is encoded into a .wav audio file. Here’s a sample in MP3 format. It sounds like plunking a rubber band.
At the core of the simulation, there is a physics function that computes the force on each particle using Hooke’s Law and updates the position using Verlet integration. Most of the execution time spends in this function. The code is shown below:
The code uses Numpy for efficient computation on arrays, but it is still not good enough. For merely 64 particles, the average framerate is just ~2 fps on my Core 2 Duo.
Speedup Simulation with Numba JIT
In the past, I would rewrite the physics function in C and use ctypes to interface it. It would take hours to get the job done. Now, I have Numba. All I need to do is to expand the Numpy array expressions into loops and decorate it with a
The Numba version of the code is shown below:
I have separated the Hooke’s Law portion and the Verlet integration portion into a new function, respectively. The simulation can now run at the maximum framerate.
We will run it in headless mode to benchmark the performance. The benchmark runs the simulation for 600 frames. The execution time is recorded for a 64 particles simulation and for a 200 particles simulation.
For 64 particles, the simulation is 24× faster with Numba. For 200 particles, the simulation is 55× faster with Numba. It is impressive!
We have seen how easy it is to use Numba to speedup Python for realtime physics simulation. The combination of Numba and PyGame makes a great tool for scientists to write and visualize realtime scientific simulation.