Numba version of inner force loop for the nsquared version of CoMD in Python

ljforce.py

import constants
import numba

class LJ_Pot(object):
    def __init__(self):
        self.sigma = 2.315
        self.epsilon = 0.167
        self.mass = 63.55 * constants.amuToInternalMass
        self.lat = 3.615
        self.lattice_type = 'FCC'
        self.cutoff = 2.5*self.sigma
        self.name = "Cu"
        self.atomic_no = 29

    def output(self):
        pass

    @numba.jit
    def computeForce(self, atoms):
        f = atoms.f
        r = atoms.r
        bounds = atoms.bounds
        nAtoms = atoms.nAtoms
        epsilon = self.epsilon

        POT_SHIFT = 1.0
        #POT_SHIFT = 0.0
        rCut2 = self.cutoff * self.cutoff

        # Options for zeroing the force array

        # slowest
        for i in range(nAtoms):
            f[i,:] = 0.0

        # works better
        #for i in range(nAtoms):
        #    for j in range(3):
        #        f[i,j] = 0.0

        # most compact
        #f[:,:] = 0.0


        ePot = 0.0

        s6 = self.sigma * self.sigma * self.sigma * self.sigma * self.sigma * self.sigma

        rCut6 = s6/(rCut2*rCut2*rCut2)
        eShift = POT_SHIFT * rCut6 * (rCut6 - 1.0)

        # loop over atoms

        iPot = 0
        for i in range(nAtoms):
            for j in range(nAtoms):
                if i < j:
                    r2 = 0.0
                    dx = r[i,0] - r[j,0]
                    if dx < -0.5*bounds[0]:
                        dx += bounds[0]
                    if dx > 0.5*bounds[0]:
                        dx -= bounds[0]
                    r2 += dx*dx

                    dy = r[i,1] - r[j,1]
                    if dy < -0.5*bounds[1]:
                        dy += bounds[1]
                    if dy > 0.5*bounds[1]:
                        dy -= bounds[1]
                    r2 += dy*dy

                    dz = r[i,2] - r[j,2]
                    if dz < -0.5*bounds[2]:
                        dz += bounds[2]
                    if dz > 0.5*bounds[2]:
                        dz -= bounds[2]
                    r2 += dz*dz

                    if r2 <= rCut2:

                        ir2 = 1.0/r2

                        r6 = s6 * (ir2*ir2*ir2)
                        eLocal = r6*(r6-1.0) - eShift

                        ePot += eLocal
                        iPot += 1


                        fr = -4.0*epsilon * r6 * ir2*(12.0*r6 - 6.0)

                        f[i,0] -= dx*fr
                        f[j,0] += dx*fr
                        f[i,1] -= dy*fr
                        f[j,1] += dy*fr
                        f[i,2] -= dz*fr
                        f[j,2] += dz*fr

        ePot = ePot*4.0*epsilon
        return ePot