From b9fd1156b2ff53995960b4c3d1b39498a1492da9 Mon Sep 17 00:00:00 2001
From: Richard Berger <richard.berger@temple.edu>
Date: Sun, 9 Jul 2017 02:10:26 -0500
Subject: [PATCH] Completed first version of fix python/integrate

This includes an example of how to implement fix NVE in Python.

The library interface was extended to provide direct access to atom data using
numpy arrays. No data copies are made and numpy operations directly manipulate
memory of the native code.

To keep this numpy dependency optional, all functions are wrapped into the
lammps.numpy sub-object which is only loaded when accessed.
---
 examples/python/README.fix_python_integrate   |  15 +++
 examples/python/in.fix_python_integrate_melt  |  23 ++++
 .../python/in.fix_python_integrate_melt_opt   |  23 ++++
 examples/python/py_integrate.py               | 110 ++++++++++++++++++
 src/PYTHON/fix_python_integrate.cpp           |  34 +++---
 5 files changed, 184 insertions(+), 21 deletions(-)
 create mode 100644 examples/python/README.fix_python_integrate
 create mode 100644 examples/python/in.fix_python_integrate_melt
 create mode 100644 examples/python/in.fix_python_integrate_melt_opt
 create mode 100644 examples/python/py_integrate.py

diff --git a/examples/python/README.fix_python_integrate b/examples/python/README.fix_python_integrate
new file mode 100644
index 0000000000..63de0910f4
--- /dev/null
+++ b/examples/python/README.fix_python_integrate
@@ -0,0 +1,15 @@
+This folder contains several LAMMPS input scripts and a python module
+file py_integrate.py to demonstrate the use of the fix style python/integrate.
+
+in.fix_python_integrate_melt:
+This is a version of the melt example which replaces the default NVE integrator
+with a Python implementation. Fix python/integrate is used to create an
+instance of the py_integrate.NVE class which implements the required interface.
+It demonstrates how to access LAMMPS data as numpy arrays. This gives direct
+access to memory owned by the C++ code, allows easy manipulation through numpy
+operations and avoids unnecessary copies.
+
+in.fix_python_integrate_melt_opt:
+This version of melt example uses NVE_Opt instead of NVE. While this Python
+implementation is still much slower than the native version, it shows that
+simple code transformations can lead to speedups.
diff --git a/examples/python/in.fix_python_integrate_melt b/examples/python/in.fix_python_integrate_melt
new file mode 100644
index 0000000000..d1103f866a
--- /dev/null
+++ b/examples/python/in.fix_python_integrate_melt
@@ -0,0 +1,23 @@
+# 3d Lennard-Jones melt
+
+units		lj
+atom_style	atomic
+
+lattice		fcc 0.8442
+region		box block 0 10 0 10 0 10
+create_box	1 box
+create_atoms	1 box
+mass		* 1.0
+
+velocity	all create 3.0 87287
+
+pair_style	lj/cut 2.5
+pair_coeff	1 1 1.0 1.0 2.5
+
+neighbor	0.3 bin
+neigh_modify	every 20 delay 0 check no
+
+fix		1 all python/integrate py_integrate.NVE
+
+thermo		50
+run		250
diff --git a/examples/python/in.fix_python_integrate_melt_opt b/examples/python/in.fix_python_integrate_melt_opt
new file mode 100644
index 0000000000..fd1d11b292
--- /dev/null
+++ b/examples/python/in.fix_python_integrate_melt_opt
@@ -0,0 +1,23 @@
+# 3d Lennard-Jones melt
+
+units		lj
+atom_style	atomic
+
+lattice		fcc 0.8442
+region		box block 0 10 0 10 0 10
+create_box	1 box
+create_atoms	1 box
+mass		* 1.0
+
+velocity	all create 3.0 87287
+
+pair_style	lj/cut 2.5
+pair_coeff	1 1 1.0 1.0 2.5
+
+neighbor	0.3 bin
+neigh_modify	every 20 delay 0 check no
+
+fix		1 all python/integrate py_integrate.NVE_Opt
+
+thermo		50
+run		250
diff --git a/examples/python/py_integrate.py b/examples/python/py_integrate.py
new file mode 100644
index 0000000000..ca45c0bcd0
--- /dev/null
+++ b/examples/python/py_integrate.py
@@ -0,0 +1,110 @@
+from __future__ import print_function
+import lammps
+import ctypes
+import traceback
+import numpy as np
+
+class LAMMPSIntegrator(object):
+    def __init__(self, ptr):
+        self.lmp = lammps.lammps(ptr=ptr)
+
+    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):
+        super(NVE, self).__init__(ptr)
+
+    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):
+    """ Tuned Python implementation of fix/nve """
+    def __init__(self, ptr):
+        super(NVE_Opt, self).__init__(ptr)
+
+    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)
+
+        for i in range(x.shape[0]):
+            vi = v[i,:] 
+            vi += dtfm[i] * f[i,:]
+            x[i,:] += dtv * vi
+
+    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)
+
+        for i in range(v.shape[0]):
+            v[i,:] += dtfm[i] * f[i,:]
diff --git a/src/PYTHON/fix_python_integrate.cpp b/src/PYTHON/fix_python_integrate.cpp
index a9c3116ed4..00f5a5544d 100644
--- a/src/PYTHON/fix_python_integrate.cpp
+++ b/src/PYTHON/fix_python_integrate.cpp
@@ -52,11 +52,11 @@ FixPythonIntegrate::FixPythonIntegrate(LAMMPS *lmp, int narg, char **arg) :
 
 
   // create integrator instance
-  char * full_cls_name = arg[2];
+  char * full_cls_name = arg[3];
   char * lastpos = strrchr(full_cls_name, '.');
 
   if (lastpos == NULL) {
-    error->all(FLERR,"Python pair style requires fully qualified class name");
+    error->all(FLERR,"Fix python/integrate requires fully qualified class name");
   }
 
   size_t module_name_length = strlen(full_cls_name) - strlen(lastpos);
@@ -91,7 +91,11 @@ FixPythonIntegrate::FixPythonIntegrate(LAMMPS *lmp, int narg, char **arg) :
   delete [] module_name;
   delete [] cls_name;
 
-  PyObject * py_integrator_obj = PyObject_CallObject(py_integrator_type, NULL);
+  PyObject * ptr = PY_VOID_POINTER(lmp);
+  PyObject * arglist = Py_BuildValue("(O)", ptr);
+  PyObject * py_integrator_obj = PyObject_CallObject(py_integrator_type, arglist);
+  Py_DECREF(arglist);
+
   if (!py_integrator_obj) {
     PyErr_Print();
     PyErr_Clear();
@@ -139,10 +143,7 @@ void FixPythonIntegrate::init()
     PyGILState_Release(gstate);
     error->all(FLERR,"Could not find 'init' method'");
   }
-  PyObject * ptr = PY_VOID_POINTER(lmp);
-  PyObject * arglist = Py_BuildValue("(O)", ptr);
-  PyObject * result = PyEval_CallObject(py_init, arglist);
-  Py_DECREF(arglist);
+  PyObject * result = PyEval_CallObject(py_init, NULL);
   PyGILState_Release(gstate);
 }
 
@@ -159,8 +160,7 @@ void FixPythonIntegrate::initial_integrate(int vflag)
     PyGILState_Release(gstate);
     error->all(FLERR,"Could not find 'initial_integrate' method'");
   }
-  PyObject * ptr = PY_VOID_POINTER(lmp);
-  PyObject * arglist = Py_BuildValue("(Oi)", ptr, vflag);
+  PyObject * arglist = Py_BuildValue("(i)", vflag);
   PyObject * result = PyEval_CallObject(py_initial_integrate, arglist);
   Py_DECREF(arglist);
   PyGILState_Release(gstate);
@@ -179,10 +179,7 @@ void FixPythonIntegrate::final_integrate()
     PyGILState_Release(gstate);
     error->all(FLERR,"Could not find 'final_integrate' method'");
   }
-  PyObject * ptr = PY_VOID_POINTER(lmp);
-  PyObject * arglist = Py_BuildValue("(O)", ptr);
-  PyObject * result = PyEval_CallObject(py_final_integrate, arglist);
-  Py_DECREF(arglist);
+  PyObject * result = PyEval_CallObject(py_final_integrate, NULL);
   PyGILState_Release(gstate);
 }
 
@@ -199,8 +196,7 @@ void FixPythonIntegrate::initial_integrate_respa(int vflag, int ilevel, int iloo
     PyGILState_Release(gstate);
     error->all(FLERR,"Could not find 'initial_integrate_respa' method'");
   }
-  PyObject * ptr = PY_VOID_POINTER(lmp);
-  PyObject * arglist = Py_BuildValue("(Oiii)", ptr, vflag, ilevel, iloop);
+  PyObject * arglist = Py_BuildValue("(iii)", vflag, ilevel, iloop);
   PyObject * result = PyEval_CallObject(py_initial_integrate_respa, arglist);
   Py_DECREF(arglist);
   PyGILState_Release(gstate);
@@ -219,8 +215,7 @@ void FixPythonIntegrate::final_integrate_respa(int ilevel, int iloop)
     PyGILState_Release(gstate);
     error->all(FLERR,"Could not find 'final_integrate_respa' method'");
   }
-  PyObject * ptr = PY_VOID_POINTER(lmp);
-  PyObject * arglist = Py_BuildValue("(Oii)", ptr, ilevel, iloop);
+  PyObject * arglist = Py_BuildValue("(ii)", ilevel, iloop);
   PyObject * result = PyEval_CallObject(py_final_integrate_respa, arglist);
   Py_DECREF(arglist);
   PyGILState_Release(gstate);
@@ -239,9 +234,6 @@ void FixPythonIntegrate::reset_dt()
     PyGILState_Release(gstate);
     error->all(FLERR,"Could not find 'reset_dt' method'");
   }
-  PyObject * ptr = PY_VOID_POINTER(lmp);
-  PyObject * arglist = Py_BuildValue("(O)", ptr);
-  PyObject * result = PyEval_CallObject(py_reset_dt, arglist);
-  Py_DECREF(arglist);
+  PyObject * result = PyEval_CallObject(py_reset_dt, NULL);
   PyGILState_Release(gstate);
 }
-- 
GitLab