Create a gist now

Instantly share code, notes, and snippets.

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