Fix python wrapper (#1532)

* Fixes to the Python interface for accessing field and particle data

* In Python wrapper, cleaned up how ngrow is handled

1 files changed, 22 insertions, 13 deletions

diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py
index 953de0698..2560da235 100755
--- a/Python/pywarpx/_libwarpx.py
+++ b/Python/pywarpx/_libwarpx.py
@@ -533,7 +533,10 @@ def get_particle_z(species_number, level=0):
 
     '''
     structs = get_particle_structs(species_number, level)
-    return [struct['z'] for struct in structs]
+    if geometry_dim == '3d':
+        return [struct['z'] for struct in structs]
+    elif geometry_dim == 'rz' or geometry_dim == '2d':
+        return [struct['y'] for struct in structs]
 
 
 def get_particle_id(species_number, level=0):
@@ -625,16 +628,16 @@ def _get_mesh_field_list(warpx_func, level, direction, include_ghosts):
     shapes = _LP_c_int()
     size = ctypes.c_int(0)
     ncomps = ctypes.c_int(0)
-    ngrow = ctypes.c_int(0)
+    ngrowvect = _LP_c_int()
     if direction is None:
         data = warpx_func(level,
                           ctypes.byref(size), ctypes.byref(ncomps),
-                          ctypes.byref(ngrow), ctypes.byref(shapes))
+                          ctypes.byref(ngrowvect), ctypes.byref(shapes))
     else:
         data = warpx_func(level, direction,
                           ctypes.byref(size), ctypes.byref(ncomps),
-                          ctypes.byref(ngrow), ctypes.byref(shapes))
-    ng = ngrow.value
+                          ctypes.byref(ngrowvect), ctypes.byref(shapes))
+    ngvect = [ngrowvect[i] for i in range(dim)]
     grid_data = []
     shapesize = dim
     if ncomps.value > 1:
@@ -651,7 +654,7 @@ def _get_mesh_field_list(warpx_func, level, direction, include_ghosts):
         if include_ghosts:
             grid_data.append(arr)
         else:
-            grid_data.append(arr[tuple([slice(ng, -ng) for _ in range(dim)])])
+            grid_data.append(arr[tuple([slice(ngvect[d], -ngvect[d]) for d in range(dim)])])
 
     _libc.free(shapes)
     _libc.free(data)
@@ -1132,15 +1135,15 @@ def get_mesh_charge_density_fp(level, include_ghosts=True):
     return _get_mesh_field_list(libwarpx.warpx_getChargeDensityFP, level, None, include_ghosts)
 
 
-def _get_mesh_array_lovects(level, direction, include_ghosts=True, getarrayfunc=None):
+def _get_mesh_array_lovects(level, direction, include_ghosts=True, getlovectsfunc=None):
     assert(0 <= level and level <= libwarpx.warpx_finestLevel())
 
     size = ctypes.c_int(0)
-    ngrow = ctypes.c_int(0)
+    ngrowvect = _LP_c_int()
     if direction is None:
-        data = getarrayfunc(level, ctypes.byref(size), ctypes.byref(ngrow))
+        data = getlovectsfunc(level, ctypes.byref(size), ctypes.byref(ngrowvect))
     else:
-        data = getarrayfunc(level, direction, ctypes.byref(size), ctypes.byref(ngrow))
+        data = getlovectsfunc(level, direction, ctypes.byref(size), ctypes.byref(ngrowvect))
 
     lovects_ref = np.ctypeslib.as_array(data, (size.value, dim))
 
@@ -1148,12 +1151,18 @@ def _get_mesh_array_lovects(level, direction, include_ghosts=True, getarrayfunc=
     # --- Also, take the transpose to give shape (dims, number of grids)
     lovects = lovects_ref.copy().T
 
-    if not include_ghosts:
-        lovects += ngrow.value
+    ng = []
+    if include_ghosts:
+        for d in range(dim):
+            ng.append(ngrowvect[d])
+    else:
+        for d in range(dim):
+            ng.append(0)
+            lovects[d,:] += ngrowvect[d]
 
     del lovects_ref
     _libc.free(data)
-    return lovects
+    return lovects, ng
 
 
 def get_mesh_electric_field_lovects(level, direction, include_ghosts=True):