submit to reddit
Siu Kwan Lam

Simple Wave Simulation with Numba and PyGame

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:

def physics(masspoints, dt, plunk, which):
ppos = masspoints[1]
cpos = masspoints[0]
N = cpos.shape[0]
# apply hooke's law
force = np.zeros((N, 2), dtype=cpos.dtype)
for i in range(1, N):
dx, dy = cpos[i] - cpos[i - 1]
dist = np.sqrt(dx**2 + dy**2)
assert dist != 0
fmag = -HOOKE_K * dist
cosine = dx / dist
sine = dy / dist
fvec = np.array([fmag * cosine, fmag * sine])
force[i - 1] -= fvec
force[i] += fvec
 
force[0] = force[-1] = 0, 0
force[which][1] += plunk
accel = force / MASS
 
# verlet integration
npos = (2 - DAMPING) * cpos - (1 - DAMPING) * ppos + accel * (dt**2)
 
masspoints[1] = cpos
masspoints[0] = npos

Python

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 jit.

The Numba version of the code is shown below:

@jit(argtypes=[f8[:,:], f8[:,:]])
def hooke(cpos, force):
N = cpos.shape[0]
for i in range(1, N):
dx = cpos[i, 0] - cpos[i - 1, 0]
dy = cpos[i, 1] - cpos[i - 1, 1]
dist = np.sqrt(dx**2 + dy**2)
fmag = -1 * HOOKE_K * dist
 
cosine = dx / dist
sine = dy / dist
fx = fmag * cosine
fy = fmag * sine
 
force[i - 1, 0] -= fx
force[i - 1, 1] -= fy
force[i, 0] += fx
force[i, 1] += fy
 
@jit(argtypes=[f8[:,:], f8[:,:], f8[:,:], f8[:,:], f8])
def verlet_integrate(npos, cpos, ppos, force, dt):
N = cpos.shape[0]
for i in range(N):
ax = force[i, 0] / MASS
ay = force[i, 1] / MASS
npos[i, 0] = (2 - DAMPING) * cpos[i, 0] - (1 - DAMPING) * ppos[i, 0] + ax * (dt**2)
npos[i, 1] = (2 - DAMPING) * cpos[i, 1] - (1 - DAMPING) * ppos[i, 1] + ay * (dt**2)
 
def physics(masspoints, dt, plunk, which):
ppos = masspoints[1]
cpos = masspoints[0]
N = cpos.shape[0]
 
# apply hooke's law
force = np.zeros((N, 2), dtype=cpos.dtype)
hooke(cpos, force)
 
force[0] = force[-1] = 0, 0
force[which][1] += plunk
 
# verlet integration
npos = np.empty((N, 2), dtype=cpos.dtype)
verlet_integrate(npos, cpos, ppos, force, dt)
 
masspoints[1] = cpos
masspoints[0] = npos

Python

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.

Benchmark

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!

Conclusion

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.

Tags: NumbaPro Accelerate
submit to reddit
comments powered by Disqus