Apply shift factors

5 files changed, 52 insertions, 24 deletions

diff --git a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp
index 95c637907..37d3b33ee 100644
--- a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp
+++ b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp
@@ -12,12 +12,12 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
     const BoxArray& ba = spectral_kspace.spectralspace_ba;
 
     // Allocate the 1D vectors
-    modified_kx_vec = spectral_kspace.AllocateAndFillModifiedKComponent( dm, 0, norder_x );
+    modified_kx_vec = spectral_kspace.getModifiedKComponent( dm, 0, norder_x );
 #if (AMREX_SPACEDIM==3)
-    modified_ky_vec = spectral_kspace.AllocateAndFillModifiedKComponent( dm, 1, norder_y );
-    modified_kz_vec = spectral_kspace.AllocateAndFillModifiedKComponent( dm, 2, norder_z );
+    modified_ky_vec = spectral_kspace.getModifiedKComponent( dm, 1, norder_y );
+    modified_kz_vec = spectral_kspace.getModifiedKComponent( dm, 2, norder_z );
 #else
-    modified_kz_vec = spectral_kspace.AllocateAndFillModifiedKComponent( dm, 1, norder_z );
+    modified_kz_vec = spectral_kspace.getModifiedKComponent( dm, 1, norder_z );
 #endif
 
     // Allocate the arrays of coefficients
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
index 052fc045e..63a7c7520 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
@@ -43,8 +43,10 @@ class SpectralFieldData
         SpectralField tmpRealField, tmpSpectralField; // Store fields before/after transform
         FFTplans forward_plan, backward_plan;
         // Factors that shift the fields between nodal and cell-centered, in spectral space
-        // (for each dimension of space)
-        amrex::Array<SpectralShiftFactor,AMREX_SPACEDIM> shift_N2C, shift_C2N;
+        SpectralShiftFactor xshift_N2C, xshift_C2N, zshift_N2C, zshift_C2N;
+#if (AMREX_SPACEDIM==3)
+        SpectralShiftFactor yshift_N2C, yshift_C2N;
+#endif
         SpectralField& getSpectralField( const int field_index );
 };
 
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
index 3fd7177e5..7d712ba03 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
@@ -25,13 +25,18 @@ SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba,
     tmpRealField = SpectralField(realspace_ba, dm, 1, 0);
     tmpSpectralField = SpectralField(spectralspace_ba, dm, 1, 0);
 
-    // Allocate the vectors that allow to shift between nodal and cell-centered
-    for (int i_dim=0; i_dim<AMREX_SPACEDIM; i_dim++) {
-        shift_C2N[i_dim] = k_space.AllocateAndFillSpectralShiftFactor(
-                            dm, i_dim, ShiftType::CenteredToNodal );
-        shift_N2C[i_dim] = k_space.AllocateAndFillSpectralShiftFactor(
-                            dm, i_dim, ShiftType::NodalToCentered );
-    }
+    // Allocate the coefficients that allow to shift between nodal and cell-centered
+    xshift_C2N = k_space.getSpectralShiftFactor(dm, 0, ShiftType::CenteredToNodal);
+    xshift_N2C = k_space.getSpectralShiftFactor(dm, 0, ShiftType::NodalToCentered);
+#if (AMREX_SPACEDIM == 3)
+    yshift_C2N = k_space.getSpectralShiftFactor(dm, 1, ShiftType::CenteredToNodal);
+    yshift_N2C = k_space.getSpectralShiftFactor(dm, 1, ShiftType::NodalToCentered);
+    zshift_C2N = k_space.getSpectralShiftFactor(dm, 2, ShiftType::CenteredToNodal);
+    zshift_N2C = k_space.getSpectralShiftFactor(dm, 2, ShiftType::NodalToCentered);
+#else
+    zshift_C2N = k_space.getSpectralShiftFactor(dm, 1, ShiftType::CenteredToNodal);
+    zshift_N2C = k_space.getSpectralShiftFactor(dm, 1, ShiftType::NodalToCentered);
+#endif
 
     // Allocate and initialize the FFT plans
     forward_plan = FFTplans(spectralspace_ba, dm);
@@ -139,6 +144,15 @@ void
 SpectralFieldData::BackwardTransform( MultiFab& mf,
                                       const int field_index, const int i_comp )
 {
+    // Check field index type, in order to apply proper shift in spectral space
+    const bool is_nodal_x = mf.is_nodal(0);
+#if (AMREX_SPACEDIM == 3)
+    const bool is_nodal_y = mf.is_nodal(1);
+    const bool is_nodal_z = mf.is_nodal(2);
+#else
+    const bool is_nodal_z = mf.is_nodal(1);
+#endif
+
     // Loop over boxes
     for ( MFIter mfi(mf); mfi.isValid(); ++mfi ){
 
@@ -148,10 +162,22 @@ SpectralFieldData::BackwardTransform( MultiFab& mf,
             SpectralField& field = getSpectralField( field_index );
             Array4<const Complex> field_arr = field[mfi].array();
             Array4<Complex> tmp_arr = tmpSpectralField[mfi].array();
+            const Complex* xshift_C2N_arr = xshift_C2N[mfi].dataPtr();
+            const Complex* yshift_C2N_arr = yshift_C2N[mfi].dataPtr();
+            const Complex* zshift_C2N_arr = zshift_C2N[mfi].dataPtr();
+            // Loop over indices within one box
             const Box spectralspace_bx = tmpSpectralField[mfi].box();
             ParallelFor( spectralspace_bx,
             [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept {
-                tmp_arr(i,j,k) = field_arr(i,j,k);
+                Complex spectral_field_value = field_arr(i,j,k);
+                // Apply proper shift in each dimension
+                if (is_nodal_x==false) spectral_field_value *= xshift_C2N_arr[i];
+#if (AMREX_SPACEDIM == 3)
+                if (is_nodal_y==false) spectral_field_value *= yshift_C2N_arr[j];
+#endif
+                if (is_nodal_z==false) spectral_field_value *= zshift_C2N_arr[k];
+                // Copy field into temporary array
+                tmp_arr(i,j,k) = spectral_field_value;
             });
         }
 
@@ -170,7 +196,8 @@ SpectralFieldData::BackwardTransform( MultiFab& mf,
         // Normalize (divide by 1/N) since the FFT result in a factor N
         {
             Box bx = mf[mfi].box();
-            const Box realspace_bx = bx.enclosedCells(); // discards last point in each nodal direction
+            const Box realspace_bx = bx.enclosedCells();
+            // `enclosedells` discards last point in each nodal direction
             AMREX_ALWAYS_ASSERT( realspace_bx == tmpRealField[mfi].box() );
             Array4<Real> mf_arr = mf[mfi].array();
             Array4<const Complex> tmp_arr = tmpRealField[mfi].array();
diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
index 90755bc0d..2f0681c3b 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
@@ -22,13 +22,12 @@ class SpectralKSpace
         SpectralKSpace( const amrex::BoxArray& realspace_ba,
                         const amrex::DistributionMapping& dm,
                         const amrex::RealVect dx );
-        KVectorComponent AllocateAndFillKComponent(
+        KVectorComponent getKComponent(
             const amrex::DistributionMapping& dm, const int i_dim ) const;
-        KVectorComponent AllocateAndFillModifiedKComponent(
+        KVectorComponent getModifiedKComponent(
             const amrex::DistributionMapping& dm, const int i_dim, const int order ) const;
-        SpectralShiftFactor AllocateAndFillSpectralShiftFactor(
-            const amrex::DistributionMapping& dm, const int i_dim, const int order ) const;
-
+        SpectralShiftFactor getSpectralShiftFactor(
+            const amrex::DistributionMapping& dm, const int i_dim, const int shift_type ) const;
 
     private:
         amrex::Array<KVectorComponent, AMREX_SPACEDIM> k_vec;
diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
index a7f726c8a..90c20b01d 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
@@ -26,12 +26,12 @@ SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba,
 
     // Allocate the components of the k vector: kx, ky (only in 3D), kz
     for (int i_dim=0; i_dim<AMREX_SPACEDIM; i_dim++) {
-        k_vec[i_dim] = AllocateAndFillKComponent( dm, i_dim );
+        k_vec[i_dim] = getKComponent( dm, i_dim );
     }
 }
 
 KVectorComponent
-SpectralKSpace::AllocateAndFillKComponent( const DistributionMapping& dm, const int i_dim ) const
+SpectralKSpace::getKComponent( const DistributionMapping& dm, const int i_dim ) const
 {
     // Initialize an empty ManagedVector in each box
     KVectorComponent k_comp = KVectorComponent(spectralspace_ba, dm);
@@ -67,7 +67,7 @@ SpectralKSpace::AllocateAndFillKComponent( const DistributionMapping& dm, const
 }
 
 KVectorComponent
-SpectralKSpace::AllocateAndFillModifiedKComponent(
+SpectralKSpace::getModifiedKComponent(
         const DistributionMapping& dm, const int i_dim, const int order ) const
 {
     // Initialize an empty ManagedVector in each box
@@ -91,7 +91,7 @@ SpectralKSpace::AllocateAndFillModifiedKComponent(
 }
 
 SpectralShiftFactor
-SpectralKSpace::AllocateAndFillSpectralShiftFactor(
+SpectralKSpace::getSpectralShiftFactor(
         const DistributionMapping& dm, const int i_dim, const int shift_type ) const
 {
     // Initialize an empty ManagedVector in each box