Switched high level Python wrapper to ctypes

1 files changed, 72 insertions, 125 deletions

diff --git a/Python/pywarpx/PGroup.py b/Python/pywarpx/PGroup.py
index 616141702..bf59d5ce5 100644
--- a/Python/pywarpx/PGroup.py
+++ b/Python/pywarpx/PGroup.py
@@ -1,15 +1,14 @@
 import numpy as np
 from . import WarpX
-from . import warpxC
+from . import _libwarpx
 
 class PGroup(object):
     """Implements a class that has the same API as a warp ParticleGroup instance.
     """
 
-    def __init__(self):
+    def __init__(self, igroup):
+        self.igroup = igroup
         self.ns = 1 # Number of species
-        self.npmax = 0 # Size of data arrays
-        self.npid = 0 # number of columns for pid.
 
         self.gallot()
 
@@ -24,10 +23,6 @@ class PGroup(object):
         self.sq = np.zeros(self.ns) # Species charge [C]
         self.sw = np.zeros(self.ns) # Species weight, (real particles per simulation particles)
 
-        self.ins = np.ones(self.ns, dtype=int) # Index of first particle in species
-        self.nps = np.zeros(self.ns, dtype=int) # Number of particles in species
-        self.ipmax = np.zeros(self.ns+1, dtype=int) # Max extent within the arrays of each species
-
         self.sid = np.arange(self.ns, dtype=int) # Global species index for each species
         self.ndts = np.ones(self.ns, dtype=int) # Stride for time step advance for each species
         self.ldts = np.ones(self.ns, dtype=int) # (logical type)
@@ -51,152 +46,104 @@ class PGroup(object):
         self.zshift = np.zeros(self.ns)
         self.gamma_ebcancel_max = np.ones(self.ns) # maximum value allowed for ExB cancellation
 
-        self.gaminv = np.ones(self.npmax) # inverse relativistic gamma factor
-        self._xp = np.zeros(self.npmax) # X-positions of particles [m]
-        self._yp = np.zeros(self.npmax) # Y-positions of particles [m]
-        self._zp = np.zeros(self.npmax) # Z-positions of particles [m]
-        self._uxp = np.zeros(self.npmax) # gamma * X-velocities of particles [m/s]
-        self._uyp = np.zeros(self.npmax) # gamma * Y-velocities of particles [m/s]
-        self._uzp = np.zeros(self.npmax) # gamma * Z-velocities of particles [m/s]
-        self._ex = np.zeros(self.npmax) # Ex of particles [V/m]
-        self._ey = np.zeros(self.npmax) # Ey of particles [V/m]
-        self._ez = np.zeros(self.npmax) # Ez of particles [V/m]
-        self._bx = np.zeros(self.npmax) # Bx of particles [T]
-        self._by = np.zeros(self.npmax) # By of particles [T]
-        self._bz = np.zeros(self.npmax) # Bz of particles [T]
-        self._pid = np.zeros((self.npmax, self.npid)) # Particle ID - used for various purposes
-
     # --- Temporary fix
     gchange = gallot
 
     def allocated(self, name):
-        return (getattr(self, name, None) is not None)
+        return True
 
     def addspecies(self):
         pass
 
-    def _updatelocations(self):
-        warpx = WarpX.warpx.warpx
-        mypc = warpx.GetPartContainer()
-
-        xplist = []
-        yplist = []
-        zplist = []
-        for ispecie in range(mypc.nSpecies()):
-            pc = mypc.GetParticleContainer(ispecie)
-            xx = pc.getLocations()
-            xplist.append(xx[0,:])
-            yplist.append(xx[1,:])
-            zplist.append(xx[2,:])
-            self.nps[ispecie] = len(xplist[-1])
-            if ispecie > 0:
-                self.ins[ispecie] = self.ins[ispecie-1] + self.nps[ispecie-1]
-                self.ipmax[ispecie+1] = self.ins[ispecie] + self.nps[ispecie] - 1
-
-        self._xp = np.concatenate(xplist)
-        self._yp = np.concatenate(yplist)
-        self._zp = np.concatenate(zplist)
-        self.npmax = len(self._xp)
-
-    def _updatevelocities(self):
-        warpx = WarpX.warpx.warpx
-        mypc = warpx.GetPartContainer()
-
-        uxplist = []
-        uyplist = []
-        uzplist = []
-        for ispecie in range(mypc.nSpecies()):
-            pc = mypc.GetParticleContainer(ispecie)
-            vv = pc.getData(0, 3)
-            uxplist.append(vv[0,:])
-            uyplist.append(vv[1,:])
-            uzplist.append(vv[2,:])
-            self.nps[ispecie] = len(uxplist[-1])
-            if ispecie > 0:
-                self.ins[ispecie] = self.ins[ispecie-1] + self.nps[ispecie-1]
-                self.ipmax[ispecie+1] = self.ins[ispecie] + self.nps[ispecie] - 1
-
-        self._uxp = np.concatenate(uxplist)
-        self._uyp = np.concatenate(uyplist)
-        self._uzp = np.concatenate(uzplist)
-        self.npmax = len(self._xp)
-
-    def _updatepids(self):
-        warpx = WarpX.warpx.warpx
-        mypc = warpx.GetPartContainer()
-
-        pidlist = []
-        for ispecie in range(mypc.nSpecies()):
-            pc = mypc.GetParticleContainer(ispecie)
-            self.npid = pc.nAttribs - 3
-            vv = pc.getData(3, self.npid)
-            pidlist.append(vv)
-            self.nps[ispecie] = len(uxplist[-1])
-            if ispecie > 0:
-                self.ins[ispecie] = self.ins[ispecie-1] + self.nps[ispecie-1]
-                self.ipmax[ispecie+1] = self.ins[ispecie] + self.nps[ispecie] - 1
-
-        self._pid = np.concatenate(pidlist.T, axis=0)
-        self.npmax = self._pid.shape[0]
-
-    def getxp(self):
-        self._updatelocations()
-        return self._xp
+
+    def getnps(self):
+        return np.array([len(self.xp)], dtype='l')
+    nps = property(getnps)
+
+    def getins(self):
+        return np.ones(self.ns, dtype='l')
+    ins = property(getins)
+
+    def getipmax(self):
+        return np.array([0, len(self.xp)], dtype='l')
+    ipmax = property(getipmax)
+
+    def getnpmax(self):
+        return self.nps.sum()
+    npmax = property(getnpmax)
+
+    def getxp(self, js=0):
+        return _libwarpx.get_particle_x(js)[self.igroup]
     xp = property(getxp)
 
