Switch and rename the shift factors

4 files changed, 83 insertions, 61 deletions

diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
index 63a7c7520..be1765ca0 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
@@ -31,7 +31,7 @@ class SpectralFieldData
                       const SpectralKSpace& k_space,
                       const amrex::DistributionMapping& dm );
         SpectralFieldData() = default; // Default constructor
-        SpectralFieldData& operator=(SpectralFieldData&& field_data) = default; // Default move assignment
+        SpectralFieldData& operator=(SpectralFieldData&& field_data) = default;
         ~SpectralFieldData();
         void ForwardTransform( const amrex::MultiFab& mf,
                                const int field_index, const int i_comp );
@@ -42,10 +42,12 @@ class SpectralFieldData
         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, backward_plan;
-        // Factors that shift the fields between nodal and cell-centered, in spectral space
-        SpectralShiftFactor xshift_N2C, xshift_C2N, zshift_N2C, zshift_C2N;
+        // Correcting "shift" factors when performing FFT from/to
+        // a cell-centered grid in real space, instead of a nodal grid
+        SpectralShiftFactor xshift_FFTfromCell, xshift_FFTtoCell,
+                            zshift_FFTfromCell, zshift_FFTtoCell;
 #if (AMREX_SPACEDIM==3)
-        SpectralShiftFactor yshift_N2C, yshift_C2N;
+        SpectralShiftFactor yshift_FFTfromCell, yshift_FFTtoCell;
 #endif
         SpectralField& getSpectralField( const int field_index );
 };
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
index 31390127e..ca5e338e5 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
@@ -21,21 +21,32 @@ SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba,
     rho_old = SpectralField(spectralspace_ba, dm, 1, 0);
     rho_new = SpectralField(spectralspace_ba, dm, 1, 0);
 
-    // Allocate temporary arrays - over different boxarrays
+    // Allocate temporary arrays - in real space and spectral space
+    // These arrays will store the data just before/after the FFT
     tmpRealField = SpectralField(realspace_ba, dm, 1, 0);
     tmpSpectralField = SpectralField(spectralspace_ba, dm, 1, 0);
 
-    // 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);
+    // By default, we assume the FFT is done from/to a nodal grid in real space
+    // It the FFT is performed from/to a cell-centered grid in real space,
+    // a correcting "shift" factor must be applied in spectral space.
+    xshift_FFTfromCell = k_space.getSpectralShiftFactor(dm, 0,
+                                    ShiftType::TransformFromCellCentered);
+    xshift_FFTtoCell = k_space.getSpectralShiftFactor(dm, 0,
+                                    ShiftType::TransformToCellCentered);
 #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);
+    yshift_FFTfromCell = k_space.getSpectralShiftFactor(dm, 1,
+                                    ShiftType::TransformFromCellCentered);
+    yshift_FFTtoCell = k_space.getSpectralShiftFactor(dm, 1,
+                                    ShiftType::TransformToCellCentered);
+    zshift_FFTfromCell = k_space.getSpectralShiftFactor(dm, 2,
+                                    ShiftType::TransformFromCellCentered);
+    zshift_FFTtoCell = k_space.getSpectralShiftFactor(dm, 2,
+                                    ShiftType::TransformToCellCentered);
 #else
-    zshift_C2N = k_space.getSpectralShiftFactor(dm, 1, ShiftType::CenteredToNodal);
-    zshift_N2C = k_space.getSpectralShiftFactor(dm, 1, ShiftType::NodalToCentered);
+    zshift_FFTfromCell = k_space.getSpectralShiftFactor(dm, 1,
+                                    ShiftType::TransformFromCellCentered);
+    zshift_FFTtoCell = k_space.getSpectralShiftFactor(dm, 1,
+                                    ShiftType::TransformToCellCentered);
 #endif
 
     // Allocate and initialize the FFT plans
@@ -131,28 +142,30 @@ SpectralFieldData::ForwardTransform( const MultiFab& mf,
         fftw_execute( forward_plan[mfi] );
 #endif
 
-        // Copy the spectral-space field `tmpSpectralField` to the appropriate field
-        // (specified by the input argument field_index )
+        // Copy the spectral-space field `tmpSpectralField` to the appropriate
+        // field (specified by the input argument field_index )
+        // and apply correcting shift factor if the real space data comes
+        // from a cell-centered grid in real space instead of a nodal grid.
         {
             SpectralField& field = getSpectralField( field_index );
             Array4<Complex> field_arr = field[mfi].array();
             Array4<const Complex> tmp_arr = tmpSpectralField[mfi].array();
-            const Complex* xshift_C2N_arr = xshift_C2N[mfi].dataPtr();
+            const Complex* xshift_arr = xshift_FFTfromCell[mfi].dataPtr();
 #if (AMREX_SPACEDIM == 3)
-            const Complex* yshift_C2N_arr = yshift_C2N[mfi].dataPtr();
+            const Complex* yshift_arr = yshift_FFTfromCell[mfi].dataPtr();
 #endif
-            const Complex* zshift_C2N_arr = zshift_C2N[mfi].dataPtr();
+            const Complex* zshift_arr = zshift_FFTfromCell[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 {
                 Complex spectral_field_value = tmp_arr(i,j,k);
                 // Apply proper shift in each dimension
-                if (is_nodal_x==false) spectral_field_value *= xshift_C2N_arr[i];
+                if (is_nodal_x==false) spectral_field_value *= xshift_arr[i];
 #if (AMREX_SPACEDIM == 3)
-                if (is_nodal_y==false) spectral_field_value *= yshift_C2N_arr[j];
+                if (is_nodal_y==false) spectral_field_value *= yshift_arr[j];
 #endif
-                if (is_nodal_z==false) spectral_field_value *= zshift_C2N_arr[k];
+                if (is_nodal_z==false) spectral_field_value *= zshift_arr[k];
                 // Copy field into temporary array
                 field_arr(i,j,k) = spectral_field_value;
             });
@@ -179,28 +192,30 @@ SpectralFieldData::BackwardTransform( MultiFab& mf,
     // Loop over boxes
     for ( MFIter mfi(mf); mfi.isValid(); ++mfi ){
 
-        // Copy the appropriate field (specified by the input argument field_index)
-        // to the spectral-space field `tmpSpectralField`
+        // Copy the spectral-space field `tmpSpectralField` to the appropriate
+        // field (specified by the input argument field_index )
+        // and apply correcting shift factor if the field is to be transformed
+        // to a cell-centered grid in real space instead of a nodal grid.
         {
             SpectralField& field = getSpectralField( field_index );
             Array4<const Complex> field_arr = field[mfi].array();
             Array4<Complex> tmp_arr = tmpSpectralField[mfi].array();
-            const Complex* xshift_N2C_arr = xshift_N2C[mfi].dataPtr();
+            const Complex* xshift_arr = xshift_FFTtoCell[mfi].dataPtr();
 #if (AMREX_SPACEDIM == 3)
-            const Complex* yshift_N2C_arr = yshift_N2C[mfi].dataPtr();
+            const Complex* yshift_arr = yshift_FFTtoCell[mfi].dataPtr();
 #endif
-            const Complex* zshift_N2C_arr = zshift_N2C[mfi].dataPtr();
+            const Complex* zshift_arr = zshift_FFTtoCell[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 {
                 Complex spectral_field_value = field_arr(i,j,k);
                 // Apply proper shift in each dimension
-                if (is_nodal_x==false) spectral_field_value *= xshift_N2C_arr[i];
+                if (is_nodal_x==false) spectral_field_value *= xshift_arr[i];
 #if (AMREX_SPACEDIM == 3)
-                if (is_nodal_y==false) spectral_field_value *= yshift_N2C_arr[j];
+                if (is_nodal_y==false) spectral_field_value *= yshift_arr[j];
 #endif
-                if (is_nodal_z==false) spectral_field_value *= zshift_N2C_arr[k];
+                if (is_nodal_z==false) spectral_field_value *= zshift_arr[k];
                 // Copy field into temporary array
                 tmp_arr(i,j,k) = spectral_field_value;
             });
diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
index fd77bb7b8..c71cf2643 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
@@ -5,11 +5,11 @@
 #include <AMReX_BoxArray.H>
 #include <AMReX_LayoutData.H>
 
-using KVectorComponent = amrex::LayoutData< amrex::Gpu::ManagedVector <amrex::Real> >;
-using SpectralShiftFactor = amrex::LayoutData< amrex::Gpu::ManagedVector <Complex> >;
+using KVectorComponent = amrex::LayoutData<amrex::Gpu::ManagedVector<amrex::Real>>;
+using SpectralShiftFactor = amrex::LayoutData<amrex::Gpu::ManagedVector<Complex>>;
 
 struct ShiftType {
-    enum{ NodalToCentered=0, CenteredToNodal=1 };
+    enum{ TransformFromCellCentered=0, TransformToCellCentered=1 };
 };
 
 /* TODO Documentation: class represent k space, and how it is distribution
@@ -28,7 +28,8 @@ class SpectralKSpace
             const amrex::DistributionMapping& dm, const int i_dim,
             const int n_order, const bool nodal ) const;
         SpectralShiftFactor getSpectralShiftFactor(
-            const amrex::DistributionMapping& dm, const int i_dim, const int shift_type ) const;
+            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 111643197..4852c7d2b 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
@@ -16,10 +16,12 @@ SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba,
     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)
+        // For local FFTs, boxes in spectral space start at 0 in each direction
+        // and have the same number of points as the real space box
+        // TODO: this will be different for the hybrid FFT scheme
         Box realspace_bx = realspace_ba[i];
-        Box bx = Box( IntVect::TheZeroVector(), realspace_bx.bigEnd() - realspace_bx.smallEnd() );
+        Box bx = Box( IntVect::TheZeroVector(),
+                      realspace_bx.bigEnd() - realspace_bx.smallEnd() );
         spectral_bl.push_back( bx );
     }
     spectralspace_ba.define( spectral_bl );
@@ -31,7 +33,8 @@ SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba,
 }
 
 KVectorComponent
-SpectralKSpace::getKComponent( 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);
@@ -58,17 +61,16 @@ SpectralKSpace::getKComponent( const DistributionMapping& dm, const int i_dim )
         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
-        // This will also be different for the real-to-complex transform
+        // TODO: this will also be different for the real-to-complex transform
+        // TODO: this will be different for the hybrid FFT scheme
     }
     return k_comp;
 }
 
 SpectralShiftFactor
-SpectralKSpace::getSpectralShiftFactor(
-        const DistributionMapping& dm, const int i_dim, const int shift_type ) const
+SpectralKSpace::getSpectralShiftFactor( const DistributionMapping& dm,
+                                        const int i_dim,
+                                        const int shift_type ) const
 {
     // Initialize an empty ManagedVector in each box
     SpectralShiftFactor shift_factor = SpectralShiftFactor( spectralspace_ba, dm );
@@ -83,8 +85,8 @@ SpectralKSpace::getSpectralShiftFactor(
         // Fill the shift coefficients
         Real sign = 0;
         switch (shift_type){
-            case ShiftType::CenteredToNodal: sign = -1.; break;
-            case ShiftType::NodalToCentered: sign = 1.;
+            case ShiftType::TransformFromCellCentered: sign = 1.; break;
+            case ShiftType::TransformToCellCentered: sign = -1.;
         }
         constexpr Complex I{0,1};
         for (int i=0; i<k.size(); i++ ){
@@ -95,9 +97,10 @@ SpectralKSpace::getSpectralShiftFactor(
 }
 
 KVectorComponent
-SpectralKSpace::getModifiedKComponent(
-        const DistributionMapping& dm, const int i_dim,
-        const int n_order, const bool nodal ) const
+SpectralKSpace::getModifiedKComponent( const DistributionMapping& dm,
+                                       const int i_dim,
+                                       const int n_order,
+                                       const bool nodal ) const
 {
     // Initialize an empty ManagedVector in each box
     KVectorComponent modified_k_comp = KVectorComponent( spectralspace_ba, dm );
@@ -118,11 +121,11 @@ SpectralKSpace::getModifiedKComponent(
         for (int i=0; i<k.size(); i++ ){
             for (int n=1; n<stencil_coef.size(); n++){
                 if (nodal){
-                    modified_k[i] = \
-                        stencil_coef[n]*std::sin( k[i]*n*delta_x )/( n*delta_x );
+                    modified_k[i] = stencil_coef[n]* \
+                        std::sin( k[i]*n*delta_x )/( n*delta_x );
                 } else {
-                    modified_k[i] = \
-                        stencil_coef[n]*std::sin( k[i]*(n-0.5)*delta_x )/( (n-0.5)*delta_x );
+                    modified_k[i] = stencil_coef[n]* \
+                        std::sin( k[i]*(n-0.5)*delta_x )/( (n-0.5)*delta_x );
                 }
             }
         }
@@ -135,16 +138,17 @@ SpectralKSpace::getModifiedKComponent(
 Vector<Real>
 getFonbergStencilCoefficients( const int n_order, const bool nodal )
 {
-    AMREX_ALWAYS_ASSERT_WITH_MESSAGE( n_order%2 == 0, "n_order should be even.");
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE( n_order%2 == 0,
+                                      "n_order should be even.");
     const int m = n_order/2;
-    Vector<Real> stencil_coef;
-    stencil_coef.resize( m+1 );
+    Vector<Real> coefs;
+    coefs.resize( m+1 );
 
     // Coefficients for nodal (a.k.a. centered) finite-difference
     if (nodal == true) {
-        stencil_coef[0] = -2.; // First coefficient
+        coefs[0] = -2.; // First coefficient
         for (int n=1; n<m+1; n++){ // Get the other coefficients by recurrence
-            stencil_coef[n] = - (m+1-n)*1./(m+n)*stencil_coef[n-1];
+            coefs[n] = - (m+1-n)*1./(m+n)*coefs[n-1];
         }
     }
     // Coefficients for staggered finite-difference
@@ -153,10 +157,10 @@ getFonbergStencilCoefficients( const int n_order, const bool nodal )
         for (int k=1; k<m+1; k++){
             prod *= (m+k)*1./(4*k);
         }
-        stencil_coef[0] = 4*m*prod*prod; // First coefficient
+        coefs[0] = 4*m*prod*prod; // First coefficient
         for (int n=1; n<m+1; n++){ // Get the other coefficients by recurrence
-            stencil_coef[n] = - ((2*n-3)*(m+1-n))*1./((2*n-1)*(m-1+n))*stencil_coef[n-1];
+            coefs[n] = - ((2*n-3)*(m+1-n))*1./((2*n-1)*(m-1+n))*coefs[n-1];
         }
     }
-    return stencil_coef;
+    return coefs;
 }