Commit b9fd1156 authored by Richard Berger's avatar Richard Berger
Browse files

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.
parent bbfe1678

examples/python/README.fix_python_integrate

0 → 100644
+15 −0
Original line number Diff line number Diff line 
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

0 → 100644
+23 −0
Original line number Diff line number Diff line 
# 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

0 → 100644
+23 −0
Original line number Diff line number Diff line 
# 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

0 → 100644
+110 −0
Original line number Diff line number Diff line 
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

+13 −21
Original line number Diff line number Diff line
@@ -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);

}

CRA Git | Maintained and supported by SUSTech CRA and CCSE