-    def getyp(self):
-        self._updatelocations()
-        return self._yp
+    def getyp(self, js=0):
+        return _libwarpx.get_particle_y(js)[self.igroup]
     yp = property(getyp)
 
-    def getzp(self):
-        self._updatelocations()
-        return self._zp
+    def getzp(self, js=0):
+        return _libwarpx.get_particle_z(js)[self.igroup]
     zp = property(getzp)
 
-    def getuxp(self):
-        self._updatevelocities()
-        return self._uxp
+    def getuxp(self, js=0):
+        return _libwarpx.get_particle_ux(js)[self.igroup]
     uxp = property(getuxp)
 
-    def getuyp(self):
-        self._updatevelocities()
-        return self._uyp
+    def getuyp(self, js=0):
+        return _libwarpx.get_particle_uy(js)[self.igroup]
     uyp = property(getuyp)
 
-    def getuzp(self):
-        self._updatevelocities()
-        return self._uzp
+    def getuzp(self, js=0):
+        return _libwarpx.get_particle_uz(js)[self.igroup]
     uzp = property(getuzp)
 
-    def getpid(self):
-        self._updatepids()
-        return self._pid
+    def getpid(self, js=0):
+        return _libwarpx.get_particle_id(js)[self.igroup]
     pid = property(getpid)
 
-    def getgaminv(self):
-        uxp = self.uxp
-        uyp = self.uyp
-        uzp = self.uzp
+    def getgaminv(self, js=0):
+        uxp = self.getuxp(js)
+        uyp = self.getuyp(js)
+        uzp = self.getuzp(js)
         return sqrt(1. - (uxp**2 + uyp**2 + uzp**2)/warpxC.c**2)
     gaminv = property(getgaminv)
 
-    def getex(self):
-        return np.zeros(self.npmax)
+    def getex(self, js=0):
+        return _libwarpx.get_particle_Ex(js)[self.igroup]
     ex = property(getex)
-    def getey(self):
-        return np.zeros(self.npmax)
+
+    def getey(self, js=0):
+        return _libwarpx.get_particle_Ey(js)[self.igroup]
     ey = property(getey)
-    def getez(self):
-        return np.zeros(self.npmax)
+
+    def getez(self, js=0):
+        return _libwarpx.get_particle_Ez(js)[self.igroup]
     ez = property(getez)
-    def getbx(self):
-        return np.zeros(self.npmax)
+
+    def getbx(self, js=0):
+        return _libwarpx.get_particle_Bx(js)[self.igroup]
     bx = property(getbx)
-    def getby(self):
-        return np.zeros(self.npmax)
+
+    def getby(self, js=0):
+        return _libwarpx.get_particle_By(js)[self.igroup]
     by = property(getby)
-    def getbz(self):
-        return np.zeros(self.npmax)
+
+    def getbz(self, js=0):
+        return _libwarpx.get_particle_Bz(js)[self.igroup]
     bz = property(getbz)
 
+class PGroups(object):
+    def __init__(self):
+        xall = _libwarpx.get_particle_x(0)
+        self.ngroups = len(xall)
+
+        self._pgroups = []
+        for igroup in range(self.ngroups):
+            self._pgroups.append(PGroup(igroup))
+        
+    def __iter__(self):
+        for igroup in range(self.ngroups):
+            yield self._pgroups[igroup]
+
+    def __getitem__(self, key):
+        return self._pgroups[key]
+