diff --git a/examples/python/README.fix_python_integrate b/examples/python/README.fix_python_integrate
new file mode 100644
index 0000000000000000000000000000000000000000..63de0910f44804a02e1f34ca78e0735d2a45d826
--- /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 0000000000000000000000000000000000000000..d1103f866aee73a93f5caab8afc19df64af038d0
--- /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 0000000000000000000000000000000000000000..fd1d11b292c4eabf33136191e5b46cc1656bc2de
--- /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 0000000000000000000000000000000000000000..ca45c0bcd0a9e847d6857530445aa024af085bc8
--- /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 a9c3116ed4cfdb6e02b0ae04a2c52335267f31ae..00f5a5544d71e2d037bb192d333842d44aa0acc5 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);
}