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 cc2c2ae05..bf9254727 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 +class ParticleSelectorEveryNth +{ +public: + template + bool operator()(const Particle& prt) + { + // return true for every `every_`th particle + return (cnt_++ % EVERY) == 0; + } + +private: + int cnt_ = 0; +}; + +class ParticleSelectorAll +{ +public: + template + bool operator()(const Particle& prt) + { + return true; + } +}; + namespace detail { // ====================================================================== // OutputParticlesHdf5 -template +template 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 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 class OutputParticlesHdf5 : OutputParticlesBase { static OutputParticlesParams adjust_params( @@ -640,7 +673,8 @@ public: return; } - detail::OutputParticlesHdf5 impl{grid, params_}; + detail::OutputParticlesHdf5 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 8c200dad0..ff68aa077 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 +using OutputParticlesTestTypes = ::testing::Types< + Config #ifdef H5_HAVE_PARALLEL - , - Config, - Config + , + Config>, + Config> #endif - >; + >; TYPED_TEST_SUITE(OutputParticlesTest, OutputParticlesTestTypes); diff --git a/src/psc_config.hxx b/src/psc_config.hxx index fcf01cd29..6401a481e 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; struct SimulationNone { @@ -252,7 +252,7 @@ struct PscConfigVpicPsc using BndParticles = BndParticlesVpic; using Checks = ChecksVpic; using Marder = MarderVpic; - using OutputParticles = OutputParticlesHdf5; + using OutputParticles = OutputParticlesHdf5; using OutputHydro = OutputHydroQ; using Dim = dim_xz; diff --git a/src/psc_flatfoil_yz.cxx b/src/psc_flatfoil_yz.cxx index a853acc05..342e52010 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>; +#else using OutputParticles = PscConfig::OutputParticles; +#endif using Moment_n = typename Moment_n_Selector::type; using Heating = typename HeatingSelector::Heating;