PSATD: add option to update E without using rho (#1128)

* Introduce option to update E with/without rho

* Clean up

* Include equations in docs

* Fix EOL whitespaces error

* Small clean-up

* Clean up

1 files changed, 79 insertions, 49 deletions

diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp
index c75e8b8c7..b8b64b80c 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp
@@ -21,10 +21,12 @@ using namespace amrex;
 PsatdAlgorithm::PsatdAlgorithm(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
-     : m_dt( dt ),
-       SpectralBaseAlgorithm( spectral_kspace, dm, norder_x, norder_y, norder_z, nodal )
+                         const int norder_z, const bool nodal, const Real dt,
+                         const bool update_with_rho)
+    // Initialize members of base class
+    : m_dt( dt ),
+      m_update_with_rho( update_with_rho ),
+      SpectralBaseAlgorithm( spectral_kspace, dm, norder_x, norder_y, norder_z, nodal )
 {
     const BoxArray& ba = spectral_kspace.spectralspace_ba;
 
@@ -45,6 +47,8 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
 void
 PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{
 
+    const bool update_with_rho = m_update_with_rho;
+
     // Loop over boxes
     for (MFIter mfi(f.fields); mfi.isValid(); ++mfi){
 
@@ -66,23 +70,24 @@ PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{
         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
+        ParallelFor( bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept
         {
-            // Record old values of the fields to be updated
             using Idx = SpectralFieldIndex;
+
             const Complex Ex_old = fields(i,j,k,Idx::Ex);
             const Complex Ey_old = fields(i,j,k,Idx::Ey);
             const Complex Ez_old = fields(i,j,k,Idx::Ez);
             const Complex Bx_old = fields(i,j,k,Idx::Bx);
             const Complex By_old = fields(i,j,k,Idx::By);
             const Complex Bz_old = fields(i,j,k,Idx::Bz);
+
             // Shortcut 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)
@@ -93,36 +98,51 @@ PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{
             const Real kz = modified_kz_arr[j];
 #endif
             constexpr Real c2 = PhysConst::c*PhysConst::c;
-            constexpr Real inv_ep0 = 1./PhysConst::ep0;
+            constexpr Real inv_eps0 = 1.0_rt/PhysConst::ep0;
+
             const Complex I = Complex{0,1};
+
             const Real C = C_arr(i,j,k);
             const Real S_ck = S_ck_arr(i,j,k);
             const Real X1 = X1_arr(i,j,k);
             const Real X2 = X2_arr(i,j,k);
             const Real X3 = X3_arr(i,j,k);
 
-
             // Update E (see WarpX online documentation: theory section)
-            fields(i,j,k,Idx::Ex) = C*Ex_old
-                        + S_ck*(c2*I*(ky*Bz_old - kz*By_old) - inv_ep0*Jx)
-                        - I*(X2*rho_new - X3*rho_old)*kx;
-            fields(i,j,k,Idx::Ey) = C*Ey_old
-                        + S_ck*(c2*I*(kz*Bx_old - kx*Bz_old) - inv_ep0*Jy)
-                        - I*(X2*rho_new - X3*rho_old)*ky;
-            fields(i,j,k,Idx::Ez) = C*Ez_old
-                        + S_ck*(c2*I*(kx*By_old - ky*Bx_old) - inv_ep0*Jz)
-                        - I*(X2*rho_new - X3*rho_old)*kz;
+
+            if (update_with_rho) {
+
+                fields(i,j,k,Idx::Ex) = C*Ex_old + S_ck*(c2*I*(ky*Bz_old-kz*By_old)-inv_eps0*Jx)
+                                        - I*(X2*rho_new-X3*rho_old)*kx;
+
+                fields(i,j,k,Idx::Ey) = C*Ey_old + S_ck*(c2*I*(kz*Bx_old-kx*Bz_old)-inv_eps0*Jy)
+                                        - I*(X2*rho_new-X3*rho_old)*ky;
+
+                fields(i,j,k,Idx::Ez) = C*Ez_old + S_ck*(c2*I*(kx*By_old-ky*Bx_old)-inv_eps0*Jz)
+                                        - I*(X2*rho_new-X3*rho_old)*kz;
+            } else {
+
+                Complex k_dot_J = kx*Jx + ky*Jy + kz*Jz;
+                Complex k_dot_E = kx*Ex_old + ky*Ey_old + kz*Ez_old;
+
+                fields(i,j,k,Idx::Ex) = C*Ex_old + S_ck*(c2*I*(ky*Bz_old-kz*By_old)-inv_eps0*Jx)
+                                        + X2*k_dot_E*kx + X3*inv_eps0*k_dot_J*kx;
+
+                fields(i,j,k,Idx::Ey) = C*Ey_old + S_ck*(c2*I*(kz*Bx_old-kx*Bz_old)-inv_eps0*Jy)
+                                        + X2*k_dot_E*ky + X3*inv_eps0*k_dot_J*ky;
+
+                fields(i,j,k,Idx::Ez) = C*Ez_old + S_ck*(c2*I*(kx*By_old-ky*Bx_old)-inv_eps0*Jz)
+                                        + X2*k_dot_E*kz + X3*inv_eps0*k_dot_J*kz;
+            }
+
             // Update B (see WarpX online documentation: theory section)
-            fields(i,j,k,Idx::Bx) = C*Bx_old
-                        - S_ck*I*(ky*Ez_old - kz*Ey_old)
-                        +   X1*I*(ky*Jz     - kz*Jy);
-            fields(i,j,k,Idx::By) = C*By_old
-                        - S_ck*I*(kz*Ex_old - kx*Ez_old)
-                        +   X1*I*(kz*Jx     - kx*Jz);
-            fields(i,j,k,Idx::Bz) = C*Bz_old
-                        - S_ck*I*(kx*Ey_old - ky*Ex_old)
-                        +   X1*I*(kx*Jy     - ky*Jx);
-        });
+
+            fields(i,j,k,Idx::Bx) = C*Bx_old - S_ck*I*(ky*Ez_old-kz*Ey_old) + X1*I*(ky*Jz-kz*Jy);
+
+            fields(i,j,k,Idx::By) = C*By_old - S_ck*I*(kz*Ex_old-kx*Ez_old) + X1*I*(kz*Jx-kx*Jz);
+
+            fields(i,j,k,Idx::Bz) = C*Bz_old - S_ck*I*(kx*Ey_old-ky*Ex_old) + X1*I*(kx*Jy-ky*Jx);
+        } );
     }
 };
 
@@ -133,8 +153,10 @@ void PsatdAlgorithm::InitializeSpectralCoefficients(const SpectralKSpace& spectr
                                     const amrex::DistributionMapping& dm,
                                     const amrex::Real dt)
 {
+    const bool update_with_rho = m_update_with_rho;
+
     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){
@@ -147,6 +169,7 @@ void PsatdAlgorithm::InitializeSpectralCoefficients(const SpectralKSpace& spectr
         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();
@@ -155,37 +178,44 @@ void PsatdAlgorithm::InitializeSpectralCoefficients(const SpectralKSpace& spectr
         Array4<Real> X3 = X3_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) +
+                std::pow(modified_kx[i],2) +
 #if (AMREX_SPACEDIM==3)
-                std::pow(modified_ky[j], 2) +
-                std::pow(modified_kz[k], 2));
+                std::pow(modified_ky[j],2) + std::pow(modified_kz[k],2));
 #else
-                std::pow(modified_kz[j], 2));
+                std::pow(modified_kz[j],2));
 #endif
