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, 173 insertions, 0 deletions

diff --git a/Source/Diagnostics/ComputeDiagFunctors/BackTransformParticleFunctor.cpp b/Source/Diagnostics/ComputeDiagFunctors/BackTransformParticleFunctor.cpp
new file mode 100644
index 000000000..b8dac11ab
--- /dev/null
+++ b/Source/Diagnostics/ComputeDiagFunctors/BackTransformParticleFunctor.cpp
@@ -0,0 +1,173 @@
+/* Copyright 2021 Revathi Jambunathan
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ */
+#include "BackTransformParticleFunctor.H"
+#include "Particles/Pusher/GetAndSetPosition.H"
+#include "Particles/WarpXParticleContainer.H"
+#include "WarpX.H"
+#include <AMReX.H>
+#include <AMReX_Print.H>
+#include <AMReX_BaseFwd.H>
+
+SelectParticles::SelectParticles (const WarpXParIter& a_pti, TmpParticles& tmp_particle_data,
+                                  amrex::Real current_z_boost, amrex::Real old_z_boost,
+                                  int a_offset)
+    : m_current_z_boost(current_z_boost), m_old_z_boost(old_z_boost)
+{
+    m_get_position = GetParticlePosition(a_pti, a_offset);
+
+    const auto lev = a_pti.GetLevel();
+    const auto index = a_pti.GetPairIndex();
+
+    zpold = tmp_particle_data[lev][index][TmpIdx::zold].dataPtr();
+}
+
+
+LorentzTransformParticles::LorentzTransformParticles ( const WarpXParIter& a_pti,
+                                TmpParticles& tmp_particle_data,
+                                amrex::Real t_boost, amrex::Real dt,
+                                amrex::Real t_lab, int a_offset)
+    : m_t_boost(t_boost), m_dt(dt), m_t_lab(t_lab)
+{
+    using namespace amrex::literals;
+
+    if (tmp_particle_data.size() == 0) return;
+    m_get_position = GetParticlePosition(a_pti, a_offset);
+
+    auto& attribs = a_pti.GetAttribs();
+    m_wpnew = attribs[PIdx::w].dataPtr();
+    m_uxpnew = attribs[PIdx::ux].dataPtr();
+    m_uypnew = attribs[PIdx::uy].dataPtr();
+    m_uzpnew = attribs[PIdx::uz].dataPtr();
+
+    const auto lev = a_pti.GetLevel();
+    const auto index = a_pti.GetPairIndex();
+
+    m_xpold = tmp_particle_data[lev][index][TmpIdx::xold].dataPtr();
+    m_ypold = tmp_particle_data[lev][index][TmpIdx::yold].dataPtr();
+    m_zpold = tmp_particle_data[lev][index][TmpIdx::zold].dataPtr();
+    m_uxpold = tmp_particle_data[lev][index][TmpIdx::uxold].dataPtr();
+    m_uypold = tmp_particle_data[lev][index][TmpIdx::uyold].dataPtr();
+    m_uzpold = tmp_particle_data[lev][index][TmpIdx::uzold].dataPtr();
+
+    m_betaboost = WarpX::beta_boost;
+    m_gammaboost = WarpX::gamma_boost;
+    m_Phys_c = PhysConst::c;
+    m_inv_c2 = 1._rt/(m_Phys_c * m_Phys_c);
+    m_uzfrm = -m_gammaboost*m_betaboost*m_Phys_c;
+}
+
+/**
+ * \brief Functor to compute Lorentz Transform and store the selected particles in existing
+ * particle buffers
+ */
+BackTransformParticleFunctor::BackTransformParticleFunctor (
+                              WarpXParticleContainer *pc_src,
+                              std::string species_name,
+                              int num_buffers)
+    : m_pc_src(pc_src), m_species_name(species_name), m_num_buffers(num_buffers)
+{
+    InitData();
+}
+
+
+void
+BackTransformParticleFunctor::operator () (ParticleContainer& pc_dst, int &totalParticleCounter, int i_buffer) const
+{
+    if (m_perform_backtransform[i_buffer] == 0) return;
+    auto &warpx = WarpX::GetInstance();
+    // get particle slice
+    const int nlevs = std::max(0, m_pc_src->finestLevel()+1);
+    auto tmp_particle_data = m_pc_src->getTmpParticleData();
+    int total_particles_added = 0;
+    for (int lev = 0; lev < nlevs; ++lev) {
+        amrex::Real t_boost = warpx.gett_new(0);
+        amrex::Real dt = warpx.getdt(0);
+
+        for (WarpXParIter pti(*m_pc_src, lev); pti.isValid(); ++pti) {
+            auto ptile_dst = pc_dst.DefineAndReturnParticleTile(lev, pti.index(), pti.LocalTileIndex() );
+        }
+
+        auto& particles = m_pc_src->GetParticles(lev);
+#ifdef AMREX_USE_OMP
+#pragma omp parallel
+#endif
+        {
+            // Temporary arrays to store copy_flag and copy_index for particles
+            // that cross the z-slice
+            amrex::Gpu::DeviceVector<int> FlagForPartCopy;
+            amrex::Gpu::DeviceVector<int> IndexForPartCopy;
+
+            for (WarpXParIter pti(*m_pc_src, lev); pti.isValid(); ++pti) {
+
+                auto index = std::make_pair(pti.index(), pti.LocalTileIndex());
+
+                const auto GetParticleFilter = SelectParticles(pti, tmp_particle_data,
+                                               m_current_z_boost[i_buffer],
+                                               m_old_z_boost[i_buffer]);
+                const auto GetParticleLorentzTransform = LorentzTransformParticles(
+                                                         pti, tmp_particle_data,
+                                                         t_boost, dt,
+                                                         m_t_lab[i_buffer]);
+
+                long const np = pti.numParticles();
+
+                FlagForPartCopy.resize(np);
+                IndexForPartCopy.resize(np);
+
+                int* const AMREX_RESTRICT Flag = FlagForPartCopy.dataPtr();
+                int* const AMREX_RESTRICT IndexLocation = IndexForPartCopy.dataPtr();
+
+                const auto& ptile_src = particles.at(index);
+                auto src_data = ptile_src.getConstParticleTileData();
+                // Flag particles that need to be copied if they cross the z-slice
+                // setting this to 1 for testing (temporarily)
+                amrex::ParallelFor(np,
+                [=] AMREX_GPU_DEVICE(int i)
+                {
+                    Flag[i] = GetParticleFilter(src_data, i);
+                });
+
+                const int total_partdiag_size = amrex::Scan::ExclusiveSum(np,Flag,IndexLocation);
+                auto& ptile_dst = pc_dst.DefineAndReturnParticleTile(lev, pti.index(), pti.LocalTileIndex() );
+                auto old_size = ptile_dst.numParticles();
+                ptile_dst.resize(old_size + total_partdiag_size);
+                auto count = amrex::filterParticles(ptile_dst, ptile_src, GetParticleFilter, 0, old_size, np);
+                auto dst_data = ptile_dst.getParticleTileData();
+                amrex::ParallelFor(np,
+                [=] AMREX_GPU_DEVICE(int i)
+                {
+                   if (Flag[i] == 1) GetParticleLorentzTransform(dst_data, src_data, i,
+                                                                 old_size + IndexLocation[i]);
+                });
+                total_particles_added += count;
+            }
+        }
+    }
+    totalParticleCounter = pc_dst.TotalNumberOfParticles();
+}
+
+
+void
+BackTransformParticleFunctor::InitData()
+{
+    m_current_z_boost.resize(m_num_buffers);
+    m_old_z_boost.resize(m_num_buffers);
+    m_t_lab.resize(m_num_buffers);
+    m_perform_backtransform.resize(m_num_buffers);
+}
+
+void
+BackTransformParticleFunctor::PrepareFunctorData ( int i_buffer, bool z_slice_in_domain,
+                              amrex::Real old_z_boost, amrex::Real current_z_boost,
+                              amrex::Real t_lab, int snapshot_full)
+{
+    m_old_z_boost.at(i_buffer) = old_z_boost;
+    m_current_z_boost.at(i_buffer) = current_z_boost;
+    m_t_lab.at(i_buffer) = t_lab;
+    m_perform_backtransform.at(i_buffer) = 0;
+    if (z_slice_in_domain == true and snapshot_full == 0) m_perform_backtransform.at(i_buffer) = 1;
+}