git-svn-id: svn://svn.icms.temple.edu/lammps-ro/trunk@8631 f3b2605a-c512-4ea7-a41b-209d697bcdaa

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)