-
-
             // Calculate coefficients
             constexpr Real c = PhysConst::c;
-            constexpr Real ep0 = PhysConst::ep0;
-            if (k_norm != 0){
-                C(i,j,k) = std::cos(c*k_norm*dt);
+            constexpr Real eps0 = PhysConst::ep0;
+
+            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);
-                X1(i,j,k) = (1. - C(i,j,k))/(ep0 * c*c * k_norm*k_norm);
-                X2(i,j,k) = (1. - S_ck(i,j,k)/dt)/(ep0 * k_norm*k_norm);
-                X3(i,j,k) = (C(i,j,k) - S_ck(i,j,k)/dt)/(ep0 * k_norm*k_norm);
-            } else { // Handle k_norm = 0, by using the analytical limit
-                C(i,j,k) = 1.;
+                X1(i,j,k) = (1.0_rt-C(i,j,k))/(eps0*c*c*k_norm*k_norm);
+                if (update_with_rho) {
+                    X2(i,j,k) = (1.0_rt-S_ck(i,j,k)/dt)/(eps0*k_norm*k_norm);
+                    X3(i,j,k) = (C(i,j,k)-S_ck(i,j,k)/dt)/(eps0*k_norm*k_norm);
+                } else {
+                    X2(i,j,k) = (1.0_rt-C(i,j,k))/(k_norm*k_norm);
+                    X3(i,j,k) = (S_ck(i,j,k)-dt)/(k_norm*k_norm);
+                }
+            } else { // Handle k_norm = 0 with analytical limit
+                C(i,j,k) = 1.0_rt;
                 S_ck(i,j,k) = dt;
-                X1(i,j,k) = 0.5 * dt*dt / ep0;
-                X2(i,j,k) = c*c * dt*dt / (6.*ep0);
-                X3(i,j,k) = - c*c * dt*dt / (3.*ep0);
+                X1(i,j,k) = 0.5_rt*dt*dt/eps0;
+                if (update_with_rho) {
+                    X2(i,j,k) = c*c*dt*dt/(6.0_rt*eps0);
+                    X3(i,j,k) = -c*c*dt*dt/(3.0_rt*eps0);
+                } else {
+                    X2(i,j,k) = 0.5_rt*dt*dt*c*c;
+                    X3(i,j,k) = -c*c*dt*dt*dt/6.0_rt;
+                }
             }
-        });
+        } );
      }
 }