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

+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.

examples/python/in.fix_python_integrate_melt

+# 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

examples/python/in.fix_python_integrate_melt_opt

+# 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

examples/python/py_integrate.py

+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,:]

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);
 }