Updated Python particles interface to allow refined levels

3 files changed, 62 insertions, 60 deletions

diff --git a/Python/pywarpx/PGroup.py b/Python/pywarpx/PGroup.py
index 48e68ceb5..f77a79004 100644
--- a/Python/pywarpx/PGroup.py
+++ b/Python/pywarpx/PGroup.py
@@ -5,9 +5,10 @@ class PGroup(object):
     """Implements a class that has the same API as a warp ParticleGroup instance.
     """
 
-    def __init__(self, igroup, ispecie):
+    def __init__(self, igroup, ispecie, level=0):
         self.igroup = igroup
         self.ispecie = ispecie
+        self.level = level
         self.ns = 1 # Number of species
 
         self.gallot()
@@ -76,35 +77,35 @@ class PGroup(object):
     npmax = property(getnpmax)
 
     def getxp(self):
-        return _libwarpx.get_particle_x(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_x(self.ispecie, self.level)[self.igroup]
     xp = property(getxp)
 
     def getyp(self):
-        return _libwarpx.get_particle_y(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_y(self.ispecie, self.level)[self.igroup]
     yp = property(getyp)
 
     def getrp(self):
-        return _libwarpx.get_particle_r(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_r(self.ispecie, self.level)[self.igroup]
     rp = property(getrp)
 
     def getzp(self):
-        return _libwarpx.get_particle_z(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_z(self.ispecie, self.level)[self.igroup]
     zp = property(getzp)
 
     def getuxp(self):
-        return _libwarpx.get_particle_ux(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_ux(self.ispecie, self.level)[self.igroup]
     uxp = property(getuxp)
 
     def getuyp(self):
-        return _libwarpx.get_particle_uy(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_uy(self.ispecie, self.level)[self.igroup]
     uyp = property(getuyp)
 
     def getuzp(self):
-        return _libwarpx.get_particle_uz(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_uz(self.ispecie, self.level)[self.igroup]
     uzp = property(getuzp)
 
     def getpid(self):
-        id = _libwarpx.get_particle_id(self.ispecie)[self.igroup]
+        id = _libwarpx.get_particle_id(self.ispecie, self.level)[self.igroup]
         pid = np.array([id]).T
         return pid
     pid = property(getpid)
@@ -117,44 +118,45 @@ class PGroup(object):
     gaminv = property(getgaminv)
 
     def getex(self):
-        return _libwarpx.get_particle_Ex(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_Ex(self.ispecie, self.level)[self.igroup]
     ex = property(getex)
 
     def getey(self):
-        return _libwarpx.get_particle_Ey(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_Ey(self.ispecie, self.level)[self.igroup]
     ey = property(getey)
 
     def getez(self):
-        return _libwarpx.get_particle_Ez(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_Ez(self.ispecie, self.level)[self.igroup]
     ez = property(getez)
 
     def getbx(self):
-        return _libwarpx.get_particle_Bx(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_Bx(self.ispecie, self.level)[self.igroup]
     bx = property(getbx)
 
     def getby(self):
-        return _libwarpx.get_particle_By(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_By(self.ispecie, self.level)[self.igroup]
     by = property(getby)
 
     def getbz(self):
-        return _libwarpx.get_particle_Bz(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_Bz(self.ispecie, self.level)[self.igroup]
     bz = property(getbz)
 
     def gettheta(self):
-        return _libwarpx.get_particle_theta(self.ispecie)[self.igroup]
+        return _libwarpx.get_particle_theta(self.ispecie, self.level)[self.igroup]
     theta = property(gettheta)
 
 class PGroups(object):
-    def __init__(self, ispecie=0):
+    def __init__(self, ispecie=0, level=0):
         self.ispecie = ispecie
+        self.level = level
 
     def setuppgroups(self):
-        xall = _libwarpx.get_particle_x(self.ispecie)
+        xall = _libwarpx.get_particle_x(self.ispecie, self.level)
         self.ngroups = len(xall)
 
         self._pgroups = []
         for igroup in range(self.ngroups):
-            self._pgroups.append(PGroup(igroup, self.ispecie))
+            self._pgroups.append(PGroup(igroup, self.ispecie, self.level))
 
     def __iter__(self):
         self.setuppgroups()
diff --git a/Python/pywarpx/WarpInterface.py b/Python/pywarpx/WarpInterface.py
index f84c22a5d..e2c9efa58 100644
--- a/Python/pywarpx/WarpInterface.py
+++ b/Python/pywarpx/WarpInterface.py
@@ -3,11 +3,11 @@ from . import fields
 from pywarpx import PGroup
 
 
-def warp_species(warp_type, picmie_species):
+def warp_species(warp_type, picmi_species, level=0):
     """Returns a Warp species that has a reference to the WarpX particles.
     """
-    elec_pgroups = PGroup.PGroups(ispecie=picmie_species.species_number)
-    return warp.Species(type=warp_type, pgroups=elec_pgroups)
+    pgroups = PGroup.PGroups(ispecie=picmi_species.species_number, level=level)
+    return warp.Species(type=warp_type, pgroups=pgroups)
 
 
 class _WarpX_FIELDtype(object):
diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py
index dc0507cb9..2b5688933 100755
--- a/Python/pywarpx/_libwarpx.py
+++ b/Python/pywarpx/_libwarpx.py
@@ -323,7 +323,7 @@ def add_particles(species_number=0,
                                  x, y, z, ux, uy, uz,
                                  attr.shape[-1], attr, unique_particles)
 
-def get_particle_structs(species_number):
+def get_particle_structs(species_number, level):
     '''
 
     This returns a list of numpy arrays containing the particle struct data
@@ -347,7 +347,7 @@ def get_particle_structs(species_number):
 
     particles_per_tile = _LP_c_int()
     num_tiles = ctypes.c_int(0)
-    data = libwarpx.warpx_getParticleStructs(species_number,
+    data = libwarpx.warpx_getParticleStructs(species_number, level,
                                              ctypes.byref(num_tiles),
                                              ctypes.byref(particles_per_tile))
 
@@ -361,7 +361,7 @@ def get_particle_structs(species_number):
     return particle_data
 
 
-def get_particle_arrays(species_number, comp):
+def get_particle_arrays(species_number, comp, level):
     '''
 
     This returns a list of numpy arrays containing the particle array data
@@ -385,7 +385,7 @@ def get_particle_arrays(species_number, comp):
 
     particles_per_tile = _LP_c_int()
     num_tiles = ctypes.c_int(0)
-    data = libwarpx.warpx_getParticleArrays(species_number, comp,
+    data = libwarpx.warpx_getParticleArrays(species_number, comp, level,
                                             ctypes.byref(num_tiles),
                                             ctypes.byref(particles_per_tile))
 
@@ -404,42 +404,42 @@ def get_particle_arrays(species_number, comp):
     return particle_data
 
 
-def get_particle_x(species_number):
+def get_particle_x(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle 'x'
     positions on each tile.
 
     '''
-    structs = get_particle_structs(species_number)
+    structs = get_particle_structs(species_number, level)
     if geometry_dim == '3d' or geometry_dim == '2d':
         return [struct['x'] for struct in structs]
     elif geometry_dim == 'rz':
         return [struct['x']*np.cos(theta) for struct, theta in zip(structs, get_particle_theta(species_number))]
 
 
-def get_particle_y(species_number):
+def get_particle_y(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle 'y'
     positions on each tile.
 
     '''
-    structs = get_particle_structs(species_number)
+    structs = get_particle_structs(species_number, level)
     if geometry_dim == '3d' or geometry_dim == '2d':
         return [struct['y'] for struct in structs]
     elif geometry_dim == 'rz':
         return [struct['x']*np.sin(theta) for struct, theta in zip(structs, get_particle_theta(species_number))]
 
 
-def get_particle_r(species_number):
+def get_particle_r(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle 'r'
     positions on each tile.
 
     '''
-    structs = get_particle_structs(species_number)
+    structs = get_particle_structs(species_number, level)
     if geometry_dim == 'rz':
         return [struct['x'] for struct in structs]
     elif geometry_dim == '3d':
@@ -448,40 +448,40 @@ def get_particle_r(species_number):
         raise Exception('get_particle_r: There is no r coordinate with 2D Cartesian')
 
 
-def get_particle_z(species_number):
+def get_particle_z(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle 'z'
     positions on each tile.
 
     '''
-    structs = get_particle_structs(species_number)
+    structs = get_particle_structs(species_number, level)
     return [struct['z'] for struct in structs]
 
 
-def get_particle_id(species_number):
+def get_particle_id(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle 'id'
     positions on each tile.
 
     '''
-    structs = get_particle_structs(species_number)
+    structs = get_particle_structs(species_number, level)
     return [struct['id'] for struct in structs]
 
 
-def get_particle_cpu(species_number):
+def get_particle_cpu(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle 'cpu'
     positions on each tile.
 
     '''
-    structs = get_particle_structs(species_number)
+    structs = get_particle_structs(species_number, level)
     return [struct['cpu'] for struct in structs]
 
 
-def get_particle_weight(species_number):
+def get_particle_weight(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle
@@ -489,10 +489,10 @@ def get_particle_weight(species_number):
 
     '''
 
-    return get_particle_arrays(species_number, 0)
+    return get_particle_arrays(species_number, 0, level)
 
 
-def get_particle_ux(species_number):
+def get_particle_ux(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle
@@ -500,10 +500,10 @@ def get_particle_ux(species_number):
 
     '''
 
-    return get_particle_arrays(species_number, 1)
+    return get_particle_arrays(species_number, 1, level)
 
 
-def get_particle_uy(species_number):
+def get_particle_uy(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle
@@ -511,10 +511,10 @@ def get_particle_uy(species_number):
 
     '''
 
-    return get_particle_arrays(species_number, 2)
+    return get_particle_arrays(species_number, 2, level)
 
 
-def get_particle_uz(species_number):
+def get_particle_uz(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle
@@ -522,10 +522,10 @@ def get_particle_uz(species_number):
 
     '''
 
-    return get_particle_arrays(species_number, 3)
+    return get_particle_arrays(species_number, 3, level)
 
 
-def get_particle_Ex(species_number):
+def get_particle_Ex(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle
@@ -533,10 +533,10 @@ def get_particle_Ex(species_number):
 
     '''
 
-    return get_particle_arrays(species_number, 4)
+    return get_particle_arrays(species_number, 4, level)
 
 
-def get_particle_Ey(species_number):
+def get_particle_Ey(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle
@@ -544,10 +544,10 @@ def get_particle_Ey(species_number):
 
     '''
 
-    return get_particle_arrays(species_number, 5)
+    return get_particle_arrays(species_number, 5, level)
 
 
-def get_particle_Ez(species_number):
+def get_particle_Ez(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle
@@ -555,10 +555,10 @@ def get_particle_Ez(species_number):
 
     '''
 
-    return get_particle_arrays(species_number, 6)
+    return get_particle_arrays(species_number, 6, level)
 
 
-def get_particle_Bx(species_number):
+def get_particle_Bx(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle
@@ -566,10 +566,10 @@ def get_particle_Bx(species_number):
 
     '''
 
-    return get_particle_arrays(species_number, 7)
+    return get_particle_arrays(species_number, 7, level)
 
 
-def get_particle_By(species_number):
+def get_particle_By(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle
@@ -577,10 +577,10 @@ def get_particle_By(species_number):
 
     '''
 
-    return get_particle_arrays(species_number, 8)
+    return get_particle_arrays(species_number, 8, level)
 
 
-def get_particle_Bz(species_number):
+def get_particle_Bz(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle
@@ -588,10 +588,10 @@ def get_particle_Bz(species_number):
 
     '''
 
-    return get_particle_arrays(species_number, 9)
+    return get_particle_arrays(species_number, 9, level)
 
 
-def get_particle_theta(species_number):
+def get_particle_theta(species_number, level=0):
     '''
 
     Return a list of numpy arrays containing the particle
@@ -600,7 +600,7 @@ def get_particle_theta(species_number):
     '''
 
     if geometry_dim == 'rz':
-        return get_particle_arrays(species_number, 10)
+        return get_particle_arrays(species_number, 10, level)
     elif geometry_dim == '3d':
         return [np.arctan2(struct['y'], struct['x']) for struct in structs]
     elif geometry_dim == '2d':