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, 79 insertions, 0 deletions

diff --git a/Source/Python/WarpXWrappers.cpp b/Source/Python/WarpXWrappers.cpp
index 1450011f9..dd4cf0e75 100644
--- a/Source/Python/WarpXWrappers.cpp
+++ b/Source/Python/WarpXWrappers.cpp
@@ -9,6 +9,7 @@
 #include "BoundaryConditions/PML.H"
 #include "Initialization/WarpXAMReXInit.H"
 #include "Particles/MultiParticleContainer.H"
+#include "Particles/ParticleBoundaryBuffer.H"
 #include "Particles/WarpXParticleContainer.H"
 #include "Utils/WarpXUtil.H"
 #include "WarpX.H"
@@ -471,6 +472,72 @@ extern "C"
         mypc.defineAllParticleTiles();
     }
 
+    int warpx_getParticleBoundaryBufferSize(const char* species_name, int boundary)
+    {
+        const std::string name(species_name);
+        auto& particle_buffers = WarpX::GetInstance().GetParticleBoundaryBuffer();
+        return particle_buffers.getNumParticlesInContainer(species_name, boundary);
+    }
+
+    int** warpx_getParticleBoundaryBufferScrapedSteps(const char* species_name, int boundary, int lev,
+                     int* num_tiles, int** particles_per_tile)
+    {
+        const std::string name(species_name);
+        auto& particle_buffers = WarpX::GetInstance().GetParticleBoundaryBuffer();
+        auto& particle_buffer = particle_buffers.getParticleBuffer(species_name, boundary);
+
+        const int comp = particle_buffer.NumIntComps() - 1;
+
+        int i = 0;
+        for (amrex::ParIter<0,0,PIdx::nattribs, 0, amrex::PinnedArenaAllocator> pti(particle_buffer, lev); pti.isValid(); ++pti, ++i) {}
+
+        // *num_tiles = myspc.numLocalTilesAtLevel(lev);
+        *num_tiles = i;
+        *particles_per_tile = (int*) malloc(*num_tiles*sizeof(int));
+
+        int** data = (int**) malloc(*num_tiles*sizeof(int*));
+        i = 0;
+        for (amrex::ParIter<0,0,PIdx::nattribs, 0, amrex::PinnedArenaAllocator> pti(particle_buffer, lev); pti.isValid(); ++pti, ++i) {
+            auto& soa = pti.GetStructOfArrays();
+            data[i] = (int*) soa.GetIntData(comp).dataPtr();
+            (*particles_per_tile)[i] = pti.numParticles();
+        }
+
+        return data;
+    }
+
+    amrex::ParticleReal** warpx_getParticleBoundaryBuffer(const char* species_name, int boundary, int lev,
+                     int* num_tiles, int** particles_per_tile, const char* comp_name)
+    {
+        const std::string name(species_name);
+        auto& particle_buffers = WarpX::GetInstance().GetParticleBoundaryBuffer();
+        auto& particle_buffer = particle_buffers.getParticleBuffer(species_name, boundary);
+
+        const int comp = warpx_getParticleCompIndex(species_name, comp_name);
+
+        int i = 0;
+        for (amrex::ParIter<0,0,PIdx::nattribs, 0, amrex::PinnedArenaAllocator> pti(particle_buffer, lev); pti.isValid(); ++pti, ++i) {}
+
+        // *num_tiles = myspc.numLocalTilesAtLevel(lev);
+        *num_tiles = i;
+        *particles_per_tile = (int*) malloc(*num_tiles*sizeof(int));
+
+        amrex::ParticleReal** data = (amrex::ParticleReal**) malloc(*num_tiles*sizeof(amrex::ParticleReal*));
+        i = 0;
+        for (amrex::ParIter<0,0,PIdx::nattribs, 0, amrex::PinnedArenaAllocator> pti(particle_buffer, lev); pti.isValid(); ++pti, ++i) {
+            auto& soa = pti.GetStructOfArrays();
+            data[i] = (amrex::ParticleReal*) soa.GetRealData(comp).dataPtr();
+            (*particles_per_tile)[i] = pti.numParticles();
+        }
+
+        return data;
+    }
+
+    void warpx_clearParticleBoundaryBuffer () {
+        auto& particle_buffers = WarpX::GetInstance().GetParticleBoundaryBuffer();
+        particle_buffers.clearParticles();
+    }
+
     void warpx_ComputeDt () {
         WarpX& warpx = WarpX::GetInstance();
         warpx.ComputeDt ();
diff --git a/Source/Python/WarpXWrappers.h b/Source/Python/WarpXWrappers.h
index 53d461e65..71eb5b4fd 100644
--- a/Source/Python/WarpXWrappers.h
+++ b/Source/Python/WarpXWrappers.h
@@ -104,6 +104,18 @@ extern "C" {
     void warpx_addRealComp(
         const char* char_species_name, const char* char_comp_name, bool comm);
 
+    int warpx_getParticleBoundaryBufferSize(const char* species_name, int boundary);
+
+    int** warpx_getParticleBoundaryBufferScrapedSteps(
+        const char* species_name, int boundary, int lev,
+        int* num_tiles, int** particles_per_tile);
+
+    amrex::ParticleReal** warpx_getParticleBoundaryBuffer(
+        const char* species_name, int boundary, int lev,
+        int* num_tiles, int** particles_per_tile, const char* comp_name);
+
+    void warpx_clearParticleBoundaryBuffer ();
+
   void warpx_ComputeDt ();
   void warpx_MoveWindow (int step, bool move_j);