Merge branch 'dev' into select_fields_in_tests

1 files changed, 188 insertions, 27 deletions

diff --git a/Source/FieldSolver/WarpXPushFieldsEM.cpp b/Source/FieldSolver/WarpXPushFieldsEM.cpp
index 4fce4717b..1df05bc0f 100644
--- a/Source/FieldSolver/WarpXPushFieldsEM.cpp
+++ b/Source/FieldSolver/WarpXPushFieldsEM.cpp
@@ -18,6 +18,40 @@
 using namespace amrex;
 
 #ifdef WARPX_USE_PSATD
+namespace {
+    void
+    PushPSATDSinglePatch (
+        SpectralSolver& solver,
+        std::array<std::unique_ptr<amrex::MultiFab>,3>& Efield,
+        std::array<std::unique_ptr<amrex::MultiFab>,3>& Bfield,
+        std::array<std::unique_ptr<amrex::MultiFab>,3>& current,
+        std::unique_ptr<amrex::MultiFab>& rho ) {
+
+        using Idx = SpectralFieldIndex;
+
+        // Perform forward Fourier transform
+        solver.ForwardTransform(*Efield[0], Idx::Ex);
+        solver.ForwardTransform(*Efield[1], Idx::Ey);
+        solver.ForwardTransform(*Efield[2], Idx::Ez);
+        solver.ForwardTransform(*Bfield[0], Idx::Bx);
+        solver.ForwardTransform(*Bfield[1], Idx::By);
+        solver.ForwardTransform(*Bfield[2], Idx::Bz);
+        solver.ForwardTransform(*current[0], Idx::Jx);
+        solver.ForwardTransform(*current[1], Idx::Jy);
+        solver.ForwardTransform(*current[2], Idx::Jz);
+        solver.ForwardTransform(*rho, Idx::rho_old, 0);
+        solver.ForwardTransform(*rho, Idx::rho_new, 1);
+        // Advance fields in spectral space
+        solver.pushSpectralFields();
+        // Perform backward Fourier Transform
+        solver.BackwardTransform(*Efield[0], Idx::Ex);
+        solver.BackwardTransform(*Efield[1], Idx::Ey);
+        solver.BackwardTransform(*Efield[2], Idx::Ez);
+        solver.BackwardTransform(*Bfield[0], Idx::Bx);
+        solver.BackwardTransform(*Bfield[1], Idx::By);
+        solver.BackwardTransform(*Bfield[2], Idx::Bz);
+    }
+}
 
 void
 WarpX::PushPSATD (amrex::Real a_dt)
@@ -31,38 +65,25 @@ WarpX::PushPSATD (amrex::Real a_dt)
         } else {
             PushPSATD_localFFT(lev, a_dt);
         }
+
+        // Evolve the fields in the PML boxes
+        if (do_pml && pml[lev]->ok()) {
+            pml[lev]->PushPSATD();
+        }
     }
 }
 
