diff options
Diffstat (limited to 'Source/FieldSolver/SpectralSolver')
| -rw-r--r-- | Source/FieldSolver/SpectralSolver/Make.package | 8 | ||||
| -rw-r--r-- | Source/FieldSolver/SpectralSolver/PsatdSolver.H | 26 | ||||
| -rw-r--r-- | Source/FieldSolver/SpectralSolver/PsatdSolver.cpp | 185 | ||||
| -rw-r--r-- | Source/FieldSolver/SpectralSolver/SpectralData.H | 43 | ||||
| -rw-r--r-- | Source/FieldSolver/SpectralSolver/SpectralData.cpp | 137 | ||||
| -rw-r--r-- | Source/FieldSolver/SpectralSolver/WarpXComplex.H | 26 |
6 files changed, 425 insertions, 0 deletions
diff --git a/Source/FieldSolver/SpectralSolver/Make.package b/Source/FieldSolver/SpectralSolver/Make.package new file mode 100644 index 000000000..a6d088038 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/Make.package @@ -0,0 +1,8 @@ +CEXE_headers += WarpXComplex.H +CEXE_headers += SpectralData.H +CEXE_sources += SpectralData.cpp +CEXE_headers += PsatdSolver.H +CEXE_sources += PsatdSolver.cpp + +INCLUDE_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver +VPATH_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver
\ No newline at end of file diff --git a/Source/FieldSolver/SpectralSolver/PsatdSolver.H b/Source/FieldSolver/SpectralSolver/PsatdSolver.H new file mode 100644 index 000000000..a2bf7e84b --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/PsatdSolver.H @@ -0,0 +1,26 @@ +#ifndef WARPX_PSATD_SOLVER_H_ +#define WARPX_PSATD_SOLVER_H_ + +#include <AMREX_Real.H> +#include <WarpXComplex.H> +#include <SpectralData.H> + +using namespace amrex; +using namespace Gpu; + +using SpectralVector = LayoutData<ManagedVector<Real>> + +class PsatdSolver +{ + using SpectralCoefficients = FabArray<BaseFab<Real>> + + public: + PsatdSolver( const BoxArray& ba, const DistributionMapping& dm, const Real* dx ); + void pushSpectralFields( SpectralData& f ) const; + + private: + SpectralVector kx, ky, kz; + SpectralCoefficients C_coef, S_ck_coef, X1_coef, X2_coef, X3_coef; +}; + +#endif // WARPX_PSATD_SOLVER_H_ diff --git a/Source/FieldSolver/SpectralSolver/PsatdSolver.cpp b/Source/FieldSolver/SpectralSolver/PsatdSolver.cpp new file mode 100644 index 000000000..2b18f7a83 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/PsatdSolver.cpp @@ -0,0 +1,185 @@ +#include <PsatdSolver.H> +#include <WarpXConst.H> +#include <cmath> + +using namespace amrex; +using namespace Gpu; + +/* + * ba: BoxArray for spectral space + * dm: DistributionMapping for spectral space + */ +PsatdSolver::PsatdSolver( const BoxArray& ba, const DistributionMapping& dm, + const Real* dx, const Real dt ) +{ + // Allocate the 1D vectors + kx = SpectralVector( ba, dm ); + ky = SpectralVector( ba, dm ); + kz = SpectralVector( ba, dm ); + for ( MFIter mfi(ba, dm); mfi.isValid(); ++mfi ){ + Box bx = ba[mfi]; + AllocateAndFillKvector( kx[mfi], bx, dx, 0 ) + AllocateAndFillKvector( ky[mfi], bx, dx, 1 ) + AllocateAndFillKvector( kz[mfi], bx, dx, 2 ) + } + + // Allocate the arrays of coefficients + C_coef = SpectralMatrix( ba, dm, 1, 0 ); + S_ck_coef = SpectralMatrix( ba, dm, 1, 0 ); + X1_coef = SpectralMatrix( ba, dm, 1, 0 ); + X2_coef = SpectralMatrix( ba, dm, 1, 0 ); + X3_coef = SpectralMatrix( ba, dm, 1, 0 ); + + // Fill them with the right values: + // Loop over boxes + for ( MFIter mfi(ba, dm); mfi.isValid(); ++mfi ){ + + const Box& bx = mfi.box(); + + // Extract pointers for the k vectors + const Real* kx = kx[mfi].dataPtr(); + const Real* ky = ky[mfi].dataPtr(); + const Real* kz = kz[mfi].dataPtr(); + // Extract arrays for the coefficients + Array4<Real> C = C_coef[mfi].array(); + Array4<Real> S_ck = S_ck_coef[mfi].array(); + Array4<Real> X1 = X1_coef[mfi].array(); + Array4<Real> X2 = X2_coef[mfi].array(); + Array4<Real> X3 = X3_coef[mfi].array(); + + // Loop over indices within one box + ParallelFor( bx, + [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept + { + // Calculate norm of vector + const Real k_norm = std::sqrt( kx[i]*kx[i] + ky[j]*ky[j] + kz[k]*kz[k] ); + + // Calculate coefficients + constexpr Real c = PhysConst::c; + constexpr Real ep0 = PhysConst::ep0; + if ( k_norm != 0 ){ + C(i,j,k) = std::cos( c*k_norm*dt ); + S_ck(i,j,k) = std::sin( c*k_norm*dt )/( c*k_norm ); + X1(i,j,k) = (1. - C(i,j,k))/(ep0 * c*c * k_norm*k_norm); + X2(i,j,k) = (1. - S_ck(i,j,k)/dt )/(ep0 * k_norm*k_norm); + X3(i,j,k) = (C(i,j,k) - S_ck(i,j,k)/dt )/(ep0 * k_norm*k_norm); + } else { // Handle k_norm = 0, by using the analytical limit + C(i,j,k) = 1.; + S_ck(i,j,k) = dt; + X1(i,j,k) = 0.5 * dt*dt / ep0; + X2(i,j,k) = c*c * dt*dt / (6.*ep0); + X3(i,j,k) = - c*c * dt*dt / (3.*ep0); + } + }); + } +} + +void +PsatdSolver::pushSpectralFields( SpectralFields& f ) const{ + + // Loop over boxes + for ( MFIter mfi(f.Ex); mfi.isValid(); ++mfi ){ + + const Box& bx = mfi.box(); + + // Extract arrays for the fields to be updated + Array4<Complex> Ex_arr = f.Ex[mfi].array(); + Array4<Complex> Ey_arr = f.Ey[mfi].array(); + Array4<Complex> Ez_arr = f.Ez[mfi].array(); + Array4<Complex> Bx_arr = f.Bx[mfi].array(); + Array4<Complex> By_arr = f.By[mfi].array(); + Array4<Complex> Bz_arr = f.Bz[mfi].array(); + // Extract arrays for J + const Array4<Complex> Jx_arr = f.Jx[mfi].array(); + const Array4<Complex> Jy_arr = f.Jy[mfi].array(); + const Array4<Complex> Jz_arr = f.Jz[mfi].array(); + const Array4<Complex> rho_old_arr = f.rho_old[mfi].array(); + const Array4<Complex> rho_new_arr = f.rho_new[mfi].array(); + // Extract arrays for the coefficients + const Array4<Real> C_arr = C_coef[mfi].array(); + const Array4<Real> S_ck_arr = S_ck_coef[mfi].array(); + const Array4<Real> inv_k2_arr = + // Extract pointers for the k vectors + const Real* kx_arr = kx[mfi].dataPtr(); + const Real* ky_arr = ky[mfi].dataPtr(); + const Real* kz_arr = kz[mfi].dataPtr(); + + // Loop over indices within one box + ParallelFor( bx, + [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept + { + // Record old values of the fields to be updated + const Complex Ex_old = Ex_arr(i,j,k); + const Complex Ey_old = Ey_arr(i,j,k); + const Complex Ez_old = Ez_arr(i,j,k); + const Complex Bx_old = Bx_arr(i,j,k); + const Complex By_old = By_arr(i,j,k); + const Complex Bz_old = Bz_arr(i,j,k); + // k vector values, and coefficients + const Real kx = kx_arr[i]; + const Real ky = ky_arr[j]; + const Real kz = kz_arr[k]; + constexpr Real c2 = PhysConst::c*PhysConst::c; + constexpr Real inv_ep0 = 1./PhysConst::ep0; + constexpr Complex I = Complex{0,1}; + const Real C = C_arr(i,j,k); + const Real S_ck = S_ck_arr(i,j,k); + const Real X1 = X1_arr(i,j,k); + const Real X2 = X2_arr(i,j,k); + const Real X3 = X3_arr(i,j,k); + // Short cut for the values of J and rho + const Complex Jx = Jx_arr(i,j,k); + const Complex Jy = Jy_arr(i,j,k); + const Complex Jz = Jz_arr(i,j,k); + + // Update E (see WarpX online documentation: theory section) + Ex_arr(i,j,k) = C*Ex_old + + S_ck*( c2*I*(ky*Bz_old - kz*By_old) - inv_ep0*Jx ) + - I*( X2*rho_new - X3*rho_old )*kx; + Ey_arr(i,j,k) = C*Ey_old + + S_ck*( c2*I*(kz*Bx_old - kx*Bz_old) - inv_ep0*Jy ) + - I*( X2*rho_new - X3*rho_old )*ky; + Ez_arr(i,j,k) = C*Ez_old + + S_ck*( c2*I*(kx*By_old - ky*Bx_old) - inv_ep0*Jz ) + - I*( X2*rho_new - X3*rho_old )*kz; + // Update B (see WarpX online documentation: theory section) + Bx_arr(i,j,k) = C*Bx_old + - S_ck*I*(ky*Ez_old - kz*Ey_old) + + X1*I*(ky*Jz_old - kz*Jy_old); + By_arr(i,j,k) = C*By_old + - S_ck*I*(kz*Ex_old - kx*Ez_old) + + X1*I*(kz*Jx_old - kx*Jz_old); + Bz_arr(i,j,k) = C*Bz_old + - S_ck*I*(kx*Ey_old - ky*Ex_old) + + X1*I*(kx*Jy_old - ky*Jx_old); + }); + } +} + +AllocateAndFillKvector( ManagedVector<Real>& k, const Box& bx, const Real* dx, const int i_dim ) +{ + // Alllocate k to the right size + int N = bx.length( i_dim ); + k.resize( N ); + + // Fill the k vector + const Real PI = std::atan(1.0)*4; + const Real dk = 2*PI/(N*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."); + // Fill positive values of k (FFT conventions: first half is positive) + for (int i=0; i<(N+1)/2; i++ ){ + k[i] = i*dk; + } + // Fill negative values of k (FFT conventions: second half is negative) + for (int i=(N+1)/2, i<N; i++){ + k[i] = (N-i)*dk; + } + // TODO: This should be quite different for the hybrid spectral code: + // In that case we should take into consideration the actual indices of the box + // and distinguish the size of the local box and that of the global FFT + // TODO: For real-to-complex, + +} diff --git a/Source/FieldSolver/SpectralSolver/SpectralData.H b/Source/FieldSolver/SpectralSolver/SpectralData.H new file mode 100644 index 000000000..e294287c5 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralData.H @@ -0,0 +1,43 @@ +#ifndef WARPX_PSATD_DATA_H_ +#define WARPX_PSATD_DATA_H_ + +#include <WarpXComplex.H> +#include <AMReX_MultiFab.H> + +using namespace amrex; + +// Declare spectral types +using SpectralField = FabArray<BaseFab<Complex>>; + +/* 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 + */ +class SpectralData +{ +#ifdef AMREX_USE_GPU +// Add cuFFT-specific code +#else + using FFTplans = LayoutData<fftw_plan>; +#endif + + public: + SpectralData( const BoxArray& realspace_ba, + const BoxArray& spectralspace_ba, + const DistributionMapping& dm ); + ~SpectralData(); + void ForwardTransform( const MultiFab& mf, const int field_index ); + void InverseTransform( MultiFab& mf, const int field_index ); + + private: + SpectralField Ex, Ey, Ez, Bx, By, Bz, Jx, Jy, Jz, rho_old, rho_new; + SpectralField tmpRealField, tmpSpectralField; // Store fields before/after transform + FFTplans forward_plan, inverse_plan; + SpectralField& getSpectralField( const int field_index ); +}; + +#endif // WARPX_PSATD_DATA_H_ diff --git a/Source/FieldSolver/SpectralSolver/SpectralData.cpp b/Source/FieldSolver/SpectralSolver/SpectralData.cpp new file mode 100644 index 000000000..316849064 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralData.cpp @@ -0,0 +1,137 @@ +#include <WarpXComplex.H> +#include <AMReX_MultiFab.H> + +using namespace amrex; + +SpectralData::SpectralData( const BoxArray& realspace_ba, + const BoxArray& spectralspace_ba, + const DistributionMapping& dm ) +{ + // 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 and initialize the FFT plans + forward_plan = FFTplans(spectralspace_ba, dm); + inverse_plan = FFTplans(spectralspace_ba, dm); + for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){ + Box bx = spectralspace_ba[mfi]; +#ifdef AMREX_USE_GPU + // Add cuFFT-specific code +#else + // Create FFTW plans + forward_plan[mfi] = fftw_plan_dft_3d( + // Swap dimensions: AMReX data is Fortran-order, but FFTW is C-order + bx.length(2), bx.length(1), bx.length(0), + reinterpret_cast<fftw_complex*>( tmpRealField[mfi].dataPtr() ), + reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ), + FFTW_FORWARD, FFTW_ESTIMATE ); + inverse_plan[mfi] = fftw_plan_dft_3d( + // Swap dimensions: AMReX data is Fortran-order, but FFTW is C-order + bx.length(2), bx.length(1), bx.length(0), + reinterpret_cast<fftw_complex*>( tmpSpectralField[mfi].dataPtr() ), + reinterpret_cast<fftw_complex*>( tmpRealField[mfi].dataPtr() ), + FFTW_BACKWARD, FFTW_ESTIMATE ); + // TODO: Add 2D code + // 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 + } +} + + +SpectralData::~SpectralData() +{ + 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( inverse_plan[mfi] ); +#endif + } +} + +/* TODO: Documentation + * Example: ForwardTransform( Efield_cp[0], SpectralFieldIndex::Ex ) + */ +void +SpectralData::ForwardTransform( const MultiFab& mf, const int field_index ) +{ + // 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<const Real> mf_arr = mf[mfi].array(); + Array4<Complex> 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); + }); + } + + // 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<Complex> field_arr = field[mfi].array(); + Array4<const Complex> 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); + }); + } + } +} + + +SpectralField& +SpectralData::getSpectralField( const int field_index ) +{ + switch(field_index) + { + case SpectralFieldIndex::Ex : return Ex; + case SpectralFieldIndex::Ey : return Ey; + case SpectralFieldIndex::Ez : return Ez; + case SpectralFieldIndex::Bx : return Bx; + case SpectralFieldIndex::By : return By; + case SpectralFieldIndex::Bz : return Bz; + case SpectralFieldIndex::Jx : return Jx; + case SpectralFieldIndex::Jy : return Jy; + case SpectralFieldIndex::Jz : return Jz; + case SpectralFieldIndex::rho_old : return rho_old; + case SpectralFieldIndex::rho_new : return rho_new; + default : return tmpSpectralField; // For synthax; should not occur in practice + } +} diff --git a/Source/FieldSolver/SpectralSolver/WarpXComplex.H b/Source/FieldSolver/SpectralSolver/WarpXComplex.H new file mode 100644 index 000000000..8e2b3977f --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/WarpXComplex.H @@ -0,0 +1,26 @@ +#ifndef WARPX_COMPLEX_H_ +#define WARPX_COMPLEX_H_ + +#include <AMReX_REAL.H> + +// Define complex type on GPU/CPU +#ifdef AMREX_USE_GPU + +#include <thrust/complex.h> +#include <cufft.h> +using Complex = thrust::complex<amrex::Real>; +static_assert( sizeof(Complex) == sizeof(cuDoubleComplex), + "The complex types in WarpX and cuFFT do not match."); + +#else + +#include <complex> +#include <fftw3.h> +using Complex = std::complex<amrex::Real>; +static_assert( sizeof(Complex) == sizeof(fftw_complex), + "The complex types in WarpX and FFTW do not match."); + +#endif +static_assert(sizeof(Complex) == sizeof(amrex::Real[2]), "Unexpected complex type."); + +#endif //WARPX_COMPLEX_H_ |