From 51688b25045cb0521d082692a04b53760a503e0f Mon Sep 17 00:00:00 2001 From: Giacomo Fiorin Date: Mon, 6 Nov 2017 15:06:54 -0500 Subject: [PATCH] Improve performance of Python integrator (NVE_Opt version) Removing the loop over atoms by using NumPy array indexing allows to recover performance close to that of plain fix nve. --- examples/python/py_integrate.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/examples/python/py_integrate.py b/examples/python/py_integrate.py index 5884ae19e8..142fe1fc13 100644 --- a/examples/python/py_integrate.py +++ b/examples/python/py_integrate.py @@ -71,7 +71,7 @@ class NVE(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"): super(NVE_Opt, self).__init__(ptr) assert(self.group_name == "all") @@ -95,11 +95,11 @@ class NVE_Opt(LAMMPSIntegrator): mass = self.mass dtfm = dtf / np.take(mass, atype) + dtfm.reshape((nlocal, 1)) - for i in range(x.shape[0]): - vi = v[i,:] - vi += dtfm[i] * f[i,:] - x[i,:] += dtv * vi + 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) @@ -112,6 +112,7 @@ class NVE_Opt(LAMMPSIntegrator): mass = self.mass dtfm = dtf / np.take(mass, atype) + dtfm.reshape((nlocal, 1)) - for i in range(v.shape[0]): - v[i,:] += dtfm[i] * f[i,:] + for d in range(v.shape[1]): + v[:,d] += dtfm[:,0] * f[:,d]