Instantly share code, notes, and snippets.

markdewing/comd.py Created Nov 13, 2015

What would you like to do?
More optimization of CoMD for Numba (using feedback from a profiler)
 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()
 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
 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
 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
 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
 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