Fix PSATD equations used with PML (#1513)

* Implement new PML PSATD equations

* Update CI test and benchmark

* Compute coefficients C1,...,C22 on the fly

* Add check on initial energy from diagnostics

4 files changed, 162 insertions, 86 deletions

diff --git a/Examples/Tests/PML/analysis_pml_psatd.py b/Examples/Tests/PML/analysis_pml_psatd.py
index a28541c53..d219fdd1d 100755
--- a/Examples/Tests/PML/analysis_pml_psatd.py
+++ b/Examples/Tests/PML/analysis_pml_psatd.py
@@ -7,7 +7,6 @@
 #
 # License: BSD-3-Clause-LBNL
 
-
 import sys
 import yt ; yt.funcs.mylog.setLevel(0)
 import numpy as np
@@ -20,7 +19,23 @@ filename = sys.argv[1]
 ############################
 ### INITIAL LASER ENERGY ###
 ############################
-energy_start = 9.1301289517e-08
+energy_start = 7.297301362346945e-08 # electromagnetic energy at iteration 50
+
+# Check consistency of field energy diagnostics with initial energy above
+ds = yt.load('pml_x_psatd_plt00050')
+all_data_level_0 = ds.covering_grid(level=0, left_edge=ds.domain_left_edge, dims=ds.domain_dimensions)
+Bx = all_data_level_0['boxlib', 'Bx'].v.squeeze()
+By = all_data_level_0['boxlib', 'By'].v.squeeze()
+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()
+energyE = np.sum(0.5 * scc.epsilon_0 * (Ex**2 + Ey**2 + Ez**2))
+energyB = np.sum(0.5 / scc.mu_0 * (Bx**2 + By**2 + Bz**2))
+energy_start_diags = energyE + energyB
+error = abs(energy_start - energy_start_diags) / energy_start
+tolerance = 1e-14
+assert (error < tolerance)
 
 ##########################
 ### FINAL LASER ENERGY ###
@@ -33,26 +48,17 @@ 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))
+energyE = np.sum(0.5 * scc.epsilon_0 * (Ex**2 + Ey**2 + Ez**2))
+energyB = np.sum(0.5 / scc.mu_0 * (Bx**2 + By**2 + Bz**2))
 energy_end = energyE + energyB
 
-Reflectivity = energy_end/energy_start
-Reflectivity_theory = 1.3806831258153887e-06
-
-error_rel = abs(Reflectivity-Reflectivity_theory) / Reflectivity_theory
-tolerance_rel = 5./100
-
-print("error_rel    : " + str(error_rel))
-print("tolerance_rel: " + str(tolerance_rel))
+reflectivity = energy_end / energy_start
+reflectivity_max = 1e-06
 
-assert( error_rel < tolerance_rel )
+print("reflectivity     = " + str(reflectivity))
+print("reflectivity_max = " + str(reflectivity_max))
 
-# 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 )
+assert(reflectivity < reflectivity_max)
 
 test_name = filename[:-9] # Could also be os.path.split(os.getcwd())[1]
 checksumAPI.evaluate_checksum(test_name, filename)
diff --git a/Regression/Checksum/benchmarks_json/pml_x_psatd.json b/Regression/Checksum/benchmarks_json/pml_x_psatd.json
index c2604f4da..82a5bd9ab 100644
--- a/Regression/Checksum/benchmarks_json/pml_x_psatd.json
+++ b/Regression/Checksum/benchmarks_json/pml_x_psatd.json
@@ -1,12 +1,12 @@
 {
   "lev=0": {
-    "Bx": 1.4955213883225615e-08,
-    "By": 2.5469838278883616e-08,
-    "Bz": 1.205728995299769e-08,
-    "Ex": 6.069794171544256,
-    "Ey": 5.355917208303502,
-    "Ez": 5.016737331847911,
-    "divE": 202005.08114094258,
-    "rho": 1.67193461204679e-06
+    "Bx": 1.1593525237393989e-08,
+    "By": 1.5110569654968376e-08,
+    "Bz": 8.985467170981363e-09,
+    "Ex": 3.9386934154767923,
+    "Ey": 4.134907100275868,
+    "Ez": 3.3085880750067016,
+    "divE": 188915.38924357053,
+    "rho": 1.6719346122590767e-06
   }
 }
