aboutsummaryrefslogtreecommitdiff
path: root/Source/FieldSolver/SpectralSolver
diff options
context:
space:
mode:
Diffstat (limited to 'Source/FieldSolver/SpectralSolver')
-rw-r--r--Source/FieldSolver/SpectralSolver/Make.package11
-rw-r--r--Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H32
-rw-r--r--Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp174
-rw-r--r--Source/FieldSolver/SpectralSolver/SpectralFieldData.H60
-rw-r--r--Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp271
-rw-r--r--Source/FieldSolver/SpectralSolver/SpectralKSpace.H54
-rw-r--r--Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp218
-rw-r--r--Source/FieldSolver/SpectralSolver/SpectralSolver.H72
-rw-r--r--Source/FieldSolver/SpectralSolver/WarpX_Complex.H27
9 files changed, 919 insertions, 0 deletions
diff --git a/Source/FieldSolver/SpectralSolver/Make.package b/Source/FieldSolver/SpectralSolver/Make.package
new file mode 100644
index 000000000..50127914d
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/Make.package
@@ -0,0 +1,11 @@
+CEXE_headers += WarpX_Complex.H
+CEXE_headers += SpectralSolver.H
+CEXE_headers += SpectralFieldData.H
+CEXE_sources += SpectralFieldData.cpp
+CEXE_headers += PsatdAlgorithm.H
+CEXE_sources += PsatdAlgorithm.cpp
+CEXE_headers += SpectralKSpace.H
+CEXE_sources += SpectralKSpace.cpp
+
+INCLUDE_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver
+VPATH_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver
diff --git a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H
new file mode 100644
index 000000000..acefcc466
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H
@@ -0,0 +1,32 @@
+#ifndef WARPX_PSATD_ALGORITHM_H_
+#define WARPX_PSATD_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.
+ */
+class PsatdAlgorithm
+{
+ using SpectralCoefficients = amrex::FabArray< amrex::BaseFab <amrex::Real> >;
+
+ 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);
+ PsatdAlgorithm() = default; // Default constructor
+ PsatdAlgorithm& operator=(PsatdAlgorithm&& algorithm) = default;
+ void pushSpectralFields(SpectralFieldData& f) const;
+
+ private:
+ // Modified finite-order vectors
+ KVectorComponent modified_kx_vec, modified_kz_vec;
+#if (AMREX_SPACEDIM==3)
+ KVectorComponent modified_ky_vec;
+#endif
+ SpectralCoefficients C_coef, S_ck_coef, X1_coef, X2_coef, X3_coef;
+};
+
+#endif // WARPX_PSATD_ALGORITHM_H_
diff --git a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp
new file mode 100644
index 000000000..56e58bcc4
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp
@@ -0,0 +1,174 @@
+#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)
+// 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
+{
+ 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) +
+#endif
+ 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);
+ 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.Ex); mfi.isValid(); ++mfi){
+
+ const Box& bx = f.Ex[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();
+ // 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
+ 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);
+ // 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);
+ // k vector values, and coefficients
+ const Real kx = modified_kx_arr[i];
+#if (AMREX_SPACEDIM==3)
+ const Real ky = modified_ky_arr[j];
+#else
+ constexpr Real ky = 0;
+#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};
+ 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)
+ 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;
+ 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;
+ 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;
+ // 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);
+ By_arr(i,j,k) = 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
+ - 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
new file mode 100644
index 000000000..c62513de6
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H
@@ -0,0 +1,60 @@
+#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 };
+};
+
+/* \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:
+ SpectralField Ex, Ey, Ez, Bx, By, Bz, Jx, Jy, Jz, rho_old, rho_new;
+ // tmpRealField and tmpSpectralField store fields
+ // right before/after the Fourier transform
+ SpectralField tmpRealField, tmpSpectralField;
+ 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
+ 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
new file mode 100644
index 000000000..291fe945e
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
@@ -0,0 +1,271 @@
+#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
+ 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);
+
+ // Allocate temporary arrays - in real space and spectral space
+ // These arrays will store the data just before/after the FFT
+ tmpRealField = SpectralField(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 ){
+ Box bx = spectralspace_ba[mfi];
+#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_3d( bx.length(2), bx.length(1), bx.length(0),
+#else
+ fftw_plan_dft_2d( bx.length(1), bx.length(0),
+#endif
+ reinterpret_cast<fftw_complex*>( tmpRealField[mfi].dataPtr() ),
+ reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ),
+ FFTW_FORWARD, FFTW_ESTIMATE );
+ backward_plan[mfi] =
+ // Swap dimensions: AMReX FAB are Fortran-order but FFTW is C-order
+#if (AMREX_SPACEDIM == 3)
+ fftw_plan_dft_3d( bx.length(2), bx.length(1), bx.length(0),
+#else
+ fftw_plan_dft_2d( bx.length(1), bx.length(0),
+#endif
+ reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ),
+ reinterpret_cast<fftw_complex*>( tmpRealField[mfi].dataPtr() ),
+ FFTW_BACKWARD, 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<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,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
+ // field (specified by the input argument 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<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];
+#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;
+ });
+ }
+ }
+}
+
+
+/* \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.
+ // 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<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();
+ // For normalization: divide by the number of points in the box
+ 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);
+ // 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];
+ // Copy field into temporary array (after normalization)
+ tmp_arr(i,j,k) = inv_N*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`
+ {
+ const Box realspace_bx = tmpRealField[mfi].box();
+ Array4<Real> mf_arr = mf[mfi].array();
+ 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,i_comp) = tmp_arr(i,j,k).real();
+ });
+ }
+ }
+}
+
+
+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.H b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
new file mode 100644
index 000000000..ad17e6423
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
@@ -0,0 +1,54 @@
+#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 int i_dim ) 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..d91891a30
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
@@ -0,0 +1,218 @@
+#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 real-to-complex FFT
+ // TODO: this will be different for the hybrid FFT scheme
+ Box realspace_bx = realspace_ba[i];
+ Print() << realspace_bx.smallEnd() << " " << realspace_bx.bigEnd() << std::endl;
+ Box bx = Box( IntVect::TheZeroVector(),
+ realspace_bx.bigEnd() - realspace_bx.smallEnd() );
+ spectral_bl.push_back( bx );
+ }
+ spectralspace_ba.define( spectral_bl );
+
+ // Allocate the components of the k vector: kx, ky (only in 3D), kz
+ for (int i_dim=0; i_dim<AMREX_SPACEDIM; i_dim++) {
+ k_vec[i_dim] = getKComponent( dm, i_dim );
+ }
+}
+
+/* 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 int i_dim ) 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
+ const Real dk = 2*MathConst::pi/(N*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.");
+ 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 real-to-complex transform
+ // 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++ ){
+ 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..7444452af
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolver.H
@@ -0,0 +1,72 @@
+#ifndef WARPX_SPECTRAL_SOLVER_H_
+#define WARPX_SPECTRAL_SOLVER_H_
+
+#include <SpectralKSpace.H>
+#include <PsatdAlgorithm.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`
+ // The body of these functions is short, since the work is done in the
+ // underlying classes `SpectralFieldData` and `PsatdAlgorithm`
+
+ /* \brief Initialize the spectral solver */
+ 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);
+ // - Initialize the algorithm (coefficients) over this space
+ 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 );
+ };
+
+ /* \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");
+ algorithm.pushSpectralFields( field_data );
+ };
+
+ private:
+ SpectralFieldData field_data; // Store field in spectral space
+ // and perform the Fourier transforms
+ PsatdAlgorithm algorithm; // Contains Psatd coefficients
+ // and field update equation
+};
+
+#endif // WARPX_SPECTRAL_SOLVER_H_
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_
generated by cgit v1.2.3 (git 2.39.1) at 2025-08-09 21:37:23 +0000