diff options
Diffstat (limited to 'Source/FieldSolver')
5 files changed, 88 insertions, 45 deletions
diff --git a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H index 16d27cab2..acefcc466 100644 --- a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H +++ b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.H @@ -4,8 +4,9 @@ #include <SpectralKSpace.H> #include <SpectralFieldData.H> -/* TODO: Write documentation -*/ +/* \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> >; @@ -16,7 +17,7 @@ class PsatdAlgorithm 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; // Default move assignment + PsatdAlgorithm& operator=(PsatdAlgorithm&& algorithm) = default; void pushSpectralFields(SpectralFieldData& f) const; private: diff --git a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp index 5ebb7144d..60e9d58c0 100644 --- a/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp +++ b/Source/FieldSolver/SpectralSolver/PsatdAlgorithm.cpp @@ -4,6 +4,7 @@ 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, @@ -28,7 +29,8 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace, X3_coef = SpectralCoefficients(ba, dm, 1, 0); // Fill them with the right values: - // Loop over boxes + // 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]; @@ -78,6 +80,8 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace, } }; +/* Advance the E and B field in spectral space (stored in `f`) + * over one time step */ void PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{ diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H index be1765ca0..c62513de6 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H +++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H @@ -8,20 +8,23 @@ // Declare type for spectral fields using SpectralField = amrex::FabArray< amrex::BaseFab <Complex> >; -/* Fields that will be stored in spectral space */ +/* 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 spectral transforms to/from real space +/* \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 + // Add cuFFT-specific code #else using FFTplans = amrex::LayoutData<fftw_plan>; #endif @@ -40,7 +43,9 @@ class SpectralFieldData private: SpectralField Ex, Ey, Ez, Bx, By, Bz, Jx, Jy, Jz, rho_old, rho_new; - SpectralField tmpRealField, tmpSpectralField; // Store fields before/after transform + // 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 diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp index ca5e338e5..15652addc 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp +++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp @@ -2,6 +2,7 @@ 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 ) @@ -52,6 +53,8 @@ SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba, // 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 @@ -59,7 +62,8 @@ SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba, #else // Create FFTW plans forward_plan[mfi] = -#if (AMREX_SPACEDIM == 3) // Swap dimensions: AMReX data is Fortran-order, but FFTW is C-order + // 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), @@ -68,7 +72,8 @@ SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba, reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ), FFTW_FORWARD, FFTW_ESTIMATE ); backward_plan[mfi] = -#if (AMREX_SPACEDIM == 3) // Swap dimensions: AMReX data is Fortran-order, but FFTW is C-order + // 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), @@ -76,8 +81,6 @@ SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba, reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ), reinterpret_cast<fftw_complex*>( tmpRealField[mfi].dataPtr() ), FFTW_BACKWARD, FFTW_ESTIMATE ); - // TODO: Do real-to-complex transform instead of complex-to-complex - // TODO: Let the user decide whether to use FFTW_ESTIMATE or FFTW_MEASURE #endif } } @@ -98,12 +101,13 @@ SpectralFieldData::~SpectralFieldData() } } -/* TODO: Documentation - * Example: ForwardTransform( Efield_cp[0], SpectralFieldIndex::Ex ) - */ +/* \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 ) + 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); @@ -123,8 +127,8 @@ SpectralFieldData::ForwardTransform( const MultiFab& mf, // As a consequence, the copy discards the *last* point of `mf` // in any direction that has *nodal* index type. { - Box bx = mf[mfi].box(); - const Box realspace_bx = bx.enclosedCells(); // discards last point in each nodal direction + 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(); @@ -174,8 +178,8 @@ SpectralFieldData::ForwardTransform( const MultiFab& mf, } -/* TODO: Documentation - */ +/* \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 ) @@ -193,9 +197,10 @@ SpectralFieldData::BackwardTransform( MultiFab& mf, for ( MFIter mfi(mf); mfi.isValid(); ++mfi ){ // Copy the spectral-space field `tmpSpectralField` to the appropriate - // field (specified by the input argument field_index ) + // 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(); @@ -207,6 +212,8 @@ SpectralFieldData::BackwardTransform( MultiFab& mf, 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); @@ -216,8 +223,8 @@ SpectralFieldData::BackwardTransform( MultiFab& mf, 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 - tmp_arr(i,j,k) = spectral_field_value; + // Copy field into temporary array (after normalization) + tmp_arr(i,j,k) = inv_N*spectral_field_value; }); } @@ -230,22 +237,15 @@ SpectralFieldData::BackwardTransform( MultiFab& mf, #endif // Copy the temporary field `tmpRealField` to the real-space field `mf` - // The copy does *not* fill the *last* point of `mf` - // in any direction that has *nodal* index type (but this point is - // in the guard cells and will be filled by guard cell exchange) - // Normalize (divide by 1/N) since the FFT result in a factor N + // in the *valid* part of the domain (guard cells are filled later, + // through guard cell exchange) { - Box bx = mf[mfi].box(); - const Box realspace_bx = bx.enclosedCells(); - // `enclosedells` discards last point in each nodal direction - AMREX_ALWAYS_ASSERT( realspace_bx == tmpRealField[mfi].box() ); + const Box realspace_valid_bx = mfi.validbox(); Array4<Real> mf_arr = mf[mfi].array(); Array4<const Complex> tmp_arr = tmpRealField[mfi].array(); - // For normalization: divide by the number of points in the box - const Real inv_N = 1./bx.numPts(); - ParallelFor( realspace_bx, + ParallelFor( realspace_valid_bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept { - mf_arr(i,j,k,i_comp) = inv_N*tmp_arr(i,j,k).real(); + mf_arr(i,j,k,i_comp) = tmp_arr(i,j,k).real(); }); } } diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp index 90866cd3b..c82e45c5e 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp +++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp @@ -45,10 +45,9 @@ SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba, } } -/* For each box, in `spectralspace_ba`, which is owned - * by the local MPI proc (as indicated by the argument `dm`), - * compute the values of the corresponding k coordinate - * along the dimension specified by `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, @@ -57,7 +56,7 @@ SpectralKSpace::getKComponent( const DistributionMapping& dm, // Initialize an empty ManagedVector in each box KVectorComponent k_comp = KVectorComponent(spectralspace_ba, dm); // Loop over boxes and allocate the corresponding ManagedVector - // for each box owned by the local MPI proc ("mfi.isValid") + // 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]; @@ -87,6 +86,15 @@ SpectralKSpace::getKComponent( const DistributionMapping& dm, 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, @@ -94,7 +102,8 @@ SpectralKSpace::getSpectralShiftFactor( const DistributionMapping& dm, { // Initialize an empty ManagedVector in each box SpectralShiftFactor shift_factor = SpectralShiftFactor( spectralspace_ba, dm ); - // Loop over boxes + // 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]; @@ -116,6 +125,19 @@ SpectralKSpace::getSpectralShiftFactor( const DistributionMapping& dm, 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, @@ -123,12 +145,13 @@ SpectralKSpace::getModifiedKComponent( const DistributionMapping& dm, const bool nodal ) const { // Initialize an empty ManagedVector in each box - KVectorComponent modified_k_comp = KVectorComponent( spectralspace_ba, dm ); + KVectorComponent modified_k_comp = KVectorComponent(spectralspace_ba, dm); // Compute real-space stencil coefficients Vector<Real> stencil_coef = getFonbergStencilCoefficients(n_order, nodal); - // Loop over boxes + // 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]; @@ -153,7 +176,12 @@ SpectralKSpace::getModifiedKComponent( const DistributionMapping& dm, return modified_k_comp; } -/* TODO: Documentation: point to Fonberg paper ; explain recurrence relation +/* 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 ) @@ -164,6 +192,11 @@ getFonbergStencilCoefficients( const int n_order, const bool nodal ) 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 |