diff --git a/python/lammps.py b/python/lammps.py index 8f4c427dff..7ae12e70bb 100644 --- a/python/lammps.py +++ b/python/lammps.py @@ -16,12 +16,6 @@ import types from ctypes import * -LMPINT = 0 -LMPDOUBLE = 1 -LMPIPTR = 2 -LMPDPTR = 3 -LMPDPTRPTR = 4 - class lammps: def __init__(self,name="",cmdargs=None): @@ -65,30 +59,26 @@ class lammps: self.lib.lammps_command(self.lmp,cmd) def extract_global(self,name,type): - if type == LMPDOUBLE: - self.lib.lammps_extract_global.restype = POINTER(c_double) - ptr = self.lib.lammps_extract_global(self.lmp,name) - return ptr[0] - if type == LMPINT: + if type == 0: self.lib.lammps_extract_global.restype = POINTER(c_int) - ptr = self.lib.lammps_extract_global(self.lmp,name) - return ptr[0] - return None + elif type == 1: + self.lib.lammps_extract_global.restype = POINTER(c_double) + else: return None + ptr = self.lib.lammps_extract_global(self.lmp,name) + return ptr[0] def extract_atom(self,name,type): - if type == LMPDPTRPTR: - self.lib.lammps_extract_atom.restype = POINTER(POINTER(c_double)) - ptr = self.lib.lammps_extract_atom(self.lmp,name) - return ptr - if type == LMPDPTR: - self.lib.lammps_extract_atom.restype = POINTER(c_double) - ptr = self.lib.lammps_extract_atom(self.lmp,name) - return ptr - if type == LMPIPTR: + if type == 0: self.lib.lammps_extract_atom.restype = POINTER(c_int) - ptr = self.lib.lammps_extract_atom(self.lmp,name) - return ptr - return None + elif type == 0: + self.lib.lammps_extract_atom.restype = POINTER(POINTER(c_int)) + elif type == 2: + self.lib.lammps_extract_atom.restype = POINTER(c_double) + elif type == 3: + self.lib.lammps_extract_atom.restype = POINTER(POINTER(c_double)) + else: return None + ptr = self.lib.lammps_extract_atom(self.lmp,name) + return ptr def extract_compute(self,id,style,type): if type == 0: @@ -150,18 +140,26 @@ class lammps: return result return None + # return total number of atoms in system + def get_natoms(self): return self.lib.lammps_get_natoms(self.lmp) - def get_coords(self): - nlen = 3 * self.lib.lammps_get_natoms(self.lmp) - coords = (c_double*nlen)() - self.lib.lammps_get_coords(self.lmp,coords) - return coords + # return vector of atom properties gathered across procs, ordered by atom ID - # assume coords is an array of c_double, as created by get_coords() - # could check if it is some other Python object and create c_double array? - # constructor for c_double array can take an arg to use to fill it? - - def put_coords(self,coords): - self.lib.lammps_put_coords(self.lmp,coords) + def gather_atoms(self,name,type,count): + natoms = self.lib.lammps_get_natoms(self.lmp) + if type == 0: + data = ((count*natoms)*c_double)() + self.lib.lammps_gather_atoms(self.lmp,name,type,count,data) + elif type == 1: + data = ((count*natoms)*c_double)() + self.lib.lammps_gather_atoms(self.lmp,name,type,count,data) + else: return None + return data + + # scatter vector of atom properties across procs, ordered by atom ID + # assume vector is of correct type and length, as created by gather_atoms() + + def scatter_atoms(self,name,type,count,data): + self.lib.lammps_scatter_atoms(self.lmp,name,type,count,data)