Particle Buffer for Backtransformed Diagnostics (#1898)

* move BTD call before Redistribute

* define particle buffer, BTD particle functor, particle filter

* select particles on the slice

* add particle functors

* add Lorentz Transform

* fix conflict

* add call to loretnz operator

* storing Particles in buffer

* This is WIP. Added particle buffers, filled them, sent them for flush with number already flushed. Does not work for multiple flushes. Crashes for OpenPMD.

* trailing endif after rebase

* adding print statements and not flushing particles in OPENPMD if numpart is 0

* last timestep flush is ensured at the end of evolve loop

* fix bug in declaring uy uy new and computing uzp

* set particle Geom, BA, and DMAP for particle flush with plotfile

* set Particle BA Geom DMAP for particle buffer and no BTD transform for force flush

* separate compute and pack from flush

* Fix Typo: resizeable -> resizable

Fixes HDF5 BTD particle output.

* new class for plotfile particles for BTD

* copy particle_H and DATA and Header. some WIP print statements

* Merge plotfile

* clean print statements

* fix warning message

* struct declaration in header, fix warning

* doxygen comments and copyright

* clean print statements

* fix eol and override function warning

* tile data

* fix output species array size bug

* fix access for particle buffer size

* clean and move time-update

* add cur_time update back

* remove cur time update which was called twice

* dont access particles flushed already for full diagnostics

* cur time must be updated for RigidInjection BTD CI test to pass

* temporarily move call to BTD

* updating time and calling BTD before movewindow

* cleanup

* reset benchmarks and analysis script

* clean and add comments

* fix particle box array, geom, dmap

* reset benchmarks for multi_J rz and ElectrostaticSphereEB_mixedBC

* wip commit

* wip commit

* add SI conversion

* abort for openpmd bp backend if species is selected. Also write particle output for BTD only if write_species is 1

* add documentation for aborting if adios is used with openpmd and add other BTD input parameters

* Apply suggestions from code review

commit Axel's suggestions from review

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* use bool instead of int

* fix doxygen format

* using h5 as backend in example test to ensure consistency with abort for particle output.

* fix doxygen comment

* reset benchmark again for comoving_2d_psatd galilean_2d_psatd multi_J_rz_psatd

* reset benchmark for background_mcc

* self-review suggestions

* reset benchmarks. Update with last snapshot full info

* Axel's PR suggestions

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* Axel's doxygen fix

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* add comments

* fix eol

* improved exception handling for stringsteam

* PR suggestions

* Axels' suggestions from code review

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* Axel's suggestions :)

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* simplify logic

* suggestions from review (Axel/Reva)

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* variable name change for clarity

* if num particles in tmp array is 0, return

* Use new BTD inputs to set up BTD for particles in the corresponding particle container

* unused var

* fix logic error

* speciesID undefined

* separate particle and field buffer calls and initialization for BTD. Data common to both are initialized separately

* rename variable so it does not ghost existing varname

* add more comments

* Assert that fields are on

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* redundant function definition

* unused variable

* unused variable zp

* 1D 2D 3D definition

* fix else

* Apply suggestions from code review

Add Prabhat's suggestion

Co-authored-by: Prabhat Kumar <89051199+prkkumar@users.noreply.github.com>

* missing semicolon and ignore xp yp for 1D

* resetting benchmarks for boosted sims and mcc sim

* temporarily changing tolerance since the relative difference for momentum_z is 3.68e-3 and the current tolerance is 2.5e-3

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Prabhat Kumar <89051199+prkkumar@users.noreply.github.com>

1 files changed, 75 insertions, 35 deletions

diff --git a/Source/Diagnostics/WarpXOpenPMD.cpp b/Source/Diagnostics/WarpXOpenPMD.cpp
index 6f5c6a36b..daf5fbefc 100644
--- a/Source/Diagnostics/WarpXOpenPMD.cpp
+++ b/Source/Diagnostics/WarpXOpenPMD.cpp
@@ -489,7 +489,7 @@ WarpXOpenPMDPlot::Init (openPMD::Access access, bool isBTD)
 
 void
 WarpXOpenPMDPlot::WriteOpenPMDParticles (const amrex::Vector<ParticleDiag>& particle_diags,
-                                         const bool isBTD)
+                  const bool isBTD, const amrex::Vector<int>& totalParticlesFlushedAlready)
 {
   WARPX_PROFILE("WarpXOpenPMDPlot::WriteOpenPMDParticles()");
 
@@ -552,28 +552,48 @@ WarpXOpenPMDPlot::WriteOpenPMDParticles (const amrex::Vector<ParticleDiag>& part
       GeometryFilter const geometry_filter(particle_diags[i].m_do_geom_filter,
                                            particle_diags[i].m_diag_domain);
 
-      using SrcData = WarpXParticleContainer::ParticleTileType::ConstParticleTileDataType;
-      tmp.copyParticles(*pc,
-                        [=] AMREX_GPU_HOST_DEVICE (const SrcData& src, int ip, const amrex::RandomEngine& engine)
-      {
-          const SuperParticleType& p = src.getSuperParticle(ip);
-          return random_filter(p, engine) * uniform_filter(p, engine)
-                 * parser_filter(p, engine) * geometry_filter(p, engine);
-      }, true);
+      if (! isBTD) {
+          using SrcData = WarpXParticleContainer::ParticleTileType::ConstParticleTileDataType;
+          tmp.copyParticles(*pc,
+                            [=] AMREX_GPU_HOST_DEVICE (const SrcData& src, int ip, const amrex::RandomEngine& engine)
+          {
+              const SuperParticleType& p = src.getSuperParticle(ip);
+              return random_filter(p, engine) * uniform_filter(p, engine)
+                     * parser_filter(p, engine) * geometry_filter(p, engine);
+          }, true);
+      } else if (isBTD) {
+          PinnedMemoryParticleContainer* pinned_pc = particle_diags[i].getPinnedParticleContainer();
+          tmp.SetParticleGeometry(0,pinned_pc->Geom(0));
+          tmp.SetParticleBoxArray(0,pinned_pc->ParticleBoxArray(0));
+          tmp.SetParticleDistributionMap(0, pinned_pc->ParticleDistributionMap(0));
+          tmp.copyParticles(*pinned_pc);
+      }
 
     // real_names contains a list of all real particle attributes.
     // real_flags is 1 or 0, whether quantity is dumped or not.
 
     {
-      DumpToFile(&tmp,
-         particle_diags[i].getSpeciesName(),
-         m_CurrentStep,
-         real_flags,
-         int_flags,
-         real_names, int_names,
-         pc->getCharge(), pc->getMass(),
-         isBTD
-      );
+      if (isBTD) {
+          DumpToFile(&tmp,
+             particle_diags[i].getSpeciesName(),
+             m_CurrentStep,
+             real_flags,
+             int_flags,
+             real_names, int_names,
+             pc->getCharge(), pc->getMass(),
+             isBTD, totalParticlesFlushedAlready[i]
+          );
+      } else {
+          DumpToFile(&tmp,
+             particle_diags[i].getSpeciesName(),
+             m_CurrentStep,
+             real_flags,
+             int_flags,
+             real_names, int_names,
+             pc->getCharge(), pc->getMass(),
+             isBTD, 0
+          );
+      }
     }
 
     // Convert momentum back to WarpX units
@@ -590,12 +610,13 @@ WarpXOpenPMDPlot::DumpToFile (ParticleContainer* pc,
                     const amrex::Vector<std::string>& real_comp_names,
                     const amrex::Vector<std::string>&  int_comp_names,
                     amrex::ParticleReal const charge,
-                    amrex::ParticleReal const mass,
-                    const bool isBTD) const
+                    amrex::ParticleReal const mass, const bool isBTD,
+                    int ParticleFlushOffset)
 {
   AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_Series != nullptr, "openPMD: series must be initialized");
 
   WarpXParticleCounter counter(pc);
+  if (counter.GetTotalNumParticles() == 0) return;
   openPMD::Iteration currIteration = GetIteration(iteration, isBTD);
 
   openPMD::ParticleSpecies currSpecies = currIteration.particles[name];
@@ -640,20 +661,29 @@ WarpXOpenPMDPlot::DumpToFile (ParticleContainer* pc,
   //
   // define positions & offsets
   //
-  SetupPos(currSpecies, counter.GetTotalNumParticles(), charge, mass);
-  SetupRealProperties(currSpecies, write_real_comp, real_comp_names, write_int_comp, int_comp_names, counter.GetTotalNumParticles());
-
+  const unsigned long long NewParticleVectorSize = counter.GetTotalNumParticles() + ParticleFlushOffset;
+  m_doParticleSetUp = false;
+  if (counter.GetTotalNumParticles() > 0 and ParticleFlushOffset == 0) {
+      // This will trigger meta-data flush for particles the first-time non-zero number of particles are flushed.
+      m_doParticleSetUp = true;
+  }
+  SetupPos(currSpecies, NewParticleVectorSize, charge, mass, isBTD);
+  SetupRealProperties(currSpecies, write_real_comp, real_comp_names, write_int_comp, int_comp_names, NewParticleVectorSize, isBTD);
   // open files from all processors, in case some will not contribute below
   m_Series->flush();
-
   for (auto currentLevel = 0; currentLevel <= pc->finestLevel(); currentLevel++)
     {
       uint64_t offset = static_cast<uint64_t>( counter.m_ParticleOffsetAtRank[currentLevel] );
-
+      // For BTD, the offset include the number of particles already flushed
+      if (isBTD) offset += ParticleFlushOffset;
       for (ParticleIter pti(*pc, currentLevel); pti.isValid(); ++pti) {
          auto const numParticleOnTile = pti.numParticles();
          uint64_t const numParticleOnTile64 = static_cast<uint64_t>( numParticleOnTile );
 
+         // Do not call storeChunk() with zero-sized particle tiles:
+         //   https://github.com/openPMD/openPMD-api/issues/1147
+         if (numParticleOnTile == 0) continue;
+
          // get position and particle ID from aos
          // note: this implementation iterates the AoS 4x...
          // if we flush late as we do now, we can also copy out the data in one go
@@ -740,11 +770,12 @@ WarpXOpenPMDPlot::SetupRealProperties (openPMD::ParticleSpecies& currSpecies,
                       const amrex::Vector<std::string>& real_comp_names,
                       const amrex::Vector<int>& write_int_comp,
                       const amrex::Vector<std::string>& int_comp_names,
-                      unsigned long long np) const
+                      const unsigned long long np, bool const isBTD) const
 {
-    auto dtype_real = openPMD::Dataset(openPMD::determineDatatype<amrex::ParticleReal>(), {np});
-    auto dtype_int  = openPMD::Dataset(openPMD::determineDatatype<int>(), {np});
-
+    std::string options = "{}";
+    if (isBTD) options = "{ \"resizable\": true }";
+    auto dtype_real = openPMD::Dataset(openPMD::determineDatatype<amrex::ParticleReal>(), {np}, options);
+    auto dtype_int  = openPMD::Dataset(openPMD::determineDatatype<int>(), {np}, options);
     //
     // the beam/input3d showed write_real_comp.size() = 16 while only 10 real comp names
     // so using the min to be safe.
@@ -767,6 +798,8 @@ WarpXOpenPMDPlot::SetupRealProperties (openPMD::ParticleSpecies& currSpecies,
         }
     }
 
+    // attributes need to be set only the first time BTD flush is called for a snapshot
+    if (isBTD and m_doParticleSetUp == false) return;
     std::set< std::string > addedRecords; // add meta-data per record only once
     for (auto idx=0; idx<m_NumSoARealAttributes; idx++) {
         auto ii = m_NumAoSRealAttributes + idx; // jump over AoS names
@@ -832,7 +865,6 @@ WarpXOpenPMDPlot::SaveRealProperty (ParticleIter& pti,
   uint64_t const numParticleOnTile64 = static_cast<uint64_t>( numParticleOnTile );
   auto const& aos = pti.GetArrayOfStructs();  // size =  numParticlesOnTile
   auto const& soa = pti.GetStructOfArrays();
-
   // first we concatinate the AoS into contiguous arrays
   {
     for( auto idx=0; idx<m_NumAoSRealAttributes; idx++ ) {
@@ -892,23 +924,31 @@ WarpXOpenPMDPlot::SetupPos (
     openPMD::ParticleSpecies& currSpecies,
     const unsigned long long& np,
     amrex::ParticleReal const charge,
-    amrex::ParticleReal const mass) const
+    amrex::ParticleReal const mass,
+    bool const isBTD)
 {
-  auto const realType = openPMD::Dataset(openPMD::determineDatatype<amrex::ParticleReal>(), {np});
-  auto const idType = openPMD::Dataset(openPMD::determineDatatype< uint64_t >(), {np});
+  std::string options = "{}";
+  if (isBTD) options = "{ \"resizable\": true }";
+  auto realType= openPMD::Dataset(openPMD::determineDatatype<amrex::ParticleReal>(), {np}, options);
+  auto idType = openPMD::Dataset(openPMD::determineDatatype< uint64_t >(), {np}, options);
 
   auto const positionComponents = detail::getParticlePositionComponentLabels();
   for( auto const& comp : positionComponents ) {
       currSpecies["positionOffset"][comp].resetDataset( realType );
-      currSpecies["positionOffset"][comp].makeConstant( 0. );
       currSpecies["position"][comp].resetDataset( realType );
   }
 
   auto const scalar = openPMD::RecordComponent::SCALAR;
   currSpecies["id"][scalar].resetDataset( idType );
   currSpecies["charge"][scalar].resetDataset( realType );
-  currSpecies["charge"][scalar].makeConstant( charge );
   currSpecies["mass"][scalar].resetDataset( realType );
+
+  if (isBTD and m_doParticleSetUp == false) return;
+  // make constant
+  for( auto const& comp : positionComponents ) {
+      currSpecies["positionOffset"][comp].makeConstant( 0. );
+  }
+  currSpecies["charge"][scalar].makeConstant( charge );
   currSpecies["mass"][scalar].makeConstant( mass );
 
   // meta data