Use Fonberg coefficients to calculate the modified k

5 files changed, 58 insertions, 43 deletions

diff --git a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H
index 2c3d8abfd..16d27cab2 100644
--- a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H
+++ b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H
@@ -11,13 +11,13 @@ class PsatdAlgorithm
     using SpectralCoefficients = amrex::FabArray< amrex::BaseFab <amrex::Real> >;
 
     public:
-        PsatdAlgorithm( const SpectralKSpace& spectral_kspace,
+        PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
                          const amrex::DistributionMapping& dm,
                          const int norder_x, const int norder_y,
-                         const int norder_z, const amrex::Real dt);
+                         const int norder_z, const bool nodal, const amrex::Real dt);
         PsatdAlgorithm() = default; // Default constructor
         PsatdAlgorithm& operator=(PsatdAlgorithm&& algorithm) = default; // Default move assignment
-        void pushSpectralFields( SpectralFieldData& f ) const;
+        void pushSpectralFields(SpectralFieldData& f) const;
 
     private:
         // Modified finite-order vectors
diff --git a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp
index 37d3b33ee..5ebb7144d 100644
--- a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp
+++ b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp
@@ -7,29 +7,29 @@ using namespace amrex;
 PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
                          const DistributionMapping& dm,
                          const int norder_x, const int norder_y,
