Implemented update without rho in RZ spectral solver (#1569)

* Implemented update without rho in RZ spectral solver

* Updated documentation for update_with_rho for RZ

* Tiny fix in GalileanPsatdAlgorithmRZ reordering declarations

* In Langmuir_multi_rz_psatd, set update_with_rho = 1

6 files changed, 66 insertions, 19 deletions

diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.H
index 3c8d3d60e..18a7f2dcd 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.H
@@ -21,7 +21,8 @@ class GalileanPsatdAlgorithmRZ : public SpectralBaseAlgorithmRZ
                                   int const n_rz_azimuthal_modes, int const norder_z,
                                   bool const nodal,
                                   const amrex::Array<amrex::Real,3>& v_galilean,
-                                  amrex::Real const dt_step);
+                                  amrex::Real const dt_step,
+                                  bool const update_with_rho);
         // Redefine functions from base class
         virtual void pushSpectralFields (SpectralFieldDataRZ & f) override final;
         virtual int getRequiredNumberOfFields () const override final {
@@ -63,9 +64,10 @@ class GalileanPsatdAlgorithmRZ : public SpectralBaseAlgorithmRZ
         // Note that dt and v_galilean are saved to use in InitializeSpectralCoefficients
         amrex::Real const m_dt;
         amrex::Array<amrex::Real,3> m_v_galilean;
+        bool m_update_with_rho;
 
         SpectralRealCoefficients C_coef, S_ck_coef;
-        SpectralComplexCoefficients Theta2_coef, X1_coef, X2_coef, X3_coef, X4_coef;
+        SpectralComplexCoefficients Theta2_coef, T_rho_coef, X1_coef, X2_coef, X3_coef, X4_coef;
 
 };
 
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.cpp
index dab4b7428..85de8ffc1 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.cpp
@@ -20,11 +20,13 @@ GalileanPsatdAlgorithmRZ::GalileanPsatdAlgorithmRZ (SpectralKSpaceRZ const & spe
                                                     int const n_rz_azimuthal_modes, int const norder_z,
                                                     bool const nodal,
                                                     const amrex::Array<amrex::Real,3>& v_galilean,
-                                                    amrex::Real const dt)
+                                                    amrex::Real const dt,
+                                                    bool const update_with_rho)
      // Initialize members of base class
      : SpectralBaseAlgorithmRZ(spectral_kspace, dm, norder_z, nodal),
        m_dt(dt),
