Fix Bugs w/ Current Correction and Vay Deposition (#3290)

* Fix Bugs w/ Current Correction and Vay Deposition

* Vay Deposition and Current Correction Cannot be Combined Together

* Add Comment for Future Implementation of Vay Deposition w/ MR

* Add Comment for Future Implementation of Vay Deposition w/ MR

* Define SyncCurrentAndRho, Clean Up

* Vay Deposition: Remove Extra FFT of Rho

* Fix Bug in RZ Geometry (Double Filtering)

* Add 2D Galilean Test w/o Periodic Single Box

* Add RZ Galilean Test w/o Periodic Single Box

* Add 3D Galilean Test w/o Periodic Single Box

1 files changed, 104 insertions, 29 deletions

diff --git a/Source/FieldSolver/WarpXPushFieldsEM.cpp b/Source/FieldSolver/WarpXPushFieldsEM.cpp
index 6d604bc15..7499ee140 100644
--- a/Source/FieldSolver/WarpXPushFieldsEM.cpp
+++ b/Source/FieldSolver/WarpXPushFieldsEM.cpp
@@ -270,7 +270,8 @@ WarpX::PSATDBackwardTransformG ()
 
 void WarpX::PSATDForwardTransformJ (
     const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_fp,
-    const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp)
+    const amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp,
+    const bool apply_kspace_filter)
 {
     SpectralFieldIndex Idx;
     int idx_jx, idx_jy, idx_jz;
@@ -299,7 +300,7 @@ void WarpX::PSATDForwardTransformJ (
 
 #ifdef WARPX_DIM_RZ
     // Apply filter in k space if needed
-    if (WarpX::use_kspace_filter)
+    if (use_kspace_filter && apply_kspace_filter)
     {
         for (int lev = 0; lev <= finest_level; ++lev)
         {
@@ -311,6 +312,8 @@ void WarpX::PSATDForwardTransformJ (
             }
         }
     }
+#else
+    amrex::ignore_unused(apply_kspace_filter);
 #endif
 }
 
@@ -349,8 +352,10 @@ void WarpX::PSATDBackwardTransformJ (
 void WarpX::PSATDForwardTransformRho (
     const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_fp,
     const amrex::Vector<std::unique_ptr<amrex::MultiFab>>& charge_cp,
-    const int icomp, const int dcomp)
+    const int icomp, const int dcomp, const bool apply_kspace_filter)
 {
+    if (charge_fp[0] == nullptr) return;
+
     const SpectralFieldIndex& Idx = spectral_solver_fp[0]->m_spectral_index;
 
     // Select index in k space
@@ -368,7 +373,7 @@ void WarpX::PSATDForwardTransformRho (
 
 #ifdef WARPX_DIM_RZ
     // Apply filter in k space if needed
-    if (WarpX::use_kspace_filter)
+    if (use_kspace_filter && apply_kspace_filter)
     {
         for (int lev = 0; lev <= finest_level; ++lev)
         {
@@ -380,6 +385,8 @@ void WarpX::PSATDForwardTransformRho (
             }
         }
     }
+#else
+    amrex::ignore_unused(apply_kspace_filter);
 #endif
 }
 
@@ -635,46 +642,114 @@ WarpX::PushPSATD ()
         "PushFieldsEM: PSATD solver selected but not built"));
 #else
 
-    PSATDForwardTransformEB(Efield_fp, Bfield_fp, Efield_cp, Bfield_cp);
-
-    amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_fp =
-        (WarpX::current_deposition_algo == CurrentDepositionAlgo::Vay) ? current_fp_vay : current_fp;
+    if (fft_periodic_single_box)
+    {
+        if (current_correction)
+        {
+            // FFT of J and rho
+            PSATDForwardTransformJ(current_fp, current_cp);
+            PSATDForwardTransformRho(rho_fp, rho_cp, 0, 0); // rho old
+            PSATDForwardTransformRho(rho_fp, rho_cp, 1, 1); // rho new
 
-    PSATDForwardTransformJ(J_fp, current_cp);
+            // Correct J in k-space
+            PSATDCurrentCorrection();
 
-    // Do rho FFTs only if needed
-    if (WarpX::update_with_rho || WarpX::current_correction || WarpX::do_dive_cleaning)
-    {
-        PSATDForwardTransformRho(rho_fp, rho_cp, 0,0); // rho old
-        PSATDForwardTransformRho(rho_fp, rho_cp, 1,1); // rho new
+            // Inverse FFT of J
+            PSATDBackwardTransformJ(current_fp, current_cp);
+        }
+        else if (current_deposition_algo == CurrentDepositionAlgo::Vay)
+        {
+            // FFT of D and rho (if used)
+            // TODO Replace current_cp with current_cp_vay once Vay deposition is implemented with MR
+            PSATDForwardTransformJ(current_fp_vay, current_cp);
+            PSATDForwardTransformRho(rho_fp, rho_cp, 0, 0); // rho old
+            PSATDForwardTransformRho(rho_fp, rho_cp, 1, 1); // rho new
+
+            // Compute J from D in k-space
+            PSATDVayDeposition();
+
+            // Inverse FFT of J, subtract cumulative sums of D
+            PSATDBackwardTransformJ(current_fp, current_cp);
+            // TODO Cumulative sums need to be fixed with periodic single box
+            PSATDSubtractCurrentPartialSumsAvg();
+
+            // FFT of J after subtraction of cumulative sums
+            PSATDForwardTransformJ(current_fp, current_cp);
+        }
+        else // no current correction, no Vay deposition
+        {
+            // FFT of J and rho (if used)
+            PSATDForwardTransformJ(current_fp, current_cp);
+            PSATDForwardTransformRho(rho_fp, rho_cp, 0, 0); // rho old
+            PSATDForwardTransformRho(rho_fp, rho_cp, 1, 1); // rho new
+        }
     }
-
-    // Correct the current in Fourier space so that the continuity equation is satisfied, and
-    // transform back to real space so that the current correction is reflected in the diagnostics
-    if (WarpX::current_correction)
+    else // no periodic single box
     {
-        PSATDCurrentCorrection();
-        PSATDBackwardTransformJ(current_fp, current_cp);
-    }
+        if (current_correction)
+        {
+            // FFT of J and rho
+#ifdef WARPX_DIM_RZ
+            // In RZ geometry, do not apply filtering here, since it is
+            // applied in the subsequent calls to these functions (below)
+            const bool apply_kspace_filter = false;
+            PSATDForwardTransformJ(current_fp, current_cp, apply_kspace_filter);
+            PSATDForwardTransformRho(rho_fp, rho_cp, 0, 0, apply_kspace_filter); // rho old
+            PSATDForwardTransformRho(rho_fp, rho_cp, 1, 1, apply_kspace_filter); // rho new
+#else
+            PSATDForwardTransformJ(current_fp, current_cp);
+            PSATDForwardTransformRho(rho_fp, rho_cp, 0, 0); // rho old
+            PSATDForwardTransformRho(rho_fp, rho_cp, 1, 1); // rho new
+#endif
 
-    // Compute the current in Fourier space according to the Vay deposition scheme, and
-    // transform back to real space so that the Vay deposition is reflected in the diagnostics
-    if (WarpX::current_deposition_algo == CurrentDepositionAlgo::Vay)
-    {
-        PSATDVayDeposition();
-        PSATDBackwardTransformJ(current_fp, current_cp);
-        PSATDSubtractCurrentPartialSumsAvg();
-        SyncCurrent(current_fp, current_cp);
+            // Correct J in k-space
+            PSATDCurrentCorrection();
+
+            // Inverse FFT of J
+            PSATDBackwardTransformJ(current_fp, current_cp);
+
+            // Synchronize J and rho
+            SyncCurrent(current_fp, current_cp);
+            SyncRho();
+        }
+        else if (current_deposition_algo == CurrentDepositionAlgo::Vay)
+        {
+            // FFT of D
+            PSATDForwardTransformJ(current_fp_vay, current_cp);
+
+            // Compute J from D in k-space
+            PSATDVayDeposition();
+
+            // Inverse FFT of J, subtract cumulative sums of D
+            PSATDBackwardTransformJ(current_fp, current_cp);
+            PSATDSubtractCurrentPartialSumsAvg();
+
+            // Synchronize J and rho (if used)
+            SyncCurrent(current_fp, current_cp);
+            SyncRho();
+        }
+
+        // FFT of J and rho (if used)
         PSATDForwardTransformJ(current_fp, current_cp);
+        PSATDForwardTransformRho(rho_fp, rho_cp, 0, 0); // rho old
+        PSATDForwardTransformRho(rho_fp, rho_cp, 1, 1); // rho new
     }
 
+    // FFT of E and B
+    PSATDForwardTransformEB(Efield_fp, Bfield_fp, Efield_cp, Bfield_cp);
+
 #ifdef WARPX_DIM_RZ
     if (pml_rz[0]) pml_rz[0]->PushPSATD(0);
 #endif
 
+    // FFT of F and G
     if (WarpX::do_dive_cleaning) PSATDForwardTransformF();
     if (WarpX::do_divb_cleaning) PSATDForwardTransformG();
+
+    // Update E, B, F, and G in k-space
     PSATDPushSpectralFields();
+
+    // Inverse FFT of E, B, F, and G
     PSATDBackwardTransformEB(Efield_fp, Bfield_fp, Efield_cp, Bfield_cp);
     if (WarpX::fft_do_time_averaging)
         PSATDBackwardTransformEBavg(Efield_avg_fp, Bfield_avg_fp, Efield_avg_cp, Bfield_avg_cp);