Added Python wrapper of charge density arrays (#783)

2 files changed, 128 insertions, 7 deletions

diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py
index 33e700668..2550cb92a 100755
--- a/Python/pywarpx/_libwarpx.py
+++ b/Python/pywarpx/_libwarpx.py
@@ -161,6 +161,10 @@ libwarpx.warpx_getCurrentDensityCP_PML.restype = _LP_LP_c_real
 libwarpx.warpx_getCurrentDensityCPLoVects_PML.restype = _LP_c_int
 libwarpx.warpx_getCurrentDensityFP_PML.restype = _LP_LP_c_real
 libwarpx.warpx_getCurrentDensityFPLoVects_PML.restype = _LP_c_int
+libwarpx.warpx_getChargeDensityCP.restype = _LP_LP_c_real
+libwarpx.warpx_getChargeDensityCPLoVects.restype = _LP_c_int
+libwarpx.warpx_getChargeDensityFP.restype = _LP_LP_c_real
+libwarpx.warpx_getChargeDensityFPLoVects.restype = _LP_c_int
 
 #libwarpx.warpx_getPMLSigma.restype = _LP_c_real
 #libwarpx.warpx_getPMLSigmaStar.restype = _LP_c_real
@@ -657,9 +661,14 @@ def _get_mesh_field_list(warpx_func, level, direction, include_ghosts):
     size = ctypes.c_int(0)
     ncomps = ctypes.c_int(0)
     ngrow = ctypes.c_int(0)
-    data = warpx_func(level, direction,
-                      ctypes.byref(size), ctypes.byref(ncomps),
-                      ctypes.byref(ngrow), ctypes.byref(shapes))
+    if direction is None:
+        data = warpx_func(level,
+                          ctypes.byref(size), ctypes.byref(ncomps),
+                          ctypes.byref(ngrow), ctypes.byref(shapes))
+    else:
+        data = warpx_func(level, direction,
+                          ctypes.byref(size), ctypes.byref(ncomps),
+                          ctypes.byref(ngrow), ctypes.byref(shapes))
     ng = ngrow.value
     grid_data = []
     shapesize = dim
@@ -1106,6 +1115,56 @@ def get_mesh_current_density_fp_pml(level, direction, include_ghosts=True):
         return _get_mesh_field_list(libwarpx.warpx_getCurrentDensityFP_PML, level, direction, include_ghosts)
     except ValueError:
         raise Exception('PML not initialized')
+def get_mesh_charge_density_cp(level, include_ghosts=True):
+    '''
+
+    This returns a list of numpy arrays containing the mesh charge density
+    data on each grid for this process. This version returns the density for
+    the coarse patch on the given level.
+
+    The data for the numpy arrays are not copied, but share the underlying
+    memory buffer with WarpX. The numpy arrays are fully writeable.
+
+    Parameters
+    ----------
+
+        level          : the AMR level to get the data for
+        include_ghosts : whether to include ghost zones or not
+
+    Returns
+    -------
+
+        A List of numpy arrays.
+
+    '''
+
+    return _get_mesh_field_list(libwarpx.warpx_getChargeDensityCP, level, None, include_ghosts)
+
+
+def get_mesh_charge_density_fp(level, include_ghosts=True):
+    '''
+
+    This returns a list of numpy arrays containing the mesh charge density
+    data on each grid for this process. This version returns the density on
+    the fine patch for the given level.
+
+    The data for the numpy arrays are not copied, but share the underlying
+    memory buffer with WarpX. The numpy arrays are fully writeable.
+
+    Parameters
+    ----------
+
+        level          : the AMR level to get the data for
+        include_ghosts : whether to include ghost zones or not
+
+    Returns
+    -------
+
+        A List of numpy arrays.
+
+    '''
+
+    return _get_mesh_field_list(libwarpx.warpx_getChargeDensityFP, level, None, include_ghosts)
 
 
 def _get_mesh_array_lovects(level, direction, include_ghosts=True, getarrayfunc=None):
@@ -1113,7 +1172,10 @@ def _get_mesh_array_lovects(level, direction, include_ghosts=True, getarrayfunc=
 
     size = ctypes.c_int(0)
     ngrow = ctypes.c_int(0)
-    data = getarrayfunc(level, direction, ctypes.byref(size), ctypes.byref(ngrow))
+    if direction is None:
+        data = getarrayfunc(level, ctypes.byref(size), ctypes.byref(ngrow))
+    else:
+        data = getarrayfunc(level, direction, ctypes.byref(size), ctypes.byref(ngrow))
 
     lovects_ref = np.ctypeslib.as_array(data, (size.value, dim))
 
@@ -1477,3 +1539,43 @@ def get_mesh_current_density_fp_lovects_pml(level, direction, include_ghosts=Tru
         return _get_mesh_array_lovects(level, direction, include_ghosts, libwarpx.warpx_getCurrentDensityFPLoVects_PML)
     except ValueError:
         raise Exception('PML not initialized')
+def get_mesh_charge_density_cp_lovects(level, include_ghosts=True):
+    '''
+
+    This returns a list of the lo vectors of the arrays containing the mesh electric field
+    data on each grid for this process.
+
+    Parameters
+    ----------
+
+        level          : the AMR level to get the data for
+        include_ghosts : whether to include ghost zones or not
+
+    Returns
+    -------
+
+        A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids)
+
+    '''
+    return _get_mesh_array_lovects(level, None, include_ghosts, libwarpx.warpx_getChargeDensityCPLoVects)
+
+def get_mesh_charge_density_fp_lovects(level, include_ghosts=True):
+    '''
+
+    This returns a list of the lo vectors of the arrays containing the mesh electric field
+    data on each grid for this process.
+
+    Parameters
+    ----------
+
+        level          : the AMR level to get the data for
+        include_ghosts : whether to include ghost zones or not
+
+    Returns
+    -------
+
+        A 2d numpy array of the lo vector for each grid with the shape (dims, number of grids)
+
+    '''
+    return _get_mesh_array_lovects(level, None, include_ghosts, libwarpx.warpx_getChargeDensityFPLoVects)
+
diff --git a/Python/pywarpx/fields.py b/Python/pywarpx/fields.py
index 916e3c887..69aa25185 100644
--- a/Python/pywarpx/fields.py
+++ b/Python/pywarpx/fields.py
@@ -28,7 +28,7 @@ class _MultiFABWrapper(object):
     """Wrapper around field arrays at level 0
     This provides a convenient way to query and set fields that are broken up into FABs.
     The indexing is based on global indices.
-     - direction: component to access, one of the values (0, 1, 2)
+     - direction: component to access, one of the values (0, 1, 2) or None
      - overlaps: is one along the axes where the grid boundaries overlap the neighboring grid
      - get_lovects: routine that returns the list of lo vectors
      - get_fabs: routine that returns the list of FABs
@@ -48,7 +48,10 @@ class _MultiFABWrapper(object):
             self.overlaps = self.overlaps[::2]
 
     def _getlovects(self):
-        lovects = self.get_lovects(self.level, self.direction, self.include_ghosts)
+        if self.direction is None:
+            lovects = self.get_lovects(self.level, self.include_ghosts)
+        else:
+            lovects = self.get_lovects(self.level, self.direction, self.include_ghosts)
         self.nghosts = -lovects.min()
         return lovects
 
@@ -63,7 +66,10 @@ class _MultiFABWrapper(object):
         return hivects
 
     def _getfields(self):
-        return self.get_fabs(self.level, self.direction, self.include_ghosts)
+        if self.direction is None:
+            return self.get_fabs(self.level, self.include_ghosts)
+        else:
+            return self.get_fabs(self.level, self.direction, self.include_ghosts)
 
     def __len__(self):
         return lend(self._getlovects())
@@ -575,6 +581,13 @@ def JzCPWrapper(level=1, include_ghosts=False):
                             get_fabs=_libwarpx.get_mesh_current_density_cp,
                             level=level, include_ghosts=include_ghosts)
 
+def RhoCPWrapper(level=1, include_ghosts=False):
+    assert level>0, Exception('Coarse patch only available on levels > 0')
+    return _MultiFABWrapper(direction=None, overlaps=[1,1,0],
+                            get_lovects=_libwarpx.get_mesh_charge_density_cp_lovects,
+                            get_fabs=_libwarpx.get_mesh_charge_density_cp,
+                            level=level, include_ghosts=include_ghosts)
+
 def ExFPWrapper(level=0, include_ghosts=False):
     return _MultiFABWrapper(direction=0, overlaps=[0,1,1],
                             get_lovects=_libwarpx.get_mesh_electric_field_fp_lovects,
@@ -628,6 +641,12 @@ def JzFPWrapper(level=0, include_ghosts=False):
                             get_lovects=_libwarpx.get_mesh_current_density_fp_lovects,
                             get_fabs=_libwarpx.get_mesh_current_density_fp,
                             level=level, include_ghosts=include_ghosts)
+
+def RhoFPWrapper(level=0, include_ghosts=False):
+    return _MultiFABWrapper(direction=None, overlaps=[1,1,1],
+                            get_lovects=_libwarpx.get_mesh_charge_density_fp_lovects,
+                            get_fabs=_libwarpx.get_mesh_charge_density_fp,
+                            level=level, include_ghosts=include_ghosts)
 def ExCPPMLWrapper(level=1, include_ghosts=False):
     assert level>0, Exception('Coarse patch only available on levels > 0')
     return _MultiFABWrapper(direction=0, overlaps=[0,1,1],