-                         const int norder_z, const Real dt)
+                         const int norder_z, const bool nodal, const Real dt)
 {
     const BoxArray& ba = spectral_kspace.spectralspace_ba;
 
     // Allocate the 1D vectors
-    modified_kx_vec = spectral_kspace.getModifiedKComponent( dm, 0, norder_x );
+    modified_kx_vec = spectral_kspace.getModifiedKComponent(dm, 0, norder_x, nodal);
 #if (AMREX_SPACEDIM==3)
-    modified_ky_vec = spectral_kspace.getModifiedKComponent( dm, 1, norder_y );
-    modified_kz_vec = spectral_kspace.getModifiedKComponent( dm, 2, norder_z );
+    modified_ky_vec = spectral_kspace.getModifiedKComponent(dm, 1, norder_y, nodal);
+    modified_kz_vec = spectral_kspace.getModifiedKComponent(dm, 2, norder_z, nodal);
 #else
-    modified_kz_vec = spectral_kspace.getModifiedKComponent( dm, 1, norder_z );
+    modified_kz_vec = spectral_kspace.getModifiedKComponent(dm, 1, norder_z, nodal);
 #endif
 
     // Allocate the arrays of coefficients
-    C_coef = SpectralCoefficients( ba, dm, 1, 0 );
-    S_ck_coef = SpectralCoefficients( ba, dm, 1, 0 );
-    X1_coef = SpectralCoefficients( ba, dm, 1, 0 );
-    X2_coef = SpectralCoefficients( ba, dm, 1, 0 );
-    X3_coef = SpectralCoefficients( ba, dm, 1, 0 );
+    C_coef = SpectralCoefficients(ba, dm, 1, 0);
+    S_ck_coef = SpectralCoefficients(ba, dm, 1, 0);
+    X1_coef = SpectralCoefficients(ba, dm, 1, 0);
+    X2_coef = SpectralCoefficients(ba, dm, 1, 0);
+    X3_coef = SpectralCoefficients(ba, dm, 1, 0);
 
     // Fill them with the right values:
     // Loop over boxes
-    for ( MFIter mfi(ba, dm); mfi.isValid(); ++mfi ){
+    for (MFIter mfi(ba, dm); mfi.isValid(); ++mfi){
 
         const Box& bx = ba[mfi];
 
@@ -47,26 +47,26 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
         Array4<Real> X3 = X3_coef[mfi].array();
 
         // Loop over indices within one box
-        ParallelFor( bx,
+        ParallelFor(bx,
         [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept
         {
             // Calculate norm of vector
             const Real k_norm = std::sqrt(
-                std::pow( modified_kx[i], 2 ) +
+                std::pow(modified_kx[i], 2) +
 #if (AMREX_SPACEDIM==3)
-                std::pow( modified_ky[j], 2 ) +
+                std::pow(modified_ky[j], 2) +
 #endif
-                std::pow( modified_kz[k], 2 ) );
+                std::pow(modified_kz[k], 2));
 
             // Calculate coefficients
             constexpr Real c = PhysConst::c;
             constexpr Real ep0 = PhysConst::ep0;
-            if ( k_norm != 0 ){
-                C(i,j,k) = std::cos( c*k_norm*dt );
-                S_ck(i,j,k) = std::sin( c*k_norm*dt )/( c*k_norm );
+            if (k_norm != 0){
+                C(i,j,k) = std::cos(c*k_norm*dt);
+                S_ck(i,j,k) = std::sin(c*k_norm*dt)/(c*k_norm);
                 X1(i,j,k) = (1. - C(i,j,k))/(ep0 * c*c * k_norm*k_norm);
-                X2(i,j,k) = (1. - S_ck(i,j,k)/dt )/(ep0 * k_norm*k_norm);
-                X3(i,j,k) = (C(i,j,k) - S_ck(i,j,k)/dt )/(ep0 * k_norm*k_norm);
+                X2(i,j,k) = (1. - S_ck(i,j,k)/dt)/(ep0 * k_norm*k_norm);
+                X3(i,j,k) = (C(i,j,k) - S_ck(i,j,k)/dt)/(ep0 * k_norm*k_norm);
             } else { // Handle k_norm = 0, by using the analytical limit
                 C(i,j,k) = 1.;
                 S_ck(i,j,k) = dt;
@@ -79,10 +79,10 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
 };
 
 void
-PsatdAlgorithm::pushSpectralFields( SpectralFieldData& f ) const{
+PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{
 
     // Loop over boxes
-    for ( MFIter mfi(f.Ex); mfi.isValid(); ++mfi ){
+    for (MFIter mfi(f.Ex); mfi.isValid(); ++mfi){
 
         const Box& bx = f.Ex[mfi].box();
 
@@ -113,7 +113,7 @@ PsatdAlgorithm::pushSpectralFields( SpectralFieldData& f ) const{
         const Real* modified_kz_arr = modified_kz_vec[mfi].dataPtr();
 
         // Loop over indices within one box
-        ParallelFor( bx,
+        ParallelFor(bx,
         [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept
         {
             // Record old values of the fields to be updated
@@ -148,24 +148,24 @@ PsatdAlgorithm::pushSpectralFields( SpectralFieldData& f ) const{
 
             // Update E (see WarpX online documentation: theory section)
             Ex_arr(i,j,k) = C*Ex_old
-                        + S_ck*( c2*I*(ky*Bz_old - kz*By_old) - inv_ep0*Jx )
-                        - I*( X2*rho_new - X3*rho_old )*kx;
+                        + S_ck*(c2*I*(ky*Bz_old - kz*By_old) - inv_ep0*Jx)
+                        - I*(X2*rho_new - X3*rho_old)*kx;
             Ey_arr(i,j,k) = C*Ey_old
-                        + S_ck*( c2*I*(kz*Bx_old - kx*Bz_old) - inv_ep0*Jy )
-                        - I*( X2*rho_new - X3*rho_old )*ky;
+                        + S_ck*(c2*I*(kz*Bx_old - kx*Bz_old) - inv_ep0*Jy)
+                        - I*(X2*rho_new - X3*rho_old)*ky;
             Ez_arr(i,j,k) = C*Ez_old
-                        + S_ck*( c2*I*(kx*By_old - ky*Bx_old) - inv_ep0*Jz )
-                        - I*( X2*rho_new - X3*rho_old )*kz;
+                        + S_ck*(c2*I*(kx*By_old - ky*Bx_old) - inv_ep0*Jz)
+                        - I*(X2*rho_new - X3*rho_old)*kz;
             // Update B (see WarpX online documentation: theory section)
             Bx_arr(i,j,k) = C*Bx_old
                         - S_ck*I*(ky*Ez_old - kz*Ey_old)
-                        +   X1*I*(ky*Jz     - kz*Jy );
+                        +   X1*I*(ky*Jz     - kz*Jy);
             By_arr(i,j,k) = C*By_old
                         - S_ck*I*(kz*Ex_old - kx*Ez_old)
-                        +   X1*I*(kz*Jx     - kx*Jz );
+                        +   X1*I*(kz*Jx     - kx*Jz);
             Bz_arr(i,j,k) = C*Bz_old
                         - S_ck*I*(kx*Ey_old - ky*Ex_old)
-                        +   X1*I*(kx*Jy     - ky*Jx );
+                        +   X1*I*(kx*Jy     - ky*Jx);
         });
     }
 };
diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
index 2f0681c3b..fd77bb7b8 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
@@ -25,7 +25,8 @@ class SpectralKSpace
         KVectorComponent getKComponent(
             const amrex::DistributionMapping& dm, const int i_dim ) const;
         KVectorComponent getModifiedKComponent(
-            const amrex::DistributionMapping& dm, const int i_dim, const int order ) const;
+            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;
 
@@ -36,4 +37,7 @@ class SpectralKSpace
         amrex::RealVect dx;
 };
 
+amrex::Vector<amrex::Real>
+getFonbergStencilCoefficients( const int n_order, const bool nodal );
+
 #endif
diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
index 77aa7342f..111643197 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
@@ -101,8 +101,13 @@ SpectralKSpace::getModifiedKComponent(
 {
     // Initialize an empty ManagedVector in each box
     KVectorComponent modified_k_comp = KVectorComponent( spectralspace_ba, dm );
+
+    // Compute real-space stencil coefficients
+    Vector<Real> stencil_coef = getFonbergStencilCoefficients(n_order, nodal);
+
     // Loop over boxes
     for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
+        Real delta_x = dx[i_dim];
         const ManagedVector<Real>& k = k_vec[i_dim][mfi];
         ManagedVector<Real>& modified_k = modified_k_comp[mfi];
 
@@ -111,9 +116,15 @@ SpectralKSpace::getModifiedKComponent(
 
         // Fill the modified k vector
         for (int i=0; i<k.size(); i++ ){
-            // For now, this simply copies the infinite order k
-            // TODO: Use the formula for finite-order modified k vector
-            modified_k[i] = k[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 );
+                } else {
+                    modified_k[i] = \
+                        stencil_coef[n]*std::sin( k[i]*(n-0.5)*delta_x )/( (n-0.5)*delta_x );
+                }
+            }
         }
     }
     return modified_k_comp;
@@ -121,12 +132,12 @@ SpectralKSpace::getModifiedKComponent(
 
 /* TODO: Documentation: point to Fonberg paper ; explain recurrence relation
  */
-Array<Real>
+Vector<Real>
 getFonbergStencilCoefficients( const int n_order, const bool nodal )
 {
     AMREX_ALWAYS_ASSERT_WITH_MESSAGE( n_order%2 == 0, "n_order should be even.");
     const int m = n_order/2;
-    Array<Real> stencil_coef;
+    Vector<Real> stencil_coef;
     stencil_coef.resize( m+1 );
 
     // Coefficients for nodal (a.k.a. centered) finite-difference
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolver.H b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
index 06d53a2a9..363ac7bd8 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolver.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
@@ -19,14 +19,14 @@ 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::RealVect dx, const amrex::Real dt )
+                        const bool nodal, 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,
             // and corresponding values of the k vectors)
             const SpectralKSpace k_space = SpectralKSpace( realspace_ba, dm, dx );
             // - Initialize algorithm (coefficients) over this space
-            algorithm = PsatdAlgorithm( k_space, dm, norder_x, norder_y, norder_z, dt );
+            algorithm = PsatdAlgorithm( k_space, dm, norder_x, norder_y, norder_z, nodal, dt );
             // - Initialize arrays for fields in Fourier space + FFT plans
             field_data = SpectralFieldData( realspace_ba, k_space, dm );
         };