Improve spectral solver on staggered grids (#1468)

* Implement Galilean PSATD equations on staggered grids

* Implement high-order interpolation in 2D/3D

* Include missing header file and small clean-up

* Fix bug for FDTD build

* Small clean-up

* Modify current correction for staggered grids

* Implement comoving PSATD scheme (formulation with rho)

* Fix single-precision builds

* Do not implement rho-free formulation for comoving PSATD yet

* Invert sign of comoving velocity to match Galilean convention

* Fix two bugs in comoving PSATD algorithm

* Update benchmark of CI test momentum-conserving-gather

* Update benchmark of CI test PlasmaAccelerationMR

* Update documentation

* Clean up comoving PSATD class

* Clean up comoving PSATD class (more)

* Clean up comoving PSATD class (more)

* Implement changes requested in PR review

* Add 2D regression test for staggered Galilean PSATD

* Add 2D regression test for staggered comoving PSATD

* Unify input files for new CI tests to avoid duplication

* Fully rebase benchmarks changes on development

* Update benchmark of Galilean hybrid test after #1536

1 files changed, 43 insertions, 5 deletions

diff --git a/Source/WarpX.cpp b/Source/WarpX.cpp
index a8afa17ec..3c057e044 100644
--- a/Source/WarpX.cpp
+++ b/Source/WarpX.cpp
@@ -81,6 +81,12 @@ long WarpX::nox = 1;
 long WarpX::noy = 1;
 long WarpX::noz = 1;
 
+// For momentum-conserving field gathering, order of interpolation from the
+// staggered positions to the grid nodes
+int WarpX::field_gathering_nox = 2;
+int WarpX::field_gathering_noy = 2;
+int WarpX::field_gathering_noz = 2;
+
 bool WarpX::use_fdtd_nci_corr = false;
 bool WarpX::galerkin_interpolation = true;
 
@@ -642,6 +648,22 @@ WarpX::ReadParameters ()
         pp.query("noy", noy);
         pp.query("noz", noz);
 
+#ifdef WARPX_USE_PSATD
+        // For momentum-conserving field gathering, read from input the order of
+        // interpolation from the staggered positions to the grid nodes
+        if (field_gathering_algo == GatheringAlgo::MomentumConserving) {
+            pp.query("field_gathering_nox", field_gathering_nox);
+            pp.query("field_gathering_noy", field_gathering_noy);
+            pp.query("field_gathering_noz", field_gathering_noz);
+        }
+
+        if (maxLevel() > 0) {
+            AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
+                field_gathering_nox == 2 && field_gathering_noy == 2 && field_gathering_noz == 2,
+                "High-order interpolation (order > 2) is not implemented with mesh refinement");
+        }
+#endif
+
         pp.query("galerkin_scheme",galerkin_interpolation);
 
         AMREX_ALWAYS_ASSERT_WITH_MESSAGE( nox == noy and nox == noz ,
@@ -688,25 +710,40 @@ WarpX::ReadParameters ()
 
         pp.query("current_correction", current_correction);
         pp.query("v_galilean", m_v_galilean);
+        pp.query("v_comoving", m_v_comoving);
         pp.query("do_time_averaging", fft_do_time_averaging);
 
-      // Scale the velocity by the speed of light
+        // Scale the velocity by the speed of light
         for (int i=0; i<3; i++) m_v_galilean[i] *= PhysConst::c;
+        for (int i=0; i<3; i++) m_v_comoving[i] *= PhysConst::c;
+
+        // The comoving PSATD algorithm is not implemented nor tested with Esirkepov current deposition
+        if (current_deposition_algo == CurrentDepositionAlgo::Esirkepov) {
+            if (m_v_comoving[0] != 0. || m_v_comoving[1] != 0. || m_v_comoving[2] != 0.) {
+                amrex::Abort("Esirkepov current deposition cannot be used with the comoving PSATD algorithm");
+            }
+        }
 
 #   ifdef WARPX_DIM_RZ
         update_with_rho = true;  // Must be true for RZ PSATD
 #   else
-        if (m_v_galilean[0] == 0. && m_v_galilean[1] == 0. && m_v_galilean[2] == 0.) {
+        if (m_v_galilean[0] == 0. && m_v_galilean[1] == 0. && m_v_galilean[2] == 0. &&
+            m_v_comoving[0] == 0. && m_v_comoving[1] == 0. && m_v_comoving[2] == 0.) {
             update_with_rho = false; // standard PSATD
         }
         else {
-            update_with_rho = true;  // Galilean PSATD
+            update_with_rho = true;  // Galilean PSATD or comoving PSATD
         }
 #   endif
 
         // Overwrite update_with_rho with value set in input file
         pp.query("update_with_rho", update_with_rho);
 
+        if (m_v_comoving[0] != 0. || m_v_comoving[1] != 0. || m_v_comoving[2] != 0.) {
+            AMREX_ALWAYS_ASSERT_WITH_MESSAGE(update_with_rho,
+                "psatd.update_with_rho must be equal to 1 for comoving PSATD");
+        }
+
 #   ifdef WARPX_DIM_RZ
         AMREX_ALWAYS_ASSERT_WITH_MESSAGE(update_with_rho,
         "psatd.update_with_rho must be equal to 1 in RZ geometry");
@@ -883,6 +920,7 @@ WarpX::AllocLevelData (int lev, const BoxArray& ba, const DistributionMapping& d
         maxwell_solver_id,
         maxLevel(),
         WarpX::m_v_galilean,
+        WarpX::m_v_comoving,
         safe_guard_cells);
 
     if (mypc->nSpeciesDepositOnMainGrid() && n_current_deposition_buffer == 0) {
@@ -1074,7 +1112,7 @@ WarpX::AllocLevelMFs (int lev, const BoxArray& ba, const DistributionMapping& dm
     }
     bool const pml_flag_false=false;
     spectral_solver_fp[lev] = std::make_unique<SpectralSolver>( realspace_ba, dm,
-        nox_fft, noy_fft, noz_fft, do_nodal, m_v_galilean, dx_vect, dt[lev],
+        nox_fft, noy_fft, noz_fft, do_nodal, m_v_galilean, m_v_comoving, dx_vect, dt[lev],
         pml_flag_false, fft_periodic_single_box, update_with_rho, fft_do_time_averaging );
 #   endif
 #endif
@@ -1192,7 +1230,7 @@ WarpX::AllocLevelMFs (int lev, const BoxArray& ba, const DistributionMapping& dm
 #   else
         c_realspace_ba.grow(ngE); // add guard cells
         spectral_solver_cp[lev] = std::make_unique<SpectralSolver>( c_realspace_ba, dm,
-            nox_fft, noy_fft, noz_fft, do_nodal, m_v_galilean, cdx_vect, dt[lev],
+            nox_fft, noy_fft, noz_fft, do_nodal, m_v_galilean, m_v_comoving, cdx_vect, dt[lev],
             pml_flag_false, fft_periodic_single_box, update_with_rho, fft_do_time_averaging );
 #   endif
 #endif