RZ spectral current correction and Galilean (#1216)

* Added stub for current correction in RZ spectral solver

* Fixed comments in RZ spectral for current correction stub

* Modified automated test for Galilean PSATD (#1033)

* Impemented current correction in RZ spectral

* Implementation Galilean version of RZ spectral solver

* For RZ spectral, do forward and backward transform with current correction

* Big fix in DivEFunctor.cpp for RZ spectral

* Added RZ rho diagnostic for saving the modes

* Implemented fft_periodic_single_box for RZ spectral

* Moved routines from SpectralSolverRZ.H to .cpp

* Added hook for VayDeposition in GalileanPsatdAlgorithmRZ

* Bug fix in DivEFunctor

* Fixes and cleanup for GalileanPsatdAlgorithmRZ

* Fix line spacing in SpectralSolverRZ.H

* Fix factor 1/2 in update of Ep_m

* Fix factor 1/2 in update of Em_m

* Fix sign error in current correction in GalileanPsatdAlgorithmRZ.cpp

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

* Add Langmuir RZ PSATD test with current correction

* Add Galilean tests with/without current correction

* For RZ psatd, simplified copy for forward transform

* Added GalileanPsatdAlgorithmRZ.cpp to CMakeLists

* Minor cleanup in RZ spectral solver

* In GalileanPsatdAlgorithmRZ.cpp use member initialization for m_v_galilean

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

* Added some _rt to GalileanPsatdAlgorithmRZ.cpp

Co-authored-by: Olga Shapoval <30510597+oshapoval@users.noreply.github.com>
Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>
Co-authored-by: Edoardo Zoni <ezoni@lbl.gov>

11 files changed, 574 insertions, 37 deletions

diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/CMakeLists.txt b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/CMakeLists.txt
index c5ebf1c3e..a050bd5c2 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/CMakeLists.txt
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/CMakeLists.txt
@@ -12,5 +12,6 @@ if(WarpX_DIMS STREQUAL RZ)
       PRIVATE
         SpectralBaseAlgorithmRZ.cpp
         PsatdAlgorithmRZ.cpp
+        GalileanPsatdAlgorithmRZ.cpp
     )
 endif()
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.H
new file mode 100644
index 000000000..3c8d3d60e
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.H
@@ -0,0 +1,72 @@
+/* Copyright 2019 David Grote
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef WARPX_GALILEAN_PSATD_ALGORITHM_RZ_H_
+#define WARPX_GALILEAN_PSATD_ALGORITHM_RZ_H_
+
+#include "SpectralBaseAlgorithmRZ.H"
+
+/* \brief Class that updates the field in spectral space
+ * and stores the coefficients of the corresponding update equation.
+ */
+class GalileanPsatdAlgorithmRZ : public SpectralBaseAlgorithmRZ
+{
+
+    public:
+        GalileanPsatdAlgorithmRZ (SpectralKSpaceRZ const & spectral_kspace,
+                                  amrex::DistributionMapping const & dm,
+                                  int const n_rz_azimuthal_modes, int const norder_z,
+                                  bool const nodal,
+                                  const amrex::Array<amrex::Real,3>& v_galilean,
+                                  amrex::Real const dt_step);
+        // Redefine functions from base class
+        virtual void pushSpectralFields (SpectralFieldDataRZ & f) override final;
+        virtual int getRequiredNumberOfFields () const override final {
+            return SpectralFieldIndex::n_fields;
+        }
+
+        void InitializeSpectralCoefficients (SpectralFieldDataRZ const & f);
+
+        /**
+         * \brief Virtual function for current correction in Fourier space
+         * This function overrides the virtual function \c CurrentCorrection in the
+         * base class \c SpectralBaseAlgorithmRZ (qualifier \c override) and cannot be
+         * overridden by further derived classes (qualifier \c final).
+         *
+         * \param[in,out] field_data all fields in Fourier space
+         * \param[in,out] current    two-dimensional array of unique pointers to MultiFab
+         *                           storing the three components of the current density
+         * \param[in]     rho        unique pointer to MultiFab storing the charge density
+         */
+        virtual void CurrentCorrection ( SpectralFieldDataRZ& field_data,
+                                         std::array<std::unique_ptr<amrex::MultiFab>,3>& current,
+                                         const std::unique_ptr<amrex::MultiFab>& rho ) override final;
+
+        /**
+         * \brief Virtual function for Vay current deposition in Fourier space
+         * This function overrides the virtual function \c VayDeposition in the
+         * base class \c SpectralBaseAlgorithmRZ and cannot be overridden by further
+         * derived classes.
+         *
+         * \param[in,out] field_data All fields in Fourier space
+         * \param[in,out] current    Array of unique pointers to \c MultiFab storing
+         *                           the three components of the current density
+         */
+        virtual void VayDeposition (SpectralFieldDataRZ& field_data,
+                                    std::array<std::unique_ptr<amrex::MultiFab>,3>& current) override final;
+
+    private:
+        bool coefficients_initialized;
+        // Note that dt and v_galilean are saved to use in InitializeSpectralCoefficients
+        amrex::Real const m_dt;
+        amrex::Array<amrex::Real,3> m_v_galilean;
+
+        SpectralRealCoefficients C_coef, S_ck_coef;
+        SpectralComplexCoefficients Theta2_coef, X1_coef, X2_coef, X3_coef, X4_coef;
+
+};
+
+#endif // WARPX_GALILEAN_PSATD_ALGORITHM_RZ_H_
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.cpp
new file mode 100644
index 000000000..ed54dd21f
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/GalileanPsatdAlgorithmRZ.cpp
@@ -0,0 +1,355 @@
+/* Copyright 2019-2020 David Grote
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ */
+#include "WarpX.H"
+#include "GalileanPsatdAlgorithmRZ.H"
+#include "Utils/WarpXConst.H"
+#include "Utils/WarpXProfilerWrapper.H"
+
+#include <cmath>
+
+using namespace amrex::literals;
+
+
+/* \brief Initialize coefficients for the update equation */
+GalileanPsatdAlgorithmRZ::GalileanPsatdAlgorithmRZ (SpectralKSpaceRZ const & spectral_kspace,
+                                                    amrex::DistributionMapping const & dm,
+                                                    int const n_rz_azimuthal_modes, int const norder_z,
+                                                    bool const nodal,
+                                                    const amrex::Array<amrex::Real,3>& v_galilean,
+                                                    amrex::Real const dt)
+     // Initialize members of base class
+     : SpectralBaseAlgorithmRZ(spectral_kspace, dm, norder_z, nodal),
+       m_dt(dt),
+       m_v_galilean(v_galilean)
+{
+
+    // Allocate the arrays of coefficients
+    amrex::BoxArray const & ba = spectral_kspace.spectralspace_ba;
+    C_coef = SpectralRealCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+    S_ck_coef = SpectralRealCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+    X1_coef = SpectralComplexCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+    X2_coef = SpectralComplexCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+    X3_coef = SpectralComplexCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+    X4_coef = SpectralComplexCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+    Theta2_coef = SpectralComplexCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+
+    coefficients_initialized = false;
+}
+
+/* Advance the E and B field in spectral space (stored in `f`)
+ * over one time step */
+void
+GalileanPsatdAlgorithmRZ::pushSpectralFields (SpectralFieldDataRZ & f)
+{
+
+    if (not coefficients_initialized) {
+        // This is called from here since it needs the kr values
+        // which can be obtained from the SpectralFieldDataRZ
+        InitializeSpectralCoefficients(f);
+        coefficients_initialized = true;
+    }
+
+    // Loop over boxes
+    for (amrex::MFIter mfi(f.fields); mfi.isValid(); ++mfi){
+
+        amrex::Box const & bx = f.fields[mfi].box();
+
+        // Extract arrays for the fields to be updated
+        amrex::Array4<Complex> const& fields = f.fields[mfi].array();
+        // Extract arrays for the coefficients
+        amrex::Array4<const amrex::Real> const& C_arr = C_coef[mfi].array();
+        amrex::Array4<const amrex::Real> const& S_ck_arr = S_ck_coef[mfi].array();
+        amrex::Array4<const Complex> const& X1_arr = X1_coef[mfi].array();
+        amrex::Array4<const Complex> const& X2_arr = X2_coef[mfi].array();
+        amrex::Array4<const Complex> const& X3_arr = X3_coef[mfi].array();
+        amrex::Array4<const Complex> const& X4_arr = X4_coef[mfi].array();
+        amrex::Array4<const Complex> const& Theta2_arr = Theta2_coef[mfi].array();
+
+        // Extract pointers for the k vectors
+        auto const & kr_modes = f.getKrArray(mfi);
+        amrex::Real const* kr_arr = kr_modes.dataPtr();
+        amrex::Real const* modified_kz_arr = modified_kz_vec[mfi].dataPtr();
+        int const nr = bx.length(0);
+
+        // Loop over indices within one box
+        // Note that k = 0
+        int const modes = f.n_rz_azimuthal_modes;
+        amrex::ParallelFor(bx, modes,
+        [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept
+        {
+
+            // All of the fields of each mode are grouped together
+            using Idx = SpectralFieldIndex;
+            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;
+
+            // Record old values of the fields to be updated
+            Complex const Ep_old = fields(i,j,k,Ep_m);
+            Complex const Em_old = fields(i,j,k,Em_m);
+            Complex const Ez_old = fields(i,j,k,Ez_m);
+            Complex const Bp_old = fields(i,j,k,Bp_m);
+            Complex const Bm_old = fields(i,j,k,Bm_m);
+            Complex const Bz_old = fields(i,j,k,Bz_m);
+            // Shortcut for the values of J and rho
+            Complex const Jp = fields(i,j,k,Jp_m);
+            Complex const Jm = fields(i,j,k,Jm_m);
+            Complex const Jz = fields(i,j,k,Jz_m);
+            Complex const rho_old = fields(i,j,k,rho_old_m);
+            Complex const rho_new = fields(i,j,k,rho_new_m);
+
+            // k vector values, and coefficients
+            // The k values for each mode are grouped together
+            int const ir = i + nr*mode;
+            amrex::Real const kr = kr_arr[ir];
+            amrex::Real const kz = modified_kz_arr[j];
+
+            constexpr amrex::Real c2 = PhysConst::c*PhysConst::c;
+            constexpr amrex::Real inv_ep0 = 1._rt/PhysConst::ep0;
+            Complex const I = Complex{0._rt,1._rt};
+            amrex::Real const C = C_arr(i,j,k,mode);
+            amrex::Real const S_ck = S_ck_arr(i,j,k,mode);
+            Complex const X1 = X1_arr(i,j,k,mode);
+            Complex const X2 = X2_arr(i,j,k,mode);
+            Complex const X3 = X3_arr(i,j,k,mode);
+            Complex const X4 = X4_arr(i,j,k,mode);
+            Complex const T2 = Theta2_arr(i,j,k,mode);
+
+            // Update E (see WarpX online documentation: theory section)
+            fields(i,j,k,Ep_m) = T2*C*Ep_old
+                        + T2*S_ck*(-c2*I*kr/2._rt*Bz_old + c2*kz*Bp_old)
+                        + X4*Jp + 0.5_rt*kr*(X2*rho_new - T2*X3*rho_old);
+            fields(i,j,k,Em_m) = T2*C*Em_old
+                        + T2*S_ck*(-c2*I*kr/2._rt*Bz_old - c2*kz*Bm_old)
+                        + X4*Jm - 0.5_rt*kr*(X2*rho_new - T2*X3*rho_old);
+            fields(i,j,k,Ez_m) = T2*C*Ez_old
+                        + T2*S_ck*(c2*I*kr*Bp_old + c2*I*kr*Bm_old)
+                        + X4*Jz - I*kz*(X2*rho_new - T2*X3*rho_old);
+            // Update B (see WarpX online documentation: theory section)
+            // Note: here X1 is T2*x1/(ep0*c*c*k_norm*k_norm), where
+            // x1 has the same definition as in the original paper
+            fields(i,j,k,Bp_m) = T2*C*Bp_old
+                        - T2*S_ck*(-I*kr/2._rt*Ez_old + kz*Ep_old)
+                        + X1*(-I*kr/2._rt*Jz + kz*Jp);
+            fields(i,j,k,Bm_m) = T2*C*Bm_old
+                        - T2*S_ck*(-I*kr/2._rt*Ez_old - kz*Em_old)
+                        + X1*(-I*kr/2._rt*Jz - kz*Jm);
+            fields(i,j,k,Bz_m) = T2*C*Bz_old
+                        - T2*S_ck*I*(kr*Ep_old + kr*Em_old)
+                        + X1*I*(kr*Jp + kr*Jm);
+        });
+    }
+};
+
+void GalileanPsatdAlgorithmRZ::InitializeSpectralCoefficients (SpectralFieldDataRZ const & f)
+{
+
+    // Fill them with the right values:
+    // Loop over boxes and allocate the corresponding coefficients
+    // for each box owned by the local MPI proc
+    for (amrex::MFIter mfi(f.fields); mfi.isValid(); ++mfi){
+
+        amrex::Box const & bx = f.fields[mfi].box();
+
+        // Extract pointers for the k vectors
+        amrex::Real const* const modified_kz = modified_kz_vec[mfi].dataPtr();
+
+        // Extract arrays for the coefficients
+        amrex::Array4<amrex::Real> const& C = C_coef[mfi].array();
+        amrex::Array4<amrex::Real> const& S_ck = S_ck_coef[mfi].array();
+        amrex::Array4<Complex> const& X1 = X1_coef[mfi].array();
+        amrex::Array4<Complex> const& X2 = X2_coef[mfi].array();
+        amrex::Array4<Complex> const& X3 = X3_coef[mfi].array();
+        amrex::Array4<Complex> const& X4 = X4_coef[mfi].array();
+        amrex::Array4<Complex> const& Theta2 = Theta2_coef[mfi].array();
+
+        // Extract real (for portability on GPU)
+        amrex::Real vz = m_v_galilean[2];
+
+        auto const & kr_modes = f.getKrArray(mfi);
+        amrex::Real const* kr_arr = kr_modes.dataPtr();
+        int const nr = bx.length(0);
+        amrex::Real const dt = m_dt;
+
+        // Loop over indices within one box
+        int const modes = f.n_rz_azimuthal_modes;
+        amrex::ParallelFor(bx, modes,
+        [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept
+        {
+            constexpr amrex::Real c = PhysConst::c;
+            constexpr amrex::Real ep0 = PhysConst::ep0;
+            Complex const I = Complex{0._rt,1._rt};
+
+            // Calculate norm of vector
+            int const ir = i + nr*mode;
+            amrex::Real const kr = kr_arr[ir];
+            amrex::Real const kz = modified_kz[j];
+            amrex::Real const k_norm = std::sqrt(kr*kr + kz*kz);
+
+            // Calculate coefficients
+            if (k_norm != 0._rt){
+
+                C(i,j,k,mode) = std::cos(c*k_norm*dt);
+                S_ck(i,j,k,mode) = std::sin(c*k_norm*dt)/(c*k_norm);
+
+                // Calculate dot product with galilean velocity
+                amrex::Real const kv = kz*vz;
+
+                amrex::Real const nu = kv/(k_norm*c);
+                Complex const theta = amrex::exp( 0.5_rt*I*kv*dt );
+                Complex const theta_star = amrex::exp( -0.5_rt*I*kv*dt );
+                Complex const e_theta = amrex::exp( I*c*k_norm*dt );
+
+                Theta2(i,j,k,mode) = theta*theta;
+
+                if ( (nu != 1._rt) && (nu != 0._rt) ) {
+
+                    // Note: the coefficients X1, X2, X do not correspond
+                    // exactly to the original Galilean paper, but the
+                    // update equation have been modified accordingly so that
+                    // the expressions/ below (with the update equations)
+                    // are mathematically equivalent to those of the paper.
+                    Complex x1 = 1._rt/(1._rt-nu*nu) *
+                        (theta_star - C(i,j,k,mode)*theta + I*kv*S_ck(i,j,k,mode)*theta);
+                    // x1, above, is identical to the original paper
+                    X1(i,j,k,mode) = theta*x1/(ep0*c*c*k_norm*k_norm);
+                    // The difference betwen X2 and X3 below, and those
+                    // from the original paper is the factor ep0*k_norm*k_norm
+                    X2(i,j,k,mode) = (x1 - theta*(1._rt - C(i,j,k,mode)))
+                                     /(theta_star-theta)/(ep0*k_norm*k_norm);
+                    X3(i,j,k,mode) = (x1 - theta_star*(1._rt - C(i,j,k,mode)))
+                                     /(theta_star-theta)/(ep0*k_norm*k_norm);
+                    X4(i,j,k,mode) = I*kv*X1(i,j,k,mode) - theta*theta*S_ck(i,j,k,mode)/ep0;
+
+                } else if (nu == 0._rt) {
+
+                    X1(i,j,k,mode) = (1._rt - C(i,j,k,mode))/(ep0 * c*c * k_norm*k_norm);
+                    X2(i,j,k,mode) = (1._rt - S_ck(i,j,k,mode)/dt)/(ep0 * k_norm*k_norm);
+                    X3(i,j,k,mode) = (C(i,j,k,mode) - S_ck(i,j,k,mode)/dt)/(ep0 * k_norm*k_norm);
+                    X4(i,j,k,mode) = -S_ck(i,j,k,mode)/ep0;
+
+                } else if ( nu == 1._rt) {
+                    X1(i,j,k,mode) = (1._rt - e_theta*e_theta + 2._rt*I*c*k_norm*dt) / (4._rt*c*c*ep0*k_norm*k_norm);
+                    X2(i,j,k,mode) = (3._rt - 4._rt*e_theta + e_theta*e_theta + 2._rt*I*c*k_norm*dt) / (4._rt*ep0*k_norm*k_norm*(1._rt - e_theta));
+                    X3(i,j,k,mode) = (3._rt - 2._rt/e_theta - 2._rt*e_theta + e_theta*e_theta - 2._rt*I*c*k_norm*dt) / (4._rt*ep0*(e_theta - 1._rt)*k_norm*k_norm);
+                    X4(i,j,k,mode) = I*(-1._rt + e_theta*e_theta + 2._rt*I*c*k_norm*dt) / (4._rt*ep0*c*k_norm);
+                }
+
+            } else { // Handle k_norm = 0, by using the analytical limit
+                C(i,j,k,mode) = 1._rt;
+                S_ck(i,j,k,mode) = dt;
+                X1(i,j,k,mode) = 0.5_rt * dt*dt / ep0;
+                X2(i,j,k,mode) = c*c * dt*dt / (6._rt*ep0);
+                X3(i,j,k,mode) = - c*c * dt*dt / (3._rt*ep0);
+                X4(i,j,k,mode) = -dt/ep0;
+                Theta2(i,j,k,mode) = 1._rt;
+            }
+        });
+     }
+}
+
+void
+GalileanPsatdAlgorithmRZ::CurrentCorrection (SpectralFieldDataRZ& field_data,
+                                             std::array<std::unique_ptr<amrex::MultiFab>,3>& current,
+                                             const std::unique_ptr<amrex::MultiFab>& rho )
+{
+    // Profiling
+    WARPX_PROFILE( "GalileanPsatdAlgorithmRZ::CurrentCorrection" );
+
+    using Idx = SpectralFieldIndex;
+
+    // Forward Fourier transform of J and rho
+    field_data.ForwardTransform( *current[0], Idx::Jx,
+                                 *current[1], Idx::Jy);
+    field_data.ForwardTransform( *current[2], Idx::Jz, 0);
+    field_data.ForwardTransform( *rho, Idx::rho_old, 0 );
+    field_data.ForwardTransform( *rho, Idx::rho_new, 1 );
+
+    // Loop over boxes
+    for (amrex::MFIter mfi(field_data.fields); mfi.isValid(); ++mfi){
+
+        amrex::Box const & bx = field_data.fields[mfi].box();
+
+        // Extract arrays for the fields to be updated
+        amrex::Array4<Complex> fields = field_data.fields[mfi].array();
+
+        // Extract pointers for the k vectors
+        auto const & kr_modes = field_data.getKrArray(mfi);
+        amrex::Real const* kr_arr = kr_modes.dataPtr();
+        amrex::Real const* modified_kz_arr = modified_kz_vec[mfi].dataPtr();
+        int const nr = bx.length(0);
+
+        // Local copy of member variables before GPU loop
+        amrex::Real vz = m_v_galilean[2];
+        amrex::Real const dt = m_dt;
+
+        // Loop over indices within one box
+        int const modes = field_data.n_rz_azimuthal_modes;
+        ParallelFor(bx, modes,
+        [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept
+        {
+
+            // All of the fields of each mode are grouped together
+            using Idx = SpectralFieldIndex;
+            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;
+
+            // Shortcuts for the values of J and rho
+            Complex const Jp = fields(i,j,k,Jp_m);
+            Complex const Jm = fields(i,j,k,Jm_m);
+            Complex const Jz = fields(i,j,k,Jz_m);
+            Complex const rho_old = fields(i,j,k,rho_old_m);
+            Complex const rho_new = fields(i,j,k,rho_new_m);
+
+            // k vector values, and coefficients
+            // The k values for each mode are grouped together
+            int const ir = i + nr*mode;
+            amrex::Real const kr = kr_arr[ir];
+            amrex::Real const kz = modified_kz_arr[j];
+            amrex::Real const k_norm2 = kr*kr + kz*kz;
+
+            constexpr Complex I = Complex{0._rt,1._rt};
+
+            // Correct J
+            if ( k_norm2 != 0._rt )
+            {
+                Complex const theta2 = amrex::exp(I*kz*vz*dt);
+                Complex const inv_1_T2 = 1._rt/(kz*vz == 0._rt ?  1._rt : 1._rt - theta2);
+                Complex const j_corr_coef = (kz == 0._rt ? 1._rt/dt : -I*kz*vz*inv_1_T2);
+                Complex const F = - (j_corr_coef*(rho_new - rho_old*theta2) + I*kz*Jz + kr*(Jp - Jm))/k_norm2;
+
+                fields(i,j,k,Jp_m) += +0.5_rt*kr*F;
+                fields(i,j,k,Jm_m) += -0.5_rt*kr*F;
+                fields(i,j,k,Jz_m) += -I*kz*F;
+            }
+        });
+    }
+
+    // Backward Fourier transform of J
+    field_data.BackwardTransform( *current[0], Idx::Jx,
+                                  *current[1], Idx::Jy);
+    field_data.BackwardTransform( *current[2], Idx::Jz, 0 );
+
+}
+
+void
+GalileanPsatdAlgorithmRZ::VayDeposition (SpectralFieldDataRZ& field_data,
+                                         std::array<std::unique_ptr<amrex::MultiFab>,3>& current)
+{
+    amrex::Abort("Vay deposition not implemented in RZ geometry");
+}
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package
index f316abefa..57f8c5e3f 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package
@@ -7,6 +7,7 @@ CEXE_sources += PMLPsatdAlgorithm.cpp
 ifeq ($(USE_RZ),TRUE)
   CEXE_sources += SpectralBaseAlgorithmRZ.cpp
   CEXE_sources += PsatdAlgorithmRZ.cpp
+  CEXE_sources += GalileanPsatdAlgorithmRZ.cpp
 endif
 
 VPATH_LOCATIONS   += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver/SpectralAlgorithms
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H
index b4c2597a8..f8df96e54 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H
@@ -61,8 +61,8 @@ class PsatdAlgorithmRZ : public SpectralBaseAlgorithmRZ
     private:
         bool coefficients_initialized;
         // Note that dt is saved to use in InitializeSpectralCoefficients
-        amrex::Real dt;
-        SpectralCoefficients C_coef, S_ck_coef, X1_coef, X2_coef, X3_coef;
+        amrex::Real m_dt;
+        SpectralRealCoefficients C_coef, S_ck_coef, X1_coef, X2_coef, X3_coef;
 };
 
 #endif // WARPX_PSATD_ALGORITHM_RZ_H_
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp
index 349f266d1..496470912 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp
@@ -4,8 +4,10 @@
  *
  * License: BSD-3-Clause-LBNL
  */
+#include "WarpX.H"
 #include "PsatdAlgorithmRZ.H"
 #include "Utils/WarpXConst.H"
+#include "Utils/WarpXProfilerWrapper.H"
 
 #include <cmath>
 
@@ -16,20 +18,20 @@ using amrex::operator""_rt;
 PsatdAlgorithmRZ::PsatdAlgorithmRZ (SpectralKSpaceRZ const & spectral_kspace,
                                     amrex::DistributionMapping const & dm,
                                     int const n_rz_azimuthal_modes, int const norder_z,
-                                    bool const nodal, amrex::Real const dt_step)
+                                    bool const nodal, amrex::Real const dt)
      // Initialize members of base class
      : SpectralBaseAlgorithmRZ(spectral_kspace, dm,
                                norder_z, nodal),
-       dt(dt_step)
+       m_dt(dt)
 {
 
     // Allocate the arrays of coefficients
     amrex::BoxArray const & ba = spectral_kspace.spectralspace_ba;
-    C_coef = SpectralCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
-    S_ck_coef = SpectralCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
-    X1_coef = SpectralCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
-    X2_coef = SpectralCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
-    X3_coef = SpectralCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+    C_coef = SpectralRealCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+    S_ck_coef = SpectralRealCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+    X1_coef = SpectralRealCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+    X2_coef = SpectralRealCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
+    X3_coef = SpectralRealCoefficients(ba, dm, n_rz_azimuthal_modes, 0);
 
     coefficients_initialized = false;
 }
@@ -164,7 +166,7 @@ void PsatdAlgorithmRZ::InitializeSpectralCoefficients (SpectralFieldDataRZ const
         auto const & kr_modes = f.getKrArray(mfi);
         amrex::Real const* kr_arr = kr_modes.dataPtr();
         int const nr = bx.length(0);
-        amrex::Real const dt_temp = dt;
+        amrex::Real const dt = m_dt;
 
         // Loop over indices within one box
         int const modes = f.n_rz_azimuthal_modes;
@@ -181,17 +183,17 @@ void PsatdAlgorithmRZ::InitializeSpectralCoefficients (SpectralFieldDataRZ const
             constexpr amrex::Real c = PhysConst::c;
             constexpr amrex::Real ep0 = PhysConst::ep0;
             if (k_norm != 0){
-                C(i,j,k,mode) = std::cos(c*k_norm*dt_temp);
-                S_ck(i,j,k,mode) = std::sin(c*k_norm*dt_temp)/(c*k_norm);
+                C(i,j,k,mode) = std::cos(c*k_norm*dt);
+                S_ck(i,j,k,mode) = std::sin(c*k_norm*dt)/(c*k_norm);
                 X1(i,j,k,mode) = (1._rt - C(i,j,k,mode))/(ep0 * c*c * k_norm*k_norm);
-                X2(i,j,k,mode) = (1._rt - S_ck(i,j,k,mode)/dt_temp)/(ep0 * k_norm*k_norm);
-                X3(i,j,k,mode) = (C(i,j,k,mode) - S_ck(i,j,k,mode)/dt_temp)/(ep0 * k_norm*k_norm);
+                X2(i,j,k,mode) = (1._rt - S_ck(i,j,k,mode)/dt)/(ep0 * k_norm*k_norm);
+                X3(i,j,k,mode) = (C(i,j,k,mode) - S_ck(i,j,k,mode)/dt)/(ep0 * k_norm*k_norm);
             } else { // Handle k_norm = 0, by using the analytical limit
                 C(i,j,k,mode) = 1._rt;
-                S_ck(i,j,k,mode) = dt_temp;
-                X1(i,j,k,mode) = 0.5_rt * dt_temp*dt_temp / ep0;
-                X2(i,j,k,mode) = c*c * dt_temp*dt_temp / (6._rt*ep0);
-                X3(i,j,k,mode) = - c*c * dt_temp*dt_temp / (3._rt*ep0);
+                S_ck(i,j,k,mode) = dt;
+                X1(i,j,k,mode) = 0.5_rt * dt*dt / ep0;
+                X2(i,j,k,mode) = c*c * dt*dt / (6._rt*ep0);
+                X3(i,j,k,mode) = - c*c * dt*dt / (3._rt*ep0);
             }
         });
      }
@@ -202,7 +204,81 @@ PsatdAlgorithmRZ::CurrentCorrection (SpectralFieldDataRZ& field_data,
                                      std::array<std::unique_ptr<amrex::MultiFab>,3>& current,
                                      const std::unique_ptr<amrex::MultiFab>& rho)
 {
-    amrex::Abort("Current correction not implemented in RZ geometry");
+    // Profiling
+    WARPX_PROFILE( "PsatdAlgorithmRZ::CurrentCorrection" );
+
+    using Idx = SpectralFieldIndex;
+
+    // Forward Fourier transform of J and rho
+    field_data.ForwardTransform( *current[0], Idx::Jx,
+                                 *current[1], Idx::Jy);
+    field_data.ForwardTransform( *current[2], Idx::Jz, 0);
+    field_data.ForwardTransform( *rho, Idx::rho_old, 0 );
+    field_data.ForwardTransform( *rho, Idx::rho_new, 1 );
+
+    // Loop over boxes
+    for (amrex::MFIter mfi(field_data.fields); mfi.isValid(); ++mfi){
+
+        amrex::Box const & bx = field_data.fields[mfi].box();
+
+        // Extract arrays for the fields to be updated
+        amrex::Array4<Complex> fields = field_data.fields[mfi].array();
+
+        // Extract pointers for the k vectors
+        auto const & kr_modes = field_data.getKrArray(mfi);
+        amrex::Real const* kr_arr = kr_modes.dataPtr();
+        amrex::Real const* modified_kz_arr = modified_kz_vec[mfi].dataPtr();
+        int const nr = bx.length(0);
+
+        // Local copy of member variables before GPU loop
+        amrex::Real const dt = m_dt;
+
+        // Loop over indices within one box
+        int const modes = field_data.n_rz_azimuthal_modes;
+        ParallelFor(bx, modes,
+        [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept
+        {
+
+            // All of the fields of each mode are grouped together
+            using Idx = SpectralFieldIndex;
+            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;
+
+            // Shortcuts for the values of J and rho
+            Complex const Jp = fields(i,j,k,Jp_m);
+            Complex const Jm = fields(i,j,k,Jm_m);
+            Complex const Jz = fields(i,j,k,Jz_m);
+            Complex const rho_old = fields(i,j,k,rho_old_m);
+            Complex const rho_new = fields(i,j,k,rho_new_m);
+
+            // k vector values, and coefficients
+            // The k values for each mode are grouped together
+            int const ir = i + nr*mode;
+            amrex::Real const kr = kr_arr[ir];
+            amrex::Real const kz = modified_kz_arr[j];
+            amrex::Real const k_norm2 = kr*kr + kz*kz;
+
+            constexpr Complex I = Complex{0._rt,1._rt};
+
+            // Correct J
+            if ( k_norm2 != 0._rt )
+            {
+                Complex const F = - ((rho_new - rho_old)/dt + I*kz*Jz + kr*(Jp - Jm))/k_norm2;
+
+                fields(i,j,k,Jp_m) += +0.5_rt*kr*F;
+                fields(i,j,k,Jm_m) += -0.5_rt*kr*F;
+                fields(i,j,k,Jz_m) += -I*kz*F;
+            }
+        });
+    }
+
+    // Backward Fourier transform of J
+    field_data.BackwardTransform( *current[0], Idx::Jx,
+                                  *current[1], Idx::Jy);
+    field_data.BackwardTransform( *current[2], Idx::Jz, 0 );
 }
 
 void
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithmRZ.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithmRZ.H
index b0a30de50..29a8d8e7f 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithmRZ.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithmRZ.H
@@ -63,7 +63,8 @@ class SpectralBaseAlgorithmRZ
 
     protected: // Meant to be used in the subclasses
 
-        using SpectralCoefficients = amrex::FabArray< amrex::BaseFab <amrex::Real> >;
+        using SpectralRealCoefficients = amrex::FabArray< amrex::BaseFab <amrex::Real> >;
+        using SpectralComplexCoefficients = amrex::FabArray< amrex::BaseFab <Complex> >;
 
         // Constructor
         SpectralBaseAlgorithmRZ(SpectralKSpaceRZ const & spectral_kspace,
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H
index a073ec483..475759b59 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H
@@ -44,12 +44,12 @@ class SpectralFieldDataRZ
         SpectralFieldDataRZ& operator=(SpectralFieldDataRZ&& field_data) = default;
         ~SpectralFieldDataRZ ();
 
-        void ForwardTransform (const amrex::MultiFab& mf,
-                               const int field_index, const int i_comp);
+        void ForwardTransform (const amrex::MultiFab& mf, const int field_index,
+                               const int i_comp);
         void ForwardTransform (const amrex::MultiFab& mf_r, const int field_index_r,
                                const amrex::MultiFab& mf_t, const int field_index_t);
-        void BackwardTransform (amrex::MultiFab& mf,
-                                const int field_index, const int i_comp);
+        void BackwardTransform (amrex::MultiFab& mf, const int field_index,
+                                const int i_comp);
         void BackwardTransform (amrex::MultiFab& mf_r, const int field_index_r,
                                 amrex::MultiFab& mf_t, const int field_index_t);
 
@@ -89,6 +89,7 @@ class SpectralFieldDataRZ
         SpectralShiftFactor zshift_FFTfromCell, zshift_FFTtoCell;
         MultiSpectralHankelTransformer multi_spectral_hankel_transformer;
         BinomialFilter binomialfilter;
+
 };
 
 #endif // WARPX_SPECTRAL_FIELD_DATA_RZ_H_
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp
index 38576a6fe..736c2dcfa 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp
@@ -301,7 +301,7 @@ SpectralFieldDataRZ::ForwardTransform (amrex::MultiFab const & field_mf, int con
 
     // Create a copy of the input multifab since the shape of field_mf
     // might not be what is needed in transform.
-    // For example, with fft_periodic_single_box, field_mf will have guard cells but
+    // For example, with periodic_single_box_fft, field_mf will have guard cells but
     // the transformed array does not.
     // Note that the copy will not include the imaginary part of mode 0 as
     // PhysicalToSpectral_Scalar expects.
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H
index 460b79366..e79d868f6 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H
@@ -29,15 +29,14 @@ class SpectralSolverRZ
                           amrex::DistributionMapping const & dm,
                           int const n_rz_azimuthal_modes,
                           int const norder_z, bool const nodal,
+                          const amrex::Array<amrex::Real,3>& v_galilean,
                           amrex::RealVect const dx, amrex::Real const dt,
                           int const lev);
 
         /* \brief Transform the component `i_comp` of MultiFab `field_mf`
          *  to spectral space, and store the corresponding result internally
          *  (in the spectral field specified by `field_index`) */
-
-        void ForwardTransform (amrex::MultiFab const & field_mf,
-                               int const field_index,
+        void ForwardTransform (amrex::MultiFab const & field_mf, int const field_index,
                                int const i_comp=0);
 
         /* \brief Transform the two MultiFabs `field_mf1` and `field_mf2`
@@ -48,8 +47,7 @@ class SpectralSolverRZ
 
         /* \brief Transform spectral field specified by `field_index` back to
          * real space, and store it in the component `i_comp` of `field_mf` */
-        void BackwardTransform (amrex::MultiFab& field_mf,
-                                int const field_index,
+        void BackwardTransform (amrex::MultiFab& field_mf, int const field_index,
                                 int const i_comp=0);
 
         /* \brief Transform spectral fields specified by `field_index1` and `field_index2`
@@ -78,6 +76,7 @@ class SpectralSolverRZ
         {
             field_data.ApplyFilter(field_index1, field_index2, field_index3);
         }
+
         /**
           * \brief Public interface to call the member function ComputeSpectralDivE
           * of the base class SpectralBaseAlgorithmRZ from objects of class SpectralSolverRZ
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp
index c6457a036..38851ef50 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp
@@ -7,6 +7,7 @@
 #include "SpectralKSpaceRZ.H"
 #include "SpectralSolverRZ.H"
 #include "SpectralAlgorithms/PsatdAlgorithmRZ.H"
+#include "SpectralAlgorithms/GalileanPsatdAlgorithmRZ.H"
 #include "WarpX.H"
 #include "Utils/WarpXProfilerWrapper.H"
 
@@ -28,6 +29,7 @@ SpectralSolverRZ::SpectralSolverRZ (amrex::BoxArray const & realspace_ba,
                                     amrex::DistributionMapping const & dm,
                                     int const n_rz_azimuthal_modes,
                                     int const norder_z, bool const nodal,
+                                    const amrex::Array<amrex::Real,3>& v_galilean,
                                     amrex::RealVect const dx, amrex::Real const dt,
                                     int const lev)
     : k_space(realspace_ba, dm, dx)
@@ -41,8 +43,15 @@ SpectralSolverRZ::SpectralSolverRZ (amrex::BoxArray const & realspace_ba,
     // - Select the algorithm depending on the input parameters
     //   Initialize the corresponding coefficients over k space
     //   PML is not supported.
-    algorithm = std::unique_ptr<PsatdAlgorithmRZ>(
-        new PsatdAlgorithmRZ(k_space, dm, n_rz_azimuthal_modes, norder_z, nodal, dt));
+    if (v_galilean[2] == 0) {
+         // v_galilean is 0: use standard PSATD algorithm
+        algorithm = std::unique_ptr<PsatdAlgorithmRZ>(
+            new PsatdAlgorithmRZ(k_space, dm, n_rz_azimuthal_modes, norder_z, nodal, dt));
+    } else {
+        // Otherwise: use the Galilean algorithm
+        algorithm = std::unique_ptr<GalileanPsatdAlgorithmRZ>(
+            new GalileanPsatdAlgorithmRZ(k_space, dm, n_rz_azimuthal_modes, norder_z, nodal, v_galilean, dt));
+    }
 
     // - Initialize arrays for fields in spectral space + FFT plans
     field_data = SpectralFieldDataRZ(realspace_ba, k_space, dm,
@@ -50,15 +59,19 @@ SpectralSolverRZ::SpectralSolverRZ (amrex::BoxArray const & realspace_ba,
                                      n_rz_azimuthal_modes, lev);
 }
 
-
+/* \brief Transform the component `i_comp` of MultiFab `field_mf`
+ *  to spectral space, and store the corresponding result internally
+ *  (in the spectral field specified by `field_index`) */
 void
-SpectralSolverRZ::ForwardTransform (amrex::MultiFab const & field_mf,
-                                    int const field_index,
+SpectralSolverRZ::ForwardTransform (amrex::MultiFab const & field_mf, int const field_index,
                                     int const i_comp) {
     WARPX_PROFILE("SpectralSolverRZ::ForwardTransform");
     field_data.ForwardTransform(field_mf, field_index, i_comp);
 }
 
+/* \brief Transform the two MultiFabs `field_mf1` and `field_mf2`
+ *  to spectral space, and store the corresponding results internally
+ *  (in the spectral field specified by `field_index1` and `field_index2`) */
 void
 SpectralSolverRZ::ForwardTransform (amrex::MultiFab const & field_mf1, int const field_index1,
                                     amrex::MultiFab const & field_mf2, int const field_index2) {
@@ -67,14 +80,17 @@ SpectralSolverRZ::ForwardTransform (amrex::MultiFab const & field_mf1, int const
                                 field_mf2, field_index2);
 }
 
+/* \brief Transform spectral field specified by `field_index` back to
+ * real space, and store it in the component `i_comp` of `field_mf` */
 void
-SpectralSolverRZ::BackwardTransform (amrex::MultiFab& field_mf,
-                                     int const field_index,
+SpectralSolverRZ::BackwardTransform (amrex::MultiFab& field_mf, int const field_index,
                                      int const i_comp) {
     WARPX_PROFILE("SpectralSolverRZ::BackwardTransform");
     field_data.BackwardTransform(field_mf, field_index, i_comp);
 }
 
+/* \brief Transform spectral fields specified by `field_index1` and `field_index2`
+ * back to real space, and store it in `field_mf1` and `field_mf2`*/
 void
 SpectralSolverRZ::BackwardTransform (amrex::MultiFab& field_mf1, int const field_index1,
                                      amrex::MultiFab& field_mf2, int const field_index2) {
@@ -83,6 +99,7 @@ SpectralSolverRZ::BackwardTransform (amrex::MultiFab& field_mf1, int const field
                                  field_mf2, field_index2);
 }
 
+/* \brief Update the fields in spectral space, over one timestep */
 void
 SpectralSolverRZ::pushSpectralFields () {
     WARPX_PROFILE("SpectralSolverRZ::pushSpectralFields");
@@ -92,12 +109,26 @@ SpectralSolverRZ::pushSpectralFields () {
     algorithm->pushSpectralFields(field_data);
 }
 
+/**
+  * \brief Public interface to call the member function ComputeSpectralDivE
+  * of the base class SpectralBaseAlgorithmRZ from objects of class SpectralSolverRZ
+  */
 void
 SpectralSolverRZ::ComputeSpectralDivE (const std::array<std::unique_ptr<amrex::MultiFab>,3>& Efield,
                                        amrex::MultiFab& divE) {
     algorithm->ComputeSpectralDivE(field_data, Efield, divE);
 }
 
+/**
+ * \brief Public interface to call the virtual function \c CurrentCorrection,
+ * defined in the base class SpectralBaseAlgorithmRZ and possibly overridden
+ * by its derived classes (e.g. PsatdAlgorithmRZ), from
+ * objects of class SpectralSolverRZ through the private unique pointer \c algorithm
+ *
+ * \param[in,out] current two-dimensional array of unique pointers to MultiFab
+ *                        storing the three components of the current density
+ * \param[in]     rho     unique pointer to MultiFab storing the charge density
+ */
 void
 SpectralSolverRZ::CurrentCorrection (std::array<std::unique_ptr<amrex::MultiFab>,3>& current,
                                       const std::unique_ptr<amrex::MultiFab>& rho) {