Make buffer of scraped particles available to Python code (#2164)

* Added wrapper to get number of particle species tracked by the scraper

Not sure if this is going to be useful, but it demonstrates a method to get information
  from the ParticleBoundaryBuffer into Python.

* Stubbed out the main wrapper functions

* Added parameters to wrapper

* Added wrapper for getting the number of particles scraped of a species on a boundary

* added picmi arguments to scrape particles at the domain boundary

* Added wrapper to get the full particle buffer into python

* rearanged the getBuffer properties code a little

* Added docstrings +other suggested changes

* Added num_particles_impacted_boundary docstring

* fixed mistake in docstring

* Changed boundary parameter to be a string for clarity

* Fixed issue with the boundary parameter for scraping

* Fixed issue with the boundary input for scraping stats wrapper

* Added demonstration of particle scraping wrapper

* Added analysis.py file

* Fix typo in one of the dimension maps

Co-authored-by: Roelof Groenewald <40245517+roelof-groenewald@users.noreply.github.com>

* Added before esolve to warpx evolve

* added test for the scraped particle buffer wrappers

* Moved python PICMI particle boundary scrape test

* Renamed test file to the correct name

* Removed old test

* added special functionality to get the timestep at which particles were scraped

* removed debug print

* added python wrapper for the clearParticles() function of the scraper buffer

* added special wrapper function to get the timesteps at which the particles in the boundary buffer were scraped

* updated test to match the non-PICMI test for the particle scraper buffer

* Fix uncaught rebase mistake

* re-activated picmi test of accessing the scraped particle buffers via python

* added documentation for the new parameters involved in the scraped particle buffer and fixed remaining issue with picmi test

* changes requested during code review

Co-authored-by: mkieburtz <michaelkieburtz@gmail.com>
Co-authored-by: Roelof <roelof.groenewald@modernelectron.com>
Co-authored-by: Roelof Groenewald <40245517+roelof-groenewald@users.noreply.github.com>

1 files changed, 126 insertions, 0 deletions

diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py
index 110de5190..ca87f4f41 100755
--- a/Python/pywarpx/_libwarpx.py
+++ b/Python/pywarpx/_libwarpx.py
@@ -120,6 +120,7 @@ class Particle(ctypes.Structure):
 # some useful typenames
 _LP_particle_p = ctypes.POINTER(ctypes.POINTER(Particle))
 _LP_c_int = ctypes.POINTER(ctypes.c_int)
+_LP_LP_c_int = ctypes.POINTER(_LP_c_int)
 _LP_c_void_p = ctypes.POINTER(ctypes.c_void_p)
 _LP_c_real = ctypes.POINTER(c_real)
 _LP_LP_c_real = ctypes.POINTER(_LP_c_real)
@@ -185,6 +186,9 @@ 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_getParticleBoundaryBufferSize.restype = ctypes.c_int
+libwarpx.warpx_getParticleBoundaryBuffer.restype = _LP_LP_c_particlereal
+libwarpx.warpx_getParticleBoundaryBufferScrapedSteps.restype = _LP_LP_c_int
 
 libwarpx.warpx_getEx_nodal_flag.restype = _LP_c_int
 libwarpx.warpx_getEy_nodal_flag.restype = _LP_c_int
@@ -763,6 +767,127 @@ def add_real_comp(species_name, pid_name, comm=True):
     )
 
 
