Fun with Numba, NumPy and F2PY

I've started using Numba to speed up MD simulations for our course in Molecular Statistics, which I teach together with +Jan Jensen and +Jimmy Charnley Kromann.

One exercise is using the Velo-Verlet algorithm to simulate a Lennard-Jones gas/liquid with periodic boundary conditions. Since we do everything in Python in this course, running the actual simulation is quite slow. What we do is we supply the students with a compiled FORTRAN module, compiled with F2PY which they can use when they've written their own Lennard-Jones gradient code.

The FORmula TRANslator module is extremely fast, compared to Python code. I rewrote the code from last year's course to a function I could use with Numba and compare directly to the F2PY module

For 100 particles:

• Numba: 0.146 ms/iteration (466x speed up)
• F2PY: 0.236 ms/iteration (289x speed up)
• Python: 68.06 ms/iteration

Now, I'm pretty sure a pure FORTRAN implementation would be faster than the F2PY, but I am very impressed with the Numba, since it's standard Python code and MUCH more simple to read than FORTRAN -- in theory all you have to do is write @autojit before a function definition. This is of course in theory ... I didn't do anything to improve on the readability of the Lennard-Jones code here. Before I posted, I had actually inserted a comment saying """This code block is unreadable""".

Now, there were a few issues I discovered:


1) @autojit vs. @jit(argtypes=[double[:,:], double[:]])

Since Numba has to decide what arguments a function takes before it's being compiled, it needs to know what possible type of arguments the function takes. If you're lucky, Numba decides on the right type and you can get away with @autojit, but sometimes autojitting makes the code SLOWER than standard, interpreted Python. In one case I got a factor of 10x slower with autojit, but explicitly stating the function argument types with @jit(argtypes = .... ) I got a speed up of 20x on the same code compared to interpreted Python.


2) Returning tuples in compiled code block

Numba does not allow a tuple to be returned inside a compiled code block. So don't do this.


3) Use Numpy properly

• Numba is NumPy aware, so code ran faster when numbers were stored in numpy.array types rather than just regular Python lists.
• At first I was numpy.round() to round to nearest integer in the periodic boundary condition code. Switching three numpy.round() calls to numpy.rint() gave a speed up of around 100x on the code execution.
• Use U[i, j] instead of U[i][j] on NumPy arrays. Not much of a difference in the vanilla Python code, but MASSIVE speed gains in compiled code.

Python Lennard-Jones code:

	import numpy
	from numba import double, jit
	rc2 = 6.25
	rc2i=1.0/rc2
	rc6i=rc2i*rc2i*rc2i
	ecut=rc6i*(rc6i-1.0)
	@jit(argtypes=[double[:,:], double[:]])
	def lennardjones(U, box):
	# Can't use (ndim, npart) = numpy.shape(U)
	# with Numba. No unpacking of tuples.
	ndim = len(U)
	npart = len(U[0])
	F = numpy.zeros((ndim, npart))
	Epot = 0.0
	Vir = 0.0
	for i in range(npart):
	for j in range(npart):
	if i > j:
	X = U[0, j] - U[0, i]
	Y = U[1, j] - U[1, i]
	Z = U[2, j] - U[2, i]
	# Periodic boundary condition
	X -= box[0] * numpy.rint(X/box[0])
	Y -= box[1] * numpy.rint(Y/box[1])
	Z -= box[2] * numpy.rint(Z/box[2])
	# Distance squared
	r2 = X*X + Y*Y + Z*Z
	if(r2 < rc2):
	r2i = 1.0 / r2
	r6i = r2i*r2i*r2i
	Epot = Epot + r6i*(r6i-1.0) - ecut
	ftmp = 48. * r6i*(r6i- 0.5) * r2i
	F[0, i] -= ftmp * X
	F[1, i] -= ftmp * Y
	F[2, i] -= ftmp * Z
	F[0, j] += ftmp * X
	F[1, j] += ftmp * Y
	F[2, j] += ftmp * Z
	Vir += ftmp
	Epot = Epot * 4.0
	return Epot, F, Vir

view raw lj.py hosted with ❤ by GitHub

F2PY Lennard-Jones code:

	C THE LENNARD-JONES POTENTIAL, GRADIENT AND THE VIRIAL
	subroutine LennardJones(U,box,Epot,F,Vir,N)
	cf2py intent(in) U
	cf2py intent(in) box
	cf2py intent(out) Epot
	cf2py intent(out) F
	cf2py intent(out) Vir
	implicit none
	double precision U,Epot,F,box,Vir
	integer N
	dimension U(3,N),F(3,N),box(3)
	C
	double precision zero,one,X,Y,Z,r2,r2i,r6i,feight,half,
	* ftmp,four,rc2,rc2i,rc6i,ecut
	integer i,j
	C
	parameter(zero=0.0D+00,one=1.0D+00,feight=48.0D+00,
	* half=0.5D+00,four=4.0D+00,rc2=(6.25D+00),
	* rc2i=one/rc2,rc6i=rc2i*rc2i*rc2i,ecut=rc6i*(rc6i-one))
	Epot = zero
	Vir = zero
	do i=1,N
	F(1,i) = zero
	F(2,i) = zero
	F(3,i) = zero
	enddo
	do i=1,N-1
	do j=i+1,N
	X = U(1,j) - U(1,i)
	Y = U(2,j) - U(2,i)
	Z = U(3,j) - U(3,i)
	X = X - box(1)*nint(X/box(1))
	Y = Y - box(2)*nint(Y/box(2))
	Z = Z - box(3)*nint(Z/box(3))
	r2 = X*X + Y*Y + Z*Z
	if(r2.lt.rc2) then
	r2i = one / r2
	r6i = r2i*r2i*r2i
	Epot = Epot + r6i*(r6i-one) - ecut
	C WE MULTIPLY BY 48 HERE AND NOT LAST.
	ftmp = feight*r6i*(r6i-half)
	F(1,i) = F(1,i) - ftmp*X*r2i
	F(2,i) = F(2,i) - ftmp*Y*r2i
	F(3,i) = F(3,i) - ftmp*Z*r2i
	F(1,j) = F(1,j) + ftmp*X*r2i
	F(2,j) = F(2,j) + ftmp*Y*r2i
	F(3,j) = F(3,j) + ftmp*Z*r2i
	Vir = Vir + ftmp
	endif
	enddo
	enddo
	Epot = Epot * four
	return
	end

view raw LJ.f hosted with ❤ by GitHub

Python Velo-Verlet solver:

	import numpy
	from forces import lennardjones
	# Initialize simulations (code not shown)
	temperature = 2.0
	n_atoms = 100
	rho = 0.01
	U = initialize_positions(n_atoms, rho)
	box = initialize_box(n_atoms, rho)
	V = initialize_velocities(U, temperature)
	# Start simulation
	n_steps = 10000
	dt = 0.001
	(epot,F,Vir) = lennardjones(U,box)
	for i in range(n_steps):
	U += V * dt + 0.5 * F * dt * dt
	F0 = F[:]
	(epot, F, Vir) = lennardjones(U, box)
	V += 0.5 * (F + F0) * dt

view raw Md_velo_verlet.py hosted with ❤ by GitHub