ICODE-ADC/lammps
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Detect correct integer type in lammps python interface

Browse Source
This commit is contained in:
Richard Berger
2017-07-15 14:52:08 -05:00
parent b9fd1156b2
commit 93be2d264e
2 changed files with 40 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
examples/python/py_integrate.py
View File

@ -4,9 +4,14 @@ import ctypes
import traceback import traceback
import numpy as np import numpy as np
class LAMMPSIntegrator(object): class LAMMPSFix(object):
def __init__(self, ptr): def __init__(self, ptr, group_name="all"):
self.lmp = lammps.lammps(ptr=ptr) self.lmp = lammps.lammps(ptr=ptr)
self.group_name = group_name
class LAMMPSIntegrator(LAMMPSFix):
def __init__(self, ptr, group_name="all"):
super(LAMMPSIntegrator, self).__init__(ptr, group_name)
def init(self): def init(self):
pass pass
@ -29,8 +34,9 @@ class LAMMPSIntegrator(object):
class NVE(LAMMPSIntegrator): class NVE(LAMMPSIntegrator):
""" Python implementation of fix/nve """ """ Python implementation of fix/nve """
def __init__(self, ptr): def __init__(self, ptr, group_name="all"):
super(NVE, self).__init__(ptr) super(NVE, self).__init__(ptr)
assert(self.group_name == "all")
def init(self): def init(self):
dt = self.lmp.extract_global("dt", 1) dt = self.lmp.extract_global("dt", 1)
@ -66,8 +72,9 @@ class NVE(LAMMPSIntegrator):
class NVE_Opt(LAMMPSIntegrator): class NVE_Opt(LAMMPSIntegrator):
""" Tuned Python implementation of fix/nve """ """ Tuned Python implementation of fix/nve """
def __init__(self, ptr): def __init__(self, ptr, group_name="all"):
super(NVE_Opt, self).__init__(ptr) super(NVE_Opt, self).__init__(ptr)
assert(self.group_name == "all")
def init(self): def init(self):
dt = self.lmp.extract_global("dt", 1) dt = self.lmp.extract_global("dt", 1)

40
python/lammps.py
View File

@ -266,25 +266,41 @@ class lammps(object):
def __init__(self, lmp): def __init__(self, lmp):
self.lmp = lmp self.lmp = lmp
def extract_atom_iarray(self, name, nelem, dim=1): def _ctype_to_numpy_int(self, ctype_int):
if dim == 1: if ctype_int == c_int32:
tmp = self.lmp.extract_atom(name, 0) return np.int32
ptr = cast(tmp, POINTER(c_int * nelem)) elif ctype_int == c_int64:
else: return np.int64
tmp = self.lmp.extract_atom(name, 1) return np.intc
ptr = cast(tmp[0], POINTER(c_int * nelem * dim))
a = np.frombuffer(ptr.contents, dtype=np.intc) def extract_atom_iarray(self, name, nelem, dim=1):
if name in ['id', 'molecule']:
c_int_type = self.lmp.c_tagint
elif name in ['image']:
c_int_type = self.lmp.c_imageint
else:
c_int_type = c_int
np_int_type = self._ctype_to_numpy_int(c_int_type)
if dim == 1:
tmp = self.lmp.extract_atom(name, 0)
ptr = cast(tmp, POINTER(c_int_type * nelem))
else:
tmp = self.lmp.extract_atom(name, 1)
ptr = cast(tmp[0], POINTER(c_int_type * nelem * dim))
a = np.frombuffer(ptr.contents, dtype=np_int_type)
a.shape = (nelem, dim) a.shape = (nelem, dim)
return a return a
def extract_atom_darray(self, name, nelem, dim=1): def extract_atom_darray(self, name, nelem, dim=1):
if dim == 1: if dim == 1:
tmp = self.lmp.extract_atom(name, 2) tmp = self.lmp.extract_atom(name, 2)
ptr = cast(tmp, POINTER(c_double * nelem)) ptr = cast(tmp, POINTER(c_double * nelem))
else: else:
tmp = self.lmp.extract_atom(name, 3) tmp = self.lmp.extract_atom(name, 3)
ptr = cast(tmp[0], POINTER(c_double * nelem * dim)) ptr = cast(tmp[0], POINTER(c_double * nelem * dim))
a = np.frombuffer(ptr.contents) a = np.frombuffer(ptr.contents)
a.shape = (nelem, dim) a.shape = (nelem, dim)