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

1 files changed, 88 insertions, 13 deletions

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;
+                }
+            }
         });
     }
 }