Implement div(E) diagnostics for spectral case (#720)

* 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.

* SpectralFieldIndex: reuse memory slot for Bx when computing divE.

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

8 files changed, 121 insertions, 24 deletions

diff --git a/Examples/Tests/PML/analysis_pml_psatd.py b/Examples/Tests/PML/analysis_pml_psatd.py
index 164c7a19e..3e5bcef04 100755
--- a/Examples/Tests/PML/analysis_pml_psatd.py
+++ b/Examples/Tests/PML/analysis_pml_psatd.py
@@ -31,6 +31,8 @@ Bz = all_data_level_0['boxlib', 'Bz'].v.squeeze()
 Ex = all_data_level_0['boxlib', 'Ex'].v.squeeze()
 Ey = all_data_level_0['boxlib', 'Ey'].v.squeeze()
 Ez = all_data_level_0['boxlib', 'Ez'].v.squeeze()
+rho = all_data_level_0['boxlib','rho'].v.squeeze()
+divE = all_data_level_0['boxlib','divE'].v.squeeze()
 energyE = np.sum(scc.epsilon_0/2*(Ex**2+Ey**2+Ez**2))
 energyB = np.sum(1./scc.mu_0/2*(Bx**2+By**2+Bz**2))
 energy_end = energyE + energyB
@@ -43,3 +45,6 @@ print("Reflectivity_theory: %s" %Reflectivity_theory)
 
 assert( abs(Reflectivity-Reflectivity_theory) < 5./100 * Reflectivity_theory )
 
+# Check relative L-infinity spatial norm of rho/epsilon_0 - div(E)
+Linf_norm = np.amax( np.abs( rho/scc.epsilon_0 - divE ) ) / np.amax( np.abs( rho/scc.epsilon_0 ) )
+assert( Linf_norm < 2.e-2 )
diff --git a/Regression/WarpX-tests.ini b/Regression/WarpX-tests.ini
index 0434dd4d0..c9fb1e19f 100644
--- a/Regression/WarpX-tests.ini
+++ b/Regression/WarpX-tests.ini
@@ -90,7 +90,7 @@ analysisRoutine = Examples/Tests/PML/analysis_pml_ckc.py
 [pml_x_psatd]
 buildDir = .
 inputFile = Examples/Tests/PML/inputs_2d
-runtime_params = warpx.do_dynamic_scheduling=0
+runtime_params = warpx.do_dynamic_scheduling=0 warpx.fields_to_plot = Ex Ey Ez Bx By Bz rho divE
 dim = 2
 addToCompileString = USE_PSATD=TRUE
 restartTest = 0
diff --git a/Source/Diagnostics/FieldIO.cpp b/Source/Diagnostics/FieldIO.cpp
index 03e5e07d5..7d6ca85e2 100644
--- a/Source/Diagnostics/FieldIO.cpp
+++ b/Source/Diagnostics/FieldIO.cpp
@@ -1,6 +1,5 @@
-/* Copyright 2019-2020 Andrew Myers, Axel Huebl, David Grote
- * Maxence Thevenet, Remi Lehe, Revathi Jambunathan
- * Weiqun Zhang
+/* Copyright 2019-2020 Andrew Myers, Axel Huebl, David Grote, Maxence Thevenet,
+ * Remi Lehe, Revathi Jambunathan, Weiqun Zhang, Edoardo Zoni
  *
  * This file is part of WarpX.
  *
@@ -13,7 +12,10 @@
 #include <AMReX_Interpolater.H>
 
 #ifdef WARPX_USE_OPENPMD
-#   include <openPMD/openPMD.hpp>
+#include <openPMD/openPMD.hpp>
+#endif
+#ifdef WARPX_USE_PSATD
+#include <SpectralSolver.H>
 #endif
 
 
@@ -643,15 +645,16 @@ WarpX::AverageAndPackFields ( Vector<std::string>& varnames,
                                     Bfield_aux[lev][2].get()},
                             WarpX::CellSize(lev) );
             } else if (fieldname == "divE"){
-                if (do_nodal) amrex::Abort("TODO: do_nodal && plot dive");
+                if (do_nodal) amrex::Abort("TODO: do_nodal && plot divE");
                 const BoxArray& ba = amrex::convert(boxArray(lev),IntVect::TheUnitVector());
-                MultiFab dive(ba,DistributionMap(lev),1,0);
-                ComputeDivE( dive, 0,
-                             {Efield_aux[lev][0].get(),
-                                     Efield_aux[lev][1].get(),
-                                     Efield_aux[lev][2].get()},
-                             WarpX::CellSize(lev) );
-                AverageAndPackScalarField( mf_avg[lev], dive, dcomp++, ngrow );
+                MultiFab divE( ba, DistributionMap(lev), 1, 0 );
+#ifdef WARPX_USE_PSATD
+                spectral_solver_fp[lev]->ComputeSpectralDivE( Efield_aux[lev], divE );
+#else
+                ComputeDivE( divE, 0, {Efield_aux[lev][0].get(), Efield_aux[lev][1].get(),
+                             Efield_aux[lev][2].get()}, WarpX::CellSize(lev) );
+#endif
+                AverageAndPackScalarField( mf_avg[lev], divE, dcomp++, ngrow );
             } else {
                 amrex::Abort("unknown field in fields_to_plot: " + fieldname);
             }
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package
index 82763f501..69758e7a9 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package
@@ -1,4 +1,5 @@
 CEXE_headers += SpectralBaseAlgorithm.H
+CEXE_sources += SpectralBaseAlgorithm.cpp
 CEXE_headers += PsatdAlgorithm.H
 CEXE_sources += PsatdAlgorithm.cpp
 
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.H
index 2487eae78..743d6d1c7 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.H
@@ -1,4 +1,4 @@
-/* Copyright 2019 Remi Lehe
+/* Copyright 2019 Remi Lehe, Edoardo Zoni
  *
  * This file is part of WarpX.
  *
@@ -31,6 +31,13 @@ class SpectralBaseAlgorithm
         // calls the subclass's destructor.
         virtual ~SpectralBaseAlgorithm() {};
 
+        /**
+         * \brief Compute spectral divergence of E
+         */
+        void ComputeSpectralDivE ( SpectralFieldData& field_data,
+                                   const std::array<std::unique_ptr<amrex::MultiFab>,3>& Efield,
+                                   amrex::MultiFab& divE );
+
     protected: // Meant to be used in the subclasses
 
         using SpectralRealCoefficients = \
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.cpp
new file mode 100644
index 000000000..c917391f9
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.cpp
@@ -0,0 +1,69 @@
+/* Copyright 2019 Edoardo Zoni
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ */
+#include <SpectralBaseAlgorithm.H>
+#include <cmath>
+
+using namespace amrex;
+
+/**
+ * \brief Compute spectral divergence of E
+ */
+void
+SpectralBaseAlgorithm::ComputeSpectralDivE (
+    SpectralFieldData& field_data,
+    const std::array<std::unique_ptr<amrex::MultiFab>,3>& Efield,
+    amrex::MultiFab& divE )
+{
+    using Idx = SpectralFieldIndex;
+
+    // Forward Fourier transform of E
+    field_data.ForwardTransform( *Efield[0], Idx::Ex, 0 );
+    field_data.ForwardTransform( *Efield[1], Idx::Ey, 0 );
+    field_data.ForwardTransform( *Efield[2], Idx::Ez, 0 );
+
+    // 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* modified_kx_arr = modified_kx_vec[mfi].dataPtr();
+#if (AMREX_SPACEDIM==3)
+        const Real* modified_ky_arr = modified_ky_vec[mfi].dataPtr();
+#endif
+        const Real* modified_kz_arr = modified_kz_vec[mfi].dataPtr();
+
+        // Loop over indices within one box
+        ParallelFor(bx,
+        [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept
+        {
+            using Idx = SpectralFieldIndex;
+            // Shortcuts for the components of E
+            const Complex Ex = fields(i,j,k,Idx::Ex);
+            const Complex Ey = fields(i,j,k,Idx::Ey);
+            const Complex Ez = fields(i,j,k,Idx::Ez);
+            // k vector values
+            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
+            const Complex I = Complex{0,1};
+
+            // div(E) in Fourier space
+            fields(i,j,k,Idx::divE) = I*(kx*Ex+ky*Ey+kz*Ez);
+        });
+    }
+
+    // Backward Fourier transform
+    field_data.BackwardTransform( divE, Idx::divE, 0 );
+};
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
index 2c5a45faa..db3979796 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
@@ -20,10 +20,12 @@
 // Declare type for spectral fields
 using SpectralField = amrex::FabArray< amrex::BaseFab <Complex> >;
 
-/** Index for the regular fields, when stored in spectral space */
+/** Index for the regular fields, when stored in spectral space:
+ *  - n_fields is automatically the total number of fields
+ *  - divE reuses the memory slot for Bx, since Bx is not used when computing divE
+ */
 struct SpectralFieldIndex {
-  enum { Ex=0, Ey, Ez, Bx, By, Bz, Jx, Jy, Jz, rho_old, rho_new, n_fields };
-  // n_fields is automatically the total number of fields
+  enum { Ex=0, Ey, Ez, Bx, By, Bz, Jx, Jy, Jz, rho_old, rho_new, n_fields, divE=3 };
 };
 
 /** Index for the PML fields, when stored in spectral space */
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolver.H b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
index 6685e489c..8b33707d0 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolver.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
@@ -1,4 +1,4 @@
-/* Copyright 2019-2020 Maxence Thevenet, Remi Lehe
+/* Copyright 2019-2020 Maxence Thevenet, Remi Lehe, Edoardo Zoni
  *
  * This file is part of WarpX.
  *
@@ -57,14 +57,24 @@ class SpectralSolver
          */
         void pushSpectralFields();
 
+        /**
+          * \brief Public interface to call the member function ComputeSpectralDivE
+          * of the base class SpectralBaseAlgorithm from objects of class SpectralSolver
+          */
+        void ComputeSpectralDivE ( const std::array<std::unique_ptr<amrex::MultiFab>,3>& Efield,
+                                   amrex::MultiFab& divE ) {
+            algorithm->ComputeSpectralDivE( field_data, Efield, divE );
+        };
+
     private:
-        SpectralFieldData field_data; // Store field in spectral space
-                                      // and perform the Fourier transforms
+
+        // Store field in spectral space and perform the Fourier transforms
+        SpectralFieldData field_data;
+
+        // Defines field update equation in spectral space and the associated coefficients.
+        // SpectralBaseAlgorithm is a base class; this pointer is meant to point
+        // to an instance of a sub-class defining a specific algorithm
         std::unique_ptr<SpectralBaseAlgorithm> algorithm;
-        // Defines field update equation in spectral space,
-        // and the associated coefficients.
-        // SpectralBaseAlgorithm is a base class ; this pointer is meant
-        // to point an instance of a *sub-class* defining a specific algorithm
 };
 #endif // WARPX_USE_PSATD
 #endif // WARPX_SPECTRAL_SOLVER_H_