Continuous injection of moving plasma (#3958)

* Continuous injection of moving plasma

* Fix const correctness

* Fix bugs in calculation of v_bulk

* Fix restart

* Use range-based for loops where possible

* Apply suggestions from code review

* Fix bug related to managed memory

* Apply suggestions from code review

* Exclude case with `moving_window_v = 0`

* Add to WarpXParticleContainer virtual function that returns pointer to plasma injector

* Add to WarpXParticleContainer member variable for current injection position

* Fix bugs

* Fix bug: use continue instead of return

10 files changed, 186 insertions, 72 deletions

diff --git a/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp b/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp
index 3ea631292..28033b3ac 100644
--- a/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp
+++ b/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp
@@ -302,7 +302,12 @@ FlushFormatPlotfile::WriteWarpXHeader(
 
         warpx.GetPartContainer().WriteHeader(HeaderFile);
 
-        HeaderFile << warpx.getcurrent_injection_position() << "\n";
+        MultiParticleContainer& mypc = warpx.GetPartContainer();
+        const int n_species = mypc.nSpecies();
+        for (int i=0; i<n_species; i++)
+        {
+             HeaderFile << mypc.GetParticleContainer(i).m_current_injection_position << "\n";
+        }
 
         HeaderFile << warpx.getdo_moving_window() << "\n";
 
diff --git a/Source/Diagnostics/WarpXIO.cpp b/Source/Diagnostics/WarpXIO.cpp
index 36f875659..aeaf85304 100644
--- a/Source/Diagnostics/WarpXIO.cpp
+++ b/Source/Diagnostics/WarpXIO.cpp
@@ -211,8 +211,12 @@ WarpX::InitFromCheckpoint ()
         }
 
         mypc->ReadHeader(is);
-        is >> current_injection_position;
-        GotoNextLine(is);
+        const int n_species = mypc->nSpecies();
+        for (int i=0; i<n_species; i++)
+        {
+             is >> mypc->GetParticleContainer(i).m_current_injection_position;
+             GotoNextLine(is);
+        }
 
         int do_moving_window_before_restart;
         is >> do_moving_window_before_restart;
diff --git a/Source/Initialization/PlasmaInjector.H b/Source/Initialization/PlasmaInjector.H
index 0f33aa062..1c2fc453d 100644
--- a/Source/Initialization/PlasmaInjector.H
+++ b/Source/Initialization/PlasmaInjector.H
@@ -133,8 +133,10 @@ public:
 
     InjectorPosition* getInjectorPosition ();
     InjectorDensity*  getInjectorDensity ();
+
     InjectorFlux*  getInjectorFlux ();
-    InjectorMomentum* getInjectorMomentum ();
+    InjectorMomentum* getInjectorMomentumDevice ();
+    InjectorMomentum* getInjectorMomentumHost ();
 
 protected:
 
diff --git a/Source/Initialization/PlasmaInjector.cpp b/Source/Initialization/PlasmaInjector.cpp
index 865e531e7..1ed868647 100644
--- a/Source/Initialization/PlasmaInjector.cpp
+++ b/Source/Initialization/PlasmaInjector.cpp
@@ -775,7 +775,13 @@ PlasmaInjector::getInjectorFlux ()
 }
 
 InjectorMomentum*
-PlasmaInjector::getInjectorMomentum ()
+PlasmaInjector::getInjectorMomentumDevice ()
 {
     return d_inj_mom;
 }
+
+InjectorMomentum*
+PlasmaInjector::getInjectorMomentumHost ()
+{
+    return h_inj_mom.get();
+}
diff --git a/Source/Particles/PhysicalParticleContainer.H b/Source/Particles/PhysicalParticleContainer.H
index 4a9beeec2..24a1b94d3 100644
--- a/Source/Particles/PhysicalParticleContainer.H
+++ b/Source/Particles/PhysicalParticleContainer.H
@@ -66,6 +66,11 @@ public:
 
     virtual void InitIonizationModule () override;
 
+    /*
+     * \brief Returns a pointer to the plasma injector.
+     */
+    virtual PlasmaInjector* GetPlasmaInjector () override;
+
     /**
      * \brief Evolve is the central function PhysicalParticleContainer that
      * advances plasma particles for a time dt (typically one timestep).
diff --git a/Source/Particles/PhysicalParticleContainer.cpp b/Source/Particles/PhysicalParticleContainer.cpp
index 0793094ff..8dff41c05 100644
--- a/Source/Particles/PhysicalParticleContainer.cpp
+++ b/Source/Particles/PhysicalParticleContainer.cpp
@@ -926,7 +926,7 @@ PhysicalParticleContainer::AddPlasma (int lev, RealBox part_realbox)
 
     InjectorPosition* inj_pos = plasma_injector->getInjectorPosition();
     InjectorDensity*  inj_rho = plasma_injector->getInjectorDensity();
-    InjectorMomentum* inj_mom = plasma_injector->getInjectorMomentum();
+    InjectorMomentum* inj_mom = plasma_injector->getInjectorMomentumDevice();
     const Real gamma_boost = WarpX::gamma_boost;
     const Real beta_boost = WarpX::beta_boost;
     const Real t = WarpX::GetInstance().gett_new(lev);
@@ -1476,7 +1476,7 @@ PhysicalParticleContainer::AddPlasmaFlux (amrex::Real dt)
 
     InjectorPosition* inj_pos = plasma_injector->getInjectorPosition();
     InjectorFlux*  inj_flux = plasma_injector->getInjectorFlux();
-    InjectorMomentum* inj_mom = plasma_injector->getInjectorMomentum();
+    InjectorMomentum* inj_mom = plasma_injector->getInjectorMomentumDevice();
     constexpr int level_zero = 0;
     const amrex::Real t = WarpX::GetInstance().gett_new(level_zero);
 
@@ -2922,6 +2922,11 @@ PhysicalParticleContainer::getIonizationFunc (const WarpXParIter& pti,
                                 ion_atomic_number);
 }
 
+PlasmaInjector* PhysicalParticleContainer::GetPlasmaInjector ()
+{
+    return plasma_injector.get();
+}
+
 void PhysicalParticleContainer::resample (const int timestep)
 {
     // In heavily load imbalanced simulations, MPI processes with few particles will spend most of
diff --git a/Source/Particles/WarpXParticleContainer.H b/Source/Particles/WarpXParticleContainer.H
index ffeead6f1..6e0d3f4f9 100644
--- a/Source/Particles/WarpXParticleContainer.H
+++ b/Source/Particles/WarpXParticleContainer.H
@@ -13,6 +13,7 @@
 #include "WarpXParticleContainer_fwd.H"
 
 #include "Evolve/WarpXDtType.H"
+#include "Initialization/PlasmaInjector.H"
 #include "Particles/ParticleBoundaries.H"
 #include "SpeciesPhysicalProperties.H"
 
@@ -121,6 +122,12 @@ public:
 
     virtual void InitIonizationModule () {}
 
+    /*
+     * \brief Virtual function that returns a pointer to the plasma injector,
+     *        for derived classes that define one (PhysicalParticleContainer).
+     */
+    virtual PlasmaInjector* GetPlasmaInjector () { return nullptr; }
+
     /**
      * Evolve is the central WarpXParticleContainer function that advances
      * particles for a time dt (typically one timestep). It is a pure virtual
@@ -241,10 +248,13 @@ public:
     virtual void ContinuousInjection(const amrex::RealBox& /*injection_box*/) {}
     // Update optional sub-class-specific injection location.
     virtual void UpdateContinuousInjectionPosition(amrex::Real /*dt*/) {}
+    bool doContinuousInjection() const {return do_continuous_injection;}
 
     // Inject a continuous flux of particles from a defined plane
     virtual void ContinuousFluxInjection(amrex::Real /*t*/, amrex::Real /*dt*/) {}
 
+    int getSpeciesId() const {return species_id;}
+
     ///
     /// This returns the total charge for all the particles in this ParticleContainer.
     /// This is needed when solving Poisson's equation with periodic boundary conditions.
@@ -306,6 +316,9 @@ public:
     int self_fields_max_iters = 200;
     int self_fields_verbosity = 2;
 
+    //! Current injection position
+    amrex::Real m_current_injection_position;
+
     // split along diagonals (0) or axes (1)
     int split_type = 0;
 
