Merge branch 'dev' of github.com:ECP-WarpX/WarpX into dev

Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp

@@ -56,8 +56,10 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace,
                 std::pow(modified_kx[i], 2) +
 #if (AMREX_SPACEDIM==3)
                 std::pow(modified_ky[j], 2) +
-#endif
                 std::pow(modified_kz[k], 2));
+#else
+                std::pow(modified_kz[j], 2));
+#endif
 
             // Calculate coefficients
             constexpr Real c = PhysConst::c;
@@ -85,23 +87,12 @@ void
 PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{
 
     // Loop over boxes
-    for (MFIter mfi(f.Ex); mfi.isValid(); ++mfi){
+    for (MFIter mfi(f.fields); mfi.isValid(); ++mfi){
 
-        const Box& bx = f.Ex[mfi].box();
+        const Box& bx = f.fields[mfi].box();
 
         // Extract arrays for the fields to be updated
-        Array4<Complex> Ex_arr = f.Ex[mfi].array();
-        Array4<Complex> Ey_arr = f.Ey[mfi].array();
-        Array4<Complex> Ez_arr = f.Ez[mfi].array();
-        Array4<Complex> Bx_arr = f.Bx[mfi].array();
-        Array4<Complex> By_arr = f.By[mfi].array();
-        Array4<Complex> Bz_arr = f.Bz[mfi].array();
-        // Extract arrays for J and rho
-        Array4<const Complex> Jx_arr = f.Jx[mfi].array();
-        Array4<const Complex> Jy_arr = f.Jy[mfi].array();
-        Array4<const Complex> Jz_arr = f.Jz[mfi].array();
-        Array4<const Complex> rho_old_arr = f.rho_old[mfi].array();
-        Array4<const Complex> rho_new_arr = f.rho_new[mfi].array();
+        Array4<Complex> fields = f.fields[mfi].array();
         // Extract arrays for the coefficients
         Array4<const Real> C_arr = C_coef[mfi].array();
         Array4<const Real> S_ck_arr = S_ck_coef[mfi].array();
@@ -120,26 +111,28 @@ PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{
         [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept
         {
             // Record old values of the fields to be updated
-            const Complex Ex_old = Ex_arr(i,j,k);
-            const Complex Ey_old = Ey_arr(i,j,k);
-            const Complex Ez_old = Ez_arr(i,j,k);
-            const Complex Bx_old = Bx_arr(i,j,k);
-            const Complex By_old = By_arr(i,j,k);
-            const Complex Bz_old = Bz_arr(i,j,k);
+            using Idx = SpectralFieldIndex;
+            const Complex Ex_old = fields(i,j,k,Idx::Ex);
+            const Complex Ey_old = fields(i,j,k,Idx::Ey);
+            const Complex Ez_old = fields(i,j,k,Idx::Ez);
+            const Complex Bx_old = fields(i,j,k,Idx::Bx);
+            const Complex By_old = fields(i,j,k,Idx::By);
+            const Complex Bz_old = fields(i,j,k,Idx::Bz);
             // Shortcut for the values of J and rho
-            const Complex Jx = Jx_arr(i,j,k);
-            const Complex Jy = Jy_arr(i,j,k);
-            const Complex Jz = Jz_arr(i,j,k);
-            const Complex rho_old = rho_old_arr(i,j,k);
-            const Complex rho_new = rho_new_arr(i,j,k);
+            const Complex Jx = fields(i,j,k,Idx::Jx);
+            const Complex Jy = fields(i,j,k,Idx::Jy);
+            const Complex Jz = fields(i,j,k,Idx::Jz);
+            const Complex rho_old = fields(i,j,k,Idx::rho_old);
+            const Complex rho_new = fields(i,j,k,Idx::rho_new);
             // k vector values, and coefficients
             const Real kx = modified_kx_arr[i];
 #if (AMREX_SPACEDIM==3)
             const Real ky = modified_ky_arr[j];
+            const Real kz = modified_kz_arr[k];
 #else
             constexpr Real ky = 0;
+            const Real kz = modified_kz_arr[j];
 #endif
-            const Real kz = modified_kz_arr[k];
             constexpr Real c2 = PhysConst::c*PhysConst::c;
             constexpr Real inv_ep0 = 1./PhysConst::ep0;
             constexpr Complex I = Complex{0,1};
@@ -150,23 +143,23 @@ PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{
             const Real X3 = X3_arr(i,j,k);
 
             // Update E (see WarpX online documentation: theory section)
-            Ex_arr(i,j,k) = C*Ex_old
+            fields(i,j,k,Idx::Ex) = C*Ex_old
                         + 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
+            fields(i,j,k,Idx::Ey) = C*Ey_old
                         + 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
+            fields(i,j,k,Idx::Ez) = C*Ez_old
                         + 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
+            fields(i,j,k,Idx::Bx) = C*Bx_old
                         - S_ck*I*(ky*Ez_old - kz*Ey_old)
                         +   X1*I*(ky*Jz     - kz*Jy);
-            By_arr(i,j,k) = C*By_old
+            fields(i,j,k,Idx::By) = C*By_old
                         - S_ck*I*(kz*Ex_old - kx*Ez_old)
                         +   X1*I*(kz*Jx     - kx*Jz);
-            Bz_arr(i,j,k) = C*Bz_old
+            fields(i,j,k,Idx::Bz) = C*Bz_old
                         - S_ck*I*(kx*Ey_old - ky*Ex_old)
                         +   X1*I*(kx*Jy     - ky*Jx);
         });

Source/FieldSolver/SpectralSolver/SpectralFieldData.H

@@ -9,8 +9,9 @@
 using SpectralField = amrex::FabArray< amrex::BaseFab <Complex> >;
 
 /* Index for the fields that will be stored in spectral space */
-struct SpectralFieldIndex{
-    enum { Ex=0, Ey, Ez, Bx, By, Bz, Jx, Jy, Jz, rho_old, rho_new };
+struct SpectralFieldIndex {
+  enum { Ex=0, Ey, Ez, Bx, By, Bz, Jx, Jy, Jz, rho_old, rho_new, n_fields };
+  // n_fields is automatically the total number of fields
 };
 
 /* \brief Class that stores the fields in spectral space, and performs the
@@ -37,12 +38,13 @@ class SpectralFieldData
         SpectralFieldData& operator=(SpectralFieldData&& field_data) = default;
         ~SpectralFieldData();
         void ForwardTransform( const amrex::MultiFab& mf,
-                               const int field_index, const int i_comp );
+                               const int field_index, const int i_comp);
         void BackwardTransform( amrex::MultiFab& mf,
-                                const int field_index, const int i_comp );
+                               const int field_index, const int i_comp);
 
     private:
-        SpectralField Ex, Ey, Ez, Bx, By, Bz, Jx, Jy, Jz, rho_old, rho_new;
+        // `fields` stores fields in spectral space, as multicomponent FabArray
+        SpectralField fields;
         // tmpRealField and tmpSpectralField store fields
         // right before/after the Fourier transform
         SpectralField tmpRealField, tmpSpectralField;
@@ -54,7 +56,6 @@ class SpectralFieldData
 #if (AMREX_SPACEDIM==3)
         SpectralShiftFactor yshift_FFTfromCell, yshift_FFTtoCell;
 #endif
-        SpectralField& getSpectralField( const int field_index );
 };
 
 #endif // WARPX_SPECTRAL_FIELD_DATA_H_

Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp

@@ -10,17 +10,9 @@ SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba,
     const BoxArray& spectralspace_ba = k_space.spectralspace_ba;
 
     // Allocate the arrays that contain the fields in spectral space
-    Ex = SpectralField(spectralspace_ba, dm, 1, 0);
-    Ey = SpectralField(spectralspace_ba, dm, 1, 0);
-    Ez = SpectralField(spectralspace_ba, dm, 1, 0);
-    Bx = SpectralField(spectralspace_ba, dm, 1, 0);
-    By = SpectralField(spectralspace_ba, dm, 1, 0);
-    Bz = SpectralField(spectralspace_ba, dm, 1, 0);
-    Jx = SpectralField(spectralspace_ba, dm, 1, 0);
-    Jy = SpectralField(spectralspace_ba, dm, 1, 0);
-    Jz = SpectralField(spectralspace_ba, dm, 1, 0);
-    rho_old = SpectralField(spectralspace_ba, dm, 1, 0);
-    rho_new = SpectralField(spectralspace_ba, dm, 1, 0);
+    // (one component per field)
+    fields = SpectralField(spectralspace_ba, dm,
+                            SpectralFieldIndex::n_fields, 0);
 
     // Allocate temporary arrays - in real space and spectral space
     // These arrays will store the data just before/after the FFT
@@ -147,12 +139,11 @@ SpectralFieldData::ForwardTransform( const MultiFab& mf,
 #endif
 
         // Copy the spectral-space field `tmpSpectralField` to the appropriate
-        // field (specified by the input argument field_index )
+        // index of the FabArray `fields` (specified by `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<Complex> fields_arr = SpectralFieldData::fields[mfi].array();
             Array4<const Complex> tmp_arr = tmpSpectralField[mfi].array();
             const Complex* xshift_arr = xshift_FFTfromCell[mfi].dataPtr();
 #if (AMREX_SPACEDIM == 3)
@@ -168,10 +159,12 @@ SpectralFieldData::ForwardTransform( const MultiFab& mf,
                 if (is_nodal_x==false) spectral_field_value *= xshift_arr[i];
 #if (AMREX_SPACEDIM == 3)
                 if (is_nodal_y==false) spectral_field_value *= yshift_arr[j];
-#endif
                 if (is_nodal_z==false) spectral_field_value *= zshift_arr[k];
-                // Copy field into temporary array
-                field_arr(i,j,k) = spectral_field_value;
+#elif (AMREX_SPACEDIM == 2)
+                if (is_nodal_z==false) spectral_field_value *= zshift_arr[j];
+#endif
+                // Copy field into the right index
+                fields_arr(i,j,k,field_index) = spectral_field_value;
             });
         }
     }
@@ -182,7 +175,8 @@ SpectralFieldData::ForwardTransform( const MultiFab& mf,
  * real space, and store it in the component `i_comp` of `mf` */
 void
 SpectralFieldData::BackwardTransform( MultiFab& mf,
-                                      const int field_index, const int i_comp )
+                                      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);
@@ -202,8 +196,7 @@ SpectralFieldData::BackwardTransform( MultiFab& mf,
         // to a cell-centered grid in real space instead of a nodal grid.
         // Normalize (divide by 1/N) since the FFT+IFFT results in a factor N
         {
-            SpectralField& field = getSpectralField( field_index );
-            Array4<const Complex> field_arr = field[mfi].array();
+            Array4<const Complex> field_arr = SpectralFieldData::fields[mfi].array();
             Array4<Complex> tmp_arr = tmpSpectralField[mfi].array();
             const Complex* xshift_arr = xshift_FFTtoCell[mfi].dataPtr();
 #if (AMREX_SPACEDIM == 3)
@@ -216,13 +209,15 @@ SpectralFieldData::BackwardTransform( MultiFab& mf,
             const Real inv_N = 1./spectralspace_bx.numPts();
             ParallelFor( spectralspace_bx,
             [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept {
-                Complex spectral_field_value = field_arr(i,j,k);
+                Complex spectral_field_value = field_arr(i,j,k,field_index);
                 // Apply proper shift in each dimension
                 if (is_nodal_x==false) spectral_field_value *= xshift_arr[i];
 #if (AMREX_SPACEDIM == 3)
                 if (is_nodal_y==false) spectral_field_value *= yshift_arr[j];
-#endif
                 if (is_nodal_z==false) spectral_field_value *= zshift_arr[k];
+#elif (AMREX_SPACEDIM == 2)
+                if (is_nodal_z==false) spectral_field_value *= zshift_arr[j];
+#endif
                 // Copy field into temporary array (after normalization)
                 tmp_arr(i,j,k) = inv_N*spectral_field_value;
             });
@@ -248,24 +243,3 @@ SpectralFieldData::BackwardTransform( MultiFab& mf,
         }
     }
 }
-
-
-SpectralField&
-SpectralFieldData::getSpectralField( const int field_index )
-{
-    switch(field_index)
-    {
-        case SpectralFieldIndex::Ex : return Ex; break;
-        case SpectralFieldIndex::Ey : return Ey; break;
-        case SpectralFieldIndex::Ez : return Ez; break;
-        case SpectralFieldIndex::Bx : return Bx; break;
-        case SpectralFieldIndex::By : return By; break;
-        case SpectralFieldIndex::Bz : return Bz; break;
-        case SpectralFieldIndex::Jx : return Jx; break;
-        case SpectralFieldIndex::Jy : return Jy; break;
-        case SpectralFieldIndex::Jz : return Jz; break;
-        case SpectralFieldIndex::rho_old : return rho_old; break;
-        case SpectralFieldIndex::rho_new : return rho_new; break;
-        default : return tmpSpectralField; // For synthax; should not occur in practice
-    }
-}

Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp

@@ -161,12 +161,13 @@ SpectralKSpace::getModifiedKComponent( const DistributionMapping& dm,
 
         // Fill the modified k vector
         for (int i=0; i<k.size(); i++ ){
+            modified_k[i] = 0;
             for (int n=1; n<stencil_coef.size(); n++){
                 if (nodal){
-                    modified_k[i] = stencil_coef[n]* \
+                    modified_k[i] += stencil_coef[n]* \
                         std::sin( k[i]*n*delta_x )/( n*delta_x );
                 } else {
-                    modified_k[i] = stencil_coef[n]* \
+                    modified_k[i] += stencil_coef[n]* \
                         std::sin( k[i]*(n-0.5)*delta_x )/( (n-0.5)*delta_x );
                 }
             }

Source/FieldSolver/WarpXFFT.cpp

@@ -130,7 +130,11 @@ WarpX::AllocLevelDataFFT (int lev)
         // Allocate and initialize objects for the spectral solver
         // (all use the same distribution mapping)
         std::array<Real,3> dx = CellSize(lev);
-        RealVect dx_vect = RealVect( AMREX_D_DECL(dx[0], dx[1], dx[2]) );
+#if (AMREX_SPACEDIM == 3)
+        RealVect dx_vect(dx[0], dx[1], dx[2]);
+#elif (AMREX_SPACEDIM == 2)
+        RealVect dx_vect(dx[0], dx[2]);
+#endif
         spectral_solver_fp[lev].reset( new SpectralSolver( ba_fp_fft, dm_fp_fft,
                  nox_fft, noy_fft, noz_fft, do_nodal, dx_vect, dt[lev] ) );
     }