Instantly share code, notes, and snippets.

markdewing/linkcell.py Created Oct 2, 2015

What would you like to do?
Numba-optimzed version of CoMD cell version.
 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): 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 @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 getNeighborBoxes(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() 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 @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
 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
to join this conversation on GitHub. Already have an account? Sign in to comment