+def _get_boundary_number(boundary):
+    '''
+
+    Utility function to find the boundary number given a boundary name.
+
+    Parameters
+    ----------
+
+        boundary       : the boundary from which to get the scraped particle data.
+                         In the form x/y/z_hi/lo or eb.
+
+    Returns
+    -------
+
+    Integer index in the boundary scraper buffer for the given boundary.
+    '''
+    if geometry_dim == '3d':
+        dimensions = {'x' : 0, 'y' : 1, 'z' : 2}
+    elif geometry_dim == '2d':
+        dimensions = {'x' : 0, 'z' : 1}
+    else:
+        raise NotImplementedError("RZ is not supported for particle scraping.")
+
+    if boundary != 'eb':
+        boundary_parts = boundary.split("_")
+        dim_num = dimensions[boundary_parts[0]]
+        if boundary_parts[1] == 'lo':
+            side = 0
+        elif boundary_parts[1] == 'hi':
+            side = 1
+        else:
+            raise RuntimeError(f'Unknown boundary specified: {boundary}')
+        boundary_num = 2 * dim_num + side
+    else:
+        boundary_num = 4 if geometry_dim == '2d' else 6
+
+    return boundary_num
+
+
+def get_particle_boundary_buffer_size(species_name, boundary):
+    '''
+
+    This returns the number of particles that have been scraped so far in the simulation
+    from the specified boundary and of the specified species.
+
+    Parameters
+    ----------
+
+        species_name   : return the number of scraped particles of this species
+        boundary       : the boundary from which to get the scraped particle data.
+                         In the form x/y/z_hi/lo
+
+    Returns
+    -------
+
+        The number of particles scraped so far from a boundary and of a species.
+
+    '''
+    return libwarpx.warpx_getParticleBoundaryBufferSize(
+        ctypes.c_char_p(species_name.encode('utf-8')),
+        _get_boundary_number(boundary)
+    )
+
+
+def get_particle_boundary_buffer(species_name, boundary, comp_name, level):
+    '''
+
+    This returns a list of numpy arrays containing the particle array data
+    for a species that has been scraped by a specific simulation boundary.
+
+    The data for the numpy arrays are not copied, but share the underlying
+    memory buffer with WarpX. The numpy arrays are fully writeable.
+
+    Parameters
+    ----------
+
+        species_name   : the species name that the data will be returned for.
+        boundary       : the boundary from which to get the scraped particle data.
+                         In the form x/y/z_hi/lo or eb.
+        comp_name      : the component of the array data that will be returned.
+                         If "step_scraped" the special attribute holding the
+                         timestep at which a particle was scraped will be
+                         returned.
+        level          : Which AMR level to retrieve scraped particle data from.
+    Returns
+    -------
+
+        A List of numpy arrays.
+
+    '''
+    particles_per_tile = _LP_c_int()
+    num_tiles = ctypes.c_int(0)
+    if comp_name == 'step_scraped':
+        data = libwarpx.warpx_getParticleBoundaryBufferScrapedSteps(
+            ctypes.c_char_p(species_name.encode('utf-8')),
+            _get_boundary_number(boundary), level,
+            ctypes.byref(num_tiles), ctypes.byref(particles_per_tile)
+        )
+    else:
+        data = libwarpx.warpx_getParticleBoundaryBuffer(
+            ctypes.c_char_p(species_name.encode('utf-8')),
+            _get_boundary_number(boundary), level,
+            ctypes.byref(num_tiles), ctypes.byref(particles_per_tile),
+            ctypes.c_char_p(comp_name.encode('utf-8'))
+        )
+
+    particle_data = []
+    for i in range(num_tiles.value):
+        arr = np.ctypeslib.as_array(data[i], (particles_per_tile[i],))
+        try:
+            # This fails on some versions of numpy
+            arr.setflags(write=1)
+        except ValueError:
+            pass
+        particle_data.append(arr)
+
+    _libc.free(particles_per_tile)
+    _libc.free(data)
+    return particle_data
+
+
 def _get_mesh_field_list(warpx_func, level, direction, include_ghosts):
     """
      Generic routine to fetch the list of field data arrays.
@@ -1677,6 +1802,7 @@ def get_mesh_charge_density_cp_lovects(level, include_ghosts=True):
     '''
     return _get_mesh_array_lovects(level, None, include_ghosts, libwarpx.warpx_getChargeDensityCPLoVects)
 
+
 def get_mesh_charge_density_fp_lovects(level, include_ghosts=True):
     '''