Vay Deposition: Separate Arrays, Correct Index Types w/ FFTs (#2965)

* Refactoring

* Separate Arrays (Fine Patch)

* Add Aborts w/ Current Centering, MR

* Cleaning

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: Remi Lehe <remi.lehe@normalesup.org>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

5 files changed, 25 insertions, 113 deletions

diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
index 4f0ad1c20..fe856c236 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
@@ -124,13 +124,7 @@ class SpectralFieldData
 
         void ForwardTransform (const int lev,
                                const amrex::MultiFab& mf, const int field_index,
-                               const int i_comp, const amrex::IntVect& stag);
-        AMREX_FORCE_INLINE
-        void ForwardTransform (const int lev,
-                               const amrex::MultiFab& mf, const int field_index, const int i_comp)
-        {
-            ForwardTransform(lev, mf, field_index, i_comp, mf.ixType().toIntVect());
-        }
+                               const int i_comp);
 
         void BackwardTransform (const int lev, amrex::MultiFab& mf, const int field_index,
                                 const int i_comp, const amrex::IntVect& fill_guards);
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
index 3926df286..460e2cc3b 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
@@ -204,22 +204,22 @@ SpectralFieldData::~SpectralFieldData()
 void
 SpectralFieldData::ForwardTransform (const int lev,
                                      const MultiFab& mf, const int field_index,
-                                     const int i_comp, const IntVect& stag)
+                                     const int i_comp)
 {
     amrex::LayoutData<amrex::Real>* cost = WarpX::getCosts(lev);
     bool do_costs = WarpXUtilLoadBalance::doCosts(cost, mf.boxArray(), mf.DistributionMap());
 
     // Check field index type, in order to apply proper shift in spectral space
 #if (AMREX_SPACEDIM >= 2)
-    const bool is_nodal_x = (stag[0] == amrex::IndexType::NODE) ? true : false;
+    const bool is_nodal_x = mf.is_nodal(0);
 #endif
 #if defined(WARPX_DIM_3D)
-    const bool is_nodal_y = (stag[1] == amrex::IndexType::NODE) ? true : false;
-    const bool is_nodal_z = (stag[2] == amrex::IndexType::NODE) ? true : false;
+    const bool is_nodal_y = mf.is_nodal(1);
+    const bool is_nodal_z = mf.is_nodal(2);
 #elif defined(WARPX_DIM_XZ) || defined(WARPX_DIM_RZ)
-    const bool is_nodal_z = (stag[1] == amrex::IndexType::NODE) ? true : false;
+    const bool is_nodal_z = mf.is_nodal(1);
 #elif defined(WARPX_DIM_1D_Z)
-    const bool is_nodal_z = (stag[0] == amrex::IndexType::NODE) ? true : false;
+    const bool is_nodal_z = mf.is_nodal(0);
 #endif
 
     // Loop over boxes
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolver.H b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
index 47ad47150..df9e9f0ff 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolver.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
@@ -93,48 +93,11 @@ class SpectralSolver
          * \param[in] mf MultiFab that is transformed to Fourier space (component i_comp)
          * \param[in] field_index index of the spectral field that stores the FFT result
          * \param[in] i_comp component of the MultiFab mf that is transformed to Fourier space
-         * \param[in] stag index type that needs to be used to perform the FFT operation
          */
         void ForwardTransform (const int lev,
                                const amrex::MultiFab& mf,
                                const int field_index,
-                               const int i_comp,
-                               const amrex::IntVect& stag);
-
-        /**
-         * \brief Overload of ForwardTransform used if stag is not passed
-         */
-        AMREX_FORCE_INLINE
-        void ForwardTransform (const int lev,
-                               const amrex::MultiFab& mf,
-                               const int field_index,
-                               const int i_comp)
-        {
-            ForwardTransform(lev, mf, field_index, i_comp, mf.ixType().toIntVect());
-        }
-
-        /**
-         * \brief Overload of ForwardTransform used if i_comp is not passed
-         */
-        AMREX_FORCE_INLINE
-        void ForwardTransform (const int lev,
-                               const amrex::MultiFab& mf,
-                               const int field_index,
-                               const amrex::IntVect& stag)
-        {
-            ForwardTransform(lev, mf, field_index, 0, stag);
-        }
-
-        /**
-         * Overload of ForwardTransform used if both i_comp and stag are not passed
-         */
-        AMREX_FORCE_INLINE
-        void ForwardTransform (const int lev,
-                               const amrex::MultiFab& mf,
-                               const int field_index)
-        {
-            ForwardTransform(lev, mf, field_index, 0, mf.ixType().toIntVect());
-        }
+                               const int i_comp = 0);
 
         /**
          * \brief Transform spectral field specified by `field_index` back to
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolver.cpp b/Source/FieldSolver/SpectralSolver/SpectralSolver.cpp
index 0f7804002..81cba183c 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolver.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolver.cpp
@@ -89,14 +89,13 @@ SpectralSolver::SpectralSolver(
 }
 
 void
-SpectralSolver::ForwardTransform( const int lev,
+SpectralSolver::ForwardTransform (const int lev,
                                   const amrex::MultiFab& mf,
                                   const int field_index,
-                                  const int i_comp,
-                                  const amrex::IntVect& stag )
+                                  const int i_comp)
 {
     WARPX_PROFILE("SpectralSolver::ForwardTransform");
-    field_data.ForwardTransform( lev, mf, field_index, i_comp, stag );
+    field_data.ForwardTransform(lev, mf, field_index, i_comp);
 }
 
 void
diff --git a/Source/FieldSolver/WarpXPushFieldsEM.cpp b/Source/FieldSolver/WarpXPushFieldsEM.cpp
index a911bd877..d02263325 100644
--- a/Source/FieldSolver/WarpXPushFieldsEM.cpp
+++ b/Source/FieldSolver/WarpXPushFieldsEM.cpp
@@ -64,35 +64,16 @@ namespace {
         SpectralSolver& solver,
 #endif
         std::array<std::unique_ptr<amrex::MultiFab>,3>& vector_field,
-        const int compx, const int compy, const int compz,
-        const amrex::IntVect& stag_x, const amrex::IntVect& stag_y, const amrex::IntVect& stag_z)
+        const int compx, const int compy, const int compz)
     {
 #ifdef WARPX_DIM_RZ
-        amrex::ignore_unused(stag_x, stag_y, stag_z);
         solver.ForwardTransform(lev, *vector_field[0], compx, *vector_field[1], compy);
         solver.ForwardTransform(lev, *vector_field[2], compz);
 #else
-        solver.ForwardTransform(lev, *vector_field[0], compx, stag_x);
-        solver.ForwardTransform(lev, *vector_field[1], compy, stag_y);
-        solver.ForwardTransform(lev, *vector_field[2], compz, stag_z);
-#endif
-    }
-
-    AMREX_FORCE_INLINE
-    void ForwardTransformVect (
-        const int lev,
-#ifdef WARPX_DIM_RZ
-        SpectralSolverRZ& solver,
-#else
-        SpectralSolver& solver,
+        solver.ForwardTransform(lev, *vector_field[0], compx);
+        solver.ForwardTransform(lev, *vector_field[1], compy);
+        solver.ForwardTransform(lev, *vector_field[2], compz);
 #endif
-        std::array<std::unique_ptr<amrex::MultiFab>,3>& vector_field,
-        const int compx, const int compy, const int compz)
-    {
-        ForwardTransformVect(lev, solver, vector_field, compx, compy, compz,
-                             (*vector_field[0]).ixType().toIntVect(),
-                             (*vector_field[1]).ixType().toIntVect(),
-                             (*vector_field[2]).ixType().toIntVect());
     }
 
     void
@@ -253,11 +234,12 @@ WarpX::PSATDBackwardTransformG ()
 }
 
 void
-WarpX::PSATDForwardTransformJ ()
+WarpX::PSATDForwardTransformJ (
+    amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_fp,
+    amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_cp)
 {
     SpectralFieldIndex Idx;
     int idx_jx, idx_jy, idx_jz;
-    amrex::IntVect jx_stag, jy_stag, jz_stag;
 
     for (int lev = 0; lev <= finest_level; ++lev)
     {
@@ -267,22 +249,7 @@ WarpX::PSATDForwardTransformJ ()
         idx_jy = (WarpX::do_multi_J) ? static_cast<int>(Idx.Jy_new) : static_cast<int>(Idx.Jy);
         idx_jz = (WarpX::do_multi_J) ? static_cast<int>(Idx.Jz_new) : static_cast<int>(Idx.Jz);
 
-        // With Vay's deposition, J stores a temporary current (D) that is later modified
-        // in Fourier space: its staggering matches that of rho and not J
-        jx_stag =
-            (WarpX::current_deposition_algo == CurrentDepositionAlgo::Vay) ?
-            WarpX::m_rho_nodal_flag : current_fp[lev][0]->ixType().toIntVect();
-
-        jy_stag =
-            (WarpX::current_deposition_algo == CurrentDepositionAlgo::Vay) ?
-            WarpX::m_rho_nodal_flag : current_fp[lev][1]->ixType().toIntVect();
-
-        jz_stag =
-            (WarpX::current_deposition_algo == CurrentDepositionAlgo::Vay) ?
-            WarpX::m_rho_nodal_flag : current_fp[lev][2]->ixType().toIntVect();
-
-        ForwardTransformVect(lev, *spectral_solver_fp[lev], current_fp[lev],
-                             idx_jx, idx_jy, idx_jz, jx_stag, jy_stag, jz_stag);
+        ForwardTransformVect(lev, *spectral_solver_fp[lev], J_fp[lev], idx_jx, idx_jy, idx_jz);
 
         if (spectral_solver_cp[lev])
         {
@@ -292,22 +259,7 @@ WarpX::PSATDForwardTransformJ ()
             idx_jy = (WarpX::do_multi_J) ? static_cast<int>(Idx.Jy_new) : static_cast<int>(Idx.Jy);
             idx_jz = (WarpX::do_multi_J) ? static_cast<int>(Idx.Jz_new) : static_cast<int>(Idx.Jz);
 
-            // With Vay's deposition, J stores a temporary current (D) that is later modified
-            // in Fourier space: its staggering matches that of rho and not J
-            jx_stag =
-                (WarpX::current_deposition_algo == CurrentDepositionAlgo::Vay) ?
-                WarpX::m_rho_nodal_flag : current_cp[lev][0]->ixType().toIntVect();
-
-            jy_stag =
-                (WarpX::current_deposition_algo == CurrentDepositionAlgo::Vay) ?
-                WarpX::m_rho_nodal_flag : current_cp[lev][1]->ixType().toIntVect();
-
-            jz_stag =
-                (WarpX::current_deposition_algo == CurrentDepositionAlgo::Vay) ?
-                WarpX::m_rho_nodal_flag : current_cp[lev][2]->ixType().toIntVect();
-
-            ForwardTransformVect(lev, *spectral_solver_cp[lev], current_cp[lev],
-                                 idx_jx, idx_jy, idx_jz, jx_stag, jy_stag, jz_stag);
+            ForwardTransformVect(lev, *spectral_solver_cp[lev], J_cp[lev], idx_jx, idx_jy, idx_jz);
         }
     }
 
@@ -528,7 +480,11 @@ WarpX::PushPSATD ()
 #else
 
     PSATDForwardTransformEB();
-    PSATDForwardTransformJ();
+
+    amrex::Vector<std::array<std::unique_ptr<amrex::MultiFab>,3>>& J_fp =
+        (WarpX::current_deposition_algo == CurrentDepositionAlgo::Vay) ? current_fp_vay : current_fp;
+
+    PSATDForwardTransformJ(J_fp, current_cp);
 
     // Do rho FFTs only if needed
     if (WarpX::update_with_rho || WarpX::current_correction || WarpX::do_dive_cleaning)