Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Implement real-to-complex FFT #98

Merged
merged 4 commits into from
May 2, 2019
Merged

Implement real-to-complex FFT #98

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e1c6d7c
Implement Real-to-complex FFT
RemiLehe Apr 24, 2019
52cf650
Fix normalization
RemiLehe Apr 24, 2019
581510f
Fix FFT size
RemiLehe Apr 24, 2019
df4a98c
Correct error in calculation of k vector
RemiLehe May 2, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion Source/FieldSolver/SpectralSolver/SpectralFieldData.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
41 changes: 23 additions & 18 deletions Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -48,31 +48,34 @@ 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
// 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),
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
}
}
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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();
Expand All @@ -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);
Expand All @@ -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;
});
}

Expand All @@ -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);
});
}
}
Expand Down
4 changes: 3 additions & 1 deletion Source/FieldSolver/SpectralSolver/SpectralKSpace.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
60 changes: 43 additions & 17 deletions Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}

Expand All @@ -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);
Expand All @@ -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;
Expand Down