#include using namespace amrex; 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 - over different boxarrays tmpRealField = SpectralField(realspace_ba, dm, 1, 0); tmpSpectralField = SpectralField(spectralspace_ba, dm, 1, 0); // Allocate the vectors that allow to shift between nodal and cell-centered for (int i_dim=0; i_dim( tmpRealField[mfi].dataPtr() ), reinterpret_cast( 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 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( tmpSpectralField[mfi].dataPtr() ), reinterpret_cast( 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 } } 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 } } } /* TODO: Documentation * Example: ForwardTransform( Efield_cp[0], SpectralFieldIndex::Ex ) */ void SpectralFieldData::ForwardTransform( const MultiFab& mf, const int field_index, const int i_comp ) { // 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 bx = mf[mfi].box(); const Box realspace_bx = bx.enclosedCells(); // discards last point in each nodal direction AMREX_ALWAYS_ASSERT( realspace_bx == tmpRealField[mfi].box() ); Array4 mf_arr = mf[mfi].array(); Array4 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 ) { SpectralField& field = getSpectralField( field_index ); Array4 field_arr = field[mfi].array(); Array4 tmp_arr = tmpSpectralField[mfi].array(); const Box spectralspace_bx = tmpSpectralField[mfi].box(); ParallelFor( spectralspace_bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept { field_arr(i,j,k) = tmp_arr(i,j,k); }); } } } /* TODO: Documentation */ void SpectralFieldData::BackwardTransform( MultiFab& mf, const int field_index, const int i_comp ) { // Loop over boxes for ( MFIter mfi(mf); mfi.isValid(); ++mfi ){ // Copy the appropriate field (specified by the input argument field_index) // to the spectral-space field `tmpSpectralField` { SpectralField& field = getSpectralField( field_index ); Array4 field_arr = field[mfi].array(); Array4 tmp_arr = tmpSpectralField[mfi].array(); const Box spectralspace_bx = tmpSpectralField[mfi].box(); ParallelFor( spectralspace_bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept { tmp_arr(i,j,k) = field_arr(i,j,k); }); } // 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` // 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 { Box bx = mf[mfi].box(); const Box realspace_bx = bx.enclosedCells(); // discards last point in each nodal direction AMREX_ALWAYS_ASSERT( realspace_bx == tmpRealField[mfi].box() ); Array4 mf_arr = mf[mfi].array(); Array4 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, [=] 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(); }); } } } 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 } }