Add Python function to extract particles scraped during the last step. (#5711)

This adds a function that extracts the particles that were scraped at
the current timestep.

This is useful in callback functions, where we often want to re-inject
particles that hit the boundary, and therefore need to select the ones
that were scraped at the current timestep.

This also avoids calling `clear_buffer`, which potentially interferes
with the `BoundaryScrapingDiagnostic`

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: RemiLehe

Examples/Tests/particle_boundary_interaction/inputs_test_rz_particle_boundary_interaction_picmi.py

@@ -128,20 +128,24 @@ def mirror_reflection():
     # STEP 1: extract the different parameters of the boundary buffer (normal, time, position)
     lev = 0  # level 0 (no mesh refinement here)
     delta_t = concat(
-        buffer.get_particle_boundary_buffer("electrons", "eb", "deltaTimeScraped", lev)
+        buffer.get_particle_scraped_this_step(
+            "electrons", "eb", "deltaTimeScraped", lev
+        )
     )
-    r = concat(buffer.get_particle_boundary_buffer("electrons", "eb", "x", lev))
-    theta = concat(buffer.get_particle_boundary_buffer("electrons", "eb", "theta", lev))
-    z = concat(buffer.get_particle_boundary_buffer("electrons", "eb", "z", lev))
+    r = concat(buffer.get_particle_scraped_this_step("electrons", "eb", "x", lev))
+    theta = concat(
+        buffer.get_particle_scraped_this_step("electrons", "eb", "theta", lev)
+    )
+    z = concat(buffer.get_particle_scraped_this_step("electrons", "eb", "z", lev))
     x = r * np.cos(theta)  # from RZ coordinates to 3D coordinates
     y = r * np.sin(theta)
-    ux = concat(buffer.get_particle_boundary_buffer("electrons", "eb", "ux", lev))
-    uy = concat(buffer.get_particle_boundary_buffer("electrons", "eb", "uy", lev))
-    uz = concat(buffer.get_particle_boundary_buffer("electrons", "eb", "uz", lev))
-    w = concat(buffer.get_particle_boundary_buffer("electrons", "eb", "w", lev))
-    nx = concat(buffer.get_particle_boundary_buffer("electrons", "eb", "nx", lev))
-    ny = concat(buffer.get_particle_boundary_buffer("electrons", "eb", "ny", lev))
-    nz = concat(buffer.get_particle_boundary_buffer("electrons", "eb", "nz", lev))
+    ux = concat(buffer.get_particle_scraped_this_step("electrons", "eb", "ux", lev))
+    uy = concat(buffer.get_particle_scraped_this_step("electrons", "eb", "uy", lev))
+    uz = concat(buffer.get_particle_scraped_this_step("electrons", "eb", "uz", lev))
+    w = concat(buffer.get_particle_scraped_this_step("electrons", "eb", "w", lev))
+    nx = concat(buffer.get_particle_scraped_this_step("electrons", "eb", "nx", lev))
+    ny = concat(buffer.get_particle_scraped_this_step("electrons", "eb", "ny", lev))
+    nz = concat(buffer.get_particle_scraped_this_step("electrons", "eb", "nz", lev))
 
     # STEP 2: use these parameters to inject particle from the same position in the plasma
     elect_pc = particle_containers.ParticleContainerWrapper(
@@ -164,8 +168,6 @@ def mirror_reflection():
     )  # adds the particle in the general particle container at the next step
     #### Can be modified depending on the model of interaction.
 
-    buffer.clear_buffer()  # reinitialise the boundary buffer
-
 
 callbacks.installafterstep(
     mirror_reflection

Examples/Tests/secondary_ion_emission/inputs_test_rz_secondary_ion_emission_picmi.py

@@ -188,26 +188,29 @@ def secondary_emission():
     elect_pc = particle_containers.ParticleContainerWrapper("electrons")
 
     if n != 0:
-        r = concat(buffer.get_particle_boundary_buffer("ions", "eb", "x", lev))
-        theta = concat(buffer.get_particle_boundary_buffer("ions", "eb", "theta", lev))
-        z = concat(buffer.get_particle_boundary_buffer("ions", "eb", "z", lev))
+        r = concat(buffer.get_particle_scraped_this_step("ions", "eb", "x", lev))
+        theta = concat(
+            buffer.get_particle_scraped_this_step("ions", "eb", "theta", lev)
+        )
+        z = concat(buffer.get_particle_scraped_this_step("ions", "eb", "z", lev))
         x = r * np.cos(theta)  # from RZ coordinates to 3D coordinates
         y = r * np.sin(theta)
-        ux = concat(buffer.get_particle_boundary_buffer("ions", "eb", "ux", lev))
-        uy = concat(buffer.get_particle_boundary_buffer("ions", "eb", "uy", lev))
-        uz = concat(buffer.get_particle_boundary_buffer("ions", "eb", "uz", lev))
-        w = concat(buffer.get_particle_boundary_buffer("ions", "eb", "w", lev))
-        nx = concat(buffer.get_particle_boundary_buffer("ions", "eb", "nx", lev))
-        ny = concat(buffer.get_particle_boundary_buffer("ions", "eb", "ny", lev))
-        nz = concat(buffer.get_particle_boundary_buffer("ions", "eb", "nz", lev))
+        ux = concat(buffer.get_particle_scraped_this_step("ions", "eb", "ux", lev))
+        uy = concat(buffer.get_particle_scraped_this_step("ions", "eb", "uy", lev))
+        uz = concat(buffer.get_particle_scraped_this_step("ions", "eb", "uz", lev))
+        w = concat(buffer.get_particle_scraped_this_step("ions", "eb", "w", lev))
+        nx = concat(buffer.get_particle_scraped_this_step("ions", "eb", "nx", lev))
+        ny = concat(buffer.get_particle_scraped_this_step("ions", "eb", "ny", lev))
+        nz = concat(buffer.get_particle_scraped_this_step("ions", "eb", "nz", lev))
         delta_t = concat(
-            buffer.get_particle_boundary_buffer("ions", "eb", "deltaTimeScraped", lev)
+            buffer.get_particle_scraped_this_step("ions", "eb", "deltaTimeScraped", lev)
         )
+
         energy_ions = 0.5 * proton_mass * w * (ux**2 + uy**2 + uz**2)
         energy_ions_in_kEv = energy_ions / (e * 1000)
         sigma_nascap_ions = sigma_nascap(energy_ions_in_kEv, delta_H, E_HMax)
-        # Loop over all ions in the EB buffer
-        for i in range(0, n):
+        # Loop over all ions that have been scraped in the last timestep
+        for i in range(0, len(w)):
             sigma = sigma_nascap_ions[i]
             # Ne_sec is number of the secondary electrons to be emitted
             Ne_sec = int(sigma + np.random.uniform())
@@ -258,7 +261,6 @@ def secondary_emission():
                     uz=uze,
                     w=we,
                 )
-        buffer.clear_buffer()  # reinitialise the boundary buffer
 
 
 # using the new particle container modified at the last step

Python/pywarpx/particle_containers.py

@@ -863,6 +863,51 @@ def get_particle_boundary_buffer(self, species_name, boundary, comp_name, level)
             raise RuntimeError("Name %s not found" % comp_name)
         return data_array
 
+    def get_particle_scraped_this_step(self, species_name, boundary, comp_name, level):
+        """
+        This returns a list of numpy or cupy arrays containing the particle array data
+        for particles that have been scraped at the current timestep,
+        for a specific species and simulation boundary.
+
+        The data for the arrays is a view of the underlying boundary buffer in WarpX ;
+        writing to these arrays will therefore also modify the underlying boundary buffer.
+
+        Parameters
+        ----------
+
+            species_name   : str
+                The species name that the data will be returned for.
+
+            boundary       : str
+                The boundary from which to get the scraped particle data in the
+                form x/y/z_hi/lo or eb.
+
+            comp_name      : str
+                The component of the array data that will be returned.
+                "x", "y", "z", "ux", "uy", "uz", "w"
+                "stepScraped","deltaTimeScraped",
+                if boundary='eb': "nx", "ny", "nz"
+
+            level          : int
+                Which AMR level to retrieve scraped particle data from.
+        """
+        # Extract the integer number of the current timestep
+        current_step = libwarpx.libwarpx_so.get_instance().getistep(level)
+
+        # Extract the data requested by the user
+        data_array = self.get_particle_boundary_buffer(
+            species_name, boundary, comp_name, level
+        )
+        step_scraped_array = self.get_particle_boundary_buffer(
+            species_name, boundary, "stepScraped", level
+        )
+
+        # Select on the particles from the previous step
+        data_array_this_step = []
+        for data, step in zip(data_array, step_scraped_array):
+            data_array_this_step.append(data[step == current_step])
+        return data_array_this_step
+
     def clear_buffer(self):
         """