-       m_v_galilean(v_galilean)
+       m_v_galilean(v_galilean),
+       m_update_with_rho(update_with_rho)
 {
 
     // Allocate the arrays of coefficients
@@ -36,16 +38,21 @@ GalileanPsatdAlgorithmRZ::GalileanPsatdAlgorithmRZ (SpectralKSpaceRZ const & spe
     X3_coef = SpectralComplexCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
     X4_coef = SpectralComplexCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
     Theta2_coef = SpectralComplexCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+    T_rho_coef = SpectralComplexCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
 
     coefficients_initialized = false;
 }
 
 /* Advance the E and B field in spectral space (stored in `f`)
- * over one time step */
+ * over one time step
+ * The algorithm is described in https://doi.org/10.1103/PhysRevE.94.053305
+ * */
 void
 GalileanPsatdAlgorithmRZ::pushSpectralFields (SpectralFieldDataRZ & f)
 {
 
+    bool const update_with_rho = m_update_with_rho;
+
     if (not coefficients_initialized) {
         // This is called from here since it needs the kr values
         // which can be obtained from the SpectralFieldDataRZ
@@ -68,6 +75,7 @@ GalileanPsatdAlgorithmRZ::pushSpectralFields (SpectralFieldDataRZ & f)
         amrex::Array4<const Complex> const& X3_arr = X3_coef[mfi].array();
         amrex::Array4<const Complex> const& X4_arr = X4_coef[mfi].array();
         amrex::Array4<const Complex> const& Theta2_arr = Theta2_coef[mfi].array();
+        amrex::Array4<const Complex> const& T_rho_arr = T_rho_coef[mfi].array();
 
         // Extract pointers for the k vectors
         auto const & kr_modes = f.getKrArray(mfi);
@@ -125,17 +133,28 @@ GalileanPsatdAlgorithmRZ::pushSpectralFields (SpectralFieldDataRZ & f)
             Complex const X3 = X3_arr(i,j,k,mode);
             Complex const X4 = X4_arr(i,j,k,mode);
             Complex const T2 = Theta2_arr(i,j,k,mode);
+            Complex const T_rho = T_rho_arr(i,j,k,mode);
+
+            Complex rho_diff;
+            if (update_with_rho) {
+                rho_diff = X2*rho_new - T2*X3*rho_old;
+            } else {
+                Complex const divE = kr*(Ep_old - Em_old) + I*kz*Ez_old;
+                Complex const divJ = kr*(Jp - Jm) + I*kz*Jz;
+
+                rho_diff = T2*(X2 - X3)*PhysConst::ep0*divE + T_rho*X2*divJ;
+            }
 
             // Update E (see WarpX online documentation: theory section)
             fields(i,j,k,Ep_m) = T2*C*Ep_old
                         + T2*S_ck*(-c2*I*kr/2._rt*Bz_old + c2*kz*Bp_old)
-                        + X4*Jp + 0.5_rt*kr*(X2*rho_new - T2*X3*rho_old);
+                        + X4*Jp + 0.5_rt*kr*rho_diff;
             fields(i,j,k,Em_m) = T2*C*Em_old
                         + T2*S_ck*(-c2*I*kr/2._rt*Bz_old - c2*kz*Bm_old)
-                        + X4*Jm - 0.5_rt*kr*(X2*rho_new - T2*X3*rho_old);
+                        + X4*Jm - 0.5_rt*kr*rho_diff;
             fields(i,j,k,Ez_m) = T2*C*Ez_old
                         + T2*S_ck*(c2*I*kr*Bp_old + c2*I*kr*Bm_old)
-                        + X4*Jz - I*kz*(X2*rho_new - T2*X3*rho_old);
+                        + X4*Jz - I*kz*rho_diff;
             // Update B (see WarpX online documentation: theory section)
             // Note: here X1 is T2*x1/(ep0*c*c*k_norm*k_norm), where
             // x1 has the same definition as in the original paper
@@ -173,6 +192,7 @@ void GalileanPsatdAlgorithmRZ::InitializeSpectralCoefficients (SpectralFieldData
         amrex::Array4<Complex> const& X3 = X3_coef[mfi].array();
         amrex::Array4<Complex> const& X4 = X4_coef[mfi].array();
         amrex::Array4<Complex> const& Theta2 = Theta2_coef[mfi].array();
+        amrex::Array4<Complex> const& T_rho = T_rho_coef[mfi].array();
 
         // Extract real (for portability on GPU)
         amrex::Real vz = m_v_galilean[2];
@@ -213,12 +233,18 @@ void GalileanPsatdAlgorithmRZ::InitializeSpectralCoefficients (SpectralFieldData
 
                 Theta2(i,j,k,mode) = theta*theta;
 
+                if (kz == 0._rt) {
+                    T_rho(i,j,k,mode) = -dt;
+                } else {
+                    T_rho(i,j,k,mode) = (1._rt - theta*theta)/(I*kz*vz);
+                }
+
                 if ( (nu != 1._rt) && (nu != 0._rt) ) {
 
                     // Note: the coefficients X1, X2, X do not correspond
                     // exactly to the original Galilean paper, but the
                     // update equation have been modified accordingly so that
-                    // the expressions/ below (with the update equations)
+                    // the expressions below (with the update equations)
                     // are mathematically equivalent to those of the paper.
                     Complex x1 = 1._rt/(1._rt-nu*nu) *
                         (theta_star - C(i,j,k,mode)*theta + I*kv*S_ck(i,j,k,mode)*theta);
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H
index f8df96e54..74bf71c29 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H
@@ -19,7 +19,8 @@ class PsatdAlgorithmRZ : public SpectralBaseAlgorithmRZ
         PsatdAlgorithmRZ(SpectralKSpaceRZ const & spectral_kspace,
                          amrex::DistributionMapping const & dm,
                          int const n_rz_azimuthal_modes, int const norder_z,
-                         bool const nodal, amrex::Real const dt_step);
+                         bool const nodal, amrex::Real const dt_step,
+                         bool const update_with_rho);
         // Redefine functions from base class
         virtual void pushSpectralFields(SpectralFieldDataRZ & f) override final;
         virtual int getRequiredNumberOfFields() const override final {
@@ -62,6 +63,7 @@ class PsatdAlgorithmRZ : public SpectralBaseAlgorithmRZ
         bool coefficients_initialized;
         // Note that dt is saved to use in InitializeSpectralCoefficients
         amrex::Real m_dt;
+        bool m_update_with_rho;
         SpectralRealCoefficients C_coef, S_ck_coef, X1_coef, X2_coef, X3_coef;
 };
 
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp
index 7cd230581..002f9e55f 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp
@@ -18,11 +18,13 @@ using amrex::operator""_rt;
 PsatdAlgorithmRZ::PsatdAlgorithmRZ (SpectralKSpaceRZ const & spectral_kspace,
                                     amrex::DistributionMapping const & dm,
                                     int const n_rz_azimuthal_modes, int const norder_z,
-                                    bool const nodal, amrex::Real const dt)
+                                    bool const nodal, amrex::Real const dt,
+                                    bool const update_with_rho)
      // Initialize members of base class
      : SpectralBaseAlgorithmRZ(spectral_kspace, dm,
                                norder_z, nodal),
-       m_dt(dt)
+       m_dt(dt),
+       m_update_with_rho(update_with_rho)
 {
 
     // Allocate the arrays of coefficients
@@ -42,6 +44,8 @@ void
 PsatdAlgorithmRZ::pushSpectralFields(SpectralFieldDataRZ & f)
 {
 
+    bool const update_with_rho = m_update_with_rho;
+
     if (not coefficients_initialized) {
         // This is called from here since it needs the kr values
         // which can be obtained from the SpectralFieldDataRZ
@@ -68,6 +72,7 @@ PsatdAlgorithmRZ::pushSpectralFields(SpectralFieldDataRZ & f)
         amrex::Real const* kr_arr = kr_modes.dataPtr();
         amrex::Real const* modified_kz_arr = modified_kz_vec[mfi].dataPtr();
         int const nr = bx.length(0);
+        amrex::Real const dt = m_dt;
 
         // Loop over indices within one box
         // Note that k = 0
@@ -119,16 +124,26 @@ PsatdAlgorithmRZ::pushSpectralFields(SpectralFieldDataRZ & f)
             amrex::Real const X2 = X2_arr(i,j,k,mode);
             amrex::Real const X3 = X3_arr(i,j,k,mode);
 
+            Complex rho_diff;
+            if (update_with_rho) {
+                rho_diff = X2*rho_new - X3*rho_old;
+            } else {
+                Complex const divE = kr*(Ep_old - Em_old) + I*kz*Ez_old;
+                Complex const divJ = kr*(Jp - Jm) + I*kz*Jz;
+
+                rho_diff = (X2 - X3)*PhysConst::ep0*divE - X2*dt*divJ;
+            }
+
             // Update E (see WarpX online documentation: theory section)
             fields(i,j,k,Ep_m) = C*Ep_old
                         + S_ck*(-c2*I*kr/2._rt*Bz_old + c2*kz*Bp_old - inv_ep0*Jp)
-                        + 0.5_rt*kr*(X2*rho_new - X3*rho_old);
+                        + 0.5_rt*kr*rho_diff;
             fields(i,j,k,Em_m) = C*Em_old
                         + S_ck*(-c2*I*kr/2._rt*Bz_old - c2*kz*Bm_old - inv_ep0*Jm)
-                        - 0.5_rt*kr*(X2*rho_new - X3*rho_old);
+                        - 0.5_rt*kr*rho_diff;
             fields(i,j,k,Ez_m) = C*Ez_old
                         + S_ck*(c2*I*kr*Bp_old + c2*I*kr*Bm_old - inv_ep0*Jz)
-                        - I*kz*(X2*rho_new - X3*rho_old);
+                        - I*kz*rho_diff;
             // Update B (see WarpX online documentation: theory section)
             fields(i,j,k,Bp_m) = C*Bp_old
                         - S_ck*(-I*kr/2._rt*Ez_old + kz*Ep_old)
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H
index e79d868f6..8b6ee51d9 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H
@@ -31,7 +31,8 @@ class SpectralSolverRZ
                           int const norder_z, bool const nodal,
                           const amrex::Array<amrex::Real,3>& v_galilean,
                           amrex::RealVect const dx, amrex::Real const dt,
-                          int const lev);
+                          int const lev,
+                          bool const update_with_rho);
 
         /* \brief Transform the component `i_comp` of MultiFab `field_mf`
          *  to spectral space, and store the corresponding result internally
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp
index 581538d6a..820db5a12 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp
@@ -31,7 +31,8 @@ SpectralSolverRZ::SpectralSolverRZ (amrex::BoxArray const & realspace_ba,
                                     int const norder_z, bool const nodal,
                                     const amrex::Array<amrex::Real,3>& v_galilean,
                                     amrex::RealVect const dx, amrex::Real const dt,
-                                    int const lev)
+                                    int const lev,
+                                    bool const update_with_rho)
     : k_space(realspace_ba, dm, dx)
 {
     // Initialize all structures using the same distribution mapping dm
@@ -46,11 +47,11 @@ SpectralSolverRZ::SpectralSolverRZ (amrex::BoxArray const & realspace_ba,
     if (v_galilean[2] == 0) {
          // v_galilean is 0: use standard PSATD algorithm
         algorithm = std::make_unique<PsatdAlgorithmRZ>(
-            k_space, dm, n_rz_azimuthal_modes, norder_z, nodal, dt);
+            k_space, dm, n_rz_azimuthal_modes, norder_z, nodal, dt, update_with_rho);
     } else {
         // Otherwise: use the Galilean algorithm
         algorithm = std::make_unique<GalileanPsatdAlgorithmRZ>(
-            k_space, dm, n_rz_azimuthal_modes, norder_z, nodal, v_galilean, dt);
+            k_space, dm, n_rz_azimuthal_modes, norder_z, nodal, v_galilean, dt, update_with_rho);
     }
 
     // - Initialize arrays for fields in spectral space + FFT plans