Current correction in Fourier space (#675)

* Start implementing PSATD push without rho.

TODO: 1) fix unit test pml_x_psatd;
      2) try new PSATD push in PML;
      3) avoid intro of new derived class?

* Correct PSATD push to fix PML test.

* Few improvements on new PSATD push:

- new class name is 'PsatdAlgorithmMixed' (both rho and J are used);
- new algorithm parameter to choose between available implementations:
  parameter name is 'psatd_push', possible values are 'standard' for
  old implementation using rho via Gauss law and continuity equation
  and 'mixed' for new implementation using rho via Gauss law and J for
  all remaining terms.

* Fix style error (tabs vs four white spaces).

* Improve comments for available PSATD algorithms.

* Correct few typos in latest comments.

* Implement first current correction:

- new member function 'CurrentCorrection' in class SpectralSolver;
- correction applied only without subcycling (in function 'OneStep_nosub');
- TODO: add correction when subcycling is used (in function 'OneStep_sub1');
- back to old implementation of PSATD push of E and B (class 'PsatdAlgorithmMixed' removed);
- TODO: PML unit test 'pml_x_psatd' does not pass.

* Small cleanup:

- remove residual option for choice of PSATD push algorithm (only one choice now);
- improve comments.

* Implement div(E) diagnostics for spectral case.

* split travis tests in bigger matrix

* split more TravisCI tests, add electrostatic, use defaults values

* typo

* Move computation of div(E) to base class SpectralBaseAlgorithm.

* need to split psatd too

* consistent variable names and use function to avoid duplication

* fix typo for qed tests

* typo

* also need to update run_tests.sg

* Update copyright tags.

* change matrix

* Add test of div(E) vs rho/epsilon_0 in PML test.

* Small clean-up.

* Small clean-up

* Remove option for current correction from input files of two new tests (not used)

* Small clean-up: remove unnecessary references

* More clean-up (minimize style changes to keep PR simple).

* Add specific 2D/3D tests for current correction:

- 'Langmuir_multi_2d_psatd_cc': same input file as 'Langmuir_multi_2d_psatd',
  except for current correction and output of divE;
- 'Langmuir_multi_psatd_cc': same input file (3D) as 'Langmuir_multi_psatd',
  except for current correction and output of divE;
- add corresponding Python scripts for analysis: same as previous ones,
  except for check on L-infinity spatial norm of rho/epsilon_0 vs div(E);
- revert changes on old tests: do not use current correction in old tests
  (benchmarks on Battra do not need to be updated).

* Improve comments.

* Start implementation of new averaging with staggering:

- face-to-cell-center and edge-to-cell-center replaced so far;
- TODO: node-to-cell-center and 1D behavior (AMREX_SPACEDIM=1).

* Remove unrelated style changes (cleaner PR).

* Avoid duplication of input files and analysis scripts for new tests

* Improve comments for Doxygen documentation.

* Improve comments for Doxygen documentation.

* Small clean-up

* first implementation of Diags base classes

* Small clean-up

* Small clean-up

* Fix erroneous non-ASCII character

* Small clean-up

* Auxiliary function for current correction in class WarpX to keep OneStep_nosub clean

* Remove unrelated style changes (cleaner PR)

* Improve comments

* Instrument virtual function 'CurrentCorrection'

* Trying to fix build error detected by LGTM analysis only

* add example, temporarily

* Continue implementation of new averaging with staggering:

- new function takes reference to single MultiFab (no vector);
- TODO: node-to-cell-center still in progress.

* Fix small bug and clean up

* Fix bug in loop over n=0,...,ncomp-1 and clean up

* add more functions

* Add Doxygen documentation and clean up

* Small clean-up in Doxygen documentation

* Compile in single precision: add _rt suffix to avoid unnecessary conversions

* Avoid accessing staggering index directly from IntVect in innermost loops

* Implement periodic-single box option for spectral

* Fix out-of-bound in the periodic, single-box case

* Replace do-while loop with for loop (default ncomp=1)

* Remove temporary pointer and pass reference to MultiFab (instead of MultiFab*)

* Replace AMREX_LAUNCH_HOST_DEVICE_LAMBDA with ParallelFor

* cleaning and initialize output mf

* use general average routine

* move flush in new class, and implemented the Plotfile derived class

* add comments

* eol

* free memory in destructor

* typo

* typo

* no need to clear MF pointers there

* though shalt not break existing tests

* FlushRaw doesnt have to be virtual for now

* The importance of being constant

* Capability to select fields in output files

* EOL

* revert to old inputs

* const in right place

* avoid brace initializer there

* oops, fix logic error in is_in

* Use old name for output image of new 2D test

* user can choose flush interval, same behavior as plot_int

* Small clean-up in Doxygen documentation

* Add option to plot raw fields

* eol

* replace ter flush with dump to avoid confusion

* add options

* Diagnostics stores a vector of functors to compute diags on the fly

* eol

* Field gather from diags to sync particle quantities

* New diagnostics handle RZ with same behavior as old ones

