Created
October 2, 2015 16:34
-
-
Save markdewing/3688c6eebc0a88081e07 to your computer and use it in GitHub Desktop.
Cython code for CoMD cell method
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment