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>

5 files changed, 111 insertions, 1 deletions

diff --git a/Source/Particles/ParticleBoundaryBuffer.H b/Source/Particles/ParticleBoundaryBuffer.H
index 452722d56..592dc574f 100644
--- a/Source/Particles/ParticleBoundaryBuffer.H
+++ b/Source/Particles/ParticleBoundaryBuffer.H
@@ -40,6 +40,10 @@ public:
 
     void printNumParticles () const;
 
+    int getNumParticlesInContainer(const std::string species_name, int boundary);
+
+    ParticleBuffer::BufferType<amrex::PinnedArenaAllocator>& getParticleBuffer(const std::string species_name, int boundary);
+
     constexpr int numBoundaries () const {
         return AMREX_SPACEDIM*2
 #ifdef AMREX_USE_EB
diff --git a/Source/Particles/ParticleBoundaryBuffer.cpp b/Source/Particles/ParticleBoundaryBuffer.cpp
index f84d05b28..16465ac76 100644
--- a/Source/Particles/ParticleBoundaryBuffer.cpp
+++ b/Source/Particles/ParticleBoundaryBuffer.cpp
@@ -97,7 +97,7 @@ void ParticleBoundaryBuffer::printNumParticles () const {
                 int np = buffer[i].isDefined() ? buffer[i].TotalNumberOfParticles(false) : 0;
                 amrex::Print() << "Species " << getSpeciesNames()[i] << " has "
                                << np << " particles in the boundary buffer "
-                               << " for side " << iside << " of dim " << idim << "\n";
+                               << "for side " << iside << " of dim " << idim << "\n";
             }
         }
     }
@@ -232,3 +232,28 @@ void ParticleBoundaryBuffer::gatherParticles (MultiParticleContainer& mypc,
     amrex::ignore_unused(distance_to_eb, dxi);
 #endif
 }
+
+int ParticleBoundaryBuffer::getNumParticlesInContainer(
+        const std::string species_name, int boundary) {
+
+    auto& buffer = m_particle_containers[boundary];
+    auto index = WarpX::GetInstance().GetPartContainer().getSpeciesID(species_name);
+
+    if (buffer[index].isDefined()) return buffer[index].TotalNumberOfParticles(false);
+    else return 0;
+}
+
+ParticleBuffer::BufferType<amrex::PinnedArenaAllocator>&
+ParticleBoundaryBuffer::getParticleBuffer(const std::string species_name, int boundary) {
+
+    auto& buffer = m_particle_containers[boundary];
+    auto index = WarpX::GetInstance().GetPartContainer().getSpeciesID(species_name);
+
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_do_boundary_buffer[boundary][index],
+                                     "Attempted to get particle buffer for boundary "
+                                     + boundary + ", which is not used!");
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(buffer[index].isDefined(),
+                                     "Tried to get a buffer that is not defined!");
+
+    return buffer[index];
+}
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);
 
diff --git a/Source/WarpX.H b/Source/WarpX.H
index dcc473ff8..3b3df6790 100644
--- a/Source/WarpX.H
+++ b/Source/WarpX.H
@@ -98,6 +98,8 @@ public:
 
     MultiParticleContainer& GetPartContainer () { return *mypc; }
 
+    ParticleBoundaryBuffer& GetParticleBoundaryBuffer () { return *m_particle_boundary_buffer; }
+
     static void shiftMF (amrex::MultiFab& mf, const amrex::Geometry& geom,
                          int num_shift, int dir, amrex::Real external_field=0.0,
                          bool useparser = false, amrex::ParserExecutor<3> const& field_parser={});