Add Spectral K space

7 files changed, 125 insertions, 63 deletions

diff --git a/Source/FieldSolver/SpectralSolver/Make.package b/Source/FieldSolver/SpectralSolver/Make.package
index 44119fa92..bb54e925c 100644
--- a/Source/FieldSolver/SpectralSolver/Make.package
+++ b/Source/FieldSolver/SpectralSolver/Make.package
@@ -1,7 +1,10 @@
 CEXE_headers += WarpX_Complex.H
 CEXE_headers += SpectralData.H
+CEXE_sources += SpectralData.cpp
 CEXE_headers += PsatdSolver.H
 CEXE_sources += PsatdSolver.cpp
+CEXE_headers += SpectralKSpace.H
+CEXE_sources += SpectralKSpace.cpp
 
 INCLUDE_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver
 VPATH_LOCATIONS   += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver
diff --git a/Source/FieldSolver/SpectralSolver/PsatdSolver.H b/Source/FieldSolver/SpectralSolver/PsatdSolver.H
index 222242d2d..bb46da670 100644
--- a/Source/FieldSolver/SpectralSolver/PsatdSolver.H
+++ b/Source/FieldSolver/SpectralSolver/PsatdSolver.H
@@ -1,15 +1,12 @@
 #ifndef WARPX_PSATD_SOLVER_H_
 #define WARPX_PSATD_SOLVER_H_
 
-#include <AMReX_REAL.H>
-#include <WarpX_Complex.H>
+#include <SpectralKSpace.H>
 #include <SpectralData.H>
 
 using namespace amrex;
 using namespace Gpu;
 
-using SpectralVector = LayoutData<ManagedVector<Real>>;
-
 /* TODO: Write documentation
 */
 class PsatdSolver
@@ -17,12 +14,13 @@ class PsatdSolver
     using SpectralCoefficients = FabArray<BaseFab<Real>>;
 
     public:
-        PsatdSolver( const BoxArray& ba, const DistributionMapping& dm,
-                     const Real* dx, const Real dt );
+        PsatdSolver( const SpectralKSpace& spectral_kspace,
+                     const DistributionMapping& dm, Real dt );
         void pushSpectralFields( SpectralData& f ) const;
 
     private:
-        SpectralVector kx_vec, ky_vec, kz_vec;
+        // Modified finite-order vectors
+        SpectralKVector modified_kx_vec, modified_ky_vec, modified_kz_vec;
         SpectralCoefficients C_coef, S_ck_coef, X1_coef, X2_coef, X3_coef;
 };
 
diff --git a/Source/FieldSolver/SpectralSolver/PsatdSolver.cpp b/Source/FieldSolver/SpectralSolver/PsatdSolver.cpp
index a2021266d..66409ca33 100644
--- a/Source/FieldSolver/SpectralSolver/PsatdSolver.cpp
+++ b/Source/FieldSolver/SpectralSolver/PsatdSolver.cpp
@@ -4,54 +4,20 @@
 
 using namespace amrex;
 
-void
-AllocateAndFillKvector( ManagedVector<Real>& k, const Box& bx, const Real* dx, const int i_dim )
-{
-    // Alllocate k to the right size
-    int N = bx.length( i_dim );
-    k.resize( N );
-
-    // Fill the k vector
-    const Real PI = std::atan(1.0)*4;
-    const Real dk = 2*PI/(N*dx[i_dim]);
-    AMREX_ALWAYS_ASSERT_WITH_MESSAGE( bx.smallEnd(i_dim) == 0,
-        "Expected box to start at 0, in spectral space.");
-    AMREX_ALWAYS_ASSERT_WITH_MESSAGE( bx.bigEnd(i_dim) == N-1,
-        "Expected different box end index in spectral space.");
-    // Fill positive values of k (FFT conventions: first half is positive)
-    for (int i=0; i<(N+1)/2; i++ ){
-        k[i] = i*dk;
-    }
-    // Fill negative values of k (FFT conventions: second half is negative)
-    for (int i=(N+1)/2; i<N; i++){
-        k[i] = (N-i)*dk;
-    }
-    // TODO: This should be quite different for the hybrid spectral code:
-    // In that case we should take into consideration the actual indices of the box
-    // and distinguish the size of the local box and that of the global FFT
-    // TODO: For real-to-complex,
-
-}
-
 /*
  * ba: BoxArray for spectral space
  * dm: DistributionMapping for spectral space
  */
-PsatdSolver::PsatdSolver( const BoxArray& ba, const DistributionMapping& dm,
-                          const Real* dx, const Real dt )
+PsatdSolver::PsatdSolver(const SpectralKSpace& spectral_kspace,
+                         const DistributionMapping& dm, const Real dt)
 {
     // Allocate the 1D vectors
-    kx_vec = SpectralVector( ba, dm );
-    ky_vec = SpectralVector( ba, dm );
-    kz_vec = SpectralVector( ba, dm );
-    for ( MFIter mfi(ba, dm); mfi.isValid(); ++mfi ){
-        Box bx = ba[mfi];
-        AllocateAndFillKvector( kx_vec[mfi], bx, dx, 0 );
-        AllocateAndFillKvector( ky_vec[mfi], bx, dx, 1 );
-        AllocateAndFillKvector( kz_vec[mfi], bx, dx, 2 );
-    }
+    modified_kx_vec = spectral_kspace.getModifiedKVector( 0 );
+    modified_ky_vec = spectral_kspace.getModifiedKVector( 1 );
+    modified_kz_vec = spectral_kspace.getModifiedKVector( 2 );
 
     // Allocate the arrays of coefficients
+    const BoxArray& ba = spectral_kspace.spectralspace_ba;
     C_coef = SpectralCoefficients( ba, dm, 1, 0 );
     S_ck_coef = SpectralCoefficients( ba, dm, 1, 0 );
     X1_coef = SpectralCoefficients( ba, dm, 1, 0 );
@@ -65,9 +31,9 @@ PsatdSolver::PsatdSolver( const BoxArray& ba, const DistributionMapping& dm,
         const Box& bx = ba[mfi];
 
         // Extract pointers for the k vectors
-        const Real* kx = kx_vec[mfi].dataPtr();
-        const Real* ky = ky_vec[mfi].dataPtr();
-        const Real* kz = kz_vec[mfi].dataPtr();
+        const Real* modified_kx = modified_kx_vec[mfi].dataPtr();
+        const Real* modified_ky = modified_ky_vec[mfi].dataPtr();
+        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();
@@ -80,7 +46,10 @@ PsatdSolver::PsatdSolver( const BoxArray& ba, const DistributionMapping& dm,
         [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept
         {
             // Calculate norm of vector
-            const Real k_norm = std::sqrt( kx[i]*kx[i] + ky[j]*ky[j] + kz[k]*kz[k] );
+            const Real k_norm = std::sqrt(
+                std::pow( modified_kx[i], 2 ) +
+                std::pow( modified_ky[j], 2 ) +
+                std::pow( modified_kz[k], 2 ) );
 
             // Calculate coefficients
             constexpr Real c = PhysConst::c;
@@ -130,9 +99,9 @@ PsatdSolver::pushSpectralFields( SpectralData& f ) const{
         Array4<const Real> X2_arr = X2_coef[mfi].array();
         Array4<const Real> X3_arr = X3_coef[mfi].array();
         // Extract pointers for the k vectors
-        const Real* kx_arr = kx_vec[mfi].dataPtr();
-        const Real* ky_arr = ky_vec[mfi].dataPtr();
-        const Real* kz_arr = kz_vec[mfi].dataPtr();
+        const Real* modified_kx_arr = modified_kx_vec[mfi].dataPtr();
+        const Real* modified_ky_arr = modified_ky_vec[mfi].dataPtr();
+        const Real* modified_kz_arr = modified_kz_vec[mfi].dataPtr();
 
         // Loop over indices within one box
         ParallelFor( bx,
@@ -152,9 +121,9 @@ PsatdSolver::pushSpectralFields( SpectralData& f ) const{
             const Complex rho_old = rho_old_arr(i,j,k);
             const Complex rho_new = rho_new_arr(i,j,k);
             // k vector values, and coefficients
-            const Real kx = kx_arr[i];
-            const Real ky = ky_arr[j];
-            const Real kz = kz_arr[k];
+            const Real kx = modified_kx_arr[i];
+            const Real ky = modified_ky_arr[j];
+            const Real kz = modified_kz_arr[k];
             constexpr Real c2 = PhysConst::c*PhysConst::c;
             constexpr Real inv_ep0 = 1./PhysConst::ep0;
             constexpr Complex I = Complex{0,1};
diff --git a/Source/FieldSolver/SpectralSolver/SpectralData.H b/Source/FieldSolver/SpectralSolver/SpectralData.H
index 87237a574..3c50b1b11 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralData.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralData.H
@@ -33,12 +33,12 @@ class SpectralData
                       const DistributionMapping& dm );
         ~SpectralData();
         void ForwardTransform( const MultiFab& mf, const int field_index );
-        void InverseTransform( MultiFab& mf, const int field_index );
+        void BackwardTransform( MultiFab& mf, const int field_index );
 
     private:
         SpectralField Ex, Ey, Ez, Bx, By, Bz, Jx, Jy, Jz, rho_old, rho_new;
         SpectralField tmpRealField, tmpSpectralField; // Store fields before/after transform
-        FFTplans forward_plan, inverse_plan;
+        FFTplans forward_plan, backward_plan;
         SpectralField& getSpectralField( const int field_index );
 };
 
diff --git a/Source/FieldSolver/SpectralSolver/SpectralData.cpp b/Source/FieldSolver/SpectralSolver/SpectralData.cpp
index 0a0b8527e..b3c902b20 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralData.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralData.cpp
@@ -25,7 +25,7 @@ SpectralData::SpectralData( const BoxArray& realspace_ba,
 
     // Allocate and initialize the FFT plans
     forward_plan = FFTplans(spectralspace_ba, dm);
-    inverse_plan = FFTplans(spectralspace_ba, dm);
+    backward_plan = FFTplans(spectralspace_ba, dm);
     for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
         Box bx = spectralspace_ba[mfi];
 #ifdef AMREX_USE_GPU
@@ -38,7 +38,7 @@ SpectralData::SpectralData( const BoxArray& realspace_ba,
             reinterpret_cast<fftw_complex*>( tmpRealField[mfi].dataPtr() ),
             reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ),
             FFTW_FORWARD, FFTW_ESTIMATE );
-        inverse_plan[mfi] = fftw_plan_dft_3d(
+        backward_plan[mfi] = fftw_plan_dft_3d(
             // Swap dimensions: AMReX data is Fortran-order, but FFTW is C-order
             bx.length(2), bx.length(1), bx.length(0),
             reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ),
@@ -60,7 +60,7 @@ SpectralData::~SpectralData()
 #else
         // Destroy FFTW plans
         fftw_destroy_plan( forward_plan[mfi] );
-        fftw_destroy_plan( inverse_plan[mfi] );
+        fftw_destroy_plan( backward_plan[mfi] );
 #endif
     }
 }
@@ -156,7 +156,7 @@ SpectralData::BackwardTransform( MultiFab& mf, const int field_index )
             Array4<const Complex> tmp_arr = tmpRealField[mfi].array();
             ParallelFor( realspace_bx,
             [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept {
-                mf_arr(i,j,k) = tmp_arr(i,j,k);
+                mf_arr(i,j,k) = tmp_arr(i,j,k).real();
             });
         }
     }
diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
new file mode 100644
index 000000000..d1aea836e
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
@@ -0,0 +1,30 @@
+#ifndef WARPX_SPECTRAL_K_SPACE_H_
+#define WARPX_SPECTRAL_K_SPACE_H_
+
+#include <AMReX_BoxArray.H>
+#include <AMReX_LayoutData.H>
+
+using namespace amrex;
+using namespace Gpu;
+
+using SpectralKVector = LayoutData<ManagedVector<Real>>;
+
+/* TODO Documentation: class represent k space, and how it is distribution
+ * for local FFT on each MPI: k spaces are not connected.
+ */
+class SpectralKSpace
+{
+    public:
+        SpectralKSpace( const BoxArray& realspace_ba, const DistributionMapping& dm, const Real* dx );
+        SpectralKVector& getModifiedKVector( const int i_dim ) const;
+        BoxArray spectralspace_ba;
+
+    private:
+        SpectralKVector kx_vec, ky_vec, kz_vec;
+        const Real* dx;
+};
+
+void
+AllocateAndFillKvector( ManagedVector<Real>& k, const Box& bx, const Real* dx, const int i_dim );
+
+#endif
diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
new file mode 100644
index 000000000..ab684444d
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
@@ -0,0 +1,62 @@
+#include <WarpXConst.H>
+#include <SpectralKSpace.H>
+
+SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba,
+                                const DistributionMapping& dm,
+                                const Real* realspace_dx )
+{
+    // Create the box array that corresponds to spectral space
+    BoxList spectral_bl; // Create empty box list
+    // Loop over boxes and fill the box list
+    for (int i=0; i < realspace_ba.size(); i++ ) {
+        // For local FFTs, each box in spectral space starts at 0 in each direction
+        // and has the same number of points as the real space box (including guard cells)
+        Box realspace_bx = realspace_ba[i];
+        Box bx = Box( IntVect::TheZeroVector(), realspace_bx.bigEnd() - realspace_bx.smallEnd() );
+        spectral_bl.push_back( bx );
+    }
+    spectralspace_ba.define( spectral_bl );
+
+    // Allocate the 1D vectors
+    kx_vec = SpectralKVector( spectralspace_ba, dm );
+    ky_vec = SpectralKVector( spectralspace_ba, dm );
+    kz_vec = SpectralKVector( spectralspace_ba, dm );
+    // Initialize the values on each box
+    for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
+        Box bx = spectralspace_ba[mfi];
+        AllocateAndFillKvector( kx_vec[mfi], bx, dx, 0 );
+        AllocateAndFillKvector( ky_vec[mfi], bx, dx, 1 );
+        AllocateAndFillKvector( kz_vec[mfi], bx, dx, 2 );
+    }
+
+    // Store the cell size
+    dx = realspace_dx;
+}
+
+void
+AllocateAndFillKvector( ManagedVector<Real>& k, const Box& bx, const Real* dx, const int i_dim )
+{
+    // Alllocate k to the right size
+    int N = bx.length( i_dim );
+    k.resize( N );
+
+    // Fill the k vector
+    const Real dk = 2*MathConst::pi/(N*dx[i_dim]);
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE( bx.smallEnd(i_dim) == 0,
+        "Expected box to start at 0, in spectral space.");
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE( bx.bigEnd(i_dim) == N-1,
+        "Expected different box end index in spectral space.");
+    // Fill positive values of k (FFT conventions: first half is positive)
+    for (int i=0; i<(N+1)/2; i++ ){
+        k[i] = i*dk;
+    }
+    // Fill negative values of k (FFT conventions: second half is negative)
+    for (int i=(N+1)/2; i<N; i++){
+        k[i] = (N-i)*dk;
+    }
+    // TODO: This should be quite different for the hybrid spectral code:
+    // In that case we should take into consideration the actual indices of the box
+    // and distinguish the size of the local box and that of the global FFT
+    // TODO: For real-to-complex,
+
+}