\ No newline at end of file
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PMLPsatdAlgorithm.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PMLPsatdAlgorithm.H
index 348282dce..a4cccbc55 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PMLPsatdAlgorithm.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PMLPsatdAlgorithm.H
@@ -65,8 +65,8 @@ class PMLPsatdAlgorithm : public SpectralBaseAlgorithm
                                     std::array<std::unique_ptr<amrex::MultiFab>,3>& current) override final;
 
     private:
-        SpectralRealCoefficients C_coef, S_ck_coef;
-
+        SpectralRealCoefficients C_coef, S_ck_coef, inv_k2_coef;
+        amrex::Real m_dt;
 };
 
 #endif // WARPX_USE_PSATD
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PMLPsatdAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PMLPsatdAlgorithm.cpp
index e1be557fe..6b9397807 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PMLPsatdAlgorithm.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PMLPsatdAlgorithm.cpp
@@ -15,20 +15,20 @@
 using namespace amrex;
 
 /* \brief Initialize coefficients for the update equation */
-PMLPsatdAlgorithm::PMLPsatdAlgorithm(
-                         const SpectralKSpace& spectral_kspace,
-                         const DistributionMapping& dm,
-                         const int norder_x, const int norder_y,
-                         const int norder_z, const bool nodal, const Real dt)
+PMLPsatdAlgorithm::PMLPsatdAlgorithm(const SpectralKSpace& spectral_kspace,
+                                     const DistributionMapping& dm,
+                                     const int norder_x, const int norder_y,
+                                     const int norder_z, const bool nodal, const Real dt)
      // Initialize members of base class
-     : SpectralBaseAlgorithm( spectral_kspace, dm,
-                              norder_x, norder_y, norder_z, nodal )
+     : SpectralBaseAlgorithm(spectral_kspace, dm, norder_x, norder_y, norder_z, nodal),
+       m_dt(dt)
 {
     const BoxArray& ba = spectral_kspace.spectralspace_ba;
 
     // Allocate the arrays of coefficients
-    C_coef = SpectralRealCoefficients(ba, dm, 1, 0);
-    S_ck_coef = SpectralRealCoefficients(ba, dm, 1, 0);
+    C_coef      = SpectralRealCoefficients(ba, dm, 1, 0);
+    S_ck_coef   = SpectralRealCoefficients(ba, dm, 1, 0);
+    inv_k2_coef = SpectralRealCoefficients(ba, dm, 1, 0);
 
     InitializeSpectralCoefficients(spectral_kspace, dm, dt);
 }
