Cython code for CoMD cell method

comd.py

import pyximport; pyximport.install()

import simflat
import initatoms
import time
import constants

import argparse

# Translation of CoMD to Python



def initValidate(sim):
    ke = initatoms.kineticEnergy(sim)
    pe = sim.ePot

    print("Initial PE (cohesive energy) : ",pe/sim.atoms.nGlobal)
    print("Initial energy : ",(ke+pe)/sim.atoms.nGlobal," atom count : ",sim.atoms.nGlobal)

firstCall = True
iStepPrev = -1
def printInfo(sim, iStep, elapsedTime):
    global firstCall, iStepPrev
    if firstCall:
        firstCall = False
        #print "# Loop  Time(fs)   Total Energy  Potential Energy   Kinetic Energy  Temperature  Performance (us/atom)  # Atoms"
        print("#{:>100s}".format("Performance"))
        print("#{:>6s} {:>10s} {:>18s} {:>18s} {:>18s} {:>12s} {:^12s} {:^12s}".\
            format("Loop", "Time(fs)", "Total Energy", "Potential Energy", "Kinetic Energy", "Temperature",  "(us/atom)", "# Atoms"))

    nEval = iStep - iStepPrev
    iStepPrev = iStep

    simtime = iStep * sim.dt
    n = sim.atoms.nGlobal
    eTotal = (sim.ePot + sim.eKinetic)/n
    eK = sim.eKinetic / n
    eU = sim.ePot / n
    Temp = (sim.eKinetic/n)/(constants.kB_eV * 1.5)
    timePerAtom = 1.0e6 * elapsedTime/(n*nEval)

    print(" {:6d} {:10.2f} {:18.12f} {:18.12f} {:18.12f} {:12.4f} {:10.4f} {:12d}".format(iStep, simtime, eTotal, eU, eK, Temp, timePerAtom, n))
    #print iStep, simtime, eTotal, eU, eK, Temp, timePerAtom, n

def printPerformanceResult(sim, elapsedTime):
    n = sim.atoms.nGlobal
    nEval = sim.nSteps
    timePerAtom = 1.0e6 * elapsedTime/(n*nEval)

    print()
    print("Average all atom update rate: %10.2f us/atom" % timePerAtom)
    print()


def parseCommandLine():
    parser = argparse.ArgumentParser()
    parser.add_argument("-x","--nx", type=int, default=10, help="number of unit cells in x")
    parser.add_argument("-y","--ny", type=int, default=10, help="number of unit cells in y")
    parser.add_argument("-z","--nz", type=int, default=10, help="number of unit cells in z")
    parser.add_argument("-N","--nSteps", type=int, default=100, help="total number of time steps")
    parser.add_argument("-n","--printRate", type=int, default=10, help="number of steps between output")
    parser.add_argument("-D","--dt", type=float, default=1, help="time step (in fs)")
    parser.add_argument("-l","--lat", type=float, default=-1, help="lattice parameter (Angstroms)")
    parser.add_argument("-T","--temp", type=float, default=600, help="initial temperature (K)")

    args = parser.parse_args()
    return args


def run_comd():
    cmd = parseCommandLine()
    # init subsystems
    sim = simflat.initSimulation(cmd)

    sim.pot.computeForce(sim.atoms, sim)
    sim.eKinetic = initatoms.kineticEnergy(sim)

    initValidate(sim)

    timestepTime = 0.0
    timestepTimeOneIteration = 0.0

    iStep = 0
    for jStep in range(0, sim.nSteps, sim.printRate):
        printInfo(sim, iStep, timestepTimeOneIteration)

        start = time.clock()
        simflat.timestep(sim, sim.printRate, sim.dt)
        end = time.clock()

        timestepTimeOneIteration = end - start
        timestepTime += timestepTimeOneIteration

        iStep += sim.printRate

    printInfo(sim, iStep, timestepTimeOneIteration)

    printPerformanceResult(sim, timestepTime)

    # validate result