* cleaning and doc

* const ref for string

* smarter for loop from Axel and typo fix from Reva

* Functors to compute some fields

* simplify code following Dave's comments

* Create subfolders and add more output options (divE etc.)

* eol

* Add documentation for periodic_single_box_fft

* For periodic, single-box, apply current correction after guard cell exchange

* rename mode_avg to convertRZmodes2cartesian

* Update CellCenterFunctor.H

* fill varnames and vector of functors at the same time

* output rho_new, not rho_old

* WarpX instance not needed here

* add const

* little bit more of reorganization

* Travis CI: force 2 MPI processes only for numprocs > 2

* Use special FFT (PR #834) and new diagnostics (PR #844) in new tests

* Improve Doxygen documentation

* Move option do_current_correction from warpx to psatd

* Fix path to output files for tests using new diagnostics

* Fix additional paths to output files for new diagnostics

* Add input paramter do_current_correction to documentation

* Fix test Langmuir_multi_psatd_hybrid: do not plot divE

* Remove input parameter amr.plot_int in tests using new diagnostics

* Trigger failing source/style checks on Travis CI

* Fix build error due to public include

* Add missing const keywords

* Change test names and corresponding analysis scripts

* Improve Python script for analysis

* Do not rename output files in old CI tests (without current correction)

* Fix output file name prefix for some tests

* Trigger Travis CI build after AMReX bug fix

* Void commit: trigger Travis CI build

* Fix some tests failing due to recent changes in master

* Use new diagnostics for particle output correctly

Co-authored-by: MaxThevenet <mthevenet@lbl.gov>

* Print tolerance and error in Python analysis

Co-authored-by: MaxThevenet <mthevenet@lbl.gov>

* Print tolerance and error in Python analysis

Co-authored-by: MaxThevenet <mthevenet@lbl.gov>

* Improve documentation

Co-authored-by: MaxThevenet <mthevenet@lbl.gov>

* Fix EOL white spaces

* Fix name of particle output variables

Co-authored-by: MaxThevenet <mthevenet@lbl.gov>
Co-authored-by: Remi Lehe <remi.lehe@normalesup.org>

3 files changed, 126 insertions, 5 deletions

diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.H
index d71c0ab18..eb440d118 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.H
@@ -1,4 +1,4 @@
-/* Copyright 2019 Maxence Thevenet, Remi Lehe, Revathi Jambunathan
+/* Copyright 2019 Maxence Thevenet, Remi Lehe, Revathi Jambunathan, Edoardo Zoni
  *
  *
  * This file is part of WarpX.
@@ -36,8 +36,25 @@ class PsatdAlgorithm : public SpectralBaseAlgorithm
                                     const amrex::DistributionMapping& dm,
                                     const amrex::Real dt);
 
+        /**
+         * \brief Virtual function for current correction in Fourier space
+         * (equation (19) of https://doi.org/10.1016/j.jcp.2013.03.010).
+         * This function overrides the virtual function \c CurrentCorrection in the
+         * base class \c SpectralBaseAlgorithm (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    three-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 ( SpectralFieldData& field_data,
+                                         std::array<std::unique_ptr<amrex::MultiFab>,3>& current,
+                                         const std::unique_ptr<amrex::MultiFab>& rho ) override final;
+
     private:
         SpectralRealCoefficients C_coef, S_ck_coef, X1_coef, X2_coef, X3_coef;
+        amrex::Real m_dt;
 };
 
 #endif // WARPX_USE_PSATD
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp
index 8fee0967d..9684fde27 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp
@@ -1,11 +1,13 @@
-/* Copyright 2019 Remi Lehe, Revathi Jambunathan
+/* Copyright 2019 Remi Lehe, Revathi Jambunathan, Edoardo Zoni
  *
  * This file is part of WarpX.
  *
  * License: BSD-3-Clause-LBNL
  */
+#include "WarpX.H"
 #include "PsatdAlgorithm.H"
 #include "Utils/WarpXConst.H"
+#include "Utils/WarpXProfilerWrapper.H"
 
 #include <cmath>
 
@@ -13,7 +15,9 @@
 #if WARPX_USE_PSATD
 using namespace amrex;
 
-/* \brief Initialize coefficients for the update equation */
+/**
+ * \brief Constructor
+ */
 PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
                          const DistributionMapping& dm,
                          const int norder_x, const int norder_y,
@@ -31,11 +35,15 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
     X2_coef = SpectralRealCoefficients(ba, dm, 1, 0);
     X3_coef = SpectralRealCoefficients(ba, dm, 1, 0);
 
+    // Initialize coefficients for update equations
     InitializeSpectralCoefficients(spectral_kspace, dm, dt);
+
+    m_dt = dt;
 }
 
-/* Advance the E and B field in spectral space (stored in `f`)
- * over one time step */
+/**
+ * \brief Advance E and B fields in spectral space (stored in `f`) over one time step
+ */
 void
 PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{
 
@@ -120,6 +128,9 @@ PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{
     }
 };
 
