diff options
Diffstat (limited to 'Source/FieldSolver')
| -rw-r--r-- | Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp | 59 | ||||
| -rw-r--r-- | Source/FieldSolver/SpectralSolver/SpectralFieldData.H | 13 | ||||
| -rw-r--r-- | Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp | 60 | ||||
| -rw-r--r-- | Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp | 5 | ||||
| -rw-r--r-- | Source/FieldSolver/WarpXFFT.cpp | 6 |
5 files changed, 58 insertions, 85 deletions
diff --git a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp index 56e58bcc4..ada7506c3 100644 --- a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp +++ b/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); }); diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H index c62513de6..9e31ac5b8 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H +++ b/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_ diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp index 291fe945e..6e6cc124f 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp +++ b/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 - } -} diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp index d91891a30..ddb2020d8 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp +++ b/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 ); } } diff --git a/Source/FieldSolver/WarpXFFT.cpp b/Source/FieldSolver/WarpXFFT.cpp index 4bccc956b..1cf5460f2 100644 --- a/Source/FieldSolver/WarpXFFT.cpp +++ b/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] ) ); } |