if __name__ == '__main__':
    run_comd()

linkcell.pyx

#from __future__ import print_function, division


import numpy as np
import math
import cython

cimport numpy as np

def initLinkCells(sim, box_size):
    gridSize = np.array([int(b/sim.pot.cutoff) for b in box_size], dtype=np.int)
    #gridSize = np.array([int(b/sim.pot.cutoff) for b in box_size])
    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:
    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

    @cython.boundscheck(False)
    def getBoxFromTuple(self, int ix, int iy, int iz):

        cdef int n, iBox
        cdef np.ndarray[dtype=np.int_t] gs
        n = self.nLocalBoxes
        gs = self.gridSize

        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

    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

    @cython.boundscheck(False)
    def getTuple(self, int iBox):
        cdef int ix,iy,iz,ink
        cdef int n
        cdef np.ndarray[dtype=np.int_t] gs

        gs = self.gridSize
        n = self.nLocalBoxes
        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

    def getNeighborBoxes(self, int iBox, np.ndarray[dtype=np.int_t] nbrBoxes):
        cdef int ix,iy,iz, i, j, k,count

        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

    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

ljforce.pyx

#from __future__ import print_function, division

import numpy as np
import constants
import cython

cimport numpy as np

class lj_pot:
    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

    @cython.boundscheck(False)
    def computeForce(self, atoms, sim):
        cdef int nLocalBoxes
        cdef int MAXATOMS
        cdef double epsilon
        cdef np.ndarray[dtype=np.double_t, ndim=2] f
        cdef np.ndarray[dtype=np.double_t, ndim=2] r
        cdef np.ndarray[dtype=np.int_t, ndim=1] gid
        cdef np.ndarray[dtype=np.int_t, ndim=1] nAtoms
        cdef np.ndarray[dtype=np.int_t, ndim=1] nbrBoxes
        cdef double r2, dx, dy, dz, rCut2, ir2, r6, s6, eShift, rCut6, fr, ePot
        cdef int iBox, iOff, ii, jj, jBox, jOff, jId, iId, nJBox, nIBox, nNbrBoxes

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

        nLocalBoxes = sim.boxes.nLocalBoxes
        r = atoms.r
        f = atoms.f
        epsilon = self.epsilon
        MAXATOMS = sim.boxes.MAXATOMS
        nAtoms = sim.boxes.nAtoms
        gid = atoms.gid

        #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

        #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)
            for jTmp in range(nNbrBoxes):
                jBox = nbrBoxes[jTmp]

                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



                        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

simflat.pyx

import ljforce
import initatoms
import linkcell
import halo
import numpy as np
import cython

cimport numpy as np

class SimFlat:
    def __init__(self, nSteps, printRate, dt):
        self.nSteps = nSteps
        self.printRate = printRate
        self.dt = dt
        self.atoms = None
        self.pot = ljforce.lj_pot()
        self.species = initSpecies(self.pot)
        self.ePot = 0.0
        self.eKinetic = 0.0
        self.boxes = None
        self.atomHalo = None


class species:
    def __init__(self, mass):
        self.mass = mass

def initSpecies(pot):
    return [species(pot.mass)]


def redistributeAtoms(sim, atoms):
    sim.boxes.updateLinkCells(atoms)
    sim.atomHalo.haloExchange(sim)

def initSimulation(cmd):
    sim = SimFlat(cmd.nSteps, cmd.printRate, cmd.dt)

    latticeConstant = cmd.lat

    if cmd.lat < 0.0:
        latticeConstant = sim.pot.lat




    box_size = np.zeros(3)
    box_size[0] = cmd.nx*latticeConstant
    box_size[1] = cmd.ny*latticeConstant
    box_size[2] = cmd.nz*latticeConstant
    sim.boxes = linkcell.initLinkCells(sim, box_size)

    sim.atoms = initatoms.initAtoms(sim.boxes.nTotalBoxes * sim.boxes.MAXATOMS)

    sim.atoms.setBounds(cmd.nx, cmd.ny, cmd.nz, latticeConstant)

    initatoms.createFccLattice(cmd.nx, cmd.ny, cmd.nz, latticeConstant, sim.atoms, sim.boxes)

    initatoms.setTemperature(sim, cmd.temp)

    sim.atomHalo = halo.Halo(sim.boxes)

    redistributeAtoms(sim, sim.atoms)



    return sim




