Implement spectral solver in 2D

6 files changed, 50 insertions, 15 deletions

diff --git a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H
index 375fa48bb..b54f5c143 100644
--- a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H
+++ b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H
@@ -21,7 +21,10 @@ class PsatdAlgorithm
 
     private:
         // Modified finite-order vectors
-        SpectralKVector modified_kx_vec, modified_ky_vec, modified_kz_vec;
+        SpectralKVector modified_kx_vec, modified_kz_vec;
+#if (AMREX_SPACEDIM==3)
+        SpectralKVector modified_ky_vec;
+#endif 
         SpectralCoefficients C_coef, S_ck_coef, X1_coef, X2_coef, X3_coef;
 };
 
diff --git a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp
index acc74148e..75df01343 100644
--- a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp
+++ b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp
@@ -13,7 +13,9 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
 
     // Allocate the 1D vectors
     modified_kx_vec = SpectralKVector( ba, dm );
+#if (AMREX_SPACEDIM==3)
     modified_ky_vec = SpectralKVector( ba, dm );
+#endif
     modified_kz_vec = SpectralKVector( ba, dm );
     // Allocate and fill them by computing the modified vector
     for ( MFIter mfi(ba, dm); mfi.isValid(); ++mfi ){
@@ -21,12 +23,18 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
         ComputeModifiedKVector(
             modified_kx_vec[mfi], spectral_kspace.kx_vec[mfi],
             bx, spectral_kspace.dx[0], norder_x );
+#if (AMREX_SPACEDIM==3)
         ComputeModifiedKVector(
             modified_ky_vec[mfi], spectral_kspace.ky_vec[mfi],
             bx, spectral_kspace.dx[1], norder_y );
         ComputeModifiedKVector(
             modified_kz_vec[mfi], spectral_kspace.kz_vec[mfi],
             bx, spectral_kspace.dx[2], norder_z );
+#else
+        ComputeModifiedKVector(
+            modified_kz_vec[mfi], spectral_kspace.kz_vec[mfi],
+            bx, spectral_kspace.dx[1], norder_z );
+#endif
     }
 
     // Allocate the arrays of coefficients
@@ -44,7 +52,9 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
 
         // Extract pointers for the k vectors
         const Real* modified_kx = modified_kx_vec[mfi].dataPtr();
+#if (AMREX_SPACEDIM==3)
         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();
@@ -60,7 +70,9 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
             // Calculate norm of vector
             const Real k_norm = std::sqrt(
                 std::pow( modified_kx[i], 2 ) +
+#if (AMREX_SPACEDIM==3)
                 std::pow( modified_ky[j], 2 ) +
+#endif
                 std::pow( modified_kz[k], 2 ) );
 
             // Calculate coefficients
@@ -112,7 +124,9 @@ PsatdAlgorithm::pushSpectralFields( SpectralFieldData& f ) const{
         Array4<const Real> X3_arr = X3_coef[mfi].array();
         // Extract pointers for the k vectors
         const Real* modified_kx_arr = modified_kx_vec[mfi].dataPtr();
+#if (AMREX_SPACEDIM==3)
         const Real* modified_ky_arr = modified_ky_vec[mfi].dataPtr();
+#endif
         const Real* modified_kz_arr = modified_kz_vec[mfi].dataPtr();
 
         // Loop over indices within one box
@@ -134,7 +148,11 @@ PsatdAlgorithm::pushSpectralFields( SpectralFieldData& f ) const{
             const Complex rho_new = rho_new_arr(i,j,k);
             // k vector values, and coefficients
             const Real kx = modified_kx_arr[i];
+#if (AMREX_SPACEDIM==3)
             const Real ky = modified_ky_arr[j];
+#else
+            constexpr Real ky = 0;
+#endif
             const Real kz = modified_kz_arr[k];
             constexpr Real c2 = PhysConst::c*PhysConst::c;
             constexpr Real inv_ep0 = 1./PhysConst::ep0;
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
index 595cbac16..25c08a367 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
@@ -34,19 +34,24 @@ SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba,
         // Add cuFFT-specific code
 #else
         // Create FFTW plans
-        forward_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),
+        forward_plan[mfi] =
+#if (AMREX_SPACEDIM == 3) // Swap dimensions: AMReX data is Fortran-order, but FFTW is C-order
+            fftw_plan_dft_3d( bx.length(2), bx.length(1), bx.length(0),
+#else
+            fftw_plan_dft_2d( bx.length(1), bx.length(0),
+#endif
             reinterpret_cast<fftw_complex*>( tmpRealField[mfi].dataPtr() ),
             reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ),
             FFTW_FORWARD, FFTW_ESTIMATE );
-        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),
+        backward_plan[mfi] =
+#if (AMREX_SPACEDIM == 3) // Swap dimensions: AMReX data is Fortran-order, but FFTW is C-order
+            fftw_plan_dft_3d( bx.length(2), bx.length(1), bx.length(0),
+#else
+            fftw_plan_dft_2d( bx.length(1), bx.length(0),
+#endif
             reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ),
             reinterpret_cast<fftw_complex*>( tmpRealField[mfi].dataPtr() ),
             FFTW_BACKWARD, FFTW_ESTIMATE );
-        // TODO: Add 2D code
         // TODO: Do real-to-complex transform instead of complex-to-complex
         // TODO: Let the user decide whether to use FFTW_ESTIMATE or FFTW_MEASURE
 #endif
diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
index a790be94f..16ec95cf5 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
@@ -14,15 +14,18 @@ class SpectralKSpace
     public:
         SpectralKSpace( const amrex::BoxArray& realspace_ba,
                         const amrex::DistributionMapping& dm,
-                        const amrex::Array<amrex::Real,3> dx );
+                        const amrex::RealVect dx );
         amrex::BoxArray spectralspace_ba;
-        SpectralKVector kx_vec, ky_vec, kz_vec;
-        amrex::Array<amrex::Real,3> dx;
+        SpectralKVector kx_vec, kz_vec;
+#if (AMREX_SPACEDIM==3)
+        SpectralKVector ky_vec;
+#endif
+        amrex::RealVect dx;
 };
 
 void
 AllocateAndFillKvector( amrex::Gpu::ManagedVector<amrex::Real>& k,
-    const amrex::Box& bx, const amrex::Array<amrex::Real,3> dx, const int i_dim );
+    const amrex::Box& bx, const amrex::RealVect dx, const int i_dim );
 
 void
 ComputeModifiedKVector( amrex::Gpu::ManagedVector<amrex::Real>& modified_k,
diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
index d05748192..76e8aef15 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
@@ -6,7 +6,7 @@ using namespace Gpu;
 
 SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba,
                                 const DistributionMapping& dm,
-                                const Array<Real,3> realspace_dx )
+                                const RealVect realspace_dx )
 {
     // Store the cell size
     dx = realspace_dx;
@@ -25,20 +25,26 @@ SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba,
 
     // Allocate the 1D vectors
     kx_vec = SpectralKVector( spectralspace_ba, dm );
+#if (AMREX_SPACEDIM == 3)
     ky_vec = SpectralKVector( spectralspace_ba, dm );
+#endif
     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 );
+#if (AMREX_SPACEDIM == 3)
         AllocateAndFillKvector( ky_vec[mfi], bx, dx, 1 );
         AllocateAndFillKvector( kz_vec[mfi], bx, dx, 2 );
+#else
+        AllocateAndFillKvector( kz_vec[mfi], bx, dx, 1 );
+#endif
     }
 }
 
 void
 AllocateAndFillKvector( ManagedVector<Real>& k, const Box& bx,
-                        const Array<Real,3> dx, const int i_dim )
+                        const RealVect dx, const int i_dim )
 {
     // Alllocate k to the right size
     int N = bx.length( i_dim );
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolver.H b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
index a471666b9..06d53a2a9 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolver.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
@@ -19,7 +19,7 @@ class SpectralSolver
         SpectralSolver( const amrex::BoxArray& realspace_ba,
                         const amrex::DistributionMapping& dm,
                         const int norder_x, const int norder_y, const int norder_z,
-                        const amrex::Array<amrex::Real,3> dx, const amrex::Real dt )
+                        const amrex::RealVect dx, const amrex::Real dt )
         {
             // Initialize all structures using the same distribution mapping dm
             // - Initialize k space (Contains size of each box in spectral space,