Fix BTD/Scrape Flush Count with Filters

ax3l · ax3l · commit 7204b2d4d358 · 2024-02-01T09:10:31.000-08:00
Move the counting of already flushed particles for writers that call the
I/O backends multiple time per data set, e.g., BTD and boundary scraping,
into the I/O backend.

Currently, filtering is done as the first step in I/O backends and thus
the previous count outside of the I/O backends was over-counting
particles that might still get filtered out.
diff --git a/Source/Diagnostics/BTDiagnostics.H b/Source/Diagnostics/BTDiagnostics.H
@@ -399,8 +399,6 @@ private:
         lab-frame data. */
     void InitializeParticleFunctors () override;
 
-    /** Update total number of particles flushed for all species for ith snapshot */
-    void UpdateTotalParticlesFlushed(int i_buffer);
     /** Reset total number of particles in the particle buffer to 0 for ith snapshot */
     void ResetTotalParticlesInBuffer(int i_buffer);
     /** Clear particle data stored in the particle buffer */
diff --git a/Source/Diagnostics/BTDiagnostics.cpp b/Source/Diagnostics/BTDiagnostics.cpp
@@ -1066,11 +1066,12 @@ BTDiagnostics::Flush (int i_buffer, bool force_flush)
     }
     m_flush_format->WriteToFile(
         m_varnames, m_mf_output[i_buffer], m_geom_output[i_buffer], warpx.getistep(),
-        labtime, m_output_species[i_buffer], nlev_output, file_name, m_file_min_digits,
+        labtime,
+        m_totalParticles_flushed_already[i_buffer],
+        m_output_species[i_buffer], nlev_output, file_name, m_file_min_digits,
         m_plot_raw_fields, m_plot_raw_fields_guards,
         use_pinned_pc, isBTD, i_buffer, m_buffer_flush_counter[i_buffer],
-        m_max_buffer_multifabs[i_buffer], m_geom_snapshot[i_buffer][0], isLastBTDFlush,
-        m_totalParticles_flushed_already[i_buffer]);
+        m_max_buffer_multifabs[i_buffer], m_geom_snapshot[i_buffer][0], isLastBTDFlush);
 
     // Rescaling the box for plotfile after WriteToFile. This is because, for plotfiles, when writing particles, amrex checks if the particles are within the bounds defined by the box. However, in BTD, particles can be (at max) 1 cell outside the bounds of the geometry. So we keep a one-cell bigger box for plotfile when writing out the particle data and rescale after.
     if (m_format == "plotfile") {
@@ -1104,7 +1105,6 @@ BTDiagnostics::Flush (int i_buffer, bool force_flush)
     NullifyFirstFlush(i_buffer);
     // if particles are selected for output then update and reset counters
     if (!m_output_species_names.empty()) {
-        UpdateTotalParticlesFlushed(i_buffer);
         ResetTotalParticlesInBuffer(i_buffer);
         ClearParticleBuffer(i_buffer);
     }
@@ -1489,15 +1489,6 @@ BTDiagnostics::PrepareParticleDataForOutput()
     }
 }
 
-void
-BTDiagnostics::UpdateTotalParticlesFlushed(int i_buffer)
-{
-    for (int isp = 0; isp < m_totalParticles_flushed_already[i_buffer].size(); ++isp) {
-        m_totalParticles_flushed_already[i_buffer][isp] += static_cast<int>(
-            m_particles_buffer[i_buffer][isp]->TotalNumberOfParticles());
-    }
-}
-
 void
 BTDiagnostics::ResetTotalParticlesInBuffer(int i_buffer)
 {
diff --git a/Source/Diagnostics/BoundaryScrapingDiagnostics.cpp b/Source/Diagnostics/BoundaryScrapingDiagnostics.cpp
@@ -158,10 +158,12 @@ BoundaryScrapingDiagnostics::Flush (int i_buffer, bool /* force_flush */)
 
     m_flush_format->WriteToFile(
         m_varnames, m_mf_output[i_buffer], m_geom_output[i_buffer], warpx.getistep(),
-        warpx.gett_new(0), m_output_species[i_buffer], nlev_output, file_prefix,
+        warpx.gett_new(0),
+        m_totalParticles_flushed_already[i_buffer],
+        m_output_species[i_buffer], nlev_output, file_prefix,
         m_file_min_digits, false, false, use_pinned_pc, isBTD,
         warpx.getistep(0), bufferID, numBTDBuffers, geom,
-        isLastBTD, m_totalParticles_flushed_already[i_buffer]);
+        isLastBTD);
 
     // Now that the data has been written out, clear out the buffer
     particle_buffer.clearParticles(i_buffer);
