Added shift function

5 files changed, 69 insertions, 20 deletions

diff --git a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp
index 0b1dbc4da..95c637907 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.AllocateAndFillModifiedKVector( dm, 0, norder_x );
+    modified_kx_vec = spectral_kspace.AllocateAndFillModifiedKComponent( dm, 0, norder_x );
 #if (AMREX_SPACEDIM==3)
-    modified_ky_vec = spectral_kspace.AllocateAndFillModifiedKVector( dm, 1, norder_y );
-    modified_kz_vec = spectral_kspace.AllocateAndFillModifiedKVector( dm, 2, norder_z );
+    modified_ky_vec = spectral_kspace.AllocateAndFillModifiedKComponent( dm, 1, norder_y );
+    modified_kz_vec = spectral_kspace.AllocateAndFillModifiedKComponent( dm, 2, norder_z );
 #else
-    modified_kz_vec = spectral_kspace.AllocateAndFillModifiedKVector( dm, 1, norder_z );
+    modified_kz_vec = spectral_kspace.AllocateAndFillModifiedKComponent( 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 2abe81889..052fc045e 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
@@ -42,6 +42,9 @@ 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
+        // (for each dimension of space)
+        amrex::Array<SpectralShiftFactor,AMREX_SPACEDIM> shift_N2C, shift_C2N;
         SpectralField& getSpectralField( const int field_index );
 };
 
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
index 6b39e9fa4..78b7d7c67 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
@@ -25,6 +25,14 @@ 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 and initialize the FFT plans
     forward_plan = FFTplans(spectralspace_ba, dm);
     backward_plan = FFTplans(spectralspace_ba, dm);
diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
index 366bc2a79..90755bc0d 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
@@ -1,10 +1,16 @@
 #ifndef WARPX_SPECTRAL_K_SPACE_H_
 #define WARPX_SPECTRAL_K_SPACE_H_
 
+#include <WarpX_Complex.H>
 #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> >;
+
+struct ShiftType {
+    enum{ NodalToCentered=0, CenteredToNodal=1 };
+};
 
 /* TODO Documentation: class represent k space, and how it is distribution
  * for local FFT on each MPI: k spaces are not connected.
@@ -12,14 +18,17 @@ using KVectorComponent = amrex::LayoutData< amrex::Gpu::ManagedVector <amrex::Re
 class SpectralKSpace
 {
     public:
+        amrex::BoxArray spectralspace_ba;
         SpectralKSpace( const amrex::BoxArray& realspace_ba,
                         const amrex::DistributionMapping& dm,
                         const amrex::RealVect dx );
-        KVectorComponent& AllocateAndFillKVector(
+        KVectorComponent AllocateAndFillKComponent(
             const amrex::DistributionMapping& dm, const int i_dim ) const;
-        KVectorComponent& AllocateAndFillModifiedKVector(
+        KVectorComponent AllocateAndFillModifiedKComponent(
             const amrex::DistributionMapping& dm, const int i_dim, const int order ) const;
-        amrex::BoxArray spectralspace_ba;
+        SpectralShiftFactor AllocateAndFillSpectralShiftFactor(
+            const amrex::DistributionMapping& dm, const int i_dim, const int order ) const;
+
 
     private:
         amrex::Array<KVectorComponent, AMREX_SPACEDIM> k_vec;
diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
index 6e6115cec..94b1486af 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
@@ -1,5 +1,6 @@
 #include <WarpXConst.H>
 #include <SpectralKSpace.H>
+#include <cmath>
 
 using namespace amrex;
 using namespace Gpu;
@@ -25,19 +26,19 @@ 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] = AllocateAndFillKVector( dm, i_dim );
+        k_vec[i_dim] = AllocateAndFillKComponent( dm, i_dim );
     }
 }
 
-KVectorComponent&
-SpectralKSpace::AllocateAndFillKVector( const DistributionMapping& dm, const int i_dim ) const
+KVectorComponent
+SpectralKSpace::AllocateAndFillKComponent( const DistributionMapping& dm, const int i_dim ) const
 {
-    // Initialize an empty vector in each box
-    KVectorComponent k_vec = KVectorComponent(spectralspace_ba, dm);
+    // Initialize an empty ManagedVector in each box
+    KVectorComponent k_comp = KVectorComponent(spectralspace_ba, dm);
     // Loop over boxes
     for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
         Box bx = spectralspace_ba[mfi];
-        ManagedVector<Real>& k = k_vec[mfi];
+        ManagedVector<Real>& k = k_comp[mfi];
 
         // Allocate k to the right size
         int N = bx.length( i_dim );
@@ -62,19 +63,19 @@ SpectralKSpace::AllocateAndFillKVector( const DistributionMapping& dm, const int
         // 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
     }
-    return k_vec;
+    return k_comp;
 }
 
-KVectorComponent&
-SpectralKSpace::AllocateAndFillModifiedKVector(
+KVectorComponent
+SpectralKSpace::AllocateAndFillModifiedKComponent(
         const DistributionMapping& dm, const int i_dim, const int order ) const
 {
-    // Initialize an empty vector in each box
-    KVectorComponent modified_k_vec = KVectorComponent( spectralspace_ba, dm );
+    // Initialize an empty ManagedVector in each box
+    KVectorComponent modified_k_comp = KVectorComponent( spectralspace_ba, dm );
     // Loop over boxes
     for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
         const ManagedVector<Real>& k = k_vec[i_dim][mfi];
-        ManagedVector<Real>& modified_k = modified_k_vec[mfi];
+        ManagedVector<Real>& modified_k = modified_k_comp[mfi];
 
         // Allocate modified_k to the same size as k
         modified_k.resize( k.size() );
@@ -86,5 +87,33 @@ SpectralKSpace::AllocateAndFillModifiedKVector(
             modified_k[i] = k[i];
         }
     }
-    return modified_k_vec;
+    return modified_k_comp;
+}
+
+SpectralShiftFactor
+SpectralKSpace::AllocateAndFillSpectralShiftFactor(
+        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 );
+    // Loop over boxes
+    for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
+        const ManagedVector<Real>& k = k_vec[i_dim][mfi];
+        ManagedVector<Complex>& shift = shift_factor[mfi];
+
+        // Allocate shift coefficients
+        shift.resize( k.size() );
+
+        // Fill the shift coefficients
+        Real sign = 0;
+        switch (shift_type){
+            case ShiftType::CenteredToNodal: sign = -1.;
+            case ShiftType::NodalToCentered: sign = 1.;
+        }
+        constexpr Complex I{0,1};
+        for (int i=0; i<k.size(); i++ ){
+            shift[i] = std::exp( I*sign*k[i]*0.5*dx[i_dim] );
+        }
+    }
+    return shift_factor;
 }