diff options
4 files changed, 56 insertions, 30 deletions
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H index 9e31ac5b8..8e58aa1d8 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H +++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H @@ -47,7 +47,8 @@ class SpectralFieldData SpectralField fields; // tmpRealField and tmpSpectralField store fields // right before/after the Fourier transform - SpectralField tmpRealField, tmpSpectralField; + 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 diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp index 6e6cc124f..4addd085c 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp +++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp @@ -16,7 +16,7 @@ SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba, // 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); + 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 @@ -56,23 +56,23 @@ SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba, 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), + fftw_plan_dft_r2c_3d( bx.length(2), bx.length(1), bx.length(0), #else - fftw_plan_dft_2d( bx.length(1), bx.length(0), + fftw_plan_dft_r2c_2d( bx.length(1), bx.length(0), #endif - reinterpret_cast<fftw_complex*>( tmpRealField[mfi].dataPtr() ), + tmpRealField[mfi].dataPtr(), reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ), - FFTW_FORWARD, FFTW_ESTIMATE ); + 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), + fftw_plan_dft_c2r_3d( bx.length(2), bx.length(1), bx.length(0), #else - fftw_plan_dft_2d( bx.length(1), bx.length(0), + fftw_plan_dft_c2r_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 ); + tmpRealField[mfi].dataPtr(), + FFTW_ESTIMATE ); #endif } } @@ -123,7 +123,7 @@ SpectralFieldData::ForwardTransform( const MultiFab& mf, 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(); + 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); @@ -235,10 +235,10 @@ SpectralFieldData::BackwardTransform( MultiFab& mf, { const Box realspace_bx = tmpRealField[mfi].box(); Array4<Real> mf_arr = mf[mfi].array(); - Array4<const Complex> tmp_arr = tmpRealField[mfi].array(); + Array4<const Real> 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(); + mf_arr(i,j,k,i_comp) = tmp_arr(i,j,k); }); } } diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H index ad17e6423..4c29728fc 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H +++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H @@ -33,7 +33,8 @@ class SpectralKSpace const amrex::DistributionMapping& dm, const amrex::RealVect realspace_dx ); KVectorComponent getKComponent( - const amrex::DistributionMapping& dm, const int i_dim ) const; + const amrex::DistributionMapping& dm, + 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; diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp index ddb2020d8..71073eb3f 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp +++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp @@ -28,19 +28,33 @@ SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba, // 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 ); + IntVect spectral_bx_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 + spectral_bx_size[0] = spectral_bx_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++) { - k_vec[i_dim] = getKComponent( dm, i_dim ); + if (i_dim==0) { + // For real-to-complex FFTs, the first axis contains + // only the positive k + only_positive_k = true; + } else { + only_positive_k = false; + } + k_vec[i_dim] = getKComponent( dm, i_dim, only_positive_k ); } } @@ -50,7 +64,8 @@ SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba, */ KVectorComponent SpectralKSpace::getKComponent( const DistributionMapping& dm, - const int i_dim ) const + const int i_dim, + const bool only_positive_k ) const { // Initialize an empty ManagedVector in each box KVectorComponent k_comp(spectralspace_ba, dm); @@ -70,16 +85,25 @@ SpectralKSpace::getKComponent( const DistributionMapping& dm, "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; + 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 real-to-complex transform // TODO: this will be different for the hybrid FFT scheme } return k_comp; |