diff --git a/Source/Utils/WarpXMovingWindow.cpp b/Source/Utils/WarpXMovingWindow.cpp
index 7e7985b98..cef4764b1 100644
--- a/Source/Utils/WarpXMovingWindow.cpp
+++ b/Source/Utils/WarpXMovingWindow.cpp
@@ -56,25 +56,81 @@ void
 WarpX::UpdatePlasmaInjectionPosition (amrex::Real a_dt)
 {
     const int dir = moving_window_dir;
-    // Continuously inject plasma in new cells (by default only on level 0)
-    if (WarpX::warpx_do_continuous_injection and (WarpX::gamma_boost > 1)){
-        // In boosted-frame simulations, the plasma has moved since the last
-        // call to this function, and injection position needs to be updated
-        current_injection_position -= WarpX::beta_boost *
+
+    // Loop over species
+    const int n_species = mypc->nSpecies();
+    for (int i=0; i<n_species; i++)
+    {
+        WarpXParticleContainer& pc = mypc->GetParticleContainer(i);
+
+        // Continuously inject plasma in new cells (by default only on level 0)
+        if (pc.doContinuousInjection())
+        {
+            PlasmaInjector* plasma_injector = pc.GetPlasmaInjector();
+            if (plasma_injector == nullptr) continue;
+
+            // Get bulk momentum and velocity of plasma
+            amrex::XDim3 u_bulk = amrex::XDim3{0._rt, 0._rt, 0._rt};
+            amrex::Real v_bulk = 0._rt;
+            if (dir == 0)
+            {
+                // dir=0 is z in 1D, x in 2D, x in 3D
+#if defined(WARPX_DIM_1D_Z)
+                u_bulk = plasma_injector->getInjectorMomentumHost()->getBulkMomentum(0._rt, 0._rt, pc.m_current_injection_position);
+                v_bulk = PhysConst::c * u_bulk.z / std::sqrt(1._rt + u_bulk.z*u_bulk.z);
+#else // 2D, 3D
+                u_bulk = plasma_injector->getInjectorMomentumHost()->getBulkMomentum(pc.m_current_injection_position, 0._rt, 0._rt);
+                v_bulk = PhysConst::c * u_bulk.x / std::sqrt(1._rt + u_bulk.x*u_bulk.x);
+#endif
+            }
+            else if (dir == 1)
+            {
+                // dir=1 is nothing in 1D, z in 2D, y in 3D
+                // (we do not expect to enter this code block in 1D)
+#if defined(WARPX_DIM_XZ) || defined(WARPX_DIM_RZ)
+                u_bulk = plasma_injector->getInjectorMomentumHost()->getBulkMomentum(0._rt, 0._rt, pc.m_current_injection_position);
+                v_bulk = PhysConst::c * u_bulk.z / std::sqrt(1._rt + u_bulk.z*u_bulk.z);
+#else // 3D
+                u_bulk = plasma_injector->getInjectorMomentumHost()->getBulkMomentum(0._rt, pc.m_current_injection_position, 0._rt);
+                v_bulk = PhysConst::c * u_bulk.y / std::sqrt(1._rt + u_bulk.y*u_bulk.y);
+#endif
+            }
+            else if (dir == 2)
+            {
+                // dir=2 is nothing in 1D, nothing in 2D, z in 3D
+                // (we do not expect to enter this code block in 1D and 2D)
+                u_bulk = plasma_injector->getInjectorMomentumHost()->getBulkMomentum(0._rt, 0._rt, pc.m_current_injection_position);
+                v_bulk = PhysConst::c * u_bulk.z / std::sqrt(1._rt + u_bulk.z*u_bulk.z);
+            }
+
+            // In boosted-frame simulations, the plasma has moved since the last
+            // call to this function, and injection position needs to be updated.
+            // Note that the bulk velocity v, obtained from getBulkMomentum, is
+            // transformed to the boosted frame velocity v' via the formula
+            // v' = (v-c*beta)/(1-v*beta/c)
+            if (WarpX::gamma_boost > 1._rt)
+            {
+                v_bulk = (v_bulk - PhysConst::c*WarpX::beta_boost)
+                         / (1._rt - v_bulk*WarpX::beta_boost/PhysConst::c);
 #if defined(WARPX_DIM_3D)
-            WarpX::boost_direction[dir] * PhysConst::c * a_dt;
+                v_bulk *= WarpX::boost_direction[dir];
 #elif defined(WARPX_DIM_XZ) || defined(WARPX_DIM_RZ)
-            // In 2D, dir=0 corresponds to x and dir=1 corresponds to z
-            // This needs to be converted in order to index `boost_direction`
-            // which has 3 components, for both 2D and 3D simulations.
-            WarpX::boost_direction[2*dir] * PhysConst::c * a_dt;
+                // In 2D, dir=0 corresponds to x and dir=1 corresponds to z.
+                // This needs to be converted to access boost_direction,
+                // which has always 3 components.
+                v_bulk *= WarpX::boost_direction[2*dir];
 #elif defined(WARPX_DIM_1D_Z)
-            // In 1D, dir=0 corresponds to z
-            // This needs to be converted in order to index `boost_direction`
-            // which has 3 components, for 1D, 2D, and 3D simulations.
-            WarpX::boost_direction[2] * PhysConst::c * a_dt;
-            amrex::ignore_unused(dir);
+                // In 1D, dir=0 corresponds to z.
+                // This needs to be converted to access boost_direction,
+                // which has always 3 components.
+                v_bulk *= WarpX::boost_direction[2];
+                amrex::ignore_unused(dir);
 #endif
+            }
+
+            // Update current injection position
+            pc.m_current_injection_position += v_bulk * a_dt;
+        }
     }
 }
 
@@ -96,15 +152,13 @@ WarpX::MoveWindow (const int step, bool move_j)
     moving_window_x += (moving_window_v - WarpX::beta_boost * PhysConst::c)/(1 - moving_window_v * WarpX::beta_boost / PhysConst::c) * dt[0];
     const int dir = moving_window_dir;
 
-    // Update warpx.current_injection_position
+    // Update warpx.current_injection_position,
     // PhysicalParticleContainer uses this injection position
     UpdatePlasmaInjectionPosition( dt[0] );
-    if (WarpX::warpx_do_continuous_injection){
-        // Update injection position for WarpXParticleContainer in mypc.
-        // Nothing to do for PhysicalParticleContainers
-        // For LaserParticleContainer, need to update the antenna position.
-        mypc->UpdateContinuousInjectionPosition( dt[0] );
-    }
+    // Update injection position for WarpXParticleContainer in mypc,
+    // nothing to do for PhysicalParticleContainer,
+    // need to update the antenna position for LaserParticleContainer.
+    mypc->UpdateContinuousInjectionPosition( dt[0] );
 
     // compute the number of cells to shift on the base level
     amrex::Real new_lo[AMREX_SPACEDIM];
@@ -312,41 +366,54 @@ WarpX::MoveWindow (const int step, bool move_j)
         }
     }
 