-void WarpX::PushPSATD_localFFT (int lev, amrex::Real /* dt */)
+void
+WarpX::PushPSATD_localFFT (int lev, amrex::Real /* dt */)
 {
-    auto& solver = *spectral_solver_fp[lev];
-
-    // Perform forward Fourier transform
-    solver.ForwardTransform(*Efield_fp[lev][0], SpectralFieldIndex::Ex);
-    solver.ForwardTransform(*Efield_fp[lev][1], SpectralFieldIndex::Ey);
-    solver.ForwardTransform(*Efield_fp[lev][2], SpectralFieldIndex::Ez);
-    solver.ForwardTransform(*Bfield_fp[lev][0], SpectralFieldIndex::Bx);
-    solver.ForwardTransform(*Bfield_fp[lev][1], SpectralFieldIndex::By);
-    solver.ForwardTransform(*Bfield_fp[lev][2], SpectralFieldIndex::Bz);
-    solver.ForwardTransform(*current_fp[lev][0], SpectralFieldIndex::Jx);
-    solver.ForwardTransform(*current_fp[lev][1], SpectralFieldIndex::Jy);
-    solver.ForwardTransform(*current_fp[lev][2], SpectralFieldIndex::Jz);
-    solver.ForwardTransform(*rho_fp[lev], SpectralFieldIndex::rho_old, 0);
-    solver.ForwardTransform(*rho_fp[lev], SpectralFieldIndex::rho_new, 1);
-
-    // Advance fields in spectral space
-    solver.pushSpectralFields();
-
-    // Perform backward Fourier Transform
-    solver.BackwardTransform(*Efield_fp[lev][0], SpectralFieldIndex::Ex);
-    solver.BackwardTransform(*Efield_fp[lev][1], SpectralFieldIndex::Ey);
-    solver.BackwardTransform(*Efield_fp[lev][2], SpectralFieldIndex::Ez);
-    solver.BackwardTransform(*Bfield_fp[lev][0], SpectralFieldIndex::Bx);
-    solver.BackwardTransform(*Bfield_fp[lev][1], SpectralFieldIndex::By);
-    solver.BackwardTransform(*Bfield_fp[lev][2], SpectralFieldIndex::Bz);
+    // Update the fields on the fine and coarse patch
+    PushPSATDSinglePatch( *spectral_solver_fp[lev],
+        Efield_fp[lev], Bfield_fp[lev], current_fp[lev], rho_fp[lev] );
+    if (spectral_solver_cp[lev]) {
+        PushPSATDSinglePatch( *spectral_solver_cp[lev],
+             Efield_cp[lev], Bfield_cp[lev], current_cp[lev], rho_cp[lev] );
+    }
 }
-
 #endif
 
 void
@@ -560,3 +581,143 @@ WarpX::EvolveF (int lev, PatchType patch_type, Real a_dt, DtType a_dt_type)
     }
 }
 