def advanceVelocity(sim, atoms, double dt):
    cdef np.ndarray[dtype=np.double_t, ndim=2] f,p 
    cdef int iBox, ii, iOff
    p = atoms.p
    f = atoms.f
    #for i in range(atoms.nAtoms):
    for iBox in range(sim.boxes.nLocalBoxes):
        for ii in range(sim.boxes.nAtoms[iBox]):
            iOff = sim.boxes.MAXATOMS * iBox + ii
            p[iOff,0] += dt*f[iOff,0]
            p[iOff,1] += dt*f[iOff,1]
            p[iOff,2] += dt*f[iOff,2]

def advancePosition(sim, atoms, double dt):
    cdef np.ndarray[dtype=np.double_t, ndim=2] r,p 
    cdef np.ndarray[dtype=np.int_t] atomSpecies
    cdef int iBox, ii, iOff
    cdef double invMass
    p = atoms.p
    r = atoms.r
    atomSpecies = atoms.species
    
    #for i in range(atoms.nAtoms):
    for iBox in range(sim.boxes.nLocalBoxes):
        for ii in range(sim.boxes.nAtoms[iBox]):
            iOff = sim.boxes.MAXATOMS * iBox + ii
            iSpecies =atomSpecies[iOff]
            invMass = 1.0/sim.species[iSpecies].mass
            r[iOff,0] += dt*p[iOff,0]*invMass
            r[iOff,1] += dt*p[iOff,1]*invMass
            r[iOff,2] += dt*p[iOff,2]*invMass


#def redistributeAtoms(sim, atoms):
#    sim.boxes.updateLinkCells(atoms)
#    sim.atomHalo.haloExchange(sim)

def dumpPos(sim):
    for iBox in range(sim.boxes.nLocalBoxes):
        for ii in range(sim.boxes.nAtoms[iBox]):
            iOff = sim.boxes.MAXATOMS * iBox + ii
            pos = sim.atoms.r[iOff,:]
            print("%d:  %20.10f %20.10f %20.10f"%(iOff+1, pos[0], pos[1], pos[2]))

def dumpVel(sim):
    for iBox in range(sim.boxes.nLocalBoxes):
        for ii in range(sim.boxes.nAtoms[iBox]):
            iOff = sim.boxes.MAXATOMS * iBox + ii
            vel = sim.atoms.p[iOff,:]
            print("%d:  %20.10g %20.10g %20.10g"%(iOff+1, vel[0], vel[1], vel[2]))

def dumpForce(sim):
    for iBox in range(sim.boxes.nLocalBoxes):
        for ii in range(sim.boxes.nAtoms[iBox]):
            iOff = sim.boxes.MAXATOMS * iBox + ii
            force = sim.atoms.f[iOff,:]
            print("%d:  %20.10g %20.10g %20.10g"%(iOff+1, force[0], force[1], force[2]))
    

def timestep(sim, int nSteps, double dt):
    cdef int ii
    for ii in range(nSteps):
        advanceVelocity(sim, sim.atoms, 0.5*dt)
        advancePosition(sim, sim.atoms, dt)
        redistributeAtoms(sim, sim.atoms)
        sim.pot.computeForce(sim.atoms, sim)
        advanceVelocity(sim, sim.atoms, 0.5*dt)
    #print('ke,pe = ',initatoms.kineticEnergy(sim)/sim.atoms.nAtoms,sim.ePot/sim.atoms.nAtoms)
    sim.eKinetic = initatoms.kineticEnergy(sim)