diff options
4 files changed, 71 insertions, 37 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..02fa2015f 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 @@ -48,7 +48,10 @@ SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba, // 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]; + // 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 @@ -56,23 +59,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( fft_size[2], fft_size[1], fft_size[0], #else - fftw_plan_dft_2d( bx.length(1), bx.length(0), + fftw_plan_dft_r2c_2d( fft_size[1], fft_size[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( fft_size[2], fft_size[1], fft_size[0], #else - fftw_plan_dft_2d( bx.length(1), bx.length(0), + fftw_plan_dft_c2r_2d( fft_size[1], fft_size[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 +126,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); @@ -194,7 +197,6 @@ SpectralFieldData::BackwardTransform( MultiFab& mf, // 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 { Array4<const Complex> field_arr = SpectralFieldData::fields[mfi].array(); Array4<Complex> tmp_arr = tmpSpectralField[mfi].array(); @@ -205,8 +207,6 @@ 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,field_index); @@ -218,8 +218,8 @@ SpectralFieldData::BackwardTransform( MultiFab& mf, #elif (AMREX_SPACEDIM == 2) if (is_nodal_z==false) spectral_field_value *= zshift_arr[j]; #endif - // Copy field into temporary array (after normalization) - tmp_arr(i,j,k) = inv_N*spectral_field_value; + // Copy field into temporary array + tmp_arr(i,j,k) = spectral_field_value; }); } @@ -232,13 +232,18 @@ SpectralFieldData::BackwardTransform( MultiFab& mf, #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 Complex> tmp_arr = tmpRealField[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 { - mf_arr(i,j,k,i_comp) = tmp_arr(i,j,k).real(); + // 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 index ad17e6423..ae7124b2e 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H +++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H @@ -33,7 +33,9 @@ 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 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; diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp index ddb2020d8..2fe78cedd 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 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++) { - k_vec[i_dim] = getKComponent( dm, 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); } } @@ -50,7 +64,9 @@ SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba, */ KVectorComponent SpectralKSpace::getKComponent( const DistributionMapping& dm, - const int i_dim ) const + 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); @@ -65,21 +81,31 @@ SpectralKSpace::getKComponent( const DistributionMapping& dm, k.resize( N ); // Fill the k vector - const Real dk = 2*MathConst::pi/(N*dx[i_dim]); + 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."); - 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; |