Refactor AddPlasma and AddPlasmaFlux (ECP-WarpX#5231)

---------

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

Source/Particles/AddPlasmaUtilities.H

@@ -0,0 +1,210 @@
+/* Copyright 2024 The WarpX Community
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ * Authors: Andrew Myers
+ */
+#ifndef WARPX_ADDPLASMAUTILITIES_H_
+#define WARPX_ADDPLASMAUTILITIES_H_
+
+#include "Initialization/PlasmaInjector.H"
+
+#ifdef WARPX_QED
+#   include "Particles/ElementaryProcess/QEDInternals/BreitWheelerEngineWrapper.H"
+#   include "Particles/ElementaryProcess/QEDInternals/QuantumSyncEngineWrapper.H"
+#endif
+
+#include <AMReX_Array.H>
+#include <AMReX_Box.H>
+#include <AMReX_GpuContainers.H>
+#include <AMReX_IntVect.H>
+#include <AMReX_REAL.H>
+#include <AMReX_RealBox.H>
+
+/*
+  Finds the overlap region between the given tile_realbox and part_realbox, returning true
+  if an overlap exists and false if otherwise. This also sets the parameters overlap_realbox,
+  overlap_box, and shifted to the appropriate values.
+ */
+bool find_overlap (const amrex::RealBox& tile_realbox, const amrex::RealBox& part_realbox,
+                   const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& dx,
+                   const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& prob_lo,
+                   amrex::RealBox& overlap_realbox, amrex::Box& overlap_box, amrex::IntVect& shifted);
+
+/*
+  Finds the overlap region between the given tile_realbox, part_realbox and the surface used for
+  flux injection, returning true if an overlap exists and false if otherwise. This also sets the
+  parameters overlap_realbox, overlap_box, and shifted to the appropriate values.
+ */
+bool find_overlap_flux (const amrex::RealBox& tile_realbox, const amrex::RealBox& part_realbox,
+                        const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& dx,
+                        const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& prob_lo,
+                        const PlasmaInjector& plasma_injector,
+                        amrex::RealBox& overlap_realbox, amrex::Box& overlap_box, amrex::IntVect& shifted);
+
+/*
+  This computes the scale_fac (used for setting the particle weights) on a volumetric basis.
+ */
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+amrex::Real compute_scale_fac_volume (const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& dx,
+                                      const amrex::Long pcount) {
+    using namespace amrex::literals;
+    return (pcount != 0) ? AMREX_D_TERM(dx[0],*dx[1],*dx[2])/pcount : 0.0_rt;
+}
+
+/*
+  Given a refinement ratio, this computes the total increase in resolution for a plane
+  defined by the normal_axis.
+ */
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+int compute_area_weights (const amrex::IntVect& iv, const int normal_axis) {
+    int r = AMREX_D_TERM(iv[0],*iv[1],*iv[2]);
+#if defined(WARPX_DIM_3D)
+    r /= iv[normal_axis];
+#elif defined(WARPX_DIM_RZ) || defined(WARPX_DIM_XZ)
+    if (normal_axis == 0) { r /= iv[0]; }
+    else if (normal_axis == 2) { r /= iv[1]; }
+#elif defined(WARPX_DIM_1D_Z)
+    if (normal_axis == 2) { r /= iv[0]; }
+#endif
+    return r;
+}
+
+/*
+  This computes the scale_fac (used for setting the particle weights) on a on area basis
+  (used for flux injection).
+ */
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+amrex::Real compute_scale_fac_area (const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& dx,
+                                    const amrex::Real num_ppc_real, const int flux_normal_axis) {
+    using namespace amrex::literals;
+    amrex::Real scale_fac = AMREX_D_TERM(dx[0],*dx[1],*dx[2])/num_ppc_real;
+    // Scale particle weight by the area of the emitting surface, within one cell
+#if defined(WARPX_DIM_3D)
+    scale_fac /= dx[flux_normal_axis];
+#elif defined(WARPX_DIM_RZ) || defined(WARPX_DIM_XZ)
+    // When emission is in the r direction, the emitting surface is a cylinder.
+    // The factor 2*pi*r is added later below.
+    if (flux_normal_axis == 0) { scale_fac /= dx[0]; }
+    // When emission is in the z direction, the emitting surface is an annulus
+    // The factor 2*pi*r is added later below.
+    if (flux_normal_axis == 2) { scale_fac /= dx[1]; }
+    // When emission is in the theta direction (flux_normal_axis == 1),
+    // the emitting surface is a rectangle, within the plane of the simulation
+#elif defined(WARPX_DIM_1D_Z)
+    if (flux_normal_axis == 2) { scale_fac /= dx[0]; }
+#endif
+    return scale_fac;
+}
+
+/*
+  These structs encapsulates several data structures needed for using the parser during plasma
+  injection.
+ */
+struct PlasmaParserWrapper
+{
+    PlasmaParserWrapper (std::size_t a_num_user_int_attribs,
+                         std::size_t a_num_user_real_attribs,
+                         const amrex::Vector< std::unique_ptr<amrex::Parser> >& a_user_int_attrib_parser,
+                         const amrex::Vector< std::unique_ptr<amrex::Parser> >& a_user_real_attrib_parser);
+
+    amrex::Gpu::PinnedVector< amrex::ParserExecutor<7> > m_user_int_attrib_parserexec_pinned;
+    amrex::Gpu::PinnedVector< amrex::ParserExecutor<7> > m_user_real_attrib_parserexec_pinned;
+};
+
+struct PlasmaParserHelper
+{
+    template <typename SoAType>
+    PlasmaParserHelper (SoAType& a_soa, std::size_t old_size,
+                        const std::vector<std::string>& a_user_int_attribs,
+                        const std::vector<std::string>& a_user_real_attribs,
+                        std::map<std::string,int>& a_particle_icomps,
+                        std::map<std::string,int>& a_particle_comps,
+                        const PlasmaParserWrapper& wrapper) :
+        m_wrapper_ptr(&wrapper) {
+        m_pa_user_int_pinned.resize(a_user_int_attribs.size());
+        m_pa_user_real_pinned.resize(a_user_real_attribs.size());
+
+#ifdef AMREX_USE_GPU
+        m_d_pa_user_int.resize(a_user_int_attribs.size());
+        m_d_pa_user_real.resize(a_user_real_attribs.size());
+        m_d_user_int_attrib_parserexec.resize(a_user_int_attribs.size());
+        m_d_user_real_attrib_parserexec.resize(a_user_real_attribs.size());
+#endif
+
+        for (std::size_t ia = 0; ia < a_user_int_attribs.size(); ++ia) {
+            m_pa_user_int_pinned[ia] = a_soa.GetIntData(a_particle_icomps[a_user_int_attribs[ia]]).data() + old_size;
+        }
+        for (std::size_t ia = 0; ia < a_user_real_attribs.size(); ++ia) {
+            m_pa_user_real_pinned[ia] = a_soa.GetRealData(a_particle_comps[a_user_real_attribs[ia]]).data() + old_size;
+        }
+
+#ifdef AMREX_USE_GPU
+        amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice, m_pa_user_int_pinned.begin(),
+                              m_pa_user_int_pinned.end(), m_d_pa_user_int.begin());
+        amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice, m_pa_user_real_pinned.begin(),
+                              m_pa_user_real_pinned.end(), m_d_pa_user_real.begin());
+        amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice, wrapper.m_user_int_attrib_parserexec_pinned.begin(),
+                              wrapper.m_user_int_attrib_parserexec_pinned.end(), m_d_user_int_attrib_parserexec.begin());
+        amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice, wrapper.m_user_real_attrib_parserexec_pinned.begin(),
+                              wrapper.m_user_real_attrib_parserexec_pinned.end(), m_d_user_real_attrib_parserexec.begin());
+#endif
+    }
+
+    int** getUserIntDataPtrs ();
+    amrex::ParticleReal** getUserRealDataPtrs ();
+    [[nodiscard]] amrex::ParserExecutor<7> const* getUserIntParserExecData () const;
+    [[nodiscard]] amrex::ParserExecutor<7> const* getUserRealParserExecData () const;
+
+    amrex::Gpu::PinnedVector<int*> m_pa_user_int_pinned;
+    amrex::Gpu::PinnedVector<amrex::ParticleReal*> m_pa_user_real_pinned;
+
+#ifdef AMREX_USE_GPU
+    // To avoid using managed memory, we first define pinned memory vector, initialize on cpu,
+    // and them memcpy to device from host
+    amrex::Gpu::DeviceVector<int*> m_d_pa_user_int;
+    amrex::Gpu::DeviceVector<amrex::ParticleReal*> m_d_pa_user_real;
+    amrex::Gpu::DeviceVector< amrex::ParserExecutor<7> > m_d_user_int_attrib_parserexec;
+    amrex::Gpu::DeviceVector< amrex::ParserExecutor<7> > m_d_user_real_attrib_parserexec;
+#endif
+    const PlasmaParserWrapper* m_wrapper_ptr;
+};
+
+#ifdef WARPX_QED
+struct QEDHelper
+{
+    template <typename SoAType>
+    QEDHelper (SoAType& a_soa, std::size_t old_size,
+               std::map<std::string,int>& a_particle_comps,
+               bool a_has_quantum_sync, bool a_has_breit_wheeler,
+               const std::shared_ptr<QuantumSynchrotronEngine>& a_shr_p_qs_engine,
+               const std::shared_ptr<BreitWheelerEngine>& a_shr_p_bw_engine)
+        : has_quantum_sync(a_has_quantum_sync), has_breit_wheeler(a_has_breit_wheeler)
+    {
+        if(has_quantum_sync){
+            quantum_sync_get_opt =
+                a_shr_p_qs_engine->build_optical_depth_functor();
+            p_optical_depth_QSR = a_soa.GetRealData(
+                a_particle_comps["opticalDepthQSR"]).data() + old_size;
+        }
+        if(has_breit_wheeler){
+            breit_wheeler_get_opt =
+                a_shr_p_bw_engine->build_optical_depth_functor();
+            p_optical_depth_BW = a_soa.GetRealData(
+                a_particle_comps["opticalDepthBW"]).data() + old_size;
+        }
+    }
+
+    amrex::ParticleReal* p_optical_depth_QSR = nullptr;
+    amrex::ParticleReal* p_optical_depth_BW  = nullptr;
+
+    const bool has_quantum_sync;
+    const bool has_breit_wheeler;
+
+    QuantumSynchrotronGetOpticalDepth quantum_sync_get_opt;
+    BreitWheelerGetOpticalDepth breit_wheeler_get_opt;
+};
+#endif
+
+#endif /*WARPX_ADDPLASMAUTILITIES_H_*/

