ICODE-ADC/lammps
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve performance of Python integrator (NVE_Opt version)

Browse Source 
Removing the loop over atoms by using NumPy array indexing allows to recover
performance close to that of plain fix nve.
This commit is contained in:
Giacomo Fiorin
2017-11-06 15:06:54 -05:00
committed by Richard Berger
parent 93be2d264e
commit 51688b2504
1 changed files with 8 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
examples/python/py_integrate.py
View File

@ -71,7 +71,7 @@ class NVE(LAMMPSIntegrator):
class NVE_Opt(LAMMPSIntegrator): class NVE_Opt(LAMMPSIntegrator):
""" Tuned Python implementation of fix/nve """ """ Performance-optimized Python implementation of fix/nve """
def __init__(self, ptr, group_name="all"): def __init__(self, ptr, group_name="all"):
super(NVE_Opt, self).__init__(ptr) super(NVE_Opt, self).__init__(ptr)
assert(self.group_name == "all") assert(self.group_name == "all")
@ -95,11 +95,11 @@ class NVE_Opt(LAMMPSIntegrator):
mass = self.mass mass = self.mass
dtfm = dtf / np.take(mass, atype) dtfm = dtf / np.take(mass, atype)
dtfm.reshape((nlocal, 1))
for i in range(x.shape[0]): for d in range(x.shape[1]):
vi = v[i,:] v[:,d] += dtfm[:,0] * f[:,d]
vi += dtfm[i] * f[i,:] x[:,d] += dtv * v[:,d]
x[i,:] += dtv * vi
def final_integrate(self): def final_integrate(self):
nlocal = self.lmp.extract_global("nlocal", 0) nlocal = self.lmp.extract_global("nlocal", 0)
@ -112,6 +112,7 @@ class NVE_Opt(LAMMPSIntegrator):
mass = self.mass mass = self.mass
dtfm = dtf / np.take(mass, atype) dtfm = dtf / np.take(mass, atype)
dtfm.reshape((nlocal, 1))
for i in range(v.shape[0]): for d in range(v.shape[1]):
v[i,:] += dtfm[i] * f[i,:] v[:,d] += dtfm[:,0] * f[:,d]