From 0b4606ff39760d1fc9471957178d962463088454 Mon Sep 17 00:00:00 2001
From: Kai Germaschewski <kai.germaschewski@gmail.com>
Date: Thu, 29 Aug 2024 12:05:04 -0400
Subject: [PATCH] outp/hdf5: add template argument for particle selector

This particle selector decides whether to include a given particle in the
output. Other than the trivial ParticleSelectorAll,
ParticleSelectorEveryNth<N> is provided ot select every nth particle.
---
 .../output_particles_hdf5_impl.hxx            | 44 ++++++++++++++++---
 src/libpsc/tests/test_output_particles.cxx    | 12 ++---
 src/psc_config.hxx                            |  4 +-
 src/psc_flatfoil_yz.cxx                       |  4 ++
 4 files changed, 51 insertions(+), 13 deletions(-)

diff --git a/src/libpsc/psc_output_particles/output_particles_hdf5_impl.hxx b/src/libpsc/psc_output_particles/output_particles_hdf5_impl.hxx
index cc2c2ae058..bf92547279 100644
--- a/src/libpsc/psc_output_particles/output_particles_hdf5_impl.hxx
+++ b/src/libpsc/psc_output_particles/output_particles_hdf5_impl.hxx
@@ -293,13 +293,38 @@ private:
   Hdf5ParticleType prt_type_;
 };
 
+template <int EVERY>
+class ParticleSelectorEveryNth
+{
+public:
+  template <typename Particle>
+  bool operator()(const Particle& prt)
+  {
+    // return true for every `every_`th particle
+    return (cnt_++ % EVERY) == 0;
+  }
+
+private:
+  int cnt_ = 0;
+};
+
+class ParticleSelectorAll
+{
+public:
+  template <typename Particle>
+  bool operator()(const Particle& prt)
+  {
+    return true;
+  }
+};
+
 namespace detail
 {
 
 // ======================================================================
 // OutputParticlesHdf5
 
-template <typename Mparticles>
+template <typename Mparticles, typename ParticleSelector>
 struct OutputParticlesHdf5
 {
   using writer_type = OutputParticlesWriterHDF5;
@@ -353,17 +378,24 @@ struct OutputParticlesHdf5
   {
     int nr_kinds = mprts.grid().kinds.size();
 
+    ParticleSelector selector;
+
     for (int p = 0; p < mprts.n_patches(); p++) {
       const int* ldims = mprts.grid().ldims;
       int nr_indices = ldims[0] * ldims[1] * ldims[2] * nr_kinds;
       off[p].resize(nr_indices + 1);
       auto&& prts = mprts[p];
       unsigned int n_prts = prts.size();
+      std::vector<int> particle_indices;
+      particle_indices.reserve(n_prts);
 
       // counting sort to get map
       for (int n = 0; n < n_prts; n++) {
-        int si = get_sort_index(prts, n);
-        off[p][si]++;
+        if (selector(prts[n])) {
+          particle_indices.push_back(n);
+          int si = get_sort_index(prts, n);
+          off[p][si]++;
+        }
       }
       // prefix sum to get offsets
       int o = 0;
@@ -377,7 +409,7 @@ struct OutputParticlesHdf5
 
       // sort a map only, not the actual particles
       map[p].resize(n_prts);
-      for (int n = 0; n < n_prts; n++) {
+      for (auto n : particle_indices) {
         int si = get_sort_index(prts, n);
         map[p][off2[si]++] = n;
       }
@@ -609,6 +641,7 @@ private:
 
 } // namespace detail
 
+template <typename ParticleSelector = ParticleSelectorAll>
 class OutputParticlesHdf5 : OutputParticlesBase
 {
   static OutputParticlesParams adjust_params(
@@ -640,7 +673,8 @@ public:
       return;
     }
 
-    detail::OutputParticlesHdf5<Mparticles> impl{grid, params_};
+    detail::OutputParticlesHdf5<Mparticles, ParticleSelector> impl{grid,
+                                                                   params_};
     impl(mprts, writer_);
   }
 
diff --git a/src/libpsc/tests/test_output_particles.cxx b/src/libpsc/tests/test_output_particles.cxx
index 8c200dad09..ff68aa0777 100644
--- a/src/libpsc/tests/test_output_particles.cxx
+++ b/src/libpsc/tests/test_output_particles.cxx
@@ -64,14 +64,14 @@ struct OutputParticlesTest : ::testing::Test
   Int3 ibn = {2, 2, 2};
 };
 
-using OutputParticlesTestTypes =
-  ::testing::Types<Config<dim_xyz, MparticlesSingle, OutputParticlesAscii>
+using OutputParticlesTestTypes = ::testing::Types<
+  Config<dim_xyz, MparticlesSingle, OutputParticlesAscii>
 #ifdef H5_HAVE_PARALLEL
-                   ,
-                   Config<dim_xyz, MparticlesSingle, OutputParticlesHdf5>,
-                   Config<dim_xyz, MparticlesDouble, OutputParticlesHdf5>
+  ,
+  Config<dim_xyz, MparticlesSingle, OutputParticlesHdf5<ParticleSelectorAll>>,
+  Config<dim_xyz, MparticlesDouble, OutputParticlesHdf5<ParticleSelectorAll>>
 #endif
-                   >;
+  >;
 
 TYPED_TEST_SUITE(OutputParticlesTest, OutputParticlesTestTypes);
 
diff --git a/src/psc_config.hxx b/src/psc_config.hxx
index fcf01cd29a..6401a481e0 100644
--- a/src/psc_config.hxx
+++ b/src/psc_config.hxx
@@ -31,7 +31,7 @@
 
 #include "../libpsc/vpic/fields_item_vpic.hxx"
 
-using OutputParticlesDefault = OutputParticlesHdf5;
+using OutputParticlesDefault = OutputParticlesHdf5<ParticleSelectorAll>;
 
 struct SimulationNone
 {
@@ -252,7 +252,7 @@ struct PscConfigVpicPsc
   using BndParticles = BndParticlesVpic<Mparticles>;
   using Checks = ChecksVpic<Mparticles, MfieldsState>;
   using Marder = MarderVpic<Mparticles, MfieldsState>;
-  using OutputParticles = OutputParticlesHdf5;
+  using OutputParticles = OutputParticlesHdf5<ParticleSelectorAll>;
   using OutputHydro = OutputHydroQ<Mparticles, MfieldsHydro,
                                    typename VpicConfig::MfieldsInterpolator>;
   using Dim = dim_xz;
diff --git a/src/psc_flatfoil_yz.cxx b/src/psc_flatfoil_yz.cxx
index a853acc059..342e520100 100644
--- a/src/psc_flatfoil_yz.cxx
+++ b/src/psc_flatfoil_yz.cxx
@@ -226,7 +226,11 @@ using Balance = PscConfig::Balance;
 using Collision = PscConfig::Collision;
 using Checks = PscConfig::Checks;
 using Marder = PscConfig::Marder;
+#if CASE == CASE_2D_SMALL
+using OutputParticles = OutputParticlesHdf5<ParticleSelectorEveryNth<10>>;
+#else
 using OutputParticles = PscConfig::OutputParticles;
+#endif
 using Moment_n = typename Moment_n_Selector<Mparticles, Dim>::type;
 using Heating = typename HeatingSelector<Mparticles>::Heating;