diff options
Diffstat (limited to 'Source/FieldSolver/SpectralSolver')
12 files changed, 996 insertions, 0 deletions
diff --git a/Source/FieldSolver/SpectralSolver/Make.package b/Source/FieldSolver/SpectralSolver/Make.package new file mode 100644 index 000000000..b0ee54095 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/Make.package @@ -0,0 +1,12 @@ +CEXE_headers += WarpX_Complex.H +CEXE_headers += SpectralSolver.H +CEXE_sources += SpectralSolver.cpp +CEXE_headers += SpectralFieldData.H +CEXE_sources += SpectralFieldData.cpp +CEXE_headers += SpectralKSpace.H +CEXE_sources += SpectralKSpace.cpp + +include $(WARPX_HOME)/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package + +INCLUDE_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver +VPATH_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package new file mode 100644 index 000000000..c62c21f44 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package @@ -0,0 +1,6 @@ +CEXE_headers += SpectralBaseAlgorithm.H +CEXE_headers += PsatdAlgorithm.H +CEXE_sources += PsatdAlgorithm.cpp + +INCLUDE_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver/SpectralAlgorithms +VPATH_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver/SpectralAlgorithms diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.H new file mode 100644 index 000000000..0487e5226 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.H @@ -0,0 +1,24 @@ +#ifndef WARPX_PSATD_ALGORITHM_H_ +#define WARPX_PSATD_ALGORITHM_H_ + +#include <SpectralBaseAlgorithm.H> + +/* \brief Class that updates the field in spectral space + * and stores the coefficients of the corresponding update equation. + */ +class PsatdAlgorithm : public SpectralBaseAlgorithm +{ + public: + PsatdAlgorithm(const SpectralKSpace& spectral_kspace, + const amrex::DistributionMapping& dm, + const int norder_x, const int norder_y, + const int norder_z, const bool nodal, + const amrex::Real dt); + // Redefine update equation from base class + virtual void pushSpectralFields(SpectralFieldData& f) const override final; + + private: + SpectralCoefficients C_coef, S_ck_coef, X1_coef, X2_coef, X3_coef; +}; + +#endif // WARPX_PSATD_ALGORITHM_H_ diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp new file mode 100644 index 000000000..37892d35a --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp @@ -0,0 +1,162 @@ +#include <PsatdAlgorithm.H> +#include <WarpXConst.H> +#include <cmath> + +using namespace amrex; + +/* \brief Initialize coefficients for the update equation */ +PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace, + const DistributionMapping& dm, + const int norder_x, const int norder_y, + const int norder_z, const bool nodal, const Real dt) + // Initialize members of base class + : SpectralBaseAlgorithm( spectral_kspace, dm, + norder_x, norder_y, norder_z, nodal ) +{ + const BoxArray& ba = spectral_kspace.spectralspace_ba; + + // 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); + + // Fill them with the right values: + // Loop over boxes and allocate the corresponding coefficients + // for each box owned by the local MPI proc + for (MFIter mfi(ba, dm); mfi.isValid(); ++mfi){ + + const Box& bx = ba[mfi]; + + // Extract pointers for the k vectors + const Real* modified_kx = modified_kx_vec[mfi].dataPtr(); +#if (AMREX_SPACEDIM==3) + const Real* modified_ky = modified_ky_vec[mfi].dataPtr(); +#endif + const Real* modified_kz = modified_kz_vec[mfi].dataPtr(); + // Extract arrays for the coefficients + Array4<Real> C = C_coef[mfi].array(); + Array4<Real> S_ck = S_ck_coef[mfi].array(); + Array4<Real> X1 = X1_coef[mfi].array(); + Array4<Real> X2 = X2_coef[mfi].array(); + Array4<Real> X3 = X3_coef[mfi].array(); + + // Loop over indices within one box + 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) + +#if (AMREX_SPACEDIM==3) + std::pow(modified_ky[j], 2) + + std::pow(modified_kz[k], 2)); +#else + std::pow(modified_kz[j], 2)); +#endif + + // 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); + 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); + } else { // Handle k_norm = 0, by using the analytical limit + C(i,j,k) = 1.; + S_ck(i,j,k) = dt; + X1(i,j,k) = 0.5 * dt*dt / ep0; + X2(i,j,k) = c*c * dt*dt / (6.*ep0); + X3(i,j,k) = - c*c * dt*dt / (3.*ep0); + } + }); + } +}; + +/* Advance the E and B field in spectral space (stored in `f`) + * over one time step */ +void +PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{ + + // Loop over boxes + for (MFIter mfi(f.fields); mfi.isValid(); ++mfi){ + + const Box& bx = f.fields[mfi].box(); + + // Extract arrays for the fields to be updated + 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(); + Array4<const Real> X1_arr = X1_coef[mfi].array(); + Array4<const Real> X2_arr = X2_coef[mfi].array(); + Array4<const Real> X3_arr = X3_coef[mfi].array(); + // Extract pointers for the k vectors + const Real* modified_kx_arr = modified_kx_vec[mfi].dataPtr(); +#if (AMREX_SPACEDIM==3) + const Real* modified_ky_arr = modified_ky_vec[mfi].dataPtr(); +#endif + const Real* modified_kz_arr = modified_kz_vec[mfi].dataPtr(); + + // Loop over indices within one box + ParallelFor(bx, + [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept + { + // Record old values of the fields to be updated + 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 = 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 + constexpr Real c2 = PhysConst::c*PhysConst::c; + constexpr Real inv_ep0 = 1./PhysConst::ep0; + constexpr Complex I = Complex{0,1}; + const Real C = C_arr(i,j,k); + const Real S_ck = S_ck_arr(i,j,k); + const Real X1 = X1_arr(i,j,k); + const Real X2 = X2_arr(i,j,k); + const Real X3 = X3_arr(i,j,k); + + // Update E (see WarpX online documentation: theory section) + 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; + 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; + 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) + fields(i,j,k,Idx::Bx) = C*Bx_old + - S_ck*I*(ky*Ez_old - kz*Ey_old) + + X1*I*(ky*Jz - kz*Jy); + fields(i,j,k,Idx::By) = C*By_old + - S_ck*I*(kz*Ex_old - kx*Ez_old) + + X1*I*(kz*Jx - kx*Jz); + 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/SpectralAlgorithms/SpectralBaseAlgorithm.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.H new file mode 100644 index 000000000..602eb2473 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithm.H @@ -0,0 +1,51 @@ +#ifndef WARPX_SPECTRAL_BASE_ALGORITHM_H_ +#define WARPX_SPECTRAL_BASE_ALGORITHM_H_ + +#include <SpectralKSpace.H> +#include <SpectralFieldData.H> + +/* \brief Class that updates the field in spectral space + * and stores the coefficients of the corresponding update equation. + * + * `SpectralBaseAlgorithm` is only a base class and cannot be used directly. + * Instead use its subclasses, which implement the specific field update + * equations for a given spectral algorithm. + */ +class SpectralBaseAlgorithm +{ + public: + // Member function that updates the fields in spectral space ; + // meant to be overridden in subclasses + virtual void pushSpectralFields(SpectralFieldData& f) const = 0; + // The destructor should also be a virtual function, so that + // a pointer to subclass of `SpectraBaseAlgorithm` actually + // calls the subclass's destructor. + virtual ~SpectralBaseAlgorithm() {}; + + protected: // Meant to be used in the subclasses + + using SpectralCoefficients = amrex::FabArray< amrex::BaseFab <amrex::Real> >; + + // Constructor + SpectralBaseAlgorithm(const SpectralKSpace& spectral_kspace, + const amrex::DistributionMapping& dm, + const int norder_x, const int norder_y, + const int norder_z, const bool nodal) + // Compute and assign the modified k vectors + : 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,nodal)), + modified_kz_vec(spectral_kspace.getModifiedKComponent(dm,2,norder_z,nodal)) +#else + modified_kz_vec(spectral_kspace.getModifiedKComponent(dm,1,norder_z,nodal)) +#endif + {}; + + // Modified finite-order vectors + KVectorComponent modified_kx_vec, modified_kz_vec; +#if (AMREX_SPACEDIM==3) + KVectorComponent modified_ky_vec; +#endif +}; + +#endif // WARPX_SPECTRAL_BASE_ALGORITHM_H_ diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H new file mode 100644 index 000000000..8e58aa1d8 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H @@ -0,0 +1,62 @@ +#ifndef WARPX_SPECTRAL_FIELD_DATA_H_ +#define WARPX_SPECTRAL_FIELD_DATA_H_ + +#include <WarpX_Complex.H> +#include <SpectralKSpace.H> +#include <AMReX_MultiFab.H> + +// Declare type for spectral fields +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, n_fields }; + // n_fields is automatically the total number of fields +}; + +/* \brief Class that stores the fields in spectral space, and performs the + * Fourier transforms between real space and spectral space + */ +class SpectralFieldData +{ + friend class PsatdAlgorithm; + + // Define the FFTplans type, which holds one fft plan per box + // (plans are only initialized for the boxes that are owned by + // the local MPI rank) +#ifdef AMREX_USE_GPU + // Add cuFFT-specific code +#else + using FFTplans = amrex::LayoutData<fftw_plan>; +#endif + + public: + SpectralFieldData( const amrex::BoxArray& realspace_ba, + const SpectralKSpace& k_space, + const amrex::DistributionMapping& dm ); + SpectralFieldData() = default; // Default constructor + SpectralFieldData& operator=(SpectralFieldData&& field_data) = default; + ~SpectralFieldData(); + 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: + // `fields` stores fields in spectral space, as multicomponent FabArray + SpectralField fields; + // tmpRealField and tmpSpectralField store fields + // right before/after the Fourier transform + SpectralField tmpSpectralField; // contains Complexs + amrex::MultiFab tmpRealField; // contains Reals + FFTplans forward_plan, backward_plan; + // Correcting "shift" factors when performing FFT from/to + // a cell-centered grid in real space, instead of a nodal grid + SpectralShiftFactor xshift_FFTfromCell, xshift_FFTtoCell, + zshift_FFTfromCell, zshift_FFTtoCell; +#if (AMREX_SPACEDIM==3) + SpectralShiftFactor yshift_FFTfromCell, yshift_FFTtoCell; +#endif +}; + +#endif // WARPX_SPECTRAL_FIELD_DATA_H_ diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp new file mode 100644 index 000000000..02fa2015f --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp @@ -0,0 +1,250 @@ +#include <SpectralFieldData.H> + +using namespace amrex; + +/* \brief Initialize fields in spectral space, and FFT plans */ +SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba, + const SpectralKSpace& k_space, + const DistributionMapping& dm ) +{ + const BoxArray& spectralspace_ba = k_space.spectralspace_ba; + + // Allocate the arrays that contain the fields in spectral space + // (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 + tmpRealField = MultiFab(realspace_ba, dm, 1, 0); + tmpSpectralField = SpectralField(spectralspace_ba, dm, 1, 0); + + // By default, we assume the FFT is done from/to a nodal grid in real space + // It the FFT is performed from/to a cell-centered grid in real space, + // a correcting "shift" factor must be applied in spectral space. + xshift_FFTfromCell = k_space.getSpectralShiftFactor(dm, 0, + ShiftType::TransformFromCellCentered); + xshift_FFTtoCell = k_space.getSpectralShiftFactor(dm, 0, + ShiftType::TransformToCellCentered); +#if (AMREX_SPACEDIM == 3) + yshift_FFTfromCell = k_space.getSpectralShiftFactor(dm, 1, + ShiftType::TransformFromCellCentered); + yshift_FFTtoCell = k_space.getSpectralShiftFactor(dm, 1, + ShiftType::TransformToCellCentered); + zshift_FFTfromCell = k_space.getSpectralShiftFactor(dm, 2, + ShiftType::TransformFromCellCentered); + zshift_FFTtoCell = k_space.getSpectralShiftFactor(dm, 2, + ShiftType::TransformToCellCentered); +#else + zshift_FFTfromCell = k_space.getSpectralShiftFactor(dm, 1, + ShiftType::TransformFromCellCentered); + zshift_FFTtoCell = k_space.getSpectralShiftFactor(dm, 1, + ShiftType::TransformToCellCentered); +#endif + + // Allocate and initialize the FFT plans + forward_plan = FFTplans(spectralspace_ba, dm); + backward_plan = FFTplans(spectralspace_ba, dm); + // Loop over boxes and allocate the corresponding plan + // for each box owned by the local MPI proc + for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){ + // Note: the size of the real-space box and spectral-space box + // differ when using real-to-complex FFT. When initializing + // the FFT plan, the valid dimensions are those of the real-space box. + IntVect fft_size = realspace_ba[mfi].length(); +#ifdef AMREX_USE_GPU + // Add cuFFT-specific code +#else + // Create FFTW plans + forward_plan[mfi] = + // Swap dimensions: AMReX FAB are Fortran-order but FFTW is C-order +#if (AMREX_SPACEDIM == 3) + fftw_plan_dft_r2c_3d( fft_size[2], fft_size[1], fft_size[0], +#else + fftw_plan_dft_r2c_2d( fft_size[1], fft_size[0], +#endif + tmpRealField[mfi].dataPtr(), + reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ), + FFTW_ESTIMATE ); + backward_plan[mfi] = + // Swap dimensions: AMReX FAB are Fortran-order but FFTW is C-order +#if (AMREX_SPACEDIM == 3) + fftw_plan_dft_c2r_3d( fft_size[2], fft_size[1], fft_size[0], +#else + fftw_plan_dft_c2r_2d( fft_size[1], fft_size[0], +#endif + reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ), + tmpRealField[mfi].dataPtr(), + FFTW_ESTIMATE ); +#endif + } +} + + +SpectralFieldData::~SpectralFieldData() +{ + if (tmpRealField.size() > 0){ + for ( MFIter mfi(tmpRealField); mfi.isValid(); ++mfi ){ +#ifdef AMREX_USE_GPU + // Add cuFFT-specific code +#else + // Destroy FFTW plans + fftw_destroy_plan( forward_plan[mfi] ); + fftw_destroy_plan( backward_plan[mfi] ); +#endif + } + } +} + +/* \brief Transform the component `i_comp` of MultiFab `mf` + * to spectral space, and store the corresponding result internally + * (in the spectral field specified by `field_index`) */ +void +SpectralFieldData::ForwardTransform( const MultiFab& mf, + 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); +#if (AMREX_SPACEDIM == 3) + const bool is_nodal_y = mf.is_nodal(1); + const bool is_nodal_z = mf.is_nodal(2); +#else + const bool is_nodal_z = mf.is_nodal(1); +#endif + + // Loop over boxes + for ( MFIter mfi(mf); mfi.isValid(); ++mfi ){ + + // Copy the real-space field `mf` to the temporary field `tmpRealField` + // This ensures that all fields have the same number of points + // before the Fourier transform. + // As a consequence, the copy discards the *last* point of `mf` + // in any direction that has *nodal* index type. + { + Box realspace_bx = mf[mfi].box(); // Copy the box + realspace_bx.enclosedCells(); // Discard last point in nodal direction + AMREX_ALWAYS_ASSERT( realspace_bx == tmpRealField[mfi].box() ); + Array4<const Real> mf_arr = mf[mfi].array(); + Array4<Real> 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,i_comp); + }); + } + + // Perform Fourier transform from `tmpRealField` to `tmpSpectralField` +#ifdef AMREX_USE_GPU + // Add cuFFT-specific code ; make sure that this is done on the same + // GPU stream as the above copy +#else + fftw_execute( forward_plan[mfi] ); +#endif + + // Copy the spectral-space field `tmpSpectralField` to the appropriate + // 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. + { + 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) + const Complex* yshift_arr = yshift_FFTfromCell[mfi].dataPtr(); +#endif + const Complex* zshift_arr = zshift_FFTfromCell[mfi].dataPtr(); + // Loop over indices within one box + const Box spectralspace_bx = tmpSpectralField[mfi].box(); + ParallelFor( spectralspace_bx, + [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept { + Complex spectral_field_value = tmp_arr(i,j,k); + // 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]; + 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 the right index + fields_arr(i,j,k,field_index) = spectral_field_value; + }); + } + } +} + + +/* \brief Transform spectral field specified by `field_index` back to + * 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 ) +{ + // Check field index type, in order to apply proper shift in spectral space + const bool is_nodal_x = mf.is_nodal(0); +#if (AMREX_SPACEDIM == 3) + const bool is_nodal_y = mf.is_nodal(1); + const bool is_nodal_z = mf.is_nodal(2); +#else + const bool is_nodal_z = mf.is_nodal(1); +#endif + + // Loop over boxes + for ( MFIter mfi(mf); mfi.isValid(); ++mfi ){ + + // Copy the spectral-space field `tmpSpectralField` to the appropriate + // field (specified by the input argument field_index) + // and apply correcting shift factor if the field is to be transformed + // to a cell-centered grid in real space instead of a nodal grid. + { + 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) + const Complex* yshift_arr = yshift_FFTtoCell[mfi].dataPtr(); +#endif + const Complex* zshift_arr = zshift_FFTtoCell[mfi].dataPtr(); + // Loop over indices within one box + const Box spectralspace_bx = tmpSpectralField[mfi].box(); + ParallelFor( spectralspace_bx, + [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept { + 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]; + 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 + tmp_arr(i,j,k) = spectral_field_value; + }); + } + + // Perform Fourier transform from `tmpSpectralField` to `tmpRealField` +#ifdef AMREX_USE_GPU + // Add cuFFT-specific code ; make sure that this is done on the same + // GPU stream as the above copy +#else + fftw_execute( backward_plan[mfi] ); +#endif + + // Copy the temporary field `tmpRealField` to the real-space field `mf` + + // Normalize (divide by 1/N) since the FFT+IFFT results in a factor N + { + const Box realspace_bx = tmpRealField[mfi].box(); + Array4<Real> mf_arr = mf[mfi].array(); + Array4<const Real> tmp_arr = tmpRealField[mfi].array(); + // Normalization: divide by the number of points in realspace + const Real inv_N = 1./realspace_bx.numPts(); + ParallelFor( realspace_bx, + [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept { + // Copy and normalize field + mf_arr(i,j,k,i_comp) = inv_N*tmp_arr(i,j,k); + }); + } + } +} diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H new file mode 100644 index 000000000..ae7124b2e --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H @@ -0,0 +1,56 @@ +#ifndef WARPX_SPECTRAL_K_SPACE_H_ +#define WARPX_SPECTRAL_K_SPACE_H_ + +#include <WarpX_Complex.H> +#include <AMReX_BoxArray.H> +#include <AMReX_LayoutData.H> + +// `KVectorComponent` and `SpectralShiftFactor` hold one 1D array +// ("ManagedVector") for each box ("LayoutData"). The arrays are +// only allocated if the corresponding box is owned by the local MPI rank. +using KVectorComponent = amrex::LayoutData< + amrex::Gpu::ManagedVector<amrex::Real> >; +using SpectralShiftFactor = amrex::LayoutData< + amrex::Gpu::ManagedVector<Complex> >; + +// Indicate the type of correction "shift" factor to apply +// when the FFT is performed from/to a cell-centered grid in real space. +struct ShiftType { + enum{ TransformFromCellCentered=0, TransformToCellCentered=1 }; +}; + +/* \brief Class that represents the spectral space. + * + * (Contains info about the size of the spectral space corresponding + * to each box in `realspace_ba`, as well as the value of the + * corresponding k coordinates) + */ +class SpectralKSpace +{ + public: + amrex::BoxArray spectralspace_ba; + SpectralKSpace( const amrex::BoxArray& realspace_ba, + const amrex::DistributionMapping& dm, + const amrex::RealVect realspace_dx ); + KVectorComponent getKComponent( + const amrex::DistributionMapping& dm, + const amrex::BoxArray& realspace_ba, + const int i_dim, const bool only_positive_k ) const; + KVectorComponent getModifiedKComponent( + 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; + + private: + amrex::Array<KVectorComponent, AMREX_SPACEDIM> k_vec; + // 3D: k_vec is an Array of 3 components, corresponding to kx, ky, kz + // 2D: k_vec is an Array of 2 components, corresponding to kx, kz + 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 new file mode 100644 index 000000000..2fe78cedd --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp @@ -0,0 +1,245 @@ +#include <WarpXConst.H> +#include <SpectralKSpace.H> +#include <cmath> + +using namespace amrex; +using namespace Gpu; + +/* \brief Initialize k space object. + * + * \param realspace_ba Box array that corresponds to the decomposition + * of the fields in real space (cell-centered ; includes guard cells) + * \param dm Indicates which MPI proc owns which box, in realspace_ba. + * \param realspace_dx Cell size of the grid in real space + */ +SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba, + const DistributionMapping& dm, + const RealVect realspace_dx ) + : dx(realspace_dx) // Store the cell size as member `dx` +{ + AMREX_ALWAYS_ASSERT_WITH_MESSAGE( + realspace_ba.ixType()==IndexType::TheCellType(), + "SpectralKSpace expects a cell-centered box."); + + // Create the box array that corresponds to spectral space + BoxList spectral_bl; // Create empty box list + // Loop over boxes and fill the box list + for (int i=0; i < realspace_ba.size(); i++ ) { + // For local FFTs, boxes in spectral space start at 0 in + // each direction and have the same number of points as the + // (cell-centered) real space box + // TODO: this will be different for the hybrid FFT scheme + Box realspace_bx = realspace_ba[i]; + IntVect fft_size = realspace_bx.length(); + // Because the spectral solver uses real-to-complex FFTs, we only + // need the positive k values along the fastest axis + // (first axis for AMReX Fortran-order arrays) in spectral space. + // This effectively reduces the size of the spectral space by half + // see e.g. the FFTW documentation for real-to-complex FFTs + IntVect spectral_bx_size = fft_size; + spectral_bx_size[0] = fft_size[0]/2 + 1; + // Define the corresponding box + Box spectral_bx = Box( IntVect::TheZeroVector(), + spectral_bx_size - IntVect::TheUnitVector() ); + spectral_bl.push_back( spectral_bx ); + } + spectralspace_ba.define( spectral_bl ); + + // Allocate the components of the k vector: kx, ky (only in 3D), kz + bool only_positive_k; + for (int i_dim=0; i_dim<AMREX_SPACEDIM; i_dim++) { + if (i_dim==0) { + // Real-to-complex FFTs: first axis contains only the positive k + only_positive_k = true; + } else { + only_positive_k = false; + } + k_vec[i_dim] = getKComponent(dm, realspace_ba, i_dim, only_positive_k); + } +} + +/* For each box, in `spectralspace_ba`, which is owned by the local MPI rank + * (as indicated by the argument `dm`), compute the values of the + * corresponding k coordinate along the dimension specified by `i_dim` + */ +KVectorComponent +SpectralKSpace::getKComponent( const DistributionMapping& dm, + const BoxArray& realspace_ba, + const int i_dim, + const bool only_positive_k ) const +{ + // Initialize an empty ManagedVector in each box + KVectorComponent k_comp(spectralspace_ba, dm); + // Loop over boxes and allocate the corresponding ManagedVector + // for each box owned by the local MPI proc + for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){ + Box bx = spectralspace_ba[mfi]; + ManagedVector<Real>& k = k_comp[mfi]; + + // Allocate k to the right size + int N = bx.length( i_dim ); + k.resize( N ); + + // Fill the k vector + IntVect fft_size = realspace_ba[mfi].length(); + const Real dk = 2*MathConst::pi/(fft_size[i_dim]*dx[i_dim]); + AMREX_ALWAYS_ASSERT_WITH_MESSAGE( bx.smallEnd(i_dim) == 0, + "Expected box to start at 0, in spectral space."); + AMREX_ALWAYS_ASSERT_WITH_MESSAGE( bx.bigEnd(i_dim) == N-1, + "Expected different box end index in spectral space."); + if (only_positive_k){ + // Fill the full axis with positive k values + // (typically: first axis, in a real-to-complex FFT) + for (int i=0; i<N; i++ ){ + k[i] = i*dk; + } + } else { + const int mid_point = (N+1)/2; + // Fill positive values of k + // (FFT conventions: first half is positive) + for (int i=0; i<mid_point; i++ ){ + k[i] = i*dk; + } + // Fill negative values of k + // (FFT conventions: second half is negative) + for (int i=mid_point; i<N; i++){ + k[i] = (i-N)*dk; + } + } + // TODO: this will be different for the hybrid FFT scheme + } + return k_comp; +} + +/* For each box, in `spectralspace_ba`, which is owned by the local MPI rank + * (as indicated by the argument `dm`), compute the values of the + * corresponding correcting "shift" factor, along the dimension + * specified by `i_dim`. + * + * (By default, we assume the FFT is done from/to a nodal grid in real space + * It the FFT is performed from/to a cell-centered grid in real space, + * a correcting "shift" factor must be applied in spectral space.) + */ +SpectralShiftFactor +SpectralKSpace::getSpectralShiftFactor( const DistributionMapping& dm, + const int i_dim, + const int shift_type ) const +{ + // Initialize an empty ManagedVector in each box + SpectralShiftFactor shift_factor( spectralspace_ba, dm ); + // Loop over boxes and allocate the corresponding ManagedVector + // for each box owned by the local MPI proc + 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::TransformFromCellCentered: sign = -1.; break; + case ShiftType::TransformToCellCentered: 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; +} + +/* \brief For each box, in `spectralspace_ba`, which is owned by the local MPI + * rank (as indicated by the argument `dm`), compute the values of the + * corresponding finite-order modified k vector, along the + * dimension specified by `i_dim` + * + * The finite-order modified k vector is the spectral-space representation + * of a finite-order stencil in real space. + * + * \param n_order Order of accuracy of the stencil, in discretizing + * a spatial derivative + * \param nodal Whether the stencil is to be applied to a nodal or + staggered set of fields + */ +KVectorComponent +SpectralKSpace::getModifiedKComponent( const DistributionMapping& dm, + const int i_dim, + const int n_order, + const bool nodal ) const +{ + // Initialize an empty ManagedVector in each box + KVectorComponent modified_k_comp(spectralspace_ba, dm); + + // Compute real-space stencil coefficients + Vector<Real> stencil_coef = getFonbergStencilCoefficients(n_order, nodal); + + // Loop over boxes and allocate the corresponding ManagedVector + // for each box owned by the local MPI proc + 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]; + + // Allocate modified_k to the same size as k + modified_k.resize( k.size() ); + + // 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]* \ + 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; +} + +/* Returns an array of coefficients, corresponding to the weight + * of each point in a finite-difference approximation (to order `n_order`) + * of a derivative. + * + * `nodal` indicates whether this finite-difference approximation is + * taken on a nodal grid or a staggered grid. + */ +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; + Vector<Real> coefs; + coefs.resize( m+1 ); + + // Note: there are closed-form formula for these coefficients, + // but they result in an overflow when evaluated numerically. + // One way to avoid the overflow is to calculate the coefficients + // by recurrence. + + // Coefficients for nodal (a.k.a. centered) finite-difference + if (nodal == true) { + coefs[0] = -2.; // First coefficient + for (int n=1; n<m+1; n++){ // Get the other coefficients by recurrence + coefs[n] = - (m+1-n)*1./(m+n)*coefs[n-1]; + } + } + // Coefficients for staggered finite-difference + else { + Real prod = 1.; + for (int k=1; k<m+1; k++){ + prod *= (m+k)*1./(4*k); + } + coefs[0] = 4*m*prod*prod; // First coefficient + for (int n=1; n<m+1; n++){ // Get the other coefficients by recurrence + coefs[n] = - ((2*n-3)*(m+1-n))*1./((2*n-1)*(m-1+n))*coefs[n-1]; + } + } + return coefs; +} diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolver.H b/Source/FieldSolver/SpectralSolver/SpectralSolver.H new file mode 100644 index 000000000..d4019a9a3 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralSolver.H @@ -0,0 +1,66 @@ +#ifndef WARPX_SPECTRAL_SOLVER_H_ +#define WARPX_SPECTRAL_SOLVER_H_ + +#include <SpectralBaseAlgorithm.H> +#include <SpectralFieldData.H> + +/* \brief Top-level class for the electromagnetic spectral solver + * + * Stores the field in spectral space, and has member functions + * to Fourier-transform the fields between real space and spectral space + * and to update fields in spectral space over one time step. + */ +class SpectralSolver +{ + public: + // Inline definition of the member functions of `SpectralSolver`, + // except the constructor (see `SpectralSolver.cpp`) + // The body of these functions is short, since the work is done in the + // underlying classes `SpectralFieldData` and `PsatdAlgorithm` + + // Constructor + SpectralSolver( const amrex::BoxArray& realspace_ba, + const amrex::DistributionMapping& dm, + const int norder_x, const int norder_y, + const int norder_z, const bool nodal, + const amrex::RealVect dx, const amrex::Real dt ); + + /* \brief Transform the component `i_comp` of MultiFab `mf` + * to spectral space, and store the corresponding result internally + * (in the spectral field specified by `field_index`) */ + void ForwardTransform( const amrex::MultiFab& mf, + const int field_index, + const int i_comp=0 ){ + BL_PROFILE("SpectralSolver::ForwardTransform"); + field_data.ForwardTransform( mf, field_index, i_comp ); + }; + + /* \brief Transform spectral field specified by `field_index` back to + * real space, and store it in the component `i_comp` of `mf` */ + void BackwardTransform( amrex::MultiFab& mf, + const int field_index, + const int i_comp=0 ){ + BL_PROFILE("SpectralSolver::BackwardTransform"); + field_data.BackwardTransform( mf, field_index, i_comp ); + }; + + /* \brief Update the fields in spectral space, over one timestep */ + void pushSpectralFields(){ + BL_PROFILE("SpectralSolver::pushSpectralFields"); + // Virtual function: the actual function used here depends + // on the sub-class of `SpectralBaseAlgorithm` that was + // initialized in the constructor of `SpectralSolver` + algorithm->pushSpectralFields( field_data ); + }; + + private: + SpectralFieldData field_data; // Store field in spectral space + // and perform the Fourier transforms + std::unique_ptr<SpectralBaseAlgorithm> algorithm; + // Defines field update equation in spectral space, + // and the associated coefficients. + // SpectralBaseAlgorithm is a base class ; this pointer is meant + // to point an instance of a *sub-class* defining a specific algorithm +}; + +#endif // WARPX_SPECTRAL_SOLVER_H_ diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolver.cpp b/Source/FieldSolver/SpectralSolver/SpectralSolver.cpp new file mode 100644 index 000000000..c21c3cfb1 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralSolver.cpp @@ -0,0 +1,35 @@ +#include <SpectralKSpace.H> +#include <SpectralSolver.H> +#include <PsatdAlgorithm.H> + +/* \brief Initialize the spectral Maxwell solver + * + * This function selects the spectral algorithm to be used, allocates the + * corresponding coefficients for the discretized field update equation, + * and prepares the structures that store the fields in spectral space. + */ +SpectralSolver::SpectralSolver( + const amrex::BoxArray& realspace_ba, + const amrex::DistributionMapping& dm, + const int norder_x, const int norder_y, + const int norder_z, const bool nodal, + const amrex::RealVect dx, const amrex::Real dt ) { + + // Initialize all structures using the same distribution mapping dm + + // - Initialize k space object (Contains info about the size of + // the spectral space corresponding to each box in `realspace_ba`, + // as well as the value of the corresponding k coordinates) + const SpectralKSpace k_space= SpectralKSpace(realspace_ba, dm, dx); + + // - Select the algorithm depending on the input parameters + // Initialize the corresponding coefficients over k space + // TODO: Add more algorithms + selection depending on input parameters + // For the moment, this only uses the standard PsatdAlgorithm + algorithm = std::unique_ptr<PsatdAlgorithm>( new PsatdAlgorithm( + k_space, dm, norder_x, norder_y, norder_z, nodal, dt ) ); + + // - Initialize arrays for fields in spectral space + FFT plans + field_data = SpectralFieldData( realspace_ba, k_space, dm ); + +}; diff --git a/Source/FieldSolver/SpectralSolver/WarpX_Complex.H b/Source/FieldSolver/SpectralSolver/WarpX_Complex.H new file mode 100644 index 000000000..c898c5baa --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/WarpX_Complex.H @@ -0,0 +1,27 @@ +#ifndef WARPX_COMPLEX_H_ +#define WARPX_COMPLEX_H_ + +#include <AMReX_REAL.H> + +// Define complex type on GPU/CPU +#ifdef AMREX_USE_GPU + +#include <thrust/complex.h> +#include <cufft.h> +using Complex = thrust::complex<amrex::Real>; +static_assert( sizeof(Complex) == sizeof(cuDoubleComplex), + "The complex types in WarpX and cuFFT do not match."); + +#else + +#include <complex> +#include <fftw3.h> +using Complex = std::complex<amrex::Real>; +static_assert( sizeof(Complex) == sizeof(fftw_complex), + "The complex types in WarpX and FFTW do not match."); + +#endif +static_assert(sizeof(Complex) == sizeof(amrex::Real[2]), + "Unexpected complex type."); + +#endif //WARPX_COMPLEX_H_ |