Create a gist now

Instantly share code, notes, and snippets.

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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment