1 files changed, 53 insertions, 0 deletions
diff --git a/Source/FieldSolver/WarpXPushFieldsEM.cpp b/Source/FieldSolver/WarpXPushFieldsEM.cpp
index 8537e89ce..2487ffa0c 100644
--- a/Source/FieldSolver/WarpXPushFieldsEM.cpp
+++ b/Source/FieldSolver/WarpXPushFieldsEM.cpp
@@ -268,6 +268,38 @@ WarpX::PSATDForwardTransformJ ()
 #endif
 }
 
+void WarpX::PSATDBackwardTransformJ ()
+{
+    SpectralFieldIndex Idx;
+    int idx_jx, idx_jy, idx_jz;
+
+    for (int lev = 0; lev <= finest_level; ++lev)
+    {
+        Idx = spectral_solver_fp[lev]->m_spectral_index;
+
+        idx_jx = static_cast<int>(Idx.Jx);
+        idx_jy = static_cast<int>(Idx.Jy);
+        idx_jz = static_cast<int>(Idx.Jz);
+
+        auto& current = (WarpX::do_current_centering) ? current_fp_nodal : current_fp;
+
+        BackwardTransformVect(lev, *spectral_solver_fp[lev], current[lev], idx_jx, idx_jy, idx_jz);
+
+        if (spectral_solver_cp[lev])
+        {
+            Idx = spectral_solver_cp[lev]->m_spectral_index;
+
+            idx_jx = static_cast<int>(Idx.Jx);
+            idx_jy = static_cast<int>(Idx.Jy);
+            idx_jz = static_cast<int>(Idx.Jz);
+
+            // Current centering is not implemented with mesh refinement, so we do not yet need to
+            // distinguish between current_cp and current_cp_nodal, as for the fine patch currents
+            BackwardTransformVect(lev, *spectral_solver_cp[lev], current_cp[lev], idx_jx, idx_jy, idx_jz);
+        }
+    }
+}
+
 void
 WarpX::PSATDForwardTransformRho (const int icomp, const int dcomp)
 {
@@ -303,6 +335,19 @@ WarpX::PSATDForwardTransformRho (const int icomp, const int dcomp)
 #endif
 }
 
+void WarpX::PSATDCurrentCorrection ()
+{
+    for (int lev = 0; lev <= finest_level; ++lev)
+    {
+        spectral_solver_fp[lev]->CurrentCorrection();
+
+        if (spectral_solver_cp[lev])
+        {
+            spectral_solver_cp[lev]->CurrentCorrection();
+        }
+    }
+}
+
 void
 WarpX::PSATDPushSpectralFields ()
 {
@@ -423,6 +468,14 @@ WarpX::PushPSATD ()
         PSATDForwardTransformRho(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)
+    {
+        PSATDCurrentCorrection();
+        PSATDBackwardTransformJ();
+    }
+
 #ifdef WARPX_DIM_RZ
     if (pml_rz[0]) pml_rz[0]->PushPSATD(0);
 #endif