RZ PSATD: Multi-J Algorithm (#2111)

* RZ PSATD: Implement Multi-J Algorithm

* Implement J_linear_in_time Option

* Reduce Style Changes

* Move Copy/Zero/Scale Functions to SpectralFieldDataRZ

* Remove Unused Member m_n_rz_azimuthal_modes from SpectralSolverRZ

* Fix CI -Werror Warnings

* Implement Same Changes of #2116, Cleaning

* Fix Bug: Pass Correct dt to SpectralSolverRZ

* Add CI Test

* CI Test: Set random_theta = 0, Update Benchmark

* Remove random_theta from Inputs

* Update Benchmark of multi_J_rz_psatd

5 files changed, 182 insertions, 23 deletions

diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H
index 4987e88d7..07c4d8142 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H
@@ -21,7 +21,11 @@ class PsatdAlgorithmRZ : public SpectralBaseAlgorithmRZ
                          const SpectralFieldIndex& spectral_index,
                          int const n_rz_azimuthal_modes, int const norder_z,
                          bool const nodal, amrex::Real const dt_step,
-                         bool const update_with_rho);
+                         bool const update_with_rho,
+                         const bool time_averaging,
+                         const bool J_linear_in_time,
+                         const bool dive_cleaning,
+                         const bool divb_cleaning);
         // Redefine functions from base class
         virtual void pushSpectralFields(SpectralFieldDataRZ & f) override final;
 
@@ -67,6 +71,10 @@ class PsatdAlgorithmRZ : public SpectralBaseAlgorithmRZ
         // Note that dt is saved to use in InitializeSpectralCoefficients
         amrex::Real m_dt;
         bool m_update_with_rho;
+        bool m_time_averaging;
+        bool m_J_linear_in_time;
+        bool m_dive_cleaning;
+        bool m_divb_cleaning;
         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 6a3c86bf6..ac8ae23e0 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp
@@ -20,12 +20,20 @@ PsatdAlgorithmRZ::PsatdAlgorithmRZ (SpectralKSpaceRZ const & spectral_kspace,
                                     const SpectralFieldIndex& spectral_index,
                                     int const n_rz_azimuthal_modes, int const norder_z,
                                     bool const nodal, amrex::Real const dt,
-                                    bool const update_with_rho)
+                                    bool const update_with_rho,
+                                    const bool time_averaging,
+                                    const bool J_linear_in_time,
+                                    const bool dive_cleaning,
+                                    const bool divb_cleaning)
      // Initialize members of base class
      : SpectralBaseAlgorithmRZ(spectral_kspace, dm, spectral_index, norder_z, nodal),
        m_spectral_index(spectral_index),
        m_dt(dt),
-       m_update_with_rho(update_with_rho)
+       m_update_with_rho(update_with_rho),
+       m_time_averaging(time_averaging),
+       m_J_linear_in_time(J_linear_in_time),
+       m_dive_cleaning(dive_cleaning),
+       m_divb_cleaning(divb_cleaning)
 {
 
     // Allocate the arrays of coefficients
@@ -37,6 +45,9 @@ PsatdAlgorithmRZ::PsatdAlgorithmRZ (SpectralKSpaceRZ const & spectral_kspace,
     X3_coef = SpectralRealCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
 
     coefficients_initialized = false;
+
+    // TODO Implement time averaging and remove this
+    amrex::ignore_unused(m_time_averaging);
 }
 
 /* Advance the E and B field in spectral space (stored in `f`)
@@ -46,6 +57,9 @@ PsatdAlgorithmRZ::pushSpectralFields(SpectralFieldDataRZ & f)
 {
 
     bool const update_with_rho = m_update_with_rho;
+    const bool J_linear_in_time = m_J_linear_in_time;
+    const bool dive_cleaning = m_dive_cleaning;
+    const bool divb_cleaning = m_divb_cleaning;
 
     if (not coefficients_initialized) {
         // This is called from here since it needs the kr values
@@ -85,17 +99,17 @@ PsatdAlgorithmRZ::pushSpectralFields(SpectralFieldDataRZ & f)
         {
 
             // All of the fields of each mode are grouped together
-            auto const Ep_m = Idx.Ex + Idx.n_fields*mode;
-            auto const Em_m = Idx.Ey + Idx.n_fields*mode;
-            auto const Ez_m = Idx.Ez + Idx.n_fields*mode;
-            auto const Bp_m = Idx.Bx + Idx.n_fields*mode;
-            auto const Bm_m = Idx.By + Idx.n_fields*mode;
-            auto const Bz_m = Idx.Bz + Idx.n_fields*mode;
-            auto const Jp_m = Idx.Jx + Idx.n_fields*mode;
-            auto const Jm_m = Idx.Jy + Idx.n_fields*mode;
-            auto const Jz_m = Idx.Jz + Idx.n_fields*mode;
-            auto const rho_old_m = Idx.rho_old + Idx.n_fields*mode;
-            auto const rho_new_m = Idx.rho_new + Idx.n_fields*mode;
+            int const Ep_m = Idx.Ex + Idx.n_fields*mode;
+            int const Em_m = Idx.Ey + Idx.n_fields*mode;
+            int const Ez_m = Idx.Ez + Idx.n_fields*mode;
+            int const Bp_m = Idx.Bx + Idx.n_fields*mode;
+            int const Bm_m = Idx.By + Idx.n_fields*mode;
+            int const Bz_m = Idx.Bz + Idx.n_fields*mode;
+            int const Jp_m = Idx.Jx + Idx.n_fields*mode;
+            int const Jm_m = Idx.Jy + Idx.n_fields*mode;
+            int const Jz_m = Idx.Jz + Idx.n_fields*mode;
+            int const rho_old_m = Idx.rho_old + Idx.n_fields*mode;
+            int const rho_new_m = Idx.rho_new + Idx.n_fields*mode;
 
             // Record old values of the fields to be updated
             Complex const Ep_old = fields(i,j,k,Ep_m);
@@ -156,6 +170,67 @@ PsatdAlgorithmRZ::pushSpectralFields(SpectralFieldDataRZ & f)
             fields(i,j,k,Bz_m) = C*Bz_old
                         - S_ck*I*(kr*Ep_old + kr*Em_old)
                         + X1*I*(kr*Jp + kr*Jm);
+
+            int F_m;
+            Complex F_old;
+            if (dive_cleaning)
+            {
+                F_m = Idx.F + Idx.n_fields*mode;
+                F_old = fields(i,j,k,F_m);
+            }
+
+            int G_m;
+            Complex G_old;
+            if (divb_cleaning)
+            {
+                G_m = Idx.G + Idx.n_fields*mode;
+                G_old = fields(i,j,k,G_m);
+            }
+
+            if (J_linear_in_time)
+            {
+                const int Jp_m_new = Idx.Jx_new + Idx.n_fields*mode;
+                const int Jm_m_new = Idx.Jy_new + Idx.n_fields*mode;
+                const int Jz_m_new = Idx.Jz_new + Idx.n_fields*mode;
+
+                const Complex Jp_new = fields(i,j,k,Jp_m_new);
+                const Complex Jm_new = fields(i,j,k,Jm_m_new);
+                const Complex Jz_new = fields(i,j,k,Jz_m_new);
+
+                fields(i,j,k,Ep_m) += -X1 * (Jp_new - Jp) / dt;
+                fields(i,j,k,Em_m) += -X1 * (Jm_new - Jm) / dt;
+                fields(i,j,k,Ez_m) += -X1 * (Jz_new - Jz) / dt;
+
+                fields(i,j,k,Bp_m) +=  X2/c2 * (kz * (Jp_new - Jp) - I * kr * 0.5_rt * (Jz_new - Jz));
+                fields(i,j,k,Bm_m) += -X2/c2 * (kz * (Jm_new - Jm) + I * kr * 0.5_rt * (Jz_new - Jz));
+                fields(i,j,k,Bz_m) += I * X2/c2 * (kr * (Jp_new - Jp) + kr * (Jm_new - Jm));
+
+                if (dive_cleaning)
+                {
+                    const Complex k_dot_J  = -I * (kr * (Jp - Jm) + I * kz * Jz);
+                    const Complex k_dot_dJ = -I * (kr * ((Jp_new - Jp) - (Jm_new - Jm))
+                        + I * kz * (Jz_new - Jz));
+                    const Complex k_dot_E = -I * (kr * (Ep_old - Em_old) + I * kz * Ez_old);
+
+                    fields(i,j,k,Ep_m) += -c2 * kr * 0.5_rt * S_ck * F_old;
+                    fields(i,j,k,Em_m) +=  c2 * kr * 0.5_rt * S_ck * F_old;
+                    fields(i,j,k,Ez_m) += I * c2 * kz * S_ck * F_old;
+
+                    fields(i,j,k,F_m) = C * F_old + S_ck * (I * k_dot_E - inv_ep0 * rho_old)
+                        - X1 * ((rho_new - rho_old) / dt + I * k_dot_J) - I * X2/c2 * k_dot_dJ;
+                }
+
+                if (divb_cleaning)
+                {
+                    const Complex k_dot_B = -I * (kr * (Bp_old - Bm_old) + I * kz * Bz_old);
+
+                    fields(i,j,k,Bp_m) += -c2 * kr * 0.5_rt * S_ck * G_old;
+                    fields(i,j,k,Bm_m) +=  c2 * kr * 0.5_rt * S_ck * G_old;
+                    fields(i,j,k,Bz_m) += I * c2 * kz * S_ck * G_old;
+
+                    fields(i,j,k,G_m) = C * G_old + I * S_ck * k_dot_B;
+                }
+            }
         });
     }
 }
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H
index bf0d6b780..ec5fa4a83 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H
@@ -64,6 +64,43 @@ class SpectralFieldDataRZ
                                     amrex::MultiFab & tempHTransformedSplit,
                                     const bool is_nodal_z);
 
+        /**
+         * \brief Copy spectral data from component \c src_comp to component \c dest_comp
+         *        of \c fields
+         *
+         * \param[in] src_comp  component of the source FabArray from which the data are copied
+         * \param[in] dest_comp component of the destination FabArray where the data are copied
+         */
+        void CopySpectralDataComp (const int src_comp, const int dest_comp)
+        {
+            // The last two arguments represent the number of components and
+            // the number of ghost cells to perform this operation
+            Copy(fields, fields, src_comp, dest_comp, n_rz_azimuthal_modes, 0);
+        }
+
+        /**
+         * \brief Set to zero the data on component \c icomp of \c fields
+         *
+         * \param[in] icomp component of the FabArray where the data are set to zero
+         */
+        void ZeroOutDataComp(const int icomp)
+        {
+            // The last argument represents the number of components to perform this operation
+            fields.setVal(0., icomp, n_rz_azimuthal_modes);
+        }
+
+        /**
+         * \brief Scale the data on component \c icomp of \c fields by a given scale factor
+         *
+         * \param[in] icomp component of the FabArray where the data are scaled
+         * \param[in] scale_factor scale factor to use for scaling
+         */
+        void ScaleDataComp(const int icomp, const amrex::Real scale_factor)
+        {
+            // The last argument represents the number of components to perform this operation
+            fields.mult(scale_factor, icomp, n_rz_azimuthal_modes);
+        }
+
         void InitFilter (amrex::IntVect const & filter_npass_each_dir, bool const compensation,
                          SpectralKSpaceRZ const & k_space);
 
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H
index 8e84f7c00..e08d5741c 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H
@@ -34,7 +34,11 @@ 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,
-                          bool const update_with_rho);
+                          bool const update_with_rho,
+                          const bool fft_do_time_averaging,
+                          const bool J_linear_in_time,
+                          const bool dive_cleaning,
+                          const bool divb_cleaning);
 
         /* \brief Transform the component `i_comp` of MultiFab `field_mf`
          *  to spectral space, and store the corresponding result internally
@@ -112,6 +116,40 @@ class SpectralSolverRZ
          */
         void VayDeposition (const int lev, std::array<std::unique_ptr<amrex::MultiFab>,3>& current);
 
+        /**
+         * \brief Copy spectral data from component \c src_comp to component \c dest_comp
+         *        of \c field_data.fields
+         *
+         * \param[in] src_comp  component of the source FabArray from which the data are copied
+         * \param[in] dest_comp component of the destination FabArray where the data are copied
+         */
+        void CopySpectralDataComp (const int src_comp, const int dest_comp)
+        {
+            field_data.CopySpectralDataComp(src_comp, dest_comp);
+        }
+
+        /**
+         * \brief Set to zero the data on component \c icomp of \c field_data.fields
+         *
+         * \param[in] icomp component of the FabArray where the data are set to zero
+         */
+        void ZeroOutDataComp (const int icomp)
+        {
+            field_data.ZeroOutDataComp(icomp);
+        }
+
+        /**
+         * \brief Scale the data on component \c icomp of \c field_data.fields
+         *        by a given scale factor
+         *
+         * \param[in] icomp component of the FabArray where the data are scaled
+         * \param[in] scale_factor scale factor to use for scaling
+         */
+        void ScaleDataComp (const int icomp, const amrex::Real scale_factor)
+        {
+            field_data.ScaleDataComp(icomp, scale_factor);
+        }
+
         SpectralFieldIndex m_spectral_index;
 
     private:
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp
index 8f46d21cd..12c2dabf2 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp
@@ -32,7 +32,11 @@ SpectralSolverRZ::SpectralSolverRZ (const int lev,
                                     int const norder_z, bool const nodal,
                                     const amrex::Array<amrex::Real,3>& v_galilean,
                                     amrex::RealVect const dx, amrex::Real const dt,
-                                    bool const update_with_rho)
+                                    bool const update_with_rho,
+                                    const bool fft_do_time_averaging,
+                                    const bool J_linear_in_time,
+                                    const bool dive_cleaning,
+                                    const bool divb_cleaning)
     : k_space(realspace_ba, dm, dx)
 {
     // Initialize all structures using the same distribution mapping dm
@@ -41,13 +45,9 @@ SpectralSolverRZ::SpectralSolverRZ (const int lev,
     //   the spectral space corresponding to each box in `realspace_ba`,
     //   as well as the value of the corresponding k coordinates.
 
-    const bool time_averaging = false;
-    const bool J_linear_in_time = false;
-    const bool dive_cleaning = false;
-    const bool divb_cleaning = false;
     const bool pml = false;
-    m_spectral_index = SpectralFieldIndex(update_with_rho, time_averaging, J_linear_in_time,
-                                          dive_cleaning, divb_cleaning, pml);
+    m_spectral_index = SpectralFieldIndex(update_with_rho, fft_do_time_averaging,
+                                          J_linear_in_time, dive_cleaning, divb_cleaning, pml);
 
     // - Select the algorithm depending on the input parameters
     //   Initialize the corresponding coefficients over k space
@@ -55,7 +55,8 @@ SpectralSolverRZ::SpectralSolverRZ (const int lev,
     if (v_galilean[2] == 0) {
          // v_galilean is 0: use standard PSATD algorithm
         algorithm = std::make_unique<PsatdAlgorithmRZ>(
-            k_space, dm, m_spectral_index, n_rz_azimuthal_modes, norder_z, nodal, dt, update_with_rho);
+            k_space, dm, m_spectral_index, n_rz_azimuthal_modes, norder_z, nodal, dt,
+            update_with_rho, fft_do_time_averaging, J_linear_in_time, dive_cleaning, divb_cleaning);
     } else {
         // Otherwise: use the Galilean algorithm
         algorithm = std::make_unique<GalileanPsatdAlgorithmRZ>(