-    // Continuously inject plasma in new cells (by default only on level 0)
-    if (WarpX::warpx_do_continuous_injection) {
-
-        const int lev = 0;
-
-        // particleBox encloses the cells where we generate particles
-        // (only injects particles in an integer number of cells,
-        // for correct particle spacing)
-        amrex::RealBox particleBox = geom[lev].ProbDomain();
-        amrex::Real new_injection_position;
-        if (moving_window_v >= 0){
-            // Forward-moving window
-            const amrex::Real dx = geom[lev].CellSize(dir);
-            new_injection_position = current_injection_position +
-                std::floor( (geom[lev].ProbHi(dir) - current_injection_position)/dx ) * dx;
-        } else {
-            // Backward-moving window
-            const amrex::Real dx = geom[lev].CellSize(dir);
-            new_injection_position = current_injection_position -
-                std::floor( (current_injection_position - geom[lev].ProbLo(dir))/dx) * dx;
-        }
-        // Modify the corresponding bounds of the particleBox
-        if (moving_window_v >= 0) {
-            particleBox.setLo( dir, current_injection_position );
-            particleBox.setHi( dir, new_injection_position );
-        } else {
-            particleBox.setLo( dir, new_injection_position );
-            particleBox.setHi( dir, current_injection_position );
-        }
+    // Loop over species
+    const int n_species = mypc->nSpecies();
+    for (int i=0; i<n_species; i++)
+    {
+        WarpXParticleContainer& pc = mypc->GetParticleContainer(i);
+
+        // Continuously inject plasma in new cells (by default only on level 0)
+        if (pc.doContinuousInjection())
+        {
+            const int lev = 0;
+
+            // particleBox encloses the cells where we generate particles
+            // (only injects particles in an integer number of cells,
+            // for correct particle spacing)
+            amrex::RealBox particleBox = geom[lev].ProbDomain();
+            amrex::Real new_injection_position = pc.m_current_injection_position;
+            if (moving_window_v > 0._rt)
+            {
+                // Forward-moving window
+                const amrex::Real dx = geom[lev].CellSize(dir);
+                new_injection_position = pc.m_current_injection_position +
+                    std::floor( (geom[lev].ProbHi(dir) - pc.m_current_injection_position)/dx ) * dx;
+            }
+            else if (moving_window_v < 0._rt)
+            {
+                // Backward-moving window
+                const amrex::Real dx = geom[lev].CellSize(dir);
+                new_injection_position = pc.m_current_injection_position -
+                    std::floor( (pc.m_current_injection_position - geom[lev].ProbLo(dir))/dx) * dx;
+            }
+            // Modify the corresponding bounds of the particleBox
+            if (moving_window_v > 0._rt)
+            {
+                particleBox.setLo( dir, pc.m_current_injection_position );
+                particleBox.setHi( dir, new_injection_position );
+            }
+            else if (moving_window_v < 0._rt)
+            {
+                particleBox.setLo( dir, new_injection_position );
+                particleBox.setHi( dir, pc.m_current_injection_position );
+            }
 
-        if (particleBox.ok() and (current_injection_position != new_injection_position)){
-            // Performs continuous injection of all WarpXParticleContainer
-            // in mypc.
-            mypc->ContinuousInjection(particleBox);
-            current_injection_position = new_injection_position;
+            if (particleBox.ok() and (pc.m_current_injection_position != new_injection_position)){
+                // Performs continuous injection of all WarpXParticleContainer
+                // in mypc.
+                pc.ContinuousInjection(particleBox);
+                pc.m_current_injection_position = new_injection_position;
+            }
         }
     }
 
diff --git a/Source/WarpX.H b/Source/WarpX.H
index fb5e403f4..dbe922812 100644
--- a/Source/WarpX.H
+++ b/Source/WarpX.H
@@ -868,7 +868,6 @@ public:
     amrex::Real getdt (int lev) const {return dt[lev];}
     int getdo_moving_window() const {return do_moving_window;}
     amrex::Real getmoving_window_x() const {return moving_window_x;}
-    amrex::Real getcurrent_injection_position () const {return current_injection_position;}
     bool getis_synchronized() const {return is_synchronized;}
 
     int maxStep () const {return max_step;}
@@ -1470,7 +1469,6 @@ private:
     amrex::Real v_particle_pml;
 
     amrex::Real moving_window_x = std::numeric_limits<amrex::Real>::max();
-    amrex::Real current_injection_position = 0;
 
     // Plasma injection parameters
     int warpx_do_continuous_injection = 0;
diff --git a/Source/WarpX.cpp b/Source/WarpX.cpp
index 17601175b..fc37ccc65 100644
--- a/Source/WarpX.cpp
+++ b/Source/WarpX.cpp
@@ -279,16 +279,25 @@ WarpX::WarpX ()
     t_old.resize(nlevs_max, std::numeric_limits<Real>::lowest());
     dt.resize(nlevs_max, std::numeric_limits<Real>::max());
 
-    // Particle Container
+    // Loop over species and set current injection position per species
     mypc = std::make_unique<MultiParticleContainer>(this);
-    warpx_do_continuous_injection = mypc->doContinuousInjection();
-    if (warpx_do_continuous_injection){
-        if (moving_window_v >= 0){
+    const int n_species = mypc->nSpecies();
+    for (int i=0; i<n_species; i++)
+    {
+        WarpXParticleContainer& pc = mypc->GetParticleContainer(i);
+
+        // Storing injection position for all species, regardless of whether
+        // they are continuously injected, since it makes looping over the
+        // elements of current_injection_position easier elsewhere in the code.
+        if (moving_window_v > 0._rt)
+        {
             // Inject particles continuously from the right end of the box
-            current_injection_position = geom[0].ProbHi(moving_window_dir);
-        } else {
+            pc.m_current_injection_position = geom[0].ProbHi(moving_window_dir);
+        }
+        else if (moving_window_v < 0._rt)
+        {
             // Inject particles continuously from the left end of the box
-            current_injection_position = geom[0].ProbLo(moving_window_dir);
+            pc.m_current_injection_position = geom[0].ProbLo(moving_window_dir);
         }
     }