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>

2 files changed, 142 insertions, 2 deletions

diff --git a/Examples/Modules/ParticleBoundaryScrape/PICMI_inputs_scrape.py b/Examples/Modules/ParticleBoundaryScrape/PICMI_inputs_scrape.py
new file mode 100644
index 000000000..79cf01a57
--- /dev/null
+++ b/Examples/Modules/ParticleBoundaryScrape/PICMI_inputs_scrape.py
@@ -0,0 +1,131 @@
+# --- Input file to test the particle scraper and the Python wrappers
+# --- to access the buffer of scraped particles.
+
+from pywarpx import picmi
+
+##########################
+# numerics parameters
+##########################
+
+# --- Number of time steps
+max_steps = 60
+diagnostic_intervals = 20
+
+# --- Grid
+nx = 64
+ny = 64
+nz = 128
+
+cfl = 0.99
+
+xmin = -125e-6
+ymin = -125e-6
+zmin = -149e-6
+xmax = 125e-6
+ymax = 125e-6
+zmax = 1e-6
+
+##########################
+# physics components
+##########################
+
+uniform_plasma_elec = picmi.UniformDistribution(
+    density = 1e23, # number of electrons per m^3
+    lower_bound = [-1e-5, -1e-5, -149e-6],
+    upper_bound = [1e-5, 1e-5, -129e-6],
+    directed_velocity = [0., 0., 2000.*picmi.constants.c] # uth the std of the (unitless) momentum
+)
+
+electrons = picmi.Species(
+    particle_type='electron', name='electrons',
+    initial_distribution=uniform_plasma_elec,
+    warpx_save_particles_at_xhi=1, warpx_save_particles_at_eb=1
+)
+
+##########################
+# numerics components
+##########################
+
+grid = picmi.Cartesian3DGrid(
+    number_of_cells = [nx, ny, nz],
+    lower_bound = [xmin, ymin, zmin],
+    upper_bound = [xmax, ymax, zmax],
+    lower_boundary_conditions=['none', 'none', 'none'],
+    upper_boundary_conditions=['none', 'none', 'none'],
+    lower_boundary_conditions_particles=['open', 'open', 'open'],
+    upper_boundary_conditions_particles=['open', 'open', 'open'],
+    warpx_max_grid_size = 128
+)
+
+solver = picmi.ElectromagneticSolver(
+    grid=grid, cfl=cfl
+)
+
+embedded_boundary = picmi.EmbeddedBoundary(
+    implicit_function="-max(max(max(x-12.5e-6,-12.5e-6-x),max(y-12.5e-6,-12.5e-6-y)),max(z-(-6.15e-5),-8.65e-5-z))"
+)
+
+##########################
+# diagnostics
+##########################
+
+field_diag = picmi.FieldDiagnostic(
+    name = 'diag1',
+    grid = grid,
+    period = diagnostic_intervals,
+    data_list = ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz'],
+    write_dir = '.',
+    warpx_file_prefix = 'Python_particle_scrape_plt'
+)
+
+##########################
+# simulation setup
+##########################
+
+sim = picmi.Simulation(
+    solver = solver,
+    max_steps = max_steps,
+    warpx_embedded_boundary=embedded_boundary,
+    verbose=True
+)
+
+sim.add_species(
+    electrons,
+    layout = picmi.GriddedLayout(
+        n_macroparticle_per_cell=[1, 1, 1], grid=grid
+    )
+)
+
+sim.add_diagnostic(field_diag)
+
+##########################
+# simulation run
+##########################
+
+# sim.write_input_file(file_name = 'inputs_from_PICMI')
+sim.step(max_steps)
+
+################################################
+# check that the wrappers to access the particle
+# buffer functions as intended
+################################################
+
+from pywarpx import _libwarpx
+
+n = _libwarpx.get_particle_boundary_buffer_size("electrons", 'eb')
+print("Number of electrons in buffer:", n)
+assert n == 612
+
+scraped_steps = _libwarpx.get_particle_boundary_buffer("electrons", 'eb', 'step_scraped', 0)
+for arr in scraped_steps:
+    assert all(arr > 40)
+
+weights = _libwarpx.get_particle_boundary_buffer("electrons", 'eb', 'w', 0)
+assert sum(len(arr) for arr in weights) == 612
+
+# clear the particle buffer
+_libwarpx.libwarpx.warpx_clearParticleBoundaryBuffer()
+# confirm that the buffer was cleared
+n = _libwarpx.get_particle_boundary_buffer_size("electrons", 'eb')
+print("Number of electrons in buffer:", n)
+assert n == 0
diff --git a/Examples/Modules/ParticleBoundaryScrape/analysis_scrape.py b/Examples/Modules/ParticleBoundaryScrape/analysis_scrape.py
index b970c4933..c325495db 100755
--- a/Examples/Modules/ParticleBoundaryScrape/analysis_scrape.py
+++ b/Examples/Modules/ParticleBoundaryScrape/analysis_scrape.py
@@ -1,6 +1,7 @@
 #! /usr/bin/env python
 
 import yt
+from pathlib import Path
 
 # This test shoots a beam of electrons at cubic embedded boundary geometry
 # At time step 40, none of the particles have hit the boundary yet. At time
@@ -9,11 +10,19 @@ import yt
 # the problem domain yet.
 
 # all particles are still there
-ds40 = yt.load("particle_scrape_plt00040")
+if Path("particle_scrape_plt00040").is_dir():
+    filename = "particle_scrape_plt00040"
+else:
+    filename = "Python_particle_scrape_plt00040"
+ds40 = yt.load(filename)
 np40 = ds40.index.particle_headers['electrons'].num_particles
 assert(np40 == 612)
 
 # all particles have been removed
-ds60 = yt.load("particle_scrape_plt00060")
+if Path("particle_scrape_plt00060").is_dir():
+    filename = "particle_scrape_plt00060"
+else:
+    filename = "Python_particle_scrape_plt00060"
+ds60 = yt.load(filename)
 np60 = ds60.index.particle_headers['electrons'].num_particles
 assert(np60 == 0)