Added optional numpy access to atom data

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.
This commit is contained in:
Richard Berger
2017-09-04 19:59:50 -04:00
parent 6b2d321d18
commit bb3ee752e5

View File

@ -162,6 +162,9 @@ class lammps(object):
pythonapi.PyCObject_AsVoidPtr.argtypes = [py_object] pythonapi.PyCObject_AsVoidPtr.argtypes = [py_object]
self.lmp = c_void_p(pythonapi.PyCObject_AsVoidPtr(ptr)) self.lmp = c_void_p(pythonapi.PyCObject_AsVoidPtr(ptr))
# optional numpy support (lazy loading)
self._numpy = None
def __del__(self): def __del__(self):
if self.lmp and self.opened: if self.lmp and self.opened:
self.lib.lammps_close(self.lmp) self.lib.lammps_close(self.lmp)
@ -236,6 +239,41 @@ class lammps(object):
ptr = self.lib.lammps_extract_atom(self.lmp,name) ptr = self.lib.lammps_extract_atom(self.lmp,name)
return ptr return ptr
@property
def numpy(self):
if not self._numpy:
import numpy as np
class LammpsNumpyWrapper:
def __init__(self, lmp):
self.lmp = lmp
def extract_atom_iarray(self, name, nelem, dim=1):
if dim == 1:
tmp = self.lmp.extract_atom(name, 0)
ptr = cast(tmp, POINTER(c_int * nelem))
else:
tmp = self.lmp.extract_atom(name, 1)
ptr = cast(tmp[0], POINTER(c_int * nelem * dim))
a = np.frombuffer(ptr.contents, dtype=np.intc)
a.shape = (nelem, dim)
return a
def extract_atom_darray(self, name, nelem, dim=1):
if dim == 1:
tmp = self.lmp.extract_atom(name, 2)
ptr = cast(tmp, POINTER(c_double * nelem))
else:
tmp = self.lmp.extract_atom(name, 3)
ptr = cast(tmp[0], POINTER(c_double * nelem * dim))
a = np.frombuffer(ptr.contents)
a.shape = (nelem, dim)
return a
self._numpy = LammpsNumpyWrapper(self)
return self._numpy
# extract compute info # extract compute info
def extract_compute(self,id,style,type): def extract_compute(self,id,style,type):