+#ifdef WARPX_DIM_RZ
+// This scales the current by the inverse volume and wraps around the depostion at negative radius.
+// It is faster to apply this on the grid than to do it particle by particle.
+// It is put here since there isn't another nice place for it.
+void
+WarpX::ApplyInverseVolumeScalingToCurrentDensity (MultiFab* Jx, MultiFab* Jy, MultiFab* Jz, int lev)
+{
+    const long ngJ = Jx->nGrow();
+    const std::array<Real,3>& dx = WarpX::CellSize(lev);
+    const Real dr = dx[0];
+
+    Box tilebox;
+
+    for ( MFIter mfi(*Jx, TilingIfNotGPU()); mfi.isValid(); ++mfi )
+    {
+
+        Array4<Real> const& Jr_arr = Jx->array(mfi);
+        Array4<Real> const& Jt_arr = Jy->array(mfi);
+        Array4<Real> const& Jz_arr = Jz->array(mfi);
+
+        tilebox = mfi.tilebox();
+        Box tbr = convert(tilebox, WarpX::jx_nodal_flag);
+        Box tbt = convert(tilebox, WarpX::jy_nodal_flag);
+        Box tbz = convert(tilebox, WarpX::jz_nodal_flag);
+
+        // Lower corner of tile box physical domain
+        // Note that this is done before the tilebox.grow so that
+        // these do not include the guard cells.
+        const std::array<Real, 3>& xyzmin = WarpX::LowerCorner(tilebox, lev);
+        const Dim3 lo = lbound(tilebox);
+        const Real rmin = xyzmin[0];
+        const int irmin = lo.x;
+
+        // Rescale current in r-z mode since the inverse volume factor was not
+        // included in the current deposition.
+        amrex::ParallelFor(tbr,
+        [=] AMREX_GPU_DEVICE (int i, int j, int k)
+        {
+            // Wrap the current density deposited in the guard cells around
+            // to the cells above the axis.
+            // Note that Jr(i==0) is at 1/2 dr.
+            if (rmin == 0. && 0 <= i && i < ngJ) {
+                Jr_arr(i,j,0) -= Jr_arr(-1-i,j,0);
+            }
+            // Apply the inverse volume scaling
+            // Since Jr is not cell centered in r, no need for distinction
+            // between on axis and off-axis factors
+            const amrex::Real r = std::abs(rmin + (i - irmin + 0.5)*dr);
+            Jr_arr(i,j,0) /= (2.*MathConst::pi*r);
+        });
+        amrex::ParallelFor(tbt,
+        [=] AMREX_GPU_DEVICE (int i, int j, int k)
+        {
+            // Wrap the current density deposited in the guard cells around
+            // to the cells above the axis.
+            // Jt is located on the boundary
+            if (rmin == 0. && 0 < i && i <= ngJ) {
+                Jt_arr(i,j,0) += Jt_arr(-i,j,0);
+            }
+
+            // Apply the inverse volume scaling
+            // Jt is forced to zero on axis.
+            const amrex::Real r = std::abs(rmin + (i - irmin)*dr);
+            if (r == 0.) {
+                Jt_arr(i,j,0) = 0.;
+            } else {
+                Jt_arr(i,j,0) /= (2.*MathConst::pi*r);
+            }
+        });
+        amrex::ParallelFor(tbz,
+        [=] AMREX_GPU_DEVICE (int i, int j, int k)
+        {
+            // Wrap the current density deposited in the guard cells around
+            // to the cells above the axis.
+            // Jz is located on the boundary
+            if (rmin == 0. && 0 < i && i <= ngJ) {
+                Jz_arr(i,j,0) += Jz_arr(-i,j,0);
+            }
+
+            // Apply the inverse volume scaling
+            const amrex::Real r = std::abs(rmin + (i - irmin)*dr);
+            if (r == 0.) {
+                // Verboncoeur JCP 164, 421-427 (2001) : corrected volume on axis
+                Jz_arr(i,j,0) /= (MathConst::pi*dr/3.);
+            } else {
+                Jz_arr(i,j,0) /= (2.*MathConst::pi*r);
+            }
+        });
+    }
+}
+
+void
+WarpX::ApplyInverseVolumeScalingToChargeDensity (MultiFab* Rho, int lev)
+{
+    const long ngRho = Rho->nGrow();
+    const std::array<Real,3>& dx = WarpX::CellSize(lev);
+    const Real dr = dx[0];
+
+    Box tilebox;
+
+    for ( MFIter mfi(*Rho, TilingIfNotGPU()); mfi.isValid(); ++mfi )
+    {
+
+        Array4<Real> const& Rho_arr = Rho->array(mfi);
+
+        tilebox = mfi.tilebox();
+        Box tb = convert(tilebox, IntVect::TheUnitVector());
+
+        // Lower corner of tile box physical domain
+        // Note that this is done before the tilebox.grow so that
+        // these do not include the guard cells.
+        const std::array<Real, 3>& xyzmin = WarpX::LowerCorner(tilebox, lev);
+        const Dim3 lo = lbound(tilebox);
+        const Real rmin = xyzmin[0];
+        const int irmin = lo.x;
+
+        // Rescale charge in r-z mode since the inverse volume factor was not
+        // included in the charge deposition.
+        amrex::ParallelFor(tb, Rho->nComp(),
+        [=] AMREX_GPU_DEVICE (int i, int j, int k, int icomp)
+        {
+            // Wrap the charge density deposited in the guard cells around
+            // to the cells above the axis.
+            // Rho is located on the boundary
+            if (rmin == 0. && 0 < i && i <= ngRho) {
+                Rho_arr(i,j,0,icomp) += Rho_arr(-i,j,0,icomp);
+            }
+
+            // Apply the inverse volume scaling
+            const amrex::Real r = std::abs(rmin + (i - irmin)*dr);
+            if (r == 0.) {
+                // Verboncoeur JCP 164, 421-427 (2001) : corrected volume on axis
+                Rho_arr(i,j,0,icomp) /= (MathConst::pi*dr/3.);
+            } else {
+                Rho_arr(i,j,0,icomp) /= (2.*MathConst::pi*r);
+            }
+        });
+    }
+}
+#endif