Periodic, single-box FFT (#881)

* Implement periodic-single box option for spectral

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

* Add documentation for periodic_single_box_fft

4 files changed, 43 insertions, 11 deletions

diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
index 37dfd84cb..d80ed9c96 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
@@ -55,7 +55,8 @@ class SpectralFieldData
         SpectralFieldData( const amrex::BoxArray& realspace_ba,
                            const SpectralKSpace& k_space,
                            const amrex::DistributionMapping& dm,
-                           const int n_field_required );
+                           const int n_field_required,
+                           const bool periodic_single_box );
         SpectralFieldData() = default; // Default constructor
         SpectralFieldData& operator=(SpectralFieldData&& field_data) = default;
         ~SpectralFieldData();
@@ -89,6 +90,8 @@ class SpectralFieldData
         */
         std::string cufftErrorToString (const cufftResult& err);
 #endif
+
+        bool m_periodic_single_box;
 };
 
 #endif // WARPX_SPECTRAL_FIELD_DATA_H_
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
index c3786d872..caa352b42 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
@@ -33,8 +33,11 @@ using fftw_precision_complex = fftw_complex;
 SpectralFieldData::SpectralFieldData( const amrex::BoxArray& realspace_ba,
                                       const SpectralKSpace& k_space,
                                       const amrex::DistributionMapping& dm,
-                                      const int n_field_required )
+                                      const int n_field_required,
+                                      const bool periodic_single_box )
 {
+    m_periodic_single_box = periodic_single_box;
+
     const BoxArray& spectralspace_ba = k_space.spectralspace_ba;
 
     // Allocate the arrays that contain the fields in spectral space
@@ -223,7 +226,12 @@ SpectralFieldData::ForwardTransform( const MultiFab& mf,
         // As a consequence, the copy discards the *last* point of `mf`
         // in any direction that has *nodal* index type.
         {
-            Box realspace_bx = mf[mfi].box(); // Copy the box
+            Box realspace_bx;
+            if (m_periodic_single_box) {
+                realspace_bx = mfi.validbox(); // Discard guard cells
+            } else {
+                realspace_bx = mf[mfi].box(); // Keep guard cells
+            }
             realspace_bx.enclosedCells(); // Discard last point in nodal direction
             AMREX_ALWAYS_ASSERT( realspace_bx == tmpRealField[mfi].box() );
             Array4<const Real> mf_arr = mf[mfi].array();
@@ -346,6 +354,7 @@ SpectralFieldData::BackwardTransform( MultiFab& mf,
                 // Copy field into temporary array
                 tmp_arr(i,j,k) = spectral_field_value;
             });
+
         }
 
         // Perform Fourier transform from `tmpSpectralField` to `tmpRealField`
@@ -389,11 +398,29 @@ SpectralFieldData::BackwardTransform( MultiFab& mf,
             const Box realspace_bx = tmpRealField[mfi].box();
             const Real inv_N = 1./realspace_bx.numPts();
 
-            ParallelFor( mfi.validbox(),
-            [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept {
-                // Copy and normalize field
-                mf_arr(i,j,k,i_comp) = inv_N*tmp_arr(i,j,k);
-            });
+            if (m_periodic_single_box) {
+                // Enforce periodicity on the nodes, by using modulo in indices
+                // This is because `tmp_arr` is cell-centered while `mf_arr` can be nodal
+                int const nx = realspace_bx.length(0);
+                int const ny = realspace_bx.length(1);
+#if (AMREX_SPACEDIM == 3)
+                int const nz = realspace_bx.length(2);
+#else
+                int const nz = 1;
+#endif
+                ParallelFor( mfi.validbox(),
+                [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept {
+                    // Copy and normalize field
+                    mf_arr(i,j,k,i_comp) = inv_N*tmp_arr(i%nx,j%ny,k%nz);
+                });
+            } else {
+                ParallelFor( mfi.validbox(),
+                [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept {
+                    // Copy and normalize field
+                    mf_arr(i,j,k,i_comp) = inv_N*tmp_arr(i,j,k);
+                });
+            }
+
         }
     }
 }
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolver.H b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
index 8b33707d0..08eef19ad 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolver.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
@@ -34,7 +34,8 @@ class SpectralSolver
                         const int norder_z, const bool nodal,
                         const amrex::Array<amrex::Real,3>& v_galilean,
                         const amrex::RealVect dx, const amrex::Real dt,
-                        const bool pml=false );
+                        const bool pml=false,
+                        const bool periodic_single_box=false );
 
         /**
          * \brief Transform the component `i_comp` of MultiFab `mf`
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolver.cpp b/Source/FieldSolver/SpectralSolver/SpectralSolver.cpp
index 9b79bc525..7b6ab83e8 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolver.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolver.cpp
@@ -28,6 +28,7 @@
  * \param dx       Cell size along each dimension
  * \param dt       Time step
  * \param pml      Whether the boxes in which the solver is applied are PML boxes
+ * \param periodic_single_box Whether the full simulation domain consists of a single periodic box (i.e. the global domain is not MPI parallelized)
  */
 SpectralSolver::SpectralSolver(
                 const amrex::BoxArray& realspace_ba,
@@ -36,7 +37,7 @@ SpectralSolver::SpectralSolver(
                 const int norder_z, const bool nodal,
                 const amrex::Array<amrex::Real,3>& v_galilean,
                 const amrex::RealVect dx, const amrex::Real dt,
-                const bool pml ) {
+                const bool pml, const bool periodic_single_box ) {
 
     // Initialize all structures using the same distribution mapping dm
 
@@ -64,7 +65,7 @@ SpectralSolver::SpectralSolver(
 
     // - Initialize arrays for fields in spectral space + FFT plans
     field_data = SpectralFieldData( realspace_ba, k_space, dm,
-            algorithm->getRequiredNumberOfFields() );
+            algorithm->getRequiredNumberOfFields(), periodic_single_box );
 
 }