@@ -36,7 +36,7 @@ PMLPsatdAlgorithm::PMLPsatdAlgorithm(
 /* Advance the E and B field in spectral space (stored in `f`)
  * over one time step */
 void
-PMLPsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{
+PMLPsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const {
 
     // Loop over boxes
     for (MFIter mfi(f.fields); mfi.isValid(); ++mfi){
@@ -45,9 +45,12 @@ PMLPsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{
 
         // Extract arrays for the fields to be updated
         Array4<Complex> fields = f.fields[mfi].array();
+
         // Extract arrays for the coefficients
         Array4<const Real> C_arr = C_coef[mfi].array();
         Array4<const Real> S_ck_arr = S_ck_coef[mfi].array();
+        Array4<const Real> inv_k2_arr = inv_k2_coef[mfi].array();
+
         // Extract pointers for the k vectors
         const Real* modified_kx_arr = modified_kx_vec[mfi].dataPtr();
 #if (AMREX_SPACEDIM==3)
@@ -55,52 +58,114 @@ PMLPsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{
 #endif
         const Real* modified_kz_arr = modified_kz_vec[mfi].dataPtr();
 
+        const amrex::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 = SpectralPMLIndex;
-            const Complex Ex_old = fields(i,j,k,Idx::Exy) \
-                                 + fields(i,j,k,Idx::Exz);
-            const Complex Ey_old = fields(i,j,k,Idx::Eyx) \
-                                 + fields(i,j,k,Idx::Eyz);
-            const Complex Ez_old = fields(i,j,k,Idx::Ezx) \
-                                 + fields(i,j,k,Idx::Ezy);
-            const Complex Bx_old = fields(i,j,k,Idx::Bxy) \
-                                 + fields(i,j,k,Idx::Bxz);
-            const Complex By_old = fields(i,j,k,Idx::Byx) \
-                                 + fields(i,j,k,Idx::Byz);
-            const Complex Bz_old = fields(i,j,k,Idx::Bzx) \
-                                 + fields(i,j,k,Idx::Bzy);
+
+            const Complex Exy = fields(i,j,k,Idx::Exy);
+            const Complex Exz = fields(i,j,k,Idx::Exz);
+            const Complex Eyx = fields(i,j,k,Idx::Eyx);
+            const Complex Eyz = fields(i,j,k,Idx::Eyz);
+            const Complex Ezx = fields(i,j,k,Idx::Ezx);
+            const Complex Ezy = fields(i,j,k,Idx::Ezy);
+
+            const Complex Bxy = fields(i,j,k,Idx::Bxy);
+            const Complex Bxz = fields(i,j,k,Idx::Bxz);
+            const Complex Byx = fields(i,j,k,Idx::Byx);
+            const Complex Byz = fields(i,j,k,Idx::Byz);
+            const Complex Bzx = fields(i,j,k,Idx::Bzx);
+            const Complex Bzy = fields(i,j,k,Idx::Bzy);
+
             // 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;
+            constexpr Real ky = 0._rt;
             const Real kz = modified_kz_arr[j];
 #endif
-            constexpr Real c2 = PhysConst::c*PhysConst::c;
-            const Complex I = Complex{0,1};
-            const Real C = C_arr(i,j,k);
-            const Real S_ck = S_ck_arr(i,j,k);
-
-            // Update E
-            fields(i,j,k,Idx::Exy) = C*fields(i,j,k,Idx::Exy) + S_ck*c2*I*ky*Bz_old;
-            fields(i,j,k,Idx::Exz) = C*fields(i,j,k,Idx::Exz) - S_ck*c2*I*kz*By_old;
-            fields(i,j,k,Idx::Eyz) = C*fields(i,j,k,Idx::Eyz) + S_ck*c2*I*kz*Bx_old;
-            fields(i,j,k,Idx::Eyx) = C*fields(i,j,k,Idx::Eyx) - S_ck*c2*I*kx*Bz_old;
-            fields(i,j,k,Idx::Ezx) = C*fields(i,j,k,Idx::Ezx) + S_ck*c2*I*kx*By_old;
-            fields(i,j,k,Idx::Ezy) = C*fields(i,j,k,Idx::Ezy) - S_ck*c2*I*ky*Bx_old;
-            // Update B
-            fields(i,j,k,Idx::Bxy) = C*fields(i,j,k,Idx::Bxy) - S_ck*I*ky*Ez_old;
-            fields(i,j,k,Idx::Bxz) = C*fields(i,j,k,Idx::Bxz) + S_ck*I*kz*Ey_old;
-            fields(i,j,k,Idx::Byz) = C*fields(i,j,k,Idx::Byz) - S_ck*I*kz*Ex_old;
-            fields(i,j,k,Idx::Byx) = C*fields(i,j,k,Idx::Byx) + S_ck*I*kx*Ez_old;
-            fields(i,j,k,Idx::Bzx) = C*fields(i,j,k,Idx::Bzx) - S_ck*I*kx*Ey_old;
-            fields(i,j,k,Idx::Bzy) = C*fields(i,j,k,Idx::Bzy) + S_ck*I*ky*Ex_old;
+            constexpr Real c2 = PhysConst::c * PhysConst::c;
+
+            const Complex I = Complex{0._rt, 1._rt};
+
+            const Real kx2 = kx*kx;
+            const Real ky2 = ky*ky;
+            const Real kz2 = kz*kz;
+            const Real k_norm = std::sqrt(kx2 + ky2 + kz2);
+
+            if (k_norm != 0._rt) {
+
+                const Real C = C_arr(i,j,k);
+                const Real S_ck = S_ck_arr(i,j,k);
+                const Real inv_k2 = inv_k2_arr(i,j,k);
+
+                const Real C1  = (kx2 * C + ky2 + kz2) * inv_k2;
+                const Real C2  = (kx2 + ky2 * C + kz2) * inv_k2;
+                const Real C3  = (kx2 + ky2 + kz2 * C) * inv_k2;
+                const Real C4  = kx2 * (C - 1._rt) * inv_k2;
+                const Real C5  = ky2 * (C - 1._rt) * inv_k2;
+                const Real C6  = kz2 * (C - 1._rt) * inv_k2;
+                const Real C7  = ky * kz * (1._rt - C) * inv_k2;
+                const Real C8  = kx * kz * (1._rt - C) * inv_k2;
+                const Real C9  = kx * ky * (1._rt - C) * inv_k2;
+                const Real C10 = c2 * kx * ky * kz * (dt - S_ck) * inv_k2;
+                const Real C11 = c2 * ky2 * kz * (dt - S_ck) * inv_k2;
+                const Real C12 = c2 * kz2 * ky * (dt - S_ck) * inv_k2;
+                const Real C13 = c2 * kz2 * kx * (dt - S_ck) * inv_k2;
+                const Real C14 = c2 * kx2 * kz * (dt - S_ck) * inv_k2;
+                const Real C15 = c2 * kx2 * ky * (dt - S_ck) * inv_k2;
+                const Real C16 = c2 * ky2 * kx * (dt - S_ck) * inv_k2;
+                const Real C17 = c2 * kx * (ky2 * dt + (kz2 + kx2) * S_ck) * inv_k2;
+                const Real C18 = c2 * kx * (kz2 * dt + (ky2 + kx2) * S_ck) * inv_k2;
+                const Real C19 = c2 * ky * (kz2 * dt + (kx2 + ky2) * S_ck) * inv_k2;
+                const Real C20 = c2 * ky * (kx2 * dt + (kz2 + ky2) * S_ck) * inv_k2;
+                const Real C21 = c2 * kz * (kx2 * dt + (ky2 + kz2) * S_ck) * inv_k2;
+                const Real C22 = c2 * kz * (ky2 * dt + (kx2 + kz2) * S_ck) * inv_k2;
+
+                // Update E
+                fields(i,j,k,Idx::Exy) = C2 * Exy + C5 * Exz + C9 * (Eyz + Eyx)
+                                       + I*C10 * (Bxy + Bxz) + I*C11 * (Byz + Byx) + I*C19 * (Bzx + Bzy);
+
+                fields(i,j,k,Idx::Exz) = C6 * Exy + C3 * Exz + C8 * (Ezx + Ezy)
+                                       - I*C10 * (Bxy + Bxz) - I*C22 * (Byz + Byx) - I*C12 * (Bzx + Bzy);
+
+                fields(i,j,k,Idx::Eyz) = C3 * Eyz + C6 * Eyx + C7 * (Ezx + Ezy)
+                                       + I*C21 * (Bxy + Bxz) + I*C10 * (Byz + Byx) + I*C13 * (Bzx + Bzy);
+
+                fields(i,j,k,Idx::Eyx) = C9 * (Exy + Exz) + C4 * Eyz + C1 * Eyx
+                                       - I*C14 * (Bxy + Bxz) - I*C10 * (Byz + Byx) - I*C18 * (Bzx + Bzy);
+
+                fields(i,j,k,Idx::Ezx) = C8 * (Exy + Exz) + C1 * Ezx + C4 * Ezy
+                                       + I*C15 * (Bxy + Bxz) + I*C17 * (Byz + Byx) + I*C10 * (Bzx + Bzy);
+
+                fields(i,j,k,Idx::Ezy) = C7 * (Eyz + Eyx) + C5 * Ezx + C2 * Ezy
+                                       - I*C20 * (Bxy + Bxz) - I*C16 * (Byz + Byx) - I*C10 * (Bzx + Bzy);
+
+                // Update B
+                fields(i,j,k,Idx::Bxy) = C2 * Bxy + C5 * Bxz + C9 * (Byz + Byx)
+                                       - I*C10/c2 * (Exy + Exz) - I*C11/c2 * (Eyz + Eyx) - I*C19/c2 * (Ezx + Ezy);
+
+                fields(i,j,k,Idx::Bxz) = C6 * Bxy + C3 * Bxz + C8 * (Bzx + Bzy)
+                                       + I*C10/c2 * (Exy + Exz) + I*C22/c2 * (Eyz + Eyx) + I*C12/c2 * (Ezx + Ezy);
+
+                fields(i,j,k,Idx::Byz) = C3 * Byz + C6 * Byx + C7 * (Bzx + Bzy)
+                                       - I*C21/c2 * (Exy + Exz) - I*C10/c2 * (Eyz + Eyx) - I*C13/c2 * (Ezx + Ezy);
+
+                fields(i,j,k,Idx::Byx) = C9 * (Bxy + Bxz) + C4 * Byz + C1 * Byx
+                                       + I*C14/c2 * (Exy + Exz) + I*C10/c2 * (Eyz + Eyx) + I*C18/c2 * (Ezx + Ezy);
+
+                fields(i,j,k,Idx::Bzx) = C8 * (Bxy + Bxz) + C1 * Bzx + C4 * Bzy
+                                       - I*C15/c2 * (Exy + Exz) - I*C17/c2 * (Eyz + Eyx) - I*C10/c2 * (Ezx + Ezy);
+
+                fields(i,j,k,Idx::Bzy) = C7 * (Byz + Byx) + C5 * Bzx + C2 * Bzy
+                                       + I*C20/c2 * (Exy + Exz) + I*C16/c2 * (Eyz + Eyx) + I*C10/c2 * (Ezx + Ezy);
+            }
         });
     }
 }
@@ -111,10 +176,11 @@ void PMLPsatdAlgorithm::InitializeSpectralCoefficients (
     const amrex::Real dt)
 {
     const BoxArray& ba = spectral_kspace.spectralspace_ba;
+
     // Fill them with the right values:
     // Loop over boxes and allocate the corresponding coefficients
     // for each box owned by the local MPI proc
-    for (MFIter mfi(ba, dm); mfi.isValid(); ++mfi){
+    for (MFIter mfi(ba, dm); mfi.isValid(); ++mfi) {
 
         const Box& bx = ba[mfi];
 
@@ -124,32 +190,36 @@ void PMLPsatdAlgorithm::InitializeSpectralCoefficients (
         const Real* modified_ky = modified_ky_vec[mfi].dataPtr();
 #endif
         const Real* modified_kz = modified_kz_vec[mfi].dataPtr();
+
         // Extract arrays for the coefficients
         Array4<Real> C = C_coef[mfi].array();
         Array4<Real> S_ck = S_ck_coef[mfi].array();
+        Array4<Real> inv_k2 = inv_k2_coef[mfi].array();
 
         // Loop over indices within one box
-        ParallelFor(bx,
-        [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept
+        ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept
         {
-            // Calculate norm of vector
-            const Real k_norm = std::sqrt(
-                std::pow(modified_kx[i], 2) +
+            const Real kx = modified_kx[i];
 #if (AMREX_SPACEDIM==3)
-                std::pow(modified_ky[j], 2) +
-                std::pow(modified_kz[k], 2));
+            const Real ky = modified_ky[j];
+            const Real kz = modified_kz[k];
 #else
-                std::pow(modified_kz[j], 2));
+            constexpr Real ky = 0._rt;
+            const     Real kz = modified_kz[j];
 #endif
 
+            // Calculate norm of vector
+            const Real k_norm = std::sqrt(kx*kx + ky*ky + kz*kz);
+            const Real k2 = k_norm * k_norm;
+
             // Calculate coefficients
-            constexpr Real c = PhysConst::c;
-            if (k_norm != 0){
-                C(i,j,k) = std::cos(c*k_norm*dt);
-                S_ck(i,j,k) = std::sin(c*k_norm*dt)/(c*k_norm);
-            } else { // Handle k_norm = 0, by using the analytical limit
-                C(i,j,k) = 1.;
-                S_ck(i,j,k) = dt;
+            constexpr Real c  = PhysConst::c;
+
+            // Coefficients for k_norm = 0 do not need to be set
+            if (k_norm != 0._rt) {
+                C(i,j,k) = std::cos(c * k_norm * dt);
+                S_ck(i,j,k) = std::sin(c * k_norm * dt) / (c * k_norm);
+                inv_k2(i,j,k) = 1._rt / k2;
             }
         });
     }