+/**
+ * \brief Initialize coefficients for update equations
+ */
 void PsatdAlgorithm::InitializeSpectralCoefficients(const SpectralKSpace& spectral_kspace,
                                     const amrex::DistributionMapping& dm,
                                     const amrex::Real dt)
@@ -179,4 +190,81 @@ void PsatdAlgorithm::InitializeSpectralCoefficients(const SpectralKSpace& spectr
         });
      }
 }
+
+void
+PsatdAlgorithm::CurrentCorrection( SpectralFieldData& field_data,
+                                   std::array<std::unique_ptr<amrex::MultiFab>,3>& current,
+                                   const std::unique_ptr<amrex::MultiFab>& rho ) {
+    // Profiling
+    WARPX_PROFILE( "PsatdAlgorithm::CurrentCorrection" );
+
+    using Idx = SpectralFieldIndex;
+
+    // Forward Fourier transform of J and rho
+    field_data.ForwardTransform( *current[0], Idx::Jx, 0 );
+    field_data.ForwardTransform( *current[1], Idx::Jy, 0 );
+    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 (MFIter mfi(field_data.fields); mfi.isValid(); ++mfi){
+
+        const Box& bx = field_data.fields[mfi].box();
+
+        // Extract arrays for the fields to be updated
+        Array4<Complex> fields = field_data.fields[mfi].array();
+        // Extract pointers for the k vectors
+        const Real* const modified_kx_arr = modified_kx_vec[mfi].dataPtr();
+#if (AMREX_SPACEDIM==3)
+        const Real* const modified_ky_arr = modified_ky_vec[mfi].dataPtr();
+#endif
+        const Real* const modified_kz_arr = modified_kz_vec[mfi].dataPtr();
+
+        // Local copy of member variables before GPU loop
+        const Real dt = m_dt;
+
+        // Loop over indices within one box
+        ParallelFor(bx,
+        [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept
+        {
+            // Record old values of the fields to be updated
+            using Idx = SpectralFieldIndex;
+            // Shortcuts for the values of J and rho
+            const Complex Jx = fields(i,j,k,Idx::Jx);
+            const Complex Jy = fields(i,j,k,Idx::Jy);
+            const Complex Jz = fields(i,j,k,Idx::Jz);
+            const Complex rho_old = fields(i,j,k,Idx::rho_old);
+            const Complex rho_new = fields(i,j,k,Idx::rho_new);
+            // k vector values, and coefficients
+            const Real kx = modified_kx_arr[i];
+#if (AMREX_SPACEDIM==3)
+            const Real ky = modified_ky_arr[j];
+            const Real kz = modified_kz_arr[k];
+#else
+            constexpr Real ky = 0;
+            const Real kz = modified_kz_arr[j];
+#endif
+            constexpr Complex I = Complex{0,1};
+
+            const Real k_norm = std::sqrt( kx*kx + ky*ky + kz*kz );
+
+            // div(J) in Fourier space
+            const Complex k_dot_J = kx*Jx + ky*Jy + kz*Jz;
+
+            // Correct J
+            if ( k_norm != 0 )
+            {
+                fields(i,j,k,Idx::Jx) = Jx - (k_dot_J-I*(rho_new-rho_old)/dt)*kx/(k_norm*k_norm);
+                fields(i,j,k,Idx::Jy) = Jy - (k_dot_J-I*(rho_new-rho_old)/dt)*ky/(k_norm*k_norm);
+                fields(i,j,k,Idx::Jz) = Jz - (k_dot_J-I*(rho_new-rho_old)/dt)*kz/(k_norm*k_norm);
+            }
+        });
+    }
+
+    // Backward Fourier transform of J
+    field_data.BackwardTransform( *current[0], Idx::Jx, 0 );
+    field_data.BackwardTransform( *current[1], Idx::Jy, 0 );
+    field_data.BackwardTransform( *current[2], Idx::Jz, 0 );
+}
 #endif // WARPX_USE_PSATD
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.H
index 7ce774ab2..7143b7aa9 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.H
@@ -32,6 +32,22 @@ class SpectralBaseAlgorithm
         virtual ~SpectralBaseAlgorithm() {};
 
         /**
+         * \brief Virtual function for current correction in Fourier space
+         * (equation (19) of https://doi.org/10.1016/j.jcp.2013.03.010).
+         * This virtual function is not pure: it can be overridden in derived
+         * classes (e.g. PsatdAlgorithm, PMLPsatdAlgorithm), but a base-class
+         * implementation is provided (empty implementation in this case).
+         *
+         * \param[in,out] field_data all fields in Fourier space
+         * \param[in,out] current    three-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 ( SpectralFieldData& field_data,
+                                         std::array<std::unique_ptr<amrex::MultiFab>,3>& current,
+                                         const std::unique_ptr<amrex::MultiFab>& rho ) {};
+
+        /**
          * \brief Compute spectral divergence of E
          */
         void ComputeSpectralDivE ( SpectralFieldData& field_data,