51688b2504
Removing the loop over atoms by using NumPy array indexing allows to recover performance close to that of plain fix nve.
119 lines
3.9 KiB
Python
119 lines
3.9 KiB
Python
from __future__ import print_function
|
|
import lammps
|
|
import ctypes
|
|
import traceback
|
|
import numpy as np
|
|
|
|
class LAMMPSFix(object):
|
|
def __init__(self, ptr, group_name="all"):
|
|
self.lmp = lammps.lammps(ptr=ptr)
|
|
self.group_name = group_name
|
|
|
|
class LAMMPSIntegrator(LAMMPSFix):
|
|
def __init__(self, ptr, group_name="all"):
|
|
super(LAMMPSIntegrator, self).__init__(ptr, group_name)
|
|
|
|
def init(self):
|
|
pass
|
|
|
|
def initial_integrate(self, vflag):
|
|
pass
|
|
|
|
def final_integrate(self):
|
|
pass
|
|
|
|
def initial_integrate_respa(self, vflag, ilevel, iloop):
|
|
pass
|
|
|
|
def final_integrate_respa(self, ilevel, iloop):
|
|
pass
|
|
|
|
def reset_dt(self):
|
|
pass
|
|
|
|
|
|
class NVE(LAMMPSIntegrator):
|
|
""" Python implementation of fix/nve """
|
|
def __init__(self, ptr, group_name="all"):
|
|
super(NVE, self).__init__(ptr)
|
|
assert(self.group_name == "all")
|
|
|
|
def init(self):
|
|
dt = self.lmp.extract_global("dt", 1)
|
|
ftm2v = self.lmp.extract_global("ftm2v", 1)
|
|
self.ntypes = self.lmp.extract_global("ntypes", 0)
|
|
self.dtv = dt
|
|
self.dtf = 0.5 * dt * ftm2v
|
|
|
|
def initial_integrate(self, vflag):
|
|
nlocal = self.lmp.extract_global("nlocal", 0)
|
|
mass = self.lmp.numpy.extract_atom_darray("mass", self.ntypes+1)
|
|
atype = self.lmp.numpy.extract_atom_iarray("type", nlocal)
|
|
x = self.lmp.numpy.extract_atom_darray("x", nlocal, dim=3)
|
|
v = self.lmp.numpy.extract_atom_darray("v", nlocal, dim=3)
|
|
f = self.lmp.numpy.extract_atom_darray("f", nlocal, dim=3)
|
|
|
|
for i in range(x.shape[0]):
|
|
dtfm = self.dtf / mass[int(atype[i])]
|
|
v[i,:]+= dtfm * f[i,:]
|
|
x[i,:] += self.dtv * v[i,:]
|
|
|
|
def final_integrate(self):
|
|
nlocal = self.lmp.extract_global("nlocal", 0)
|
|
mass = self.lmp.numpy.extract_atom_darray("mass", self.ntypes+1)
|
|
atype = self.lmp.numpy.extract_atom_iarray("type", nlocal)
|
|
v = self.lmp.numpy.extract_atom_darray("v", nlocal, dim=3)
|
|
f = self.lmp.numpy.extract_atom_darray("f", nlocal, dim=3)
|
|
|
|
for i in range(v.shape[0]):
|
|
dtfm = self.dtf / mass[int(atype[i])]
|
|
v[i,:] += dtfm * f[i,:]
|
|
|
|
|
|
class NVE_Opt(LAMMPSIntegrator):
|
|
""" Performance-optimized Python implementation of fix/nve """
|
|
def __init__(self, ptr, group_name="all"):
|
|
super(NVE_Opt, self).__init__(ptr)
|
|
assert(self.group_name == "all")
|
|
|
|
def init(self):
|
|
dt = self.lmp.extract_global("dt", 1)
|
|
ftm2v = self.lmp.extract_global("ftm2v", 1)
|
|
self.ntypes = self.lmp.extract_global("ntypes", 0)
|
|
self.dtv = dt
|
|
self.dtf = 0.5 * dt * ftm2v
|
|
self.mass = self.lmp.numpy.extract_atom_darray("mass", self.ntypes+1)
|
|
|
|
def initial_integrate(self, vflag):
|
|
nlocal = self.lmp.extract_global("nlocal", 0)
|
|
atype = self.lmp.numpy.extract_atom_iarray("type", nlocal)
|
|
x = self.lmp.numpy.extract_atom_darray("x", nlocal, dim=3)
|
|
v = self.lmp.numpy.extract_atom_darray("v", nlocal, dim=3)
|
|
f = self.lmp.numpy.extract_atom_darray("f", nlocal, dim=3)
|
|
dtf = self.dtf
|
|
dtv = self.dtv
|
|
mass = self.mass
|
|
|
|
dtfm = dtf / np.take(mass, atype)
|
|
dtfm.reshape((nlocal, 1))
|
|
|
|
for d in range(x.shape[1]):
|
|
v[:,d] += dtfm[:,0] * f[:,d]
|
|
x[:,d] += dtv * v[:,d]
|
|
|
|
def final_integrate(self):
|
|
nlocal = self.lmp.extract_global("nlocal", 0)
|
|
mass = self.lmp.numpy.extract_atom_darray("mass", self.ntypes+1)
|
|
atype = self.lmp.numpy.extract_atom_iarray("type", nlocal)
|
|
v = self.lmp.numpy.extract_atom_darray("v", nlocal, dim=3)
|
|
f = self.lmp.numpy.extract_atom_darray("f", nlocal, dim=3)
|
|
dtf = self.dtf
|
|
dtv = self.dtv
|
|
mass = self.mass
|
|
|
|
dtfm = dtf / np.take(mass, atype)
|
|
dtfm.reshape((nlocal, 1))
|
|
|
|
for d in range(v.shape[1]):
|
|
v[:,d] += dtfm[:,0] * f[:,d]
|