Instantly share code, notes, and snippets.

# markdewing/comd.py Created Oct 2, 2015

What would you like to do?
Cython code for CoMD cell method
 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()
 #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
 #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
 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)