Transform rho_old and rho_new

3 files changed, 16 insertions, 10 deletions

diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
index a74bfe22b..2abe81889 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
@@ -33,8 +33,10 @@ class SpectralFieldData
         SpectralFieldData() = default; // Default constructor
         SpectralFieldData& operator=(SpectralFieldData&& field_data) = default; // Default move assignment
         ~SpectralFieldData();
-        void ForwardTransform( const amrex::MultiFab& mf, const int field_index );
-        void BackwardTransform( amrex::MultiFab& mf, const int field_index );
+        void ForwardTransform( const amrex::MultiFab& mf,
+                               const int field_index, const int i_comp );
+        void BackwardTransform( amrex::MultiFab& mf,
+                                const int field_index, const int i_comp );
 
     private:
         SpectralField Ex, Ey, Ez, Bx, By, Bz, Jx, Jy, Jz, rho_old, rho_new;
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
index 66e1e3470..ee6dbff6a 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
@@ -73,7 +73,8 @@ SpectralFieldData::~SpectralFieldData()
  * Example: ForwardTransform( Efield_cp[0], SpectralFieldIndex::Ex )
  */
 void
-SpectralFieldData::ForwardTransform( const MultiFab& mf, const int field_index )
+SpectralFieldData::ForwardTransform( const MultiFab& mf,
+                                     const int field_index, const int i_comp )
 {
     // Loop over boxes
     for ( MFIter mfi(mf); mfi.isValid(); ++mfi ){
@@ -91,7 +92,7 @@ SpectralFieldData::ForwardTransform( const MultiFab& mf, const int field_index )
             Array4<Complex> tmp_arr = tmpRealField[mfi].array();
             ParallelFor( realspace_bx,
             [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept {
-                tmp_arr(i,j,k) = mf_arr(i,j,k);
+                tmp_arr(i,j,k) = mf_arr(i,j,k,i_comp);
             });
         }
 
@@ -122,7 +123,8 @@ SpectralFieldData::ForwardTransform( const MultiFab& mf, const int field_index )
 /* TODO: Documentation
  */
 void
-SpectralFieldData::BackwardTransform( MultiFab& mf, const int field_index )
+SpectralFieldData::BackwardTransform( MultiFab& mf,
+                                      const int field_index, const int i_comp )
 {
     // Loop over boxes
     for ( MFIter mfi(mf); mfi.isValid(); ++mfi ){
@@ -160,7 +162,7 @@ SpectralFieldData::BackwardTransform( MultiFab& mf, const int field_index )
             Array4<const Complex> tmp_arr = tmpRealField[mfi].array();
             ParallelFor( realspace_bx,
             [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept {
-                mf_arr(i,j,k) = tmp_arr(i,j,k).real();
+                mf_arr(i,j,k,i_comp) = tmp_arr(i,j,k).real();
             });
         }
     }
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolver.H b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
index 521e558ba..a471666b9 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolver.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
@@ -30,11 +30,13 @@ class SpectralSolver
             // - Initialize arrays for fields in Fourier space + FFT plans
             field_data = SpectralFieldData( realspace_ba, k_space, dm );
         };
-        void ForwardTransform( const amrex::MultiFab& mf, const int field_index ){
-            field_data.ForwardTransform( mf, field_index );
+        void ForwardTransform( const amrex::MultiFab& mf,
+                               const int field_index, const int i_comp=0 ){
+            field_data.ForwardTransform( mf, field_index, i_comp );
         };
-        void BackwardTransform( amrex::MultiFab& mf, const int field_index ){
-            field_data.BackwardTransform( mf, field_index );
+        void BackwardTransform( amrex::MultiFab& mf,
+                                const int field_index, const int i_comp=0 ){
+            field_data.BackwardTransform( mf, field_index, i_comp );
         };
         void pushSpectralFields(){
             algorithm.pushSpectralFields( field_data );