diff --git a/Source/Diagnostics/Diagnostics.cpp b/Source/Diagnostics/Diagnostics.cpp
@@ -516,6 +516,16 @@ Diagnostics::InitBaseData ()
 
     // allocate vector of particle buffers
     m_output_species.resize(m_num_buffers);
+
+    // Initialize total number of particles flushed
+    m_totalParticles_flushed_already.resize(m_num_buffers);
+    for (int i_buffer = 0; i_buffer < m_num_buffers; ++i_buffer) {
+        int const n_species = static_cast<int>(m_output_species_names.size());
+        m_totalParticles_flushed_already[i_buffer].resize(n_species);
+        for (int i_species=0; i_species<n_species; i_species++) {
+            m_totalParticles_flushed_already[i_buffer][i_species] = 0;
+        }
+    }
 }
 
 void
diff --git a/Source/Diagnostics/FlushFormats/FlushFormat.H b/Source/Diagnostics/FlushFormats/FlushFormat.H
@@ -16,6 +16,7 @@ public:
         const amrex::Vector<amrex::MultiFab>& mf,
         amrex::Vector<amrex::Geometry>& geom,
         amrex::Vector<int> iteration, double time,
+        amrex::Vector<int>& totalParticlesFlushedAlready,
         const amrex::Vector<ParticleDiag>& particle_diags, int nlev,
         std::string prefix, int file_min_digits,
         bool plot_raw_fields,
@@ -24,8 +25,7 @@ public:
         bool isBTD = false, int snapshotID = -1,
         int bufferID = 1, int numBuffers = 1,
         const amrex::Geometry& full_BTD_snapshot = amrex::Geometry(),
-        bool isLastBTDFlush = false,
-        const amrex::Vector<int>& totalParticlesFlushedAlready = amrex::Vector<int>() ) const = 0;
+        bool isLastBTDFlush = false) const = 0;
 
     FlushFormat () = default;
     virtual ~FlushFormat() = default;
diff --git a/Source/Diagnostics/FlushFormats/FlushFormatAscent.H b/Source/Diagnostics/FlushFormats/FlushFormatAscent.H
@@ -33,6 +33,7 @@ public:
         const amrex::Vector<amrex::MultiFab>& mf,
         amrex::Vector<amrex::Geometry>& geom,
         amrex::Vector<int> iteration, double time,
+        amrex::Vector<int>& totalParticlesFlushedAlready,
         const amrex::Vector<ParticleDiag>& particle_diags, int nlev,
         std::string prefix, int file_min_digits,
         bool plot_raw_fields,
@@ -41,8 +42,7 @@ public:
         bool isBTD = false, int snapshotID = -1,
         int bufferID = 1, int numBuffers = 1,
         const amrex::Geometry& full_BTD_snapshot = amrex::Geometry(),
-        bool isLastBTDFlush = false,
-        const amrex::Vector<int>& totalParticlesFlushedAlready = amrex::Vector<int>() ) const override;
+        bool isLastBTDFlush = false ) const override;
 
 #ifdef AMREX_USE_ASCENT
     /** \brief Do in-situ visualization for particle data.
diff --git a/Source/Diagnostics/FlushFormats/FlushFormatAscent.cpp b/Source/Diagnostics/FlushFormats/FlushFormatAscent.cpp
@@ -15,13 +15,14 @@ FlushFormatAscent::WriteToFile (
     const amrex::Vector<amrex::MultiFab>& mf,
     amrex::Vector<amrex::Geometry>& geom,
     const amrex::Vector<int> iteration, const double time,
+    amrex::Vector<int>& /* totalParticlesFlushedAlready*/,
     const amrex::Vector<ParticleDiag>& particle_diags, int nlev,
     const std::string prefix, int file_min_digits, bool plot_raw_fields,
     bool plot_raw_fields_guards,
     const bool /*use_pinned_pc*/,
     bool isBTD, int /*snapshotID*/, int /*bufferID*/, int /*numBuffers*/,
     const amrex::Geometry& /*full_BTD_snapshot*/,
-    bool /*isLastBTDFlush*/, const amrex::Vector<int>& /* totalParticlesFlushedAlready*/) const
+    bool /*isLastBTDFlush*/) const
 {
 #ifdef AMREX_USE_ASCENT
     WARPX_PROFILE("FlushFormatAscent::WriteToFile()");
diff --git a/Source/Diagnostics/FlushFormats/FlushFormatCheckpoint.H b/Source/Diagnostics/FlushFormats/FlushFormatCheckpoint.H
@@ -20,6 +20,7 @@ class FlushFormatCheckpoint final : public FlushFormatPlotfile
         const amrex::Vector<amrex::MultiFab>& mf,
         amrex::Vector<amrex::Geometry>& geom,
         amrex::Vector<int> iteration, double time,
+        amrex::Vector<int>& totalParticlesFlushedAlready,
         const amrex::Vector<ParticleDiag>& particle_diags, int nlev,
         std::string prefix, int file_min_digits,
         bool plot_raw_fields,
@@ -28,8 +29,7 @@ class FlushFormatCheckpoint final : public FlushFormatPlotfile
         bool isBTD = false, int snapshotID = -1,
         int bufferID = 1, int numBuffers = 1,
         const amrex::Geometry& full_BTD_snapshot = amrex::Geometry(),
-        bool isLastBTDFlush = false,
-        const amrex::Vector<int>& totalParticlesFlushedAlready = amrex::Vector<int>() ) const final;
+        bool isLastBTDFlush = false) const final;
 
     void CheckpointParticles (const std::string& dir,
                               const amrex::Vector<ParticleDiag>& particle_diags) const;
diff --git a/Source/Diagnostics/FlushFormats/FlushFormatCheckpoint.cpp b/Source/Diagnostics/FlushFormats/FlushFormatCheckpoint.cpp
@@ -31,6 +31,7 @@ FlushFormatCheckpoint::WriteToFile (
         const amrex::Vector<amrex::MultiFab>& /*mf*/,
         amrex::Vector<amrex::Geometry>& geom,
         const amrex::Vector<int> iteration, const double /*time*/,
+        amrex::Vector<int>& /* totalParticlesFlushedAlready*/,
         const amrex::Vector<ParticleDiag>& particle_diags, int nlev,
         const std::string prefix, int file_min_digits,
         bool /*plot_raw_fields*/,
@@ -39,7 +40,7 @@ FlushFormatCheckpoint::WriteToFile (
         bool /*isBTD*/, int /*snapshotID*/,
         int /*bufferID*/, int /*numBuffers*/,
         const amrex::Geometry& /*full_BTD_snapshot*/,
-        bool /*isLastBTDFlush*/, const amrex::Vector<int>& /* totalParticlesFlushedAlready*/) const
+        bool /*isLastBTDFlush*/) const
 {
     WARPX_PROFILE("FlushFormatCheckpoint::WriteToFile()");
 
diff --git a/Source/Diagnostics/FlushFormats/FlushFormatOpenPMD.H b/Source/Diagnostics/FlushFormats/FlushFormatOpenPMD.H
@@ -32,6 +32,7 @@ public:
         const amrex::Vector<amrex::MultiFab>& mf,
         amrex::Vector<amrex::Geometry>& geom,
         amrex::Vector<int> iteration, double time,
+        amrex::Vector<int>& totalParticlesFlushedAlready,
         const amrex::Vector<ParticleDiag>& particle_diags, int output_levels,
         std::string prefix, int file_min_digits,
         bool plot_raw_fields,
@@ -40,8 +41,7 @@ public:
         bool isBTD = false, int snapshotID = -1,
         int bufferID = 1, int numBuffers = 1,
         const amrex::Geometry& full_BTD_snapshot = amrex::Geometry(),
-        bool isLastBTDFlush = false,
-        const amrex::Vector<int>& totalParticlesFlushedAlready = amrex::Vector<int>() ) const override;
+        bool isLastBTDFlush = false ) const override;
 
     ~FlushFormatOpenPMD () override = default;
 
diff --git a/Source/Diagnostics/FlushFormats/FlushFormatOpenPMD.cpp b/Source/Diagnostics/FlushFormats/FlushFormatOpenPMD.cpp
@@ -120,13 +120,14 @@ FlushFormatOpenPMD::WriteToFile (
     const amrex::Vector<amrex::MultiFab>& mf,
     amrex::Vector<amrex::Geometry>& geom,
     const amrex::Vector<int> iteration, const double time,
+    amrex::Vector<int>& totalParticlesFlushedAlready,
     const amrex::Vector<ParticleDiag>& particle_diags, int output_levels,
     const std::string prefix, int file_min_digits, bool plot_raw_fields,
     bool plot_raw_fields_guards,
     const bool use_pinned_pc,
     bool isBTD, int snapshotID, int bufferID, int numBuffers,
     const amrex::Geometry& full_BTD_snapshot,
-    bool isLastBTDFlush, const amrex::Vector<int>& totalParticlesFlushedAlready) const
+    bool isLastBTDFlush) const
 {
     WARPX_PROFILE("FlushFormatOpenPMD::WriteToFile()");
     const std::string& filename = amrex::Concatenate(prefix, iteration[0], file_min_digits);
@@ -164,7 +165,7 @@ FlushFormatOpenPMD::WriteToFile (
 
     // particles: all (reside only on locally finest level)
     m_OpenPMDPlotWriter->WriteOpenPMDParticles(
-        particle_diags, static_cast<amrex::Real>(time), use_pinned_pc, isBTD, isLastBTDFlush, totalParticlesFlushedAlready);
+        particle_diags, static_cast<amrex::Real>(time), totalParticlesFlushedAlready, use_pinned_pc, isBTD, isLastBTDFlush);
 
     // signal that no further updates will be written to this iteration
     m_OpenPMDPlotWriter->CloseStep(isBTD, isLastBTDFlush);
diff --git a/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.H b/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.H
@@ -27,6 +27,7 @@ public:
         const amrex::Vector<amrex::MultiFab>& mf,
         amrex::Vector<amrex::Geometry>& geom,
         amrex::Vector<int> iteration, double time,
+        amrex::Vector<int>& totalParticlesFlushedAlready,
         const amrex::Vector<ParticleDiag>& particle_diags, int nlev,
         std::string prefix, int file_min_digits,
         bool plot_raw_fields,
@@ -35,8 +36,7 @@ public:
         bool isBTD = false, int snapshotID = -1,
         int bufferID = 1, int numBuffers = 1,
         const amrex::Geometry& full_BTD_snapshot = amrex::Geometry(),
-        bool isLastBTDFlush = false,
-        const amrex::Vector<int>& totalParticlesFlushedAlready = amrex::Vector<int>() ) const override;
+        bool isLastBTDFlush = false) const override;
 
     /** Write general info of the run into the plotfile */
     void WriteJobInfo(const std::string& dir) const;
@@ -49,11 +49,13 @@ public:
      * \param[in] dir name of output directory
      * \param[in] particle_diags Each element of this vector handles output of 1 species.
      * \param[in] time the simulation time on the coarsest level
+     * \param[inout] totalParticlesFlushedAlready already flushed particles per species
      * \param[in] isBTD whether this is a back-transformed diagnostic
      */
     void WriteParticles(const std::string& dir,
                         const amrex::Vector<ParticleDiag>& particle_diags,
                         amrex::Real time,
+                        amrex::Vector<int>& totalParticlesFlushedAlready,
                         bool isBTD = false) const;
 
     FlushFormatPlotfile () = default;
diff --git a/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp b/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp
@@ -59,13 +59,14 @@ FlushFormatPlotfile::WriteToFile (
     const amrex::Vector<amrex::MultiFab>& mf,
     amrex::Vector<amrex::Geometry>& geom,
     const amrex::Vector<int> iteration, const double time,
+    amrex::Vector<int>& totalParticlesFlushedAlready,
     const amrex::Vector<ParticleDiag>& particle_diags, int nlev,
     const std::string prefix, int file_min_digits, bool plot_raw_fields,
     bool plot_raw_fields_guards,
     const bool /*use_pinned_pc*/,
     bool isBTD, int snapshotID,  int bufferID, int numBuffers,
     const amrex::Geometry& /*full_BTD_snapshot*/,
-    bool isLastBTDFlush, const amrex::Vector<int>& /* totalParticlesFlushedAlready*/) const
+    bool isLastBTDFlush) const
 {
     WARPX_PROFILE("FlushFormatPlotfile::WriteToFile()");
     auto & warpx = WarpX::GetInstance();
@@ -99,7 +100,7 @@ FlushFormatPlotfile::WriteToFile (
 
     WriteAllRawFields(plot_raw_fields, nlev, filename, plot_raw_fields_guards);
 
-    WriteParticles(filename, particle_diags, static_cast<amrex::Real>(time), isBTD);
+    WriteParticles(filename, particle_diags, static_cast<amrex::Real>(time), totalParticlesFlushedAlready, isBTD);
 
     WriteJobInfo(filename);
 
@@ -340,9 +341,11 @@ FlushFormatPlotfile::WriteWarpXHeader(
 void
 FlushFormatPlotfile::WriteParticles(const std::string& dir,
                                     const amrex::Vector<ParticleDiag>& particle_diags,
-                                    const amrex::Real time, bool isBTD) const
+                                    const amrex::Real time,
+                                    amrex::Vector<int>& totalParticlesFlushedAlready,
+                                    bool isBTD) const
 {
-
+    int i = 0;
     for (const auto& part_diag : particle_diags) {
         WarpXParticleContainer* pc = part_diag.getParticleContainer();
         PinnedMemoryParticleContainer* pinned_pc = part_diag.getPinnedParticleContainer();
@@ -418,6 +421,10 @@ FlushFormatPlotfile::WriteParticles(const std::string& dir,
             dir, part_diag.getSpeciesName(),
             real_flags, int_flags,
             real_names, int_names);
+
+        // keep book of filtered-and-written particles
+        totalParticlesFlushedAlready[i] += tmp.TotalNumberOfParticles(false, true);
+        i++;
     }
 }
 
diff --git a/Source/Diagnostics/FlushFormats/FlushFormatSensei.H b/Source/Diagnostics/FlushFormats/FlushFormatSensei.H
@@ -53,6 +53,7 @@ public:
         const amrex::Vector<amrex::MultiFab>& mf,
         amrex::Vector<amrex::Geometry>& geom,
         amrex::Vector<int> iteration, double time,
+        amrex::Vector<int>& totalParticlesFlushedAlready,
         const amrex::Vector<ParticleDiag>& particle_diags, int nlev,
         std::string prefix, int file_min_digits,
         bool plot_raw_fields,
@@ -61,8 +62,7 @@ public:
         bool isBTD = false, int snapshotID = -1,
         int bufferID = 1, int numBuffers = 1,
         const amrex::Geometry& full_BTD_snapshot = amrex::Geometry(),
-        bool isLastBTDFlush = false,
-        const amrex::Vector<int>& totalParticlesFlushedAlready = amrex::Vector<int>() ) const override;
+        bool isLastBTDFlush = false) const override;
 
     /** \brief Do in-situ visualization for particle data.
      * \param[in] particle_diags Each element of this vector handles output of 1 species.
diff --git a/Source/Diagnostics/FlushFormats/FlushFormatSensei.cpp b/Source/Diagnostics/FlushFormats/FlushFormatSensei.cpp
@@ -48,13 +48,13 @@ FlushFormatSensei::WriteToFile (
     const amrex::Vector<amrex::MultiFab>& mf,
     amrex::Vector<amrex::Geometry>& geom,
     const amrex::Vector<int> iteration, const double time,
+    amrex::Vector<int>& totalParticlesFlushedAlready,
     const amrex::Vector<ParticleDiag>& particle_diags,
     int nlev, const std::string prefix, int file_min_digits,
     bool plot_raw_fields, bool plot_raw_fields_guards,
     const bool use_pinned_pc,
     bool isBTD, int /*snapshotID*/, int /*bufferID*/, int /*numBuffers*/,
-    const amrex::Geometry& /*full_BTD_snapshot*/, bool /*isLastBTDFlush*/,
-    const amrex::Vector<int>& totalParticlesFlushedAlready) const
+    const amrex::Geometry& /*full_BTD_snapshot*/, bool /*isLastBTDFlush*/) const
 {
     amrex::ignore_unused(
         geom, nlev, prefix, file_min_digits,
diff --git a/Source/Diagnostics/FullDiagnostics.cpp b/Source/Diagnostics/FullDiagnostics.cpp
@@ -75,6 +75,16 @@ FullDiagnostics::InitializeParticleBuffer ()
             m_output_species[i_buffer].push_back(ParticleDiag(m_diag_name, species, mpc.GetParticleContainerPtr(idx)));
         }
     }
+
+    // Initialize total number of particles flushed
+    m_totalParticles_flushed_already.resize(m_num_buffers);
+    for (int i_buffer = 0; i_buffer < m_num_buffers; ++i_buffer) {
+        int const n_species = static_cast<int>(m_output_species_names.size());
+        m_totalParticles_flushed_already[i_buffer].resize(n_species);
+        for (int i_species=0; i_species<n_species; i_species++) {
+            m_totalParticles_flushed_already[i_buffer][i_species] = 0;
+        }
+    }
 }
 
 void
@@ -134,7 +144,9 @@ FullDiagnostics::Flush ( int i_buffer, bool /* force_flush */ )
 
     m_flush_format->WriteToFile(
         m_varnames, m_mf_output[i_buffer], m_geom_output[i_buffer], warpx.getistep(),
-        warpx.gett_new(0), m_output_species[i_buffer], nlev_output, m_file_prefix,
+        warpx.gett_new(0),
+        m_totalParticles_flushed_already[i_buffer],
+        m_output_species[i_buffer], nlev_output, m_file_prefix,
         m_file_min_digits, m_plot_raw_fields, m_plot_raw_fields_guards);
 
     FlushRaw();
diff --git a/Source/Diagnostics/WarpXOpenPMD.H b/Source/Diagnostics/WarpXOpenPMD.H
@@ -123,10 +123,10 @@ public:
   void WriteOpenPMDParticles (
               const amrex::Vector<ParticleDiag>& particle_diags,
               amrex::Real time,
+              amrex::Vector<int>& totalParticlesFlushedAlready,
               bool use_pinned_pc = false,
               bool isBTD = false,
-              bool isLastBTDFlush = false,
-              const amrex::Vector<int>& totalParticlesFlushedAlready = amrex::Vector<int>());
+              bool isLastBTDFlush = false);
 
   /** Write out all openPMD fields for all active MR levels
    *
@@ -290,9 +290,9 @@ private:
    * @param[in] int_comp_names The int attribute names, from WarpX
    * @param[in] charge         Charge of the particles (note: fix for ions)
    * @param[in] mass           Mass of the particles
+   * @param[inout] ParticleFlushOffset previously flushed number of particles in BTD
    * @param[in] isBTD is this a backtransformed diagnostics (BTD) write?
    * @param[in] isLastBTDFlush is this the last time we will flush this BTD station?
-   * @param[in] ParticleFlushOffset previously flushed number of particles in BTD
    */
   void DumpToFile (ParticleContainer* pc,
             const std::string& name,
@@ -303,9 +303,9 @@ private:
             const amrex::Vector<std::string>&  int_comp_names,
             amrex::ParticleReal charge,
             amrex::ParticleReal mass,
+            int & ParticleFlushOffset,
             bool isBTD = false,
-            bool isLastBTDFlush = false,
-            int ParticleFlushOffset = 0);
+            bool isLastBTDFlush = false);
 
   /** Get the openPMD-api filename for openPMD::Series
    *
diff --git a/Source/Diagnostics/WarpXOpenPMD.cpp b/Source/Diagnostics/WarpXOpenPMD.cpp
