Skip to content

Instantly share code, notes, and snippets.

@markdewing

markdewing/comd.py

Created November 13, 2015 16:22
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save markdewing/8bd6bd8dbef8613004fe to your computer and use it in GitHub Desktop.
Save markdewing/8bd6bd8dbef8613004fe to your computer and use it in GitHub Desktop.
Download ZIP
More optimization of CoMD for Numba (using feedback from a profiler)
Raw
comd.py
from __future__ import print_function, division
import simflat
import initatoms
import time
import constants
import argparse
import numba
# Translation of CoMD to Python
def advanceVelocity(sim, atoms, dt):
#for i in range(atoms.nAtoms):
MAXATOMS = sim.boxes.MAXATOMS
p = atoms.p
f = atoms.f
for iBox in range(sim.boxes.nLocalBoxes):
for ii in range(sim.boxes.nAtoms[iBox]):
iOff = MAXATOMS * iBox + ii
p[iOff,0] += dt*f[iOff,0]
p[iOff,1] += dt*f[iOff,1]
p[iOff,2] += dt*f[iOff,2]
@numba.njit
def advanceVelocityInner(MAXATOMS, nLocalBoxes, nAtoms, p, f, dt):
#for i in range(atoms.nAtoms):
#MAXATOMS = sim.boxes.MAXATOMS
#p = atoms.p
#f = atoms.f
for iBox in range(nLocalBoxes):
for ii in range(nAtoms[iBox]):
iOff = 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, dt):
#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 = atoms.species[iOff]
invMass = 1.0/sim.species[iSpecies].mass
atoms.r[iOff,0] += dt*atoms.p[iOff,0]*invMass
atoms.r[iOff,1] += dt*atoms.p[iOff,1]*invMass
atoms.r[iOff,2] += dt*atoms.p[iOff,2]*invMass
@numba.njit
def advancePositionInner(MAXATOMS, nLocalBoxes, nAtoms, species, r, p, invMassArray, dt):
#for i in range(atoms.nAtoms):
for iBox in range(nLocalBoxes):
for ii in range(nAtoms[iBox]):
iOff = MAXATOMS * iBox + ii
iSpecies = species[iOff]
#invMass = 1.0/sim.species[iSpecies].mass
invMass = invMassArray[iSpecies]
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, nSteps, dt):
for ii in range(nSteps):
#advanceVelocity(sim, sim.atoms, 0.5*dt)
advanceVelocityInner(sim.boxes.MAXATOMS,sim.boxes.nLocalBoxes, sim.boxes.nAtoms, sim.atoms.p, sim.atoms.f, 0.5*dt)
#advancePosition(sim, sim.atoms, dt)
advancePositionInner(sim.boxes.MAXATOMS, sim.boxes.nLocalBoxes, sim.boxes.nAtoms, sim.atoms.species, sim.atoms.r, sim.atoms.p, sim.invMass, dt)
simflat.redistributeAtoms(sim, sim.atoms)
sim.ePot = sim.pot.computeForce(sim.atoms, sim)
#advanceVelocity(sim, sim.atoms, 0.5*dt)
advanceVelocityInner(sim.boxes.MAXATOMS,sim.boxes.nLocalBoxes, sim.boxes.nAtoms, sim.atoms.p, sim.atoms.f, 0.5*dt)
#print('ke,pe = ',initatoms.kineticEnergy(sim)/sim.atoms.nAtoms,sim.ePot/sim.atoms.nAtoms)
#sim.eKinetic = initatoms.kineticEnergy(sim)
boxes = sim.boxes
atoms = sim.atoms
sim.eKinetic = initatoms.kineticEnergyInner(boxes.nLocalBoxes, boxes.MAXATOMS, boxes.nAtoms, atoms.species, sim.invMass, atoms.p)
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 - sim.nSkip * sim.printRate
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("--skip", type=int, default=1, help="number of blocks to skip in averaged 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.ePot = 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()
timestep(sim, sim.printRate, sim.dt)
end = time.clock()
timestepTimeOneIteration = end - start
if jStep >= sim.nSkip * sim.printRate:
timestepTime += timestepTimeOneIteration
iStep += sim.printRate
printInfo(sim, iStep, timestepTimeOneIteration)
printPerformanceResult(sim, timestepTime)
# validate result
if __name__ == '__main__':
run_comd()
Raw
halo.py
from __future__ import print_function, division
import numpy as np
import numba
import linkcell
class Halo(object):
HALO_X_MINUS = 0
HALO_X_PLUS = 1
HALO_Y_MINUS = 2
HALO_Y_PLUS = 3
HALO_Z_MINUS = 4
HALO_Z_PLUS = 5
HALO_X_AXIS = 0
HALO_Y_AXIS = 1
HALO_Z_AXIS = 2
def __init__(self, boxes):
self.bufCapacity = 0
self.nCells = np.zeros(6, dtype=np.int)
self.pbcFactor = np.zeros((6,3))
self.nCells[self.HALO_X_MINUS] = 2*(boxes.gridSize[1] + 2) * (boxes.gridSize[2] + 2)
self.nCells[self.HALO_Y_MINUS] = 2*(boxes.gridSize[0] + 2) * (boxes.gridSize[2] + 2)
self.nCells[self.HALO_Z_MINUS] = 2*(boxes.gridSize[0] + 2) * (boxes.gridSize[1] + 2)
self.nCells[self.HALO_X_PLUS] = self.nCells[self.HALO_X_MINUS]
self.nCells[self.HALO_Y_PLUS] = self.nCells[self.HALO_Y_MINUS]
self.nCells[self.HALO_Z_PLUS] = self.nCells[self.HALO_Z_MINUS]
#self.cellList = []
self.cellList = np.zeros( (6, 128), dtype=np.int)
for ii in range(6):
#self.cellList.append(self.mkAtomCellList(boxes, ii))
self.mkAtomCellList(boxes, ii, self.cellList)
self.pbcFactor[self.HALO_X_MINUS, self.HALO_X_AXIS] = 1.0
self.pbcFactor[self.HALO_X_PLUS, self.HALO_X_AXIS] = -1.0
self.pbcFactor[self.HALO_Y_MINUS, self.HALO_Y_AXIS] = 1.0
self.pbcFactor[self.HALO_Y_PLUS, self.HALO_Y_AXIS] = -1.0
self.pbcFactor[self.HALO_Z_MINUS, self.HALO_Z_AXIS] = 1.0
self.pbcFactor[self.HALO_Z_PLUS, self.HALO_Z_AXIS] = -1.0
def haloExchange(self, sim):
boxes = sim.boxes
atoms = sim.atoms
for iAxis in range(3):
self.exchangeData(iAxis, sim)
#atoms.nLocal = exchangeDataInner(iAxis,
# self.nCells, self.cellList, self.pbcFactor,
# boxes.nAtoms, boxes.MAXATOMS, boxes.invBoxSize, boxes.nLocalBoxes, boxes.gridSize, #boxes
# atoms.bounds, atoms.gid, atoms.species, atoms.r, atoms.p, atoms.nLocal) #atoms
def exchangeData(self, iAxis, sim):
faceM = 2*iAxis
faceP = faceM + 1
boxes = sim.boxes
atoms = sim.atoms
int_bufM, fp_bufM = self.createBuffer(sim ,faceM)
int_bufP, fp_bufP = self.createBuffer(sim ,faceP)
shiftM = self.pbcFactor[faceM, :] * atoms.bounds
shiftP = self.pbcFactor[faceP, :] * atoms.bounds
loadAtomsBufferInner(int_bufM, fp_bufM, shiftM, self.nCells[faceM], self.cellList[faceM, :],
boxes.nAtoms, boxes.MAXATOMS, atoms.gid, atoms.species, atoms.r, atoms.p)
loadAtomsBufferInner(int_bufP, fp_bufP, shiftP, self.nCells[faceP], self.cellList[faceP, :],
boxes.nAtoms, boxes.MAXATOMS, atoms.gid, atoms.species, atoms.r, atoms.p)
atoms.nLocal = unloadAtomsBufferInner(int_bufP, fp_bufP, boxes.nAtoms, boxes.MAXATOMS, boxes.invBoxSize,
boxes.nLocalBoxes, boxes.gridSize, atoms.nLocal,
atoms.gid, atoms.species, atoms.r, atoms.p)
atoms.nLocal = unloadAtomsBufferInner(int_bufM, fp_bufM, boxes.nAtoms, boxes.MAXATOMS, boxes.invBoxSize,
boxes.nLocalBoxes, boxes.gridSize, atoms.nLocal,
atoms.gid, atoms.species, atoms.r, atoms.p)
def createBuffer(self, sim, face):
nEntry = 0
for iCell in range(self.nCells[face]):
iBox = self.cellList[face][iCell]
nEntry += sim.boxes.nAtoms[iBox]
int_buf = np.zeros((nEntry, 2), dtype=np.int)
fp_buf = np.zeros((nEntry, 6))
return int_buf, fp_buf
def loadAtomsBuffer(self, sim, face, int_buf, fp_buf):
shift = self.pbcFactor[face,:] * sim.atoms.bounds
nCells = self.nCells[face]
cellList = self.cellList[face]
idx = 0
for iCell in range(nCells):
iBox = cellList[iCell]
for ii in range(sim.boxes.nAtoms[iBox]):
iOff = iBox*sim.boxes.MAXATOMS + ii
int_buf[idx, 0] = sim.atoms.gid[iOff]
int_buf[idx, 1] = sim.atoms.species[iOff]
fp_buf[idx, 0] = sim.atoms.r[iOff][0] + shift[0]
fp_buf[idx, 1] = sim.atoms.r[iOff][1] + shift[1]
fp_buf[idx, 2] = sim.atoms.r[iOff][2] + shift[2]
fp_buf[idx, 3] = sim.atoms.p[iOff][0]
fp_buf[idx, 4] = sim.atoms.p[iOff][1]
fp_buf[idx, 5] = sim.atoms.p[iOff][2]
idx += 1
def unloadAtomsBuffer(self, sim, face, int_buf, fp_buf):
size = int_buf.shape[0]
for idx in range(size):
gid = int_buf[idx, 0]
species = int_buf[idx, 1]
rx = fp_buf[idx, 0]
ry = fp_buf[idx, 1]
rz = fp_buf[idx, 2]
px = fp_buf[idx, 3]
py = fp_buf[idx, 4]
pz = fp_buf[idx, 5]
sim.boxes.putAtomInBox(sim.atoms, gid, species, rx, ry, rz, px, py, pz)
def mkAtomCellList(self, boxes, iFace, cellList):
xBegin = -1
xEnd = boxes.gridSize[0] + 1
yBegin = -1
yEnd = boxes.gridSize[1] + 1
zBegin = -1
zEnd = boxes.gridSize[2] + 1
if iFace == self.HALO_X_MINUS: xEnd = xBegin + 2
if iFace == self.HALO_X_PLUS: xBegin = xEnd - 2
if iFace == self.HALO_Y_MINUS: yEnd = yBegin + 2
if iFace == self.HALO_Y_PLUS: yBegin = yEnd - 2
if iFace == self.HALO_Z_MINUS: zEnd = zBegin + 2
if iFace == self.HALO_Z_PLUS: zBegin = zEnd - 2
#cells = []
idx = 0
for ix in range(xBegin,xEnd):
for iy in range(yBegin,yEnd):
for iz in range(zBegin,zEnd):
#cells.append(boxes.getBoxFromTuple(ix,iy,iz))
cellList[iFace, idx] = boxes.getBoxFromTuple(ix,iy,iz)
idx += 1
#return cells
@numba.njit
def loadAtomsBufferInner(int_buf, fp_buf, shift, nCells, cellList, nAtoms, MAXATOMS, gid, species, r, p):
#shift = self.pbcFactor[face,:] * sim.atoms.bounds
#nCells = self.nCells[face]
#cellList = self.cellList[face]
#nAtoms = sim.boxes.nAtoms
idx = 0
for iCell in range(nCells):
iBox = cellList[iCell]
for ii in range(nAtoms[iBox]):
iOff = iBox*MAXATOMS + ii
int_buf[idx, 0] = gid[iOff]
int_buf[idx, 1] = species[iOff]
fp_buf[idx, 0] = r[iOff][0] + shift[0]
fp_buf[idx, 1] = r[iOff][1] + shift[1]
fp_buf[idx, 2] = r[iOff][2] + shift[2]
fp_buf[idx, 3] = p[iOff][0]
fp_buf[idx, 4] = p[iOff][1]
fp_buf[idx, 5] = p[iOff][2]
idx += 1
@numba.njit
def unloadAtomsBufferInner(int_buf, fp_buf, nAtoms, MAXATOMS, invBoxSize, nLocalBoxes, gridSize, nLocal, agid, aspecies,r, p):
size = int_buf.shape[0]
for idx in range(size):
gid = int_buf[idx, 0]
species = int_buf[idx, 1]
rx = fp_buf[idx, 0]
ry = fp_buf[idx, 1]
rz = fp_buf[idx, 2]
px = fp_buf[idx, 3]
py = fp_buf[idx, 4]
pz = fp_buf[idx, 5]
nLocal += linkcell.putAtomInBoxInner(nAtoms, MAXATOMS, invBoxSize, nLocalBoxes, gridSize, nLocal, agid, aspecies, r, p, gid, species, rx, ry, rz, px, py, pz)
return nLocal
#def linkcell.putAtomInBoxInner(nAtoms, MAXATOMS, invBoxSize, nLocalBoxes, gs, nLocal, gid, species, r, p, self_gid, iType, x, y, z, px, py, pz):
@numba.njit
def createBufferInner(nCells, cellList, face, nAtoms):
nEntry = 0
for iCell in range(nCells[face]):
iBox = cellList[face][iCell]
nEntry += nAtoms[iBox]
# 11/13/2014, with Numba 0.22
# Had trouble getting Numba to accept a 2-d array
# initialization .
# Eventually I tried collapsing it to a 1-d array,
# but that ran slower
#shape_int = np.array( [nEntry, 2], dtype=np.int)
#int_buf = np.zeros(shape_int, dtype=np.int)
#shape_fp = np.array( [nEntry, 6], dtype=np.int)
#fp_buf = np.zeros(shape_fp)
#int_buf = np.zeros(nEntry*2, dtype=np.int)
#int_buf = np.zeros(nEntry*2)
#fp_buf = np.zeros(nEntry*6)
return int_buf, fp_buf
# This didn't help performance. Didn't investigate why.
@numba.njit
def exchangeDataInner(iAxis,
nCells, cellList, pbcFactor, #self
nAtoms, MAXATOMS, invBoxSize, nLocalBoxes, gridSize, #boxes
bounds, gid, species, r, p, nLocal): #atoms
faceM = 2*iAxis
faceP = faceM + 1
int_bufM, fp_bufM = createBufferInner(nCells, cellList, faceM, nAtoms)
int_bufP, fp_bufP = createBufferInner(nCells, cellList, faceP, nAtoms)
shiftM = pbcFactor[faceM, :] * bounds
shiftP = pbcFactor[faceP, :] * bounds
loadAtomsBufferInner(int_bufM, fp_bufM, shiftM, nCells[faceM], cellList[faceM, :],
nAtoms, MAXATOMS, gid, species, r, p)
loadAtomsBufferInner(int_bufP, fp_bufP, shiftP, nCells[faceP], cellList[faceP, :],
nAtoms, MAXATOMS, gid, species, r, p)
nLocal = unloadAtomsBufferInner(int_bufP, fp_bufP, nAtoms, MAXATOMS, invBoxSize,
nLocalBoxes, gridSize, nLocal,
gid, species, r, p)
nLocal = unloadAtomsBufferInner(int_bufM, fp_bufM, nAtoms, MAXATOMS, invBoxSize,
nLocalBoxes, gridSize, nLocal,
gid, species, r, p)
return nLocal
Raw
initatoms.py
from __future__ import print_function, division
import constants
import numpy as np
import math
import random
import numba
class Atoms(object):
def __init__(self, n):
self.r = np.zeros( (n,3) )
self.f = np.zeros( (n,3) )
self.p = np.zeros( (n,3) )
self.species = np.zeros(n, dtype=np.int)
self.gid = np.zeros(n, dtype=np.int)
self.bounds = np.zeros(3)
self.maxTotalAtoms = n
self.nLocal = 0
self.nGlobal = 0
def setBounds(self, nx, ny, nz, lat):
self.bounds[0] = nx*lat
self.bounds[1] = ny*lat
self.bounds[2] = nz*lat
def initAtoms(maxatoms):
return Atoms(maxatoms)
def createFccLattice(nx, ny, nz, lat, atoms, boxes):
"""Create atom positions on a face centered cubic (FCC) lattice
with nx * ny * nz unit cells and lattice constance 'lat'
"""
nb = 4 # number of atoms in this basis
basis = [ (0.25, 0.25, 0.25),
(0.25, 0.75, 0.75),
(0.75, 0.25, 0.75),
(0.75, 0.75, 0.25)
]
idx = 0
# loop over ix,iy,iz
for ix in range(nx):
for iy in range(ny):
for iz in range(nz):
for ib in range(nb):
rx = (ix+basis[ib][0])*lat
ry = (iy+basis[ib][1])*lat
rz = (iz+basis[ib][2])*lat
#atoms.r[idx,0] = rx
#atoms.r[idx,1] = ry
#atoms.r[idx,2] = rz
gid = ib + nb*(iz + nz*(iy + ny*ix))
boxes.putAtomInBox(atoms, gid, 0, rx, ry, rz)
idx += 1
atoms.nGlobal = idx
assert idx == nb*nx*ny*nz
#idx should equal nx*ny*nz*nb
def kineticEnergy(sim):
ke = 0.0
#for i in range(sim.atoms.nAtoms):
for iBox in range(sim.boxes.nLocalBoxes):
iOff = sim.boxes.MAXATOMS * iBox
for ii in range(sim.boxes.nAtoms[iBox]):
iSpecies = sim.atoms.species[iOff]
invMass = 0.5/sim.species[iSpecies].mass
ke += invMass * (sim.atoms.p[iOff,0]*sim.atoms.p[iOff,0] + \
sim.atoms.p[iOff,1]*sim.atoms.p[iOff,1] + sim.atoms.p[iOff,2]*sim.atoms.p[iOff,2])
iOff += 1
return ke
@numba.njit
def kineticEnergyInner(nLocalBoxes, MAXATOMS, nAtoms, species, invMassArray, p):
ke = 0.0
#for i in range(sim.atoms.nAtoms):
for iBox in range(nLocalBoxes):
iOff = MAXATOMS * iBox
for ii in range(nAtoms[iBox]):
iSpecies = species[iOff]
#invMass = 0.5/sim.species[iSpecies].mass
invMass = 0.5*invMassArray[iSpecies]
ke += invMass * (p[iOff,0]*p[iOff,0] + \
p[iOff,1]*p[iOff,1] + p[iOff,2]*p[iOff,2])
iOff += 1
return ke
def computeVcm(sim):
vcm = np.zeros(3)
totalMass = 0.0
#for i in range(sim.atoms.nAtoms):
for iBox in range(sim.boxes.nLocalBoxes):
iOff = sim.boxes.MAXATOMS * iBox
for ii in range(sim.boxes.nAtoms[iBox]):
vcm += sim.atoms.p[iOff,:]
iSpecies = sim.atoms.species[iOff]
totalMass += sim.species[iSpecies].mass
iOff += 1
vcm = vcm/totalMass
return vcm
def setVcm(sim, newVcm):
oldVcm = computeVcm(sim)
shift = newVcm - oldVcm
#for i in range(sim.atoms.nAtoms):
for iBox in range(sim.boxes.nLocalBoxes):
iOff = sim.boxes.MAXATOMS * iBox
for ii in range(sim.boxes.nAtoms[iBox]):
iSpecies = sim.atoms.species[iOff]
mass = sim.species[iSpecies].mass
sim.atoms.p[iOff,:] += mass*shift
iOff += 1
def setTemperature(sim, temperature):
#for i in range(sim.atoms.nAtoms):
for iBox in range(sim.boxes.nLocalBoxes):
iOff = sim.boxes.MAXATOMS * iBox
for ii in range(sim.boxes.nAtoms[iBox]):
iSpecies = sim.atoms.species[iOff]
mass = sim.species[iSpecies].mass
sigma = math.sqrt(constants.kB_eV * temperature/mass)
sim.atoms.p[iOff,0] = mass * sigma * random.gauss(1.0, 1.0)
sim.atoms.p[iOff,1] = mass * sigma * random.gauss(1.0, 1.0)
sim.atoms.p[iOff,2] = mass * sigma * random.gauss(1.0, 1.0)
iOff += 1
if temperature == 0.0:
return
vZero = np.zeros(3)
setVcm(sim, vZero)
ke = kineticEnergy(sim)
temp = (ke/sim.atoms.nGlobal)/constants.kB_eV/1.5
scaleFactor = math.sqrt(temperature/temp)
#for i in range(sim.atoms.nAtoms):
for iBox in range(sim.boxes.nLocalBoxes):
iOff = sim.boxes.MAXATOMS * iBox
for ii in range(sim.boxes.nAtoms[iBox]):
sim.atoms.p[iOff,:] *= scaleFactor
iOff += 1
Raw
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):
atoms.nLocal = putAtomInBoxInner(self.nAtoms, self.MAXATOMS, self.invBoxSize, self.nLocalBoxes,
self.gridSize, atoms.nLocal, atoms.gid, atoms.species, atoms.r, atoms.p,
gid, iType, x, y, z, px, py, pz)
#iBox = self.getBoxFromCoord([x, y, z])
#iBox = getBoxFromCoordInner(x, y, z, self.invBoxSize, self.nLocalBoxes, self.gridSize)
#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 getNeighborBoxesInner(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()
emptyHaloCellsInner(self.nLocalBoxes, self.nTotalBoxes, self.nAtoms)
for iBox in range(self.nLocalBoxes):
iOff = iBox * self.MAXATOMS
ii = 0
while ii < self.nAtoms[iBox]:
#jBox = self.getBoxFromCoord(atoms.r[iOff + ii])
r = atoms.r[iOff + ii]
jBox = getBoxFromCoordInner(r[0], r[1], r[2], self.invBoxSize, self.nLocalBoxes, self.gridSize)
if jBox != iBox:
#self.moveAtom(atoms, ii, iBox, jBox)
atoms.nLocal = moveAtomInner(ii, iBox, jBox, self.MAXATOMS, atoms.gid, atoms.species, atoms.r, atoms.p, atoms.f, self.nAtoms, self.nLocalBoxes, atoms.nLocal)
else:
ii += 1
@numba.njit
def emptyHaloCellsInner(nLocalBoxes, nTotalBoxes, nAtoms):
for ii in range(nLocalBoxes, nTotalBoxes):
nAtoms[ii] = 0
@numba.njit
def copyAtomInner(iAtom, iBox, jAtom, jBox, MAXATOMS, gid, species, r, p, f):
iOff = MAXATOMS*iBox + iAtom
jOff = MAXATOMS*jBox + jAtom
gid[jOff] = gid[iOff]
species[jOff] = species[iOff]
r[jOff] = r[iOff]
p[jOff] = p[iOff]
f[jOff] = f[iOff]
@numba.njit
def moveAtomInner(iId, iBox, jBox, MAXATOMS, gid, species, r, p, f, nAtoms, nLocalBoxes, nLocal):
nj = nAtoms[jBox]
#self.copyAtom(atoms, iId, iBox, nj, jBox)
copyAtomInner(iId, iBox, nj, jBox, MAXATOMS, gid, species, r, p, f)
nAtoms[jBox] += 1
nAtoms[iBox] -= 1
ni = nAtoms[iBox]
if ni:
#self.copyAtom(atoms, ni, iBox, iId, iBox)
copyAtomInner(ni, iBox, iId, iBox, MAXATOMS, gid, species, r, p, f)
if jBox > nLocalBoxes:
nLocal -= 1
return nLocal
@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
@numba.njit
def getBoxFromCoordInner(x, y, z, invBoxSize, nLocalBoxes, gs):
ix = int(math.floor(x*invBoxSize[0]))
iy = int(math.floor(y*invBoxSize[1]))
iz = int(math.floor(z*invBoxSize[2]))
return getBoxFromTupleInner(ix, iy, iz, nLocalBoxes, gs)
@numba.njit
def putAtomInBoxInner(nAtoms, MAXATOMS, invBoxSize, nLocalBoxes, gs, nLocal, gid, species, r, p, self_gid, iType, x, y, z, px, py, pz):
iBox = getBoxFromCoordInner(x, y, z, invBoxSize, nLocalBoxes, gs)
iOff = iBox*MAXATOMS
iOff += nAtoms[iBox]
if iBox < nLocalBoxes:
nLocal += 1
nAtoms[iBox] += 1
gid[iOff] = self_gid
species[iOff] = iType
r[iOff][0] = x
r[iOff][1] = y
r[iOff][2] = z
p[iOff][0] = px
p[iOff][1] = py
p[iOff][2] = pz
return nLocal
@numba.njit
def updateLinkCellsInner(nLocalBoxes, nTotalBoxes, nAtoms, MAXATOMS, invBoxSize, gridSize, nLocal, gid, species, r, p, f):
# empty halo cells
#self.emptyHaloCells()
emptyHaloCellsInner(nLocalBoxes, nTotalBoxes, nAtoms)
for iBox in range(nLocalBoxes):
iOff = iBox * MAXATOMS
ii = 0
while ii < nAtoms[iBox]:
#jBox = self.getBoxFromCoord(atoms.r[iOff + ii])
rr = r[iOff + ii]
jBox = getBoxFromCoordInner(rr[0], rr[1], rr[2], invBoxSize, nLocalBoxes, gridSize)
if jBox != iBox:
#self.moveAtom(atoms, ii, iBox, jBox)
nLocal = moveAtomInner(ii, iBox, jBox, MAXATOMS, gid, species, r, p, f, nAtoms, nLocalBoxes, nLocal)
else:
ii += 1
return nLocal
Raw
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
Raw
simflat.py
from __future__ import print_function, division
import ljforce
import initatoms
import linkcell
import halo
import numpy as np
class SimFlat(object):
def __init__(self, nSteps, printRate, dt, nSkip):
self.nSteps = nSteps
self.printRate = printRate
self.nSkip = nSkip
self.dt = dt
self.atoms = None
self.pot = ljforce.LJ_Pot()
self.species = initSpecies(self.pot)
self.invMass = np.array([1.0/s.mass for s in self.species])
self.ePot = 0.0
self.eKinetic = 0.0
self.boxes = None
self.atomHalo = None
class Species(object):
def __init__(self, mass):
self.mass = mass
def initSpecies(pot):
return [Species(pot.mass)]
def redistributeAtoms(sim, atoms):
boxes = sim.boxes
atoms.nLocal = linkcell.updateLinkCellsInner(boxes.nLocalBoxes, boxes.nTotalBoxes, boxes.nAtoms, boxes.MAXATOMS,
boxes.invBoxSize, boxes.gridSize,
atoms.nLocal, atoms.gid, atoms.species, atoms.r, atoms.p, atoms.f)
sim.atomHalo.haloExchange(sim)
def initSimulation(cmd):
sim = SimFlat(cmd.nSteps, cmd.printRate, cmd.dt, cmd.skip)
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment