Numba-optimzed version of CoMD cell version.

linkcell.py

from __future__ import print_function, division

import numpy as np
import math

import numba

def initLinkCells(sim, box_size):
    gridSize = np.array([int(b/sim.pot.cutoff) for b in box_size], dtype=np.int)
    boxSize = box_size*1.0/gridSize
    nLocalBoxes = gridSize[0]*gridSize[1]*gridSize[2]

    nHaloBoxes = 2*((gridSize[0] + 2) * (gridSize[1] + gridSize[2] + 2) + (gridSize[1] * gridSize[2]))

    nTotalBoxes = nLocalBoxes + nHaloBoxes

    return LinkCell(nLocalBoxes, nTotalBoxes, gridSize, 1.0/boxSize)

class LinkCell(object):
    def __init__(self, nLocalBoxes, nTotalBoxes, gridSize, invBoxSize):
        self.MAXATOMS = 64
        self.nLocalBoxes = nLocalBoxes
        self.nTotalBoxes = nTotalBoxes
        self.nAtoms = np.zeros(nTotalBoxes, dtype=np.int)
        self.gridSize = gridSize
        self.invBoxSize = invBoxSize

    @numba.jit
    def getBoxFromTuple(self, ix, iy, iz):
        iBox = 0
        # Halo in Z+
        if iz == self.gridSize[2]:
            iBox = self.nLocalBoxes + 2*self.gridSize[2]*self.gridSize[1] + 2*self.gridSize[2]*(self.gridSize[0]+2)  \
                   + (self.gridSize[0]+2)*(self.gridSize[1]+2) + (self.gridSize[0] + 2)*(iy+1) + (ix+1)

        # Halo in Z-
        elif iz == -1:
            iBox = self.nLocalBoxes + 2*self.gridSize[2]*self.gridSize[1] + 2*self.gridSize[2]*(self.gridSize[0]+2) \
                    + (self.gridSize[0]+2)*(iy+1) + (ix+1)

        # Halo in Y+
        elif iy == self.gridSize[1]:
            iBox = self.nLocalBoxes + 2*self.gridSize[2]*self.gridSize[1] + self.gridSize[2]*(self.gridSize[0]+2) \
                    + (self.gridSize[0]+2)*iz + (ix+1)

        # Halo in Y-
        elif iy == -1:
            iBox = self.nLocalBoxes + 2*self.gridSize[2]*self.gridSize[1] + iz*(self.gridSize[0] + 2) + (ix+1)

        # Halo in X+
        elif ix == self.gridSize[0]:
            iBox = self.nLocalBoxes + self.gridSize[1]*self.gridSize[2] + iz*self.gridSize[1] + iy

        # Halo in X-
        elif ix == -1:
            iBox = self.nLocalBoxes + iz*self.gridSize[1] + iy

        # Local link cell
        else:
            iBox = ix + self.gridSize[0]*iy + self.gridSize[0]*self.gridSize[1]*iz

        if iBox >= self.nTotalBoxes:
            print('iBox too large',iBox,self.nTotalBoxes)
            print('ix,iy,iz',ix,iy,iz)
        assert iBox >=0
        assert iBox < self.nTotalBoxes
        return iBox

    def getBoxFromCoord(self, r):
        ix = int(math.floor(r[0]*self.invBoxSize[0]))
        iy = int(math.floor(r[1]*self.invBoxSize[1]))
        iz = int(math.floor(r[2]*self.invBoxSize[2]))
        return self.getBoxFromTuple(ix, iy, iz)

    def putAtomInBox(self, atoms, gid, iType, x, y, z, px=0.0, py=0.0, pz=0.0):
        iBox = self.getBoxFromCoord([x, y, z])
        iOff = iBox*self.MAXATOMS
        iOff += self.nAtoms[iBox]

        if iBox < self.nLocalBoxes:
            atoms.nLocal += 1

        self.nAtoms[iBox] += 1
        atoms.gid[iOff] = gid
        atoms.species[iOff] = iType
        atoms.r[iOff][0] = x
        atoms.r[iOff][1] = y
        atoms.r[iOff][2] = z
        atoms.p[iOff][0] = px
        atoms.p[iOff][1] = py
        atoms.p[iOff][2] = pz

    @numba.jit
    def getTuple(self, iBox):
        ix,iy,iz = 0,0,0
        # local box
        if iBox < self.nLocalBoxes:
            ix = iBox % self.gridSize[0]
            iBox //= self.gridSize[0]
            iy = iBox % self.gridSize[1]
            iz = iBox // self.gridSize[1]
        else:   # halo box
            ink = iBox - self.nLocalBoxes
            if ink < 2*self.gridSize[1]*self.gridSize[2]:
                if ink < self.gridSize[1]*self.gridSize[2]:
                    ix = 0
                else:
                    ink -= self.gridSize[1]*self.gridSize[2]
                    ix = self.gridSize[0] + 1
                iy = 1 + ink % self.gridSize[1]
                iz = 1 + ink // self.gridSize[1]

            elif ink < 2*self.gridSize[2]*(self.gridSize[1] + self.gridSize[0] + 2):
                ink -= 2 * self.gridSize[2] * self.gridSize[1]
                if ink < (self.gridSize[0] + 2)*self.gridSize[2]:
                    iy = 0
                else:
                    ink -= (self.gridSize[0] + 2)*self.gridSize[2]
                    iy = self.gridSize[1] + 1
                ix = ink % (self.gridSize[0] + 2)
                iz = 1 + ink // (self.gridSize[0] + 2)
            else:
                ink -= 2*self.gridSize[2]*(self.gridSize[1] + self.gridSize[0] + 2)
                if ink < (self.gridSize[0] + 2)*(self.gridSize[1] + 2):
                    iz = 0
                else:
                    ink -= (self.gridSize[0] + 2)*(self.gridSize[1] + 2)
                    iz = self.gridSize[2] + 1
                ix = ink % (self.gridSize[0] + 2)
                iy = ink // (self.gridSize[0] + 2)

            ix -= 1
            iy -= 1
            iz -= 1

        return ix,iy,iz

    @numba.jit
    def getNeighborBoxes(self, iBox, nbrBoxes):
        #ix,iy,iz = self.getTuple(iBox)
        #count = 0
        #for i in range(ix-1, ix+2):
        #    for j in range(iy-1, iy+2):
        #        for k in range(iz-1, iz+2):
        #            nbrBoxes[count] = self.getBoxFromTuple(i, j, k)
        #            count += 1
        #return count
        gs = self.gridSize
        n = self.nLocalBoxes
        return getNeighborBoxes(iBox, nbrBoxes, n, gs)

    def copyAtom(self, atoms, iAtom, iBox, jAtom, jBox):
        iOff = self.MAXATOMS*iBox + iAtom
        jOff = self.MAXATOMS*jBox + jAtom
        atoms.gid[jOff] = atoms.gid[iOff]
        atoms.species[jOff] = atoms.species[iOff]
        atoms.r[jOff] = atoms.r[iOff]
        atoms.p[jOff] = atoms.p[iOff]
        atoms.f[jOff] = atoms.f[iOff]

    def moveAtom(self, atoms, iId, iBox, jBox):
        nj = self.nAtoms[jBox]
        self.copyAtom(atoms, iId, iBox, nj, jBox)
        self.nAtoms[jBox] += 1

        self.nAtoms[iBox] -= 1
        ni = self.nAtoms[iBox]
        if ni:
            self.copyAtom(atoms, ni, iBox, iId, iBox)
        if jBox > self.nLocalBoxes:
            atoms.nLocal -= 1

    def emptyHaloCells(self):
        for ii in range(self.nLocalBoxes, self.nTotalBoxes):
            self.nAtoms[ii] = 0

    def updateLinkCells(self, atoms):
        # empty halo cells
        self.emptyHaloCells()
        for iBox in range(self.nLocalBoxes):
            iOff = iBox * self.MAXATOMS
            ii = 0
            while ii < self.nAtoms[iBox]:
                jBox = self.getBoxFromCoord(atoms.r[iOff + ii])
                if jBox != iBox:
                    self.moveAtom(atoms, ii, iBox, jBox)
                else:
                    ii += 1

@numba.jit
def getBoxFromTupleInner(ix, iy, iz, n, gs):
    iBox = 0
    # Halo in Z+

    if iz == gs[2]:
        iBox = n + 2*gs[2]*gs[1] + 2*gs[2]*(gs[0]+2)  \
               + (gs[0]+2)*(gs[1]+2) + (gs[0] + 2)*(iy+1) + (ix+1)

    # Halo in Z-
    elif iz == -1:
        iBox = n + 2*gs[2]*gs[1] + 2*gs[2]*(gs[0]+2) \
                + (gs[0]+2)*(iy+1) + (ix+1)

    # Halo in Y+
    elif iy == gs[1]:
        iBox = n + 2*gs[2]*gs[1] + gs[2]*(gs[0]+2) \
                + (gs[0]+2)*iz + (ix+1)

    # Halo in Y-
    elif iy == -1:
        iBox = n + 2*gs[2]*gs[1] + iz*(gs[0] + 2) + (ix+1)

    # Halo in X+
    elif ix == gs[0]:
        iBox = n + gs[1]*gs[2] + iz*gs[1] + iy

    # Halo in X-
    elif ix == -1:
        iBox = n + iz*gs[1] + iy

    # Local link cell
    else:
        iBox = ix + gs[0]*iy + gs[0]*gs[1]*iz

    #if iBox >= self.nTotalBoxes:
    #    print('iBox too large',iBox,self.nTotalBoxes)
    #    print('ix,iy,iz',ix,iy,iz)
    #assert iBox >=0
    #assert iBox < self.nTotalBoxes
    return iBox

@numba.jit
def getTupleInner(iBox, n, gs):
    ix,iy,iz = 0,0,0
    # local box
    if iBox < n:
        ix = iBox % gs[0]
        iBox //= gs[0]
        iy = iBox % gs[1]
        iz = iBox // gs[1]
    else:   # halo box
        ink = iBox - n
        if ink < 2*gs[1]*gs[2]:
            if ink < gs[1]*gs[2]:
                ix = 0
            else:
                ink -= gs[1]*gs[2]
                ix = gs[0] + 1
            iy = 1 + ink % gs[1]
            iz = 1 + ink // gs[1]

        elif ink < 2*gs[2]*(gs[1] + gs[0] + 2):
            ink -= 2 * gs[2] * gs[1]
            if ink < (gs[0] + 2)*gs[2]:
                iy = 0
            else:
                ink -= (gs[0] + 2)*gs[2]
                iy = gs[1] + 1
            ix = ink % (gs[0] + 2)
            iz = 1 + ink // (gs[0] + 2)
        else:
            ink -= 2*gs[2]*(gs[1] + gs[0] + 2)
            if ink < (gs[0] + 2)*(gs[1] + 2):
                iz = 0
            else:
                ink -= (gs[0] + 2)*(gs[1] + 2)
                iz = gs[2] + 1
            ix = ink % (gs[0] + 2)
            iy = ink // (gs[0] + 2)

        ix -= 1
        iy -= 1
        iz -= 1

    return ix,iy,iz

@numba.njit
def getNeighborBoxesInner(iBox, nbrBoxes, n, gs):
    ix,iy,iz = getTupleInner(iBox, n, gs)
    count = 0

    for i in range(ix-1, ix+2):
        for j in range(iy-1, iy+2):
            for k in range(iz-1, iz+2):
                nbrBoxes[count] = getBoxFromTupleInner(i, j, k, n, gs)
                count += 1
    return count

ljforce.py

from __future__ import print_function, division

import numpy as np
import constants
import linkcell

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, sim):
        POT_SHIFT = 1.0
        #POT_SHIFT = 0.0
        rCut2 = self.cutoff * self.cutoff
        nLocalBoxes = sim.boxes.nLocalBoxes
        MAXATOMS = sim.boxes.MAXATOMS
        nAtoms = sim.boxes.nAtoms
        gridSize = sim.boxes.gridSize
        f = atoms.f
        r = atoms.r
        #getNeighborBoxes = sim.boxes.getNeighborBoxes
        #getNeighborBoxesInner = linkcell.getNeighborBoxesInner
        gid = atoms.gid
        epsilon = self.epsilon


        #for i in range(atoms.nAtoms):
        for iBox in range(nLocalBoxes):
            iOff = MAXATOMS * iBox
            for ii in range(nAtoms[iBox]):
                f[iOff,:] = 0.0
                iOff += 1

        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)

        nbrBoxes = np.zeros(27, dtype=np.int)
        # loop over atoms

        iPot = 0
        #for i in range(atoms.nAtoms):
        #    for j in range(atoms.nAtoms):
        for iBox in range(nLocalBoxes):
            nIBox = nAtoms[iBox]
            #nNbrBoxes = sim.boxes.getNeighborBoxes(iBox, nbrBoxes)
            #nNbrBoxes = getNeighborBoxes(iBox, nbrBoxes)
            nNbrBoxes = linkcell.getNeighborBoxesInner(iBox, nbrBoxes, nLocalBoxes, gridSize)
            for jTmp in range(nNbrBoxes):
                jBox = nbrBoxes[jTmp]
                #assert jBox >= 0

                nJBox = nAtoms[jBox]
                if nJBox == 0:
                    continue

                for ii in range(nIBox):
                    iOff = MAXATOMS * iBox + ii
                    iId = gid[iOff]
                    for jj in range(nJBox):
                        jOff = MAXATOMS * jBox + jj
                        jId = gid[jOff]
                        if jBox < nLocalBoxes and jId <= iId:
                            continue


                        r2 = 0.0
                        dx = r[iOff,0] - r[jOff,0]
                        r2 += dx*dx

                        dy = r[iOff,1] - r[jOff,1]
                        r2 += dy*dy

                        dz = r[iOff,2] - r[jOff,2]
                        r2 += dz*dz

                        if r2 > rCut2:
                            continue

                        ir2 = 1.0/r2

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

                        if jBox < nLocalBoxes:
                            ePot += eLocal
                        else:
                            ePot += 0.5 * eLocal

                        iPot += 1


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

                        f[iOff,0] -= dx*fr
                        f[jOff,0] += dx*fr
                        f[iOff,1] -= dy*fr
                        f[jOff,1] += dy*fr
                        f[iOff,2] -= dz*fr
                        f[jOff,2] += dz*fr

        ePot = ePot*4.0*epsilon
        return ePot