Source/Particles/AddPlasmaUtilities.cpp

@@ -0,0 +1,158 @@
+/* Copyright 2024 The WarpX Community
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ * Authors: Andrew Myers
+ */
+#include "AddPlasmaUtilities.H"
+
+#include <cmath>
+
+bool find_overlap (const amrex::RealBox& tile_realbox, const amrex::RealBox& part_realbox,
+                   const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& dx,
+                   const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& prob_lo,
+                   amrex::RealBox& overlap_realbox, amrex::Box& overlap_box, amrex::IntVect& shifted)
+{
+    using namespace amrex::literals;
+
+    bool no_overlap = false;
+    for (int dir=0; dir<AMREX_SPACEDIM; dir++) {
+        if ( tile_realbox.lo(dir) <= part_realbox.hi(dir) ) {
+            const amrex::Real ncells_adjust = std::floor( (tile_realbox.lo(dir) - part_realbox.lo(dir))/dx[dir] );
+            overlap_realbox.setLo( dir, part_realbox.lo(dir) + std::max(ncells_adjust, 0._rt) * dx[dir]);
+        } else {
+            no_overlap = true; break;
+        }
+        if ( tile_realbox.hi(dir) >= part_realbox.lo(dir) ) {
+            const amrex::Real ncells_adjust = std::floor( (part_realbox.hi(dir) - tile_realbox.hi(dir))/dx[dir] );
+            overlap_realbox.setHi( dir, part_realbox.hi(dir) - std::max(ncells_adjust, 0._rt) * dx[dir]);
+        } else {
+            no_overlap = true; break;
+        }
+        // Count the number of cells in this direction in overlap_realbox
+        overlap_box.setSmall( dir, 0 );
+        overlap_box.setBig( dir,
+                            int( std::round((overlap_realbox.hi(dir)-overlap_realbox.lo(dir))
+                                            /dx[dir] )) - 1);
+        shifted[dir] =
+            static_cast<int>(std::round((overlap_realbox.lo(dir)-prob_lo[dir])/dx[dir]));
+        // shifted is exact in non-moving-window direction.  That's all we care.
+    }
+    return no_overlap;
+}
+
+bool find_overlap_flux (const amrex::RealBox& tile_realbox, const amrex::RealBox& part_realbox,
+                        const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& dx,
+                        const amrex::GpuArray<amrex::Real, AMREX_SPACEDIM>& prob_lo,
+                        const PlasmaInjector& plasma_injector,
+                        amrex::RealBox& overlap_realbox, amrex::Box& overlap_box, amrex::IntVect& shifted)
+{
+    using namespace amrex::literals;
+
+    bool no_overlap = false;
+    for (int dir=0; dir<AMREX_SPACEDIM; dir++) {
+#if (defined(WARPX_DIM_3D))
+        if (dir == plasma_injector.flux_normal_axis) {
+#elif defined(WARPX_DIM_XZ) || defined(WARPX_DIM_RZ)
+        if (2*dir == plasma_injector.flux_normal_axis) {
+            // The above formula captures the following cases:
+            // - flux_normal_axis=0 (emission along x/r) and dir=0
+            // - flux_normal_axis=2 (emission along z) and dir=1
+#elif defined(WARPX_DIM_1D_Z)
+        if ( (dir==0) && (plasma_injector.flux_normal_axis==2) ) {
+#endif
+            if (plasma_injector.flux_direction > 0) {
+                if (plasma_injector.surface_flux_pos <  tile_realbox.lo(dir) ||
+                    plasma_injector.surface_flux_pos >= tile_realbox.hi(dir)) {
+                    no_overlap = true;
+                    break;
+                }
+            } else {
+                if (plasma_injector.surface_flux_pos <= tile_realbox.lo(dir) ||
+                    plasma_injector.surface_flux_pos >  tile_realbox.hi(dir)) {
+                    no_overlap = true;
+                    break;
+                }
+            }
+            overlap_realbox.setLo( dir, plasma_injector.surface_flux_pos );
+            overlap_realbox.setHi( dir, plasma_injector.surface_flux_pos );
+            overlap_box.setSmall( dir, 0 );
+            overlap_box.setBig( dir, 0 );
+            shifted[dir] =
+                static_cast<int>(std::round((overlap_realbox.lo(dir)-prob_lo[dir])/dx[dir]));
+        } else {
+            if ( tile_realbox.lo(dir) <= part_realbox.hi(dir) ) {
+                const amrex::Real ncells_adjust = std::floor( (tile_realbox.lo(dir) - part_realbox.lo(dir))/dx[dir] );
+                overlap_realbox.setLo( dir, part_realbox.lo(dir) + std::max(ncells_adjust, 0._rt) * dx[dir]);
+            } else {
+                no_overlap = true; break;
+            }
+            if ( tile_realbox.hi(dir) >= part_realbox.lo(dir) ) {
+                const amrex::Real ncells_adjust = std::floor( (part_realbox.hi(dir) - tile_realbox.hi(dir))/dx[dir] );
+                overlap_realbox.setHi( dir, part_realbox.hi(dir) - std::max(ncells_adjust, 0._rt) * dx[dir]);
+            } else {
+                no_overlap = true; break;
+            }
+            // Count the number of cells in this direction in overlap_realbox
+            overlap_box.setSmall( dir, 0 );
+            overlap_box.setBig( dir,
+                                int( std::round((overlap_realbox.hi(dir)-overlap_realbox.lo(dir))
+                                                /dx[dir] )) - 1);
+            shifted[dir] =
+                static_cast<int>(std::round((overlap_realbox.lo(dir)-prob_lo[dir])/dx[dir]));
+            // shifted is exact in non-moving-window direction.  That's all we care.
+        }
+    }
+
+    return no_overlap;
+}
+
+PlasmaParserWrapper::PlasmaParserWrapper (const std::size_t a_num_user_int_attribs,
+                                          const std::size_t a_num_user_real_attribs,
+                                          const amrex::Vector< std::unique_ptr<amrex::Parser> >& a_user_int_attrib_parser,
+                                          const amrex::Vector< std::unique_ptr<amrex::Parser> >& a_user_real_attrib_parser)
+
+{
+    m_user_int_attrib_parserexec_pinned.resize(a_num_user_int_attribs);
+    m_user_real_attrib_parserexec_pinned.resize(a_num_user_real_attribs);
+
+    for (std::size_t ia = 0; ia < a_num_user_int_attribs; ++ia) {
+        m_user_int_attrib_parserexec_pinned[ia] = a_user_int_attrib_parser[ia]->compile<7>();
+    }
+    for (std::size_t ia = 0; ia < a_num_user_real_attribs; ++ia) {
+        m_user_real_attrib_parserexec_pinned[ia] = a_user_real_attrib_parser[ia]->compile<7>();
+    }
+}
+
+int** PlasmaParserHelper::getUserIntDataPtrs () {
+#ifdef AMREX_USE_GPU
+    return m_d_pa_user_int.dataPtr();
+#else
+    return m_pa_user_int_pinned.dataPtr();
+#endif
+}
+
+amrex::ParticleReal** PlasmaParserHelper::getUserRealDataPtrs () {
+#ifdef AMREX_USE_GPU
+    return m_d_pa_user_real.dataPtr();
+#else
+    return m_pa_user_real_pinned.dataPtr();
+#endif
+}
+
+amrex::ParserExecutor<7> const* PlasmaParserHelper::getUserIntParserExecData () const {
+#ifdef AMREX_USE_GPU
+    return m_d_user_int_attrib_parserexec.dataPtr();
+#else
+    return m_wrapper_ptr->m_user_int_attrib_parserexec_pinned.dataPtr();
+#endif
+}
+
+amrex::ParserExecutor<7> const* PlasmaParserHelper::getUserRealParserExecData () const {
+#ifdef AMREX_USE_GPU
+    return m_d_user_real_attrib_parserexec.dataPtr();
+#else
+    return m_wrapper_ptr->m_user_real_attrib_parserexec_pinned.dataPtr();
+#endif
+}

Source/Particles/CMakeLists.txt

@@ -2,6 +2,7 @@ foreach(D IN LISTS WarpX_DIMS)
     warpx_set_suffix_dims(SD ${D})
     target_sources(lib_${SD}
       PRIVATE
+        AddPlasmaUtilities.cpp
         MultiParticleContainer.cpp
         ParticleBoundaries.cpp
         PhotonParticleContainer.cpp

Source/Particles/Make.package

@@ -1,3 +1,4 @@
+CEXE_sources += AddPlasmaUtilities.cpp
 CEXE_sources += MultiParticleContainer.cpp
 CEXE_sources += WarpXParticleContainer.cpp
 CEXE_sources += RigidInjectedParticleContainer.cpp

Source/Particles/PhysicalParticleContainer.H

@@ -392,6 +392,14 @@ public:
     }
 
 protected:
+
+    /*
+       Finds the box defining the region where refine injection should be used, if that
+       option is enabled. Currently this only works for numLevels() == 2 and static mesh
+       refinement.
+    */
+    bool findRefinedInjectionBox (amrex::Box& fine_injection_box, amrex::IntVect& rrfac);
+
     std::string species_name;
     std::vector<std::unique_ptr<PlasmaInjector>> plasma_injectors;