diff options
| author |  David Grote <grote1@llnl.gov>
| 2020-04-14 15:51:10 -0700 |
|---|---|---|
| committer |  GitHub <noreply@github.com>
| 2020-04-14 15:51:10 -0700 |
| commit | d61cdd0fbef823f5135f7e9c4251ec09ee586651 (patch) | |
| tree | 7fc39e980f69013840402c9419c7e2134b3d3f6c /Source/FieldSolver/SpectralSolver | |
| parent | b28820fe545779411bfb32ca3c574ad89a2a611a (diff) | |
| download | WarpX-d61cdd0fbef823f5135f7e9c4251ec09ee586651.tar.gz WarpX-d61cdd0fbef823f5135f7e9c4251ec09ee586651.tar.zst WarpX-d61cdd0fbef823f5135f7e9c4251ec09ee586651.zip | |
Implementation of the RZ spectral solver (#816)
* For diagnostics, added RZ modes of scalars, allowed different centerings
* For RZ, generalized the centering of the inverse volume scaling of J and rho
* Fixed spacing in ConstructTotalRZScalarField
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Added Python wrapper of charge density arrays
* Add assert ensuring that Jr is never node-centered
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Small fixes to Python to better handle particle weights
* Implementation of the RZ spectral solver
* Removed k-space filtering code
* Removed more k-space filter code from RZ spectral solver
* For RZ spectral, added _rt for literals and cleaned up namespace use
* In RZ spectral solver, cleaned up some member names
* Update Docs/source/building/rzgeometry.rst
Small fix for clarity.
Co-Authored-By: Remi Lehe <remi.lehe@normalesup.org>
* Update Source/Evolve/WarpXEvolve.cpp
Fix macro indentation
Co-Authored-By: Remi Lehe <remi.lehe@normalesup.org>
* Update Source/Evolve/WarpXEvolve.cpp
Fix more macro indentation
Co-Authored-By: Remi Lehe <remi.lehe@normalesup.org>
* Update Source/Evolve/WarpXEvolve.cpp
Fix another macro indentation
Co-Authored-By: Remi Lehe <remi.lehe@normalesup.org>
* New diagnostics support RZ (#836)
* Start implementation of new averaging with staggering:
- face-to-cell-center and edge-to-cell-center replaced so far;
- TODO: node-to-cell-center and 1D behavior (AMREX_SPACEDIM=1).
* first implementation of Diags base classes
* add example, temporarily
* Continue implementation of new averaging with staggering:
- new function takes reference to single MultiFab (no vector);
- TODO: node-to-cell-center still in progress.
* Fix small bug and clean up
* Fix bug in loop over n=0,...,ncomp-1 and clean up
* add more functions
* Add Doxygen documentation and clean up
* Small clean-up in Doxygen documentation
* Compile in single precision: add _rt suffix to avoid unnecessary conversions
* Avoid accessing staggering index directly from IntVect in innermost loops
* Replace do-while loop with for loop (default ncomp=1)
* Remove temporary pointer and pass reference to MultiFab (instead of MultiFab*)
* Replace AMREX_LAUNCH_HOST_DEVICE_LAMBDA with ParallelFor
* cleaning and initialize output mf
* use general average routine
* move flush in new class, and implemented the Plotfile derived class
* add comments
* eol
* free memory in destructor
* typo
* typo
* no need to clear MF pointers there
* though shalt not break existing tests
* FlushRaw doesnt have to be virtual for now
* The importance of being constant
* Capability to select fields in output files
* EOL
* revert to old inputs
* const in right place
* avoid brace initializer there
* oops, fix logic error in is_in
* user can choose flush interval, same behavior as plot_int
* Add option to plot raw fields
* eol
* replace ter flush with dump to avoid confusion
* add options
* Diagnostics stores a vector of functors to compute diags on the fly
* eol
* Field gather from diags to sync particle quantities
* New diagnostics handle RZ with same behavior as old ones
* cleaning and doc
* const ref for string
* smarter for loop from Axel and typo fix from Reva
* simplify code following Dave's comments
* rename mode_avg to convertRZmodes2cartesian
* Update CellCenterFunctor.H
* fill varnames and vector of functors at the same time
* WarpX instance not needed here
* add const
Co-authored-by: Edoardo Zoni <ezoni@lbl.gov>
* Add load balance options documentation (#842)
* Add load balance options documentation
* Add load balance options documentations
* EOL
* Replace tilebox by growntilebox (#849)
* Updated Profiling information in running_cpp (#776)
* Fixed link that was pointing to 404 error page
* Added motivation for profiling and TINYPROFILERS explanation
* Moved section on NERSC profiling to developers Docs
* Update Docs/source/running_cpp/profiling.rst
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Update tiny profilers suggestion Docs/source/running_cpp/profiling.rst
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Fix typo Docs/source/running_cpp/profiling.rst
Co-authored-by: MaxThevenet <mthevenet@lbl.gov>
* Add a few additional diags (divE etc.) (#844)
* Start implementation of new averaging with staggering:
- face-to-cell-center and edge-to-cell-center replaced so far;
- TODO: node-to-cell-center and 1D behavior (AMREX_SPACEDIM=1).
* first implementation of Diags base classes
* add example, temporarily
* Continue implementation of new averaging with staggering:
- new function takes reference to single MultiFab (no vector);
- TODO: node-to-cell-center still in progress.
* Fix small bug and clean up
* Fix bug in loop over n=0,...,ncomp-1 and clean up
* add more functions
* Add Doxygen documentation and clean up
* Small clean-up in Doxygen documentation
* Compile in single precision: add _rt suffix to avoid unnecessary conversions
* Avoid accessing staggering index directly from IntVect in innermost loops
* Replace do-while loop with for loop (default ncomp=1)
* Remove temporary pointer and pass reference to MultiFab (instead of MultiFab*)
* Replace AMREX_LAUNCH_HOST_DEVICE_LAMBDA with ParallelFor
* cleaning and initialize output mf
* use general average routine
* move flush in new class, and implemented the Plotfile derived class
* add comments
* eol
* free memory in destructor
* typo
* typo
* no need to clear MF pointers there
* though shalt not break existing tests
* FlushRaw doesnt have to be virtual for now
* The importance of being constant
* Capability to select fields in output files
* EOL
* revert to old inputs
* const in right place
* avoid brace initializer there
* oops, fix logic error in is_in
* user can choose flush interval, same behavior as plot_int
* Add option to plot raw fields
* eol
* replace ter flush with dump to avoid confusion
* add options
* Diagnostics stores a vector of functors to compute diags on the fly
* eol
* Field gather from diags to sync particle quantities
* New diagnostics handle RZ with same behavior as old ones
* cleaning and doc
* const ref for string
* smarter for loop from Axel and typo fix from Reva
* Functors to compute some fields
* simplify code following Dave's comments
* Create subfolders and add more output options (divE etc.)
* eol
* rename mode_avg to convertRZmodes2cartesian
* Update CellCenterFunctor.H
* fill varnames and vector of functors at the same time
* output rho_new, not rho_old
* WarpX instance not needed here
* add const
* little bit more of reorganization
* Apply suggestions from code review
Co-Authored-By: Axel Huebl <axel.huebl@plasma.ninja>
* add a bunch of const
* make derived classes final
Co-authored-by: Edoardo Zoni <ezoni@lbl.gov>
Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>
* Add Initial Distribution Test (#735)
* Add Histogram
* Add normalization
* Add doc
* Minor
* Minor
* Fix a bug
* Add gaussian distribution test
* Fix alert and change amr.plot_int
* Add maxwell-boltzmann distribution test
* Add maxwell-boltzmann distribution test
* Add maxwell-boltzmann distribution test
* Add maxwell-juttner
* Minor
* Typo
* Minor
* Minor
* Add const
* Apply suggestions from code review
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Modify based on suggestions.
* Add histogram name
* Add bin values
* Don't add histogram name
* Modify read_raw_data.py
* Add doc
* Change ux,uy,uz units
* Change ux,uy,uz units
* Change if format
* Save some variables
* Change more
* Minor
* Fix a bug on GPU
* Fix a bug on GPU
* Add wrong species name abort
* Minor doc
* Change #include format
* Apply suggestions from code review
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Add const
* Change to member variables
* revert
* Change units based on changes of PR#727
* merge
* Add Gaussian position distribution test
* Minor
* Change based on suggestions
* Use read_raw_data.py
* Minor
* Change to no normalization
* Add more in doc
* doc
* doc
* Use relative error
* Don't divide by bin_size
* Change based on suggestions
Co-authored-by: MaxThevenet <mthevenet@lbl.gov>
* Tests: Fix Bool Switch Typo OMP (#854)
useOMP is 0 (False) or 1 (True)
* Costs vector of (pointer to) vector (#829)
* [WIP] costs mf --> costs vector
* [WIP] costs vector
* [WIP] vector costs
* formatting
* makefile
* [WIP] costs vector
* [WIP] *= costs
* wts do not need to divide by num cells
* Tiling safety on CPU
* Add tests
* EOL
* Remove unneeded input
* Update Source/WarpX.H costs documentation
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Update timers with times only if user Timers update
* warpx.-->WarpX::
* warpx.-->WarpX::
* warpx.-->WarpX::
* warpx.-->WarpX::
* warpx.-->WarpX::
* add dev synch
* Update Regression/WarpX-tests.ini
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Delete inputs_loadbalance_costs_heuristic
* Update and rename inputs_loadbalancecosts_timers to inputs_loadbalancecosts
* Update WarpX-tests.ini
Co-authored-by: MaxThevenet <mthevenet@lbl.gov>
* [mini-PR] Read species distribution from OPMD file (#847)
* Added <species>.profile=external_file and .profile_file
* Added description of input parameters to Docs
* Changed from profile to injection option for external file
* Fix typo in amrex abort message (due to copy paste)
* Added the OpenPMD use amrex abort message
* Minor fix - not sure how to remove EOL issue
* Tried to add AddExternalFileBeam functon to PhysicalParticleContainer
* Trued to fix EOL white space issue
* Added read/print species name from OPMD file
* Update Source/Initialization/PlasmaInjector.cpp
Co-Authored-By: Axel Huebl <axel.huebl@plasma.ninja>
* Update Source/Particles/PhysicalParticleContainer.cpp
Co-Authored-By: Axel Huebl <axel.huebl@plasma.ninja>
* Update Source/Particles/PhysicalParticleContainer.cpp
Co-Authored-By: Axel Huebl <axel.huebl@plasma.ninja>
* Update Source/Particles/PhysicalParticleContainer.H
Co-Authored-By: Axel Huebl <axel.huebl@plasma.ninja>
* Update Source/Particles/PhysicalParticleContainer.cpp
Co-Authored-By: Axel Huebl <axel.huebl@plasma.ninja>
* No need to include openPMD header yet
* Fix EOL according to @ax3l's recommendation in #845
* Remove commented out AbortMessage
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Removed commented out part initialization (used only in branch for next PR)
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Added warning that this is WIP
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Changed function name to AddPlasmaFromFile
* Removed AMReX warning from loop
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>
Co-authored-by: MaxThevenet <mthevenet@lbl.gov>
* Ignore python build/dist and egg folders (#850)
* Travis CI: set max numprocs=2 and do not overwrite (#860)
* [mini-PR] Fix bug in Breit-Wheeler engine (#852)
* fixed bug in BW engine
* fixed bug
* fixed bug
* fixed bug
* fixed bug
* fixed bug
* eliminate useless variable
* updated test
* updated inputfile
* Updated tests
* increase tolerance from .04 to .07 in QED 3D BW test
* do plot pos_bw and ele_bw
Co-authored-by: MaxThevenet <mthevenet@lbl.gov>
* Documentation update - towards full SI (#301)
* Added blank line after list. Changed characters in link to Q. H. Liu paper so hyoerlink works with sphinx-build 2.1.2.
* Added line cut unintentionally.
* Removed line added unintentionally.
* Same as before.
* Same as before, but hopefully successfully
* Same as before, but hopefully successfully
* Minor changes to description of PWFA example run
* Revert "Minor changes to description of PWFA example run"
This reverts commit a4d7fa969c906959b018efe683a3e585cbd741f9.
* Revert "Profiler wrapper to allow for cudaDeviceSynchronize (#738)"
This reverts commit bbefc3dad687f4370afd5bc85386d983201cb321.
* Revert "Revert "Minor changes to description of PWFA example run""
This reverts commit 965982d35361ff54d0ad10101ffc31605117e5ac.
* Revert "Minor changes to description of PWFA example run"
This reverts commit a4d7fa969c906959b018efe683a3e585cbd741f9.
* I am making a huge mess with merging
* Minor changes to description of PWFA example run
* Added explanation PWFA simulation section
* Re-structuring. Adding sections for each choice.
* Minor fix to note
* Minor changes to text
* Time step description + fixed line length
* Added FDTD and CKC selection
* Added max time step calculations
* Trying to fix EOL issue
* Added mesh refinement and moving window
* Fixed minor issues
* Fix EOL issues again
* Fixed typo - auxiliary
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
Co-authored-by: Diana Amorim <diana@henrivincenti.dhcp.lbl.gov>
Co-authored-by: MaxThevenet <mthevenet@lbl.gov>
* Remove compiler warnings (#843)
* Fix compiler warnings with DIM=2
* Fix compiler warnings with USE_RZ=TRUE
* Fix compiler warnings with USE_PSATD=TRUE and DIM=2
* Fix compiler warnings with USE_PSATD=TRUE and DIM=3
* Fix bug: discard only return value when calling DefineAndReturnParticleTile
* Remove unused variables not triggering warnings
* [mini-PR] Fix energy calculation for photons in reduced diagnostics (#861)
* fix energy calculation for photons
* fixed typo and mada calculation clearer
* added photon particles in reduced diags test
* Update Source/Diagnostics/ReducedDiags/ParticleEnergy.cpp
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Update Source/Diagnostics/ReducedDiags/ParticleEnergy.cpp
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
Co-authored-by: MaxThevenet <mthevenet@lbl.gov>
* Port rigid injection to the gpu (#862)
* port rigid injection to the gpu
* eol
* Apply suggestions from code review
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* define csqi
Co-authored-by: MaxThevenet <mthevenet@lbl.gov>
* Added blocking factor to 2d and RZ geometries (#864)
* doc: fix formatting for ascent yaml examples (#865)
* [mini-PR] Clarifying ionizable particle charge (#863)
* Added documentation note on ionization particle charge
* Added correct charge of ion to be ionized
* Corrected multiplication symbol
Co-Authored-By: MaxThevenet <mthevenet@lbl.gov>
* Testing doxygen issue
* Charge correction only to ionizable species
* Trying to fix doxygen url fetch issue
Co-authored-by: MaxThevenet <mthevenet@lbl.gov>
* In HankelTransform, added explicit matrix multiply for GPU
* In RZ spectral solver, update setval to be on device
* Removed CEXE_headers from FieldSolver/SpectralSolver/Make.package
* In HankelTransform, added check of the Bessel root finder
* Updated includes to RZ spectral solver
* Added comments on how Hankel transform matrix is calculated
* Added more comments to Hankel transform calculation
* For RZ spectral solver, cleaned up naming and add subdirectory for Hankel transform files
* Cleaned up code in PsatdAlgorithmRZ.cpp
* Updated comment for fields in SpectralFieldDataRZ.cpp
* Changed HankelTransformer to MultiSpectralHankelTransformer
* Updates comments in RZ spectral solver
* Removed code for k-space filtering that will be added in a later PR
* For RZ spectral solver, passed lev in
* Fixed comment in SpectralFieldDataRZ.cpp
Co-authored-by: MaxThevenet <mthevenet@lbl.gov>
Co-authored-by: Remi Lehe <remi.lehe@normalesup.org>
Co-authored-by: Edoardo Zoni <ezoni@lbl.gov>
Co-authored-by: Michael E Rowan <38045958+mrowan137@users.noreply.github.com>
Co-authored-by: NeilZaim <49716072+NeilZaim@users.noreply.github.com>
Co-authored-by: L. Diana Amorim <LDianaAmorim@lbl.gov>
Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>
Co-authored-by: Yinjian Zhao <yinjianzhao@lbl.gov>
Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>
Co-authored-by: Luca Fedeli <luca.fedeli@cea.fr>
Co-authored-by: Diana Amorim <diana@henrivincenti.dhcp.lbl.gov>
Co-authored-by: Andrew Myers <atmyers@lbl.gov>
Co-authored-by: Cyrus Harrison <cyrush@llnl.gov>
Diffstat (limited to 'Source/FieldSolver/SpectralSolver')
19 files changed, 1909 insertions, 1 deletions
diff --git a/Source/FieldSolver/SpectralSolver/Make.package b/Source/FieldSolver/SpectralSolver/Make.package index 2d7aafc8f..549347135 100644 --- a/Source/FieldSolver/SpectralSolver/Make.package +++ b/Source/FieldSolver/SpectralSolver/Make.package @@ -2,6 +2,13 @@ CEXE_sources += SpectralSolver.cpp CEXE_sources += SpectralFieldData.cpp CEXE_sources += SpectralKSpace.cpp +ifeq ($(USE_RZ),TRUE) + CEXE_sources += SpectralSolverRZ.cpp + CEXE_sources += SpectralFieldDataRZ.cpp + CEXE_sources += SpectralKSpaceRZ.cpp + include $(WARPX_HOME)/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/Make.package +endif + include $(WARPX_HOME)/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package VPATH_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package index 93cefcb66..436d99adf 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package +++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/Make.package @@ -3,4 +3,9 @@ CEXE_sources += PsatdAlgorithm.cpp CEXE_sources += GalileanAlgorithm.cpp CEXE_sources += PMLPsatdAlgorithm.cpp +ifeq ($(USE_RZ),TRUE) + CEXE_sources += SpectralBaseAlgorithmRZ.cpp + CEXE_sources += PsatdAlgorithmRZ.cpp +endif + VPATH_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver/SpectralAlgorithms diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H new file mode 100644 index 000000000..309bb050f --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.H @@ -0,0 +1,38 @@ +/* Copyright 2019 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#ifndef WARPX_PSATD_ALGORITHM_RZ_H_ +#define WARPX_PSATD_ALGORITHM_RZ_H_ + +#include "SpectralBaseAlgorithmRZ.H" + +/* \brief Class that updates the field in spectral space + * and stores the coefficients of the corresponding update equation. + */ +class PsatdAlgorithmRZ : public SpectralBaseAlgorithmRZ +{ + + public: + PsatdAlgorithmRZ(SpectralKSpaceRZ const & spectral_kspace, + amrex::DistributionMapping const & dm, + int const n_rz_azimuthal_modes, int const norder_z, + bool const nodal, amrex::Real const dt_step); + // Redefine functions from base class + virtual void pushSpectralFields(SpectralFieldDataRZ & f) override final; + virtual int getRequiredNumberOfFields() const override final { + return SpectralFieldIndex::n_fields; + } + + void InitializeSpectralCoefficients(SpectralFieldDataRZ const & f); + + private: + bool coefficients_initialized; + // Note that dt is saved to use in InitializeSpectralCoefficients + amrex::Real dt; + SpectralCoefficients C_coef, S_ck_coef, X1_coef, X2_coef, X3_coef; +}; + +#endif // WARPX_PSATD_ALGORITHM_RZ_H_ diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp new file mode 100644 index 000000000..6ec4c523a --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithmRZ.cpp @@ -0,0 +1,198 @@ +/* Copyright 2019-2020 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#include "PsatdAlgorithmRZ.H" +#include "Utils/WarpXConst.H" + +#include <cmath> + +using amrex::operator""_rt; + + +/* \brief Initialize coefficients for the update equation */ +PsatdAlgorithmRZ::PsatdAlgorithmRZ (SpectralKSpaceRZ const & spectral_kspace, + amrex::DistributionMapping const & dm, + int const n_rz_azimuthal_modes, int const norder_z, + bool const nodal, amrex::Real const dt_step) + // Initialize members of base class + : SpectralBaseAlgorithmRZ(spectral_kspace, dm, + norder_z, nodal), + dt(dt_step) +{ + + // Allocate the arrays of coefficients + amrex::BoxArray const & ba = spectral_kspace.spectralspace_ba; + C_coef = SpectralCoefficients(ba, dm, n_rz_azimuthal_modes, 0); + S_ck_coef = SpectralCoefficients(ba, dm, n_rz_azimuthal_modes, 0); + X1_coef = SpectralCoefficients(ba, dm, n_rz_azimuthal_modes, 0); + X2_coef = SpectralCoefficients(ba, dm, n_rz_azimuthal_modes, 0); + X3_coef = SpectralCoefficients(ba, dm, n_rz_azimuthal_modes, 0); + + coefficients_initialized = false; +} + +/* Advance the E and B field in spectral space (stored in `f`) + * over one time step */ +void +PsatdAlgorithmRZ::pushSpectralFields(SpectralFieldDataRZ & f) +{ + + if (not coefficients_initialized) { + // This is called from here since it needs the kr values + // which can be obtained from the SpectralFieldDataRZ + InitializeSpectralCoefficients(f); + coefficients_initialized = true; + } + + // Loop over boxes + for (amrex::MFIter mfi(f.fields); mfi.isValid(); ++mfi){ + + amrex::Box const & bx = f.fields[mfi].box(); + + // Extract arrays for the fields to be updated + amrex::Array4<Complex> const& fields = f.fields[mfi].array(); + // Extract arrays for the coefficients + amrex::Array4<const amrex::Real> const& C_arr = C_coef[mfi].array(); + amrex::Array4<const amrex::Real> const& S_ck_arr = S_ck_coef[mfi].array(); + amrex::Array4<const amrex::Real> const& X1_arr = X1_coef[mfi].array(); + amrex::Array4<const amrex::Real> const& X2_arr = X2_coef[mfi].array(); + amrex::Array4<const amrex::Real> const& X3_arr = X3_coef[mfi].array(); + + // Extract pointers for the k vectors + auto const & kr_modes = f.getKrArray(mfi); + amrex::Real const* kr_arr = kr_modes.dataPtr(); + amrex::Real const* modified_kz_arr = modified_kz_vec[mfi].dataPtr(); + int const nr = bx.length(0); + + // Loop over indices within one box + // Note that k = 0 + int const modes = f.n_rz_azimuthal_modes; + amrex::ParallelFor(bx, modes, + [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept + { + + // All of the fields of each mode are grouped together + using Idx = SpectralFieldIndex; + auto const Ep_m = Idx::Ex + Idx::n_fields*mode; + auto const Em_m = Idx::Ey + Idx::n_fields*mode; + auto const Ez_m = Idx::Ez + Idx::n_fields*mode; + auto const Bp_m = Idx::Bx + Idx::n_fields*mode; + auto const Bm_m = Idx::By + Idx::n_fields*mode; + auto const Bz_m = Idx::Bz + Idx::n_fields*mode; + auto const Jp_m = Idx::Jx + Idx::n_fields*mode; + auto const Jm_m = Idx::Jy + Idx::n_fields*mode; + auto const Jz_m = Idx::Jz + Idx::n_fields*mode; + auto const rho_old_m = Idx::rho_old + Idx::n_fields*mode; + auto const rho_new_m = Idx::rho_new + Idx::n_fields*mode; + + // Record old values of the fields to be updated + Complex const Ep_old = fields(i,j,k,Ep_m); + Complex const Em_old = fields(i,j,k,Em_m); + Complex const Ez_old = fields(i,j,k,Ez_m); + Complex const Bp_old = fields(i,j,k,Bp_m); + Complex const Bm_old = fields(i,j,k,Bm_m); + Complex const Bz_old = fields(i,j,k,Bz_m); + // Shortcut for the values of J and rho + Complex const Jp = fields(i,j,k,Jp_m); + Complex const Jm = fields(i,j,k,Jm_m); + Complex const Jz = fields(i,j,k,Jz_m); + Complex const rho_old = fields(i,j,k,rho_old_m); + Complex const rho_new = fields(i,j,k,rho_new_m); + + // k vector values, and coefficients + // The k values for each mode are grouped together + int const ir = i + nr*mode; + amrex::Real const kr = kr_arr[ir]; + amrex::Real const kz = modified_kz_arr[j]; + + constexpr amrex::Real c2 = PhysConst::c*PhysConst::c; + constexpr amrex::Real inv_ep0 = 1._rt/PhysConst::ep0; + Complex const I = Complex{0._rt,1._rt}; + amrex::Real const C = C_arr(i,j,k,mode); + amrex::Real const S_ck = S_ck_arr(i,j,k,mode); + amrex::Real const X1 = X1_arr(i,j,k,mode); + amrex::Real const X2 = X2_arr(i,j,k,mode); + amrex::Real const X3 = X3_arr(i,j,k,mode); + + // Update E (see WarpX online documentation: theory section) + fields(i,j,k,Ep_m) = C*Ep_old + + S_ck*(-c2*I*kr/2._rt*Bz_old + c2*kz*Bp_old - inv_ep0*Jp) + + kr*(X2*rho_new - X3*rho_old); + fields(i,j,k,Em_m) = C*Em_old + + S_ck*(-c2*I*kr/2._rt*Bz_old - c2*kz*Bm_old - inv_ep0*Jm) + - kr*(X2*rho_new - X3*rho_old); + fields(i,j,k,Ez_m) = C*Ez_old + + S_ck*(c2*I*kr*Bp_old + c2*I*kr*Bm_old - inv_ep0*Jz) + - I*kz*(X2*rho_new - X3*rho_old); + // Update B (see WarpX online documentation: theory section) + fields(i,j,k,Bp_m) = C*Bp_old + - S_ck*(-I*kr/2._rt*Ez_old + kz*Ep_old) + + X1*(-I*kr/2._rt*Jz + kz*Jp); + fields(i,j,k,Bm_m) = C*Bm_old + - S_ck*(-I*kr/2._rt*Ez_old - kz*Em_old) + + X1*(-I*kr/2._rt*Jz - kz*Jm); + fields(i,j,k,Bz_m) = C*Bz_old + - S_ck*I*(kr*Ep_old + kr*Em_old) + + X1*I*(kr*Jp + kr*Jm); + }); + } +}; + +void PsatdAlgorithmRZ::InitializeSpectralCoefficients (SpectralFieldDataRZ const & f) +{ + + // Fill them with the right values: + // Loop over boxes and allocate the corresponding coefficients + // for each box owned by the local MPI proc + for (amrex::MFIter mfi(f.fields); mfi.isValid(); ++mfi){ + + amrex::Box const & bx = f.fields[mfi].box(); + + // Extract pointers for the k vectors + amrex::Real const* const modified_kz = modified_kz_vec[mfi].dataPtr(); + + // Extract arrays for the coefficients + amrex::Array4<amrex::Real> const& C = C_coef[mfi].array(); + amrex::Array4<amrex::Real> const& S_ck = S_ck_coef[mfi].array(); + amrex::Array4<amrex::Real> const& X1 = X1_coef[mfi].array(); + amrex::Array4<amrex::Real> const& X2 = X2_coef[mfi].array(); + amrex::Array4<amrex::Real> const& X3 = X3_coef[mfi].array(); + + auto const & kr_modes = f.getKrArray(mfi); + amrex::Real const* kr_arr = kr_modes.dataPtr(); + int const nr = bx.length(0); + amrex::Real const dt_temp = dt; + + // Loop over indices within one box + int const modes = f.n_rz_azimuthal_modes; + amrex::ParallelFor(bx, modes, + [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept + { + // Calculate norm of vector + int const ir = i + nr*mode; + amrex::Real const kr = kr_arr[ir]; + amrex::Real const kz = modified_kz[j]; + amrex::Real const k_norm = std::sqrt(kr*kr + kz*kz); + + // Calculate coefficients + constexpr amrex::Real c = PhysConst::c; + constexpr amrex::Real ep0 = PhysConst::ep0; + if (k_norm != 0){ + C(i,j,k,mode) = std::cos(c*k_norm*dt_temp); + S_ck(i,j,k,mode) = std::sin(c*k_norm*dt_temp)/(c*k_norm); + X1(i,j,k,mode) = (1._rt - C(i,j,k,mode))/(ep0 * c*c * k_norm*k_norm); + X2(i,j,k,mode) = (1._rt - S_ck(i,j,k,mode)/dt_temp)/(ep0 * k_norm*k_norm); + X3(i,j,k,mode) = (C(i,j,k,mode) - S_ck(i,j,k,mode)/dt_temp)/(ep0 * k_norm*k_norm); + } else { // Handle k_norm = 0, by using the analytical limit + C(i,j,k,mode) = 1._rt; + S_ck(i,j,k,mode) = dt_temp; + X1(i,j,k,mode) = 0.5_rt * dt_temp*dt_temp / ep0; + X2(i,j,k,mode) = c*c * dt_temp*dt_temp / (6._rt*ep0); + X3(i,j,k,mode) = - c*c * dt_temp*dt_temp / (3._rt*ep0); + } + }); + } +} diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithmRZ.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithmRZ.H new file mode 100644 index 000000000..aae17f387 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithmRZ.H @@ -0,0 +1,54 @@ +/* Copyright 2019 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#ifndef WARPX_SPECTRAL_BASE_ALGORITHM_RZ_H_ +#define WARPX_SPECTRAL_BASE_ALGORITHM_RZ_H_ + +#include "FieldSolver/SpectralSolver/SpectralKSpaceRZ.H" +#include "FieldSolver/SpectralSolver/SpectralFieldDataRZ.H" + +/* \brief Class that updates the field in spectral space + * and stores the coefficients of the corresponding update equation. + * + * `SpectralBaseAlgorithmRZ` is only a base class and cannot be used directly. + * Instead use its subclasses, which implement the specific field update + * equations for a given spectral algorithm. + */ +class SpectralBaseAlgorithmRZ +{ + public: + // Virtual member function ; meant to be overridden in subclasses + virtual void pushSpectralFields(SpectralFieldDataRZ & f) = 0; + virtual int getRequiredNumberOfFields() const = 0; + // The destructor should also be a virtual function, so that + // a pointer to subclass of `SpectraBaseAlgorithm` actually + // calls the subclass's destructor. + virtual ~SpectralBaseAlgorithmRZ() {}; + + /** + * \brief Compute spectral divergence of E + */ + void ComputeSpectralDivE ( SpectralFieldDataRZ& field_data, + const std::array<std::unique_ptr<amrex::MultiFab>,3>& Efield, + amrex::MultiFab& divE ); + + protected: // Meant to be used in the subclasses + + using SpectralCoefficients = amrex::FabArray< amrex::BaseFab <amrex::Real> >; + + // Constructor + SpectralBaseAlgorithmRZ(SpectralKSpaceRZ const & spectral_kspace, + amrex::DistributionMapping const & dm, + int const norder_z, bool const nodal) + // Compute and assign the modified k vectors + : modified_kz_vec(spectral_kspace.getModifiedKComponent(dm, 1, norder_z, nodal)) + {}; + + // Modified finite-order vectors + KVectorComponent modified_kz_vec; +}; + +#endif // WARPX_SPECTRAL_BASE_ALGORITHM_RZ_H_ diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithmRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithmRZ.cpp new file mode 100644 index 000000000..679fc5fba --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/SpectralBaseAlgorithmRZ.cpp @@ -0,0 +1,74 @@ +/* Copyright 2019 Edoardo Zoni + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#include "SpectralBaseAlgorithmRZ.H" + +#include <cmath> + +using namespace amrex; + +/** + * \brief Compute spectral divergence of E + */ +void +SpectralBaseAlgorithmRZ::ComputeSpectralDivE ( + SpectralFieldDataRZ& field_data, + const std::array<std::unique_ptr<amrex::MultiFab>,3>& Efield, + amrex::MultiFab& divE ) +{ + using amrex::operator""_rt; + using Idx = SpectralFieldIndex; + + // Forward Fourier transform of E + field_data.ForwardTransform( *Efield[0], Idx::Ex, + *Efield[1], Idx::Ey ); + field_data.ForwardTransform( *Efield[2], Idx::Ez, 0 ); + + // Loop over boxes + for (MFIter mfi(field_data.fields); mfi.isValid(); ++mfi){ + + Box const & bx = field_data.fields[mfi].box(); + + // Extract arrays for the fields to be updated + Array4<Complex> fields = field_data.fields[mfi].array(); + + // Extract pointers for the k vectors + auto const & kr = field_data.getKrArray(mfi); + Real const * kr_arr = kr.dataPtr(); + Real const * modified_kz_arr = modified_kz_vec[mfi].dataPtr(); + + int const nr = bx.length(0); + int const modes = field_data.n_rz_azimuthal_modes; + constexpr int n_fields = SpectralFieldIndex::n_fields; + + // Loop over indices within one box + ParallelFor(bx, modes, + [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept + { + int const ic1 = Idx::Ex + mode*n_fields; + int const ic2 = Idx::Ey + mode*n_fields; + int const ic3 = Idx::Ez + mode*n_fields; + + // Shortcuts for the components of E + Complex const Ep = fields(i,j,0,ic1); + Complex const Em = fields(i,j,0,ic2); + Complex const Ez = fields(i,j,0,ic3); + + // k vector values + int const ir = i + nr*mode; + Real const kr = kr_arr[ir]; + Real const kz = modified_kz_arr[j]; + Complex const I = Complex{0._rt,1._rt}; + + // div(E) in Fourier space + int const ic = Idx::divE + mode*n_fields; + fields(i,j,0,ic) = kr*(Ep - Em) + I*kz*Ez; + }); + } + + // Backward Fourier transform + field_data.BackwardTransform( divE, Idx::divE, 0 ); +}; diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H new file mode 100644 index 000000000..d0e29d070 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H @@ -0,0 +1,84 @@ +/* Copyright 2019-2020 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#ifndef WARPX_SPECTRAL_FIELD_DATA_RZ_H_ +#define WARPX_SPECTRAL_FIELD_DATA_RZ_H_ + +#include "SpectralKSpaceRZ.H" +#include "SpectralFieldData.H" +#include "SpectralHankelTransform/SpectralHankelTransformer.H" +#include <AMReX_MultiFab.H> + +/* \brief Class that stores the fields in spectral space, and performs the + * Fourier transforms between real space and spectral space + */ +class SpectralFieldDataRZ +{ + + public: + + // Define the FFTplans type, which holds one fft plan per box + // (plans are only initialized for the boxes that are owned by + // the local MPI rank) +#ifdef AMREX_USE_GPU + using FFTplans = amrex::LayoutData<cufftHandle>; +#else + using FFTplans = amrex::LayoutData<fftw_plan>; +#endif + // Similarly, define the Hankel transformers for each box. + using MultiSpectralHankelTransformer = amrex::LayoutData<SpectralHankelTransformer>; + + SpectralFieldDataRZ(const amrex::BoxArray& realspace_ba, + const SpectralKSpaceRZ& k_space, + const amrex::DistributionMapping& dm, + const int n_field_required, + const int n_modes, + const int lev); + SpectralFieldDataRZ() = default; // Default constructor + SpectralFieldDataRZ& operator=(SpectralFieldDataRZ&& field_data) = default; + ~SpectralFieldDataRZ(); + + void ForwardTransform(const amrex::MultiFab& mf, + const int field_index, const int i_comp); + void ForwardTransform(const amrex::MultiFab& mf_r, const int field_index_r, + const amrex::MultiFab& mf_t, const int field_index_t); + void BackwardTransform(amrex::MultiFab& mf, + const int field_index, const int i_comp); + void BackwardTransform(amrex::MultiFab& mf_r, const int field_index_r, + amrex::MultiFab& mf_t, const int field_index_t); + + void FABZForwardTransform(amrex::MFIter const & mfi, + amrex::MultiFab const & tempHTransformedSplit, + int field_index, const bool is_nodal_z); + void FABZBackwardTransform(amrex::MFIter const & mfi, const int field_index, + amrex::MultiFab & tempHTransformedSplit, + const bool is_nodal_z); + + // Returns an array that holds the kr for all of the modes + HankelTransform::RealVector const & getKrArray(amrex::MFIter const & mfi) const { + return multi_spectral_hankel_transformer[mfi].getKrArray(); + } + + // `fields` stores fields in spectral space, as multicomponent FabArray + SpectralField fields; + + int n_rz_azimuthal_modes; + + private: + + // tempHTransformed and tmpSpectralField store fields + // right before/after the z Fourier transform + SpectralField tempHTransformed; // contains Complexes + SpectralField tmpSpectralField; // contains Complexes + 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 + SpectralShiftFactor zshift_FFTfromCell, zshift_FFTtoCell; + MultiSpectralHankelTransformer multi_spectral_hankel_transformer; + +}; + +#endif // WARPX_SPECTRAL_FIELD_DATA_RZ_H_ diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp new file mode 100644 index 000000000..4a66d924a --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp @@ -0,0 +1,444 @@ +/* Copyright 2019-2020 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#include "SpectralFieldDataRZ.H" + +#include "WarpX.H" + +using amrex::operator""_rt; + +/* \brief Initialize fields in spectral space, and FFT plans + * + * \param realspace_ba Box array that corresponds to the decomposition + * * of the fields in real space (cell-centered ; includes guard cells only in z) + * \param k_space Defined the domain of the k space + * \param dm Indicates which MPI proc owns which box, in realspace_ba + * \param n_field_required Specifies the number of fields that will be transformed + * \param n_modes Number of cylindrical modes + * */ +SpectralFieldDataRZ::SpectralFieldDataRZ (amrex::BoxArray const & realspace_ba, + SpectralKSpaceRZ const & k_space, + amrex::DistributionMapping const & dm, + int const n_field_required, + int const n_modes, + int const lev) + : n_rz_azimuthal_modes(n_modes) +{ + amrex::BoxArray const & spectralspace_ba = k_space.spectralspace_ba; + + // Allocate the arrays that contain the fields in spectral space. + // SpectralField is comparable to a MultiFab but stores complex numbers. + // This stores all of the transformed fields in one place, with the last dimension + // being the list of fields, defined by SpectralFieldIndex, for all of the modes. + // The fields of each mode are grouped together, so that the index of a + // field for a specific mode is given by field_index + mode*n_fields. + fields = SpectralField(spectralspace_ba, dm, n_rz_azimuthal_modes*n_field_required, 0); + + // Allocate temporary arrays - in real space and spectral space. + // These complex arrays will store the data just before/after the z FFT. + // Note that the realspace_ba should not include the radial guard cells. + tempHTransformed = SpectralField(realspace_ba, dm, n_rz_azimuthal_modes, 0); + tmpSpectralField = SpectralField(spectralspace_ba, dm, n_rz_azimuthal_modes, 0); + + // By default, we assume the z FFT is done from/to a nodal grid in real space. + // It the FFT is performed from/to a cell-centered grid in real space, + // a correcting "shift" factor must be applied in spectral space. + zshift_FFTfromCell = k_space.getSpectralShiftFactor(dm, 1, + ShiftType::TransformFromCellCentered); + zshift_FFTtoCell = k_space.getSpectralShiftFactor(dm, 1, + ShiftType::TransformToCellCentered); + + // Allocate and initialize the FFT plans and Hankel transformer. + forward_plan = FFTplans(spectralspace_ba, dm); +#ifndef AMREX_USE_GPU + // The backward plan is not needed with GPU since it would be the same + // as the forward plan anyway. + backward_plan = FFTplans(spectralspace_ba, dm); +#endif + multi_spectral_hankel_transformer = MultiSpectralHankelTransformer(spectralspace_ba, dm); + + // Loop over boxes and allocate the corresponding plan + // for each box owned by the local MPI proc. + for (amrex::MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi){ + amrex::IntVect grid_size = realspace_ba[mfi].length(); +#ifdef AMREX_USE_GPU + // Create cuFFT plan. + // This is alway complex to complex. + // This plan is for one azimuthal mode only. + cufftResult result; + int fft_length[] = {grid_size[1]}; + int inembed[] = {grid_size[1]}; + int istride = grid_size[0]; + int idist = 1; + int onembed[] = {grid_size[1]}; + int ostride = grid_size[0]; + int odist = 1; + int batch = grid_size[0]; // number of ffts + result = cufftPlanMany(&forward_plan[mfi], 1, fft_length, inembed, istride, idist, + onembed, ostride, odist, CUFFT_Z2Z, batch); + if (result != CUFFT_SUCCESS) { + amrex::Print() << " cufftPlanMany failed! \n"; + } + // The backward plane is the same as the forward since the direction is passed when executed. + +#else + // Create FFTW plans. + fftw_iodim dims[1]; + fftw_iodim howmany_dims[2]; + dims[0].n = grid_size[1]; + dims[0].is = grid_size[0]; + dims[0].os = grid_size[0]; + howmany_dims[0].n = n_rz_azimuthal_modes; + howmany_dims[0].is = grid_size[0]*grid_size[1]; + howmany_dims[0].os = grid_size[0]*grid_size[1]; + howmany_dims[1].n = grid_size[0]; + howmany_dims[1].is = 1; + howmany_dims[1].os = 1; + forward_plan[mfi] = + // Note that AMReX FAB are Fortran-order. + fftw_plan_guru_dft(1, // int rank + dims, + 2, // int howmany_rank, + howmany_dims, + reinterpret_cast<fftw_complex*>(tempHTransformed[mfi].dataPtr()), // fftw_complex *in + reinterpret_cast<fftw_complex*>(tmpSpectralField[mfi].dataPtr()), // fftw_complex *out + FFTW_FORWARD, // int sign + FFTW_ESTIMATE); // unsigned flags + backward_plan[mfi] = + fftw_plan_guru_dft(1, // int rank + dims, + 2, // int howmany_rank, + howmany_dims, + reinterpret_cast<fftw_complex*>(tmpSpectralField[mfi].dataPtr()), // fftw_complex *in + reinterpret_cast<fftw_complex*>(tempHTransformed[mfi].dataPtr()), // fftw_complex *out + FFTW_BACKWARD, // int sign + FFTW_ESTIMATE); // unsigned flags +#endif + + // Create the Hankel transformer for each box. + std::array<amrex::Real,3> xmax = WarpX::UpperCorner(mfi.tilebox(), lev); + multi_spectral_hankel_transformer[mfi] = SpectralHankelTransformer(grid_size[0], n_rz_azimuthal_modes, xmax[0]); + } +} + + +SpectralFieldDataRZ::~SpectralFieldDataRZ() +{ + if (fields.size() > 0){ + for (amrex::MFIter mfi(fields); mfi.isValid(); ++mfi){ +#ifdef AMREX_USE_GPU + // Destroy cuFFT plans. + cufftDestroy(forward_plan[mfi]); + // cufftDestroy(backward_plan[mfi]); // This was never allocated. +#else + // Destroy FFTW plans. + fftw_destroy_plan(forward_plan[mfi]); + fftw_destroy_plan(backward_plan[mfi]); +#endif + } + } +} + +/* \brief Z Transform the FAB to spectral space, + * and store the corresponding result internally + * (in the spectral field specified by `field_index`) + * The input, tempHTransformedSplit, is the complex, Hankel transformed + * data, which is stored wih the real and imaginary parts split. + * The input should include the imaginary component of mode 0 + * (even though it is all zeros). */ +void +SpectralFieldDataRZ::FABZForwardTransform (amrex::MFIter const & mfi, + amrex::MultiFab const & tempHTransformedSplit, + int const field_index, const bool is_nodal_z) +{ + // Copy the split complex to the interleaved complex. + + amrex::Box const& realspace_bx = tempHTransformed[mfi].box(); + + amrex::Array4<const amrex::Real> const& split_arr = tempHTransformedSplit[mfi].array(); + amrex::Array4<Complex> const& complex_arr = tempHTransformed[mfi].array(); + + int const modes = n_rz_azimuthal_modes; + ParallelFor(realspace_bx, modes, + [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept { + int const mode_r = 2*mode; + int const mode_i = 2*mode + 1; + complex_arr(i,j,k,mode) = Complex{split_arr(i,j,k,mode_r), split_arr(i,j,k,mode_i)}; + }); + + // Perform Fourier transform from `tempHTransformed` to `tmpSpectralField`. +#ifdef AMREX_USE_GPU + // Perform Fast Fourier Transform on GPU using cuFFT. + // Make sure that this is done on the same + // GPU stream as the above copy. + cufftResult result; + cudaStream_t stream = amrex::Gpu::Device::cudaStream(); + cufftSetStream(forward_plan[mfi], stream); + for (int mode=0 ; mode < n_rz_azimuthal_modes ; mode++) { + result = cufftExecZ2Z(forward_plan[mfi], + reinterpret_cast<cuDoubleComplex*>(tempHTransformed[mfi].dataPtr(mode)), // cuDoubleComplex *in + reinterpret_cast<cuDoubleComplex*>(tmpSpectralField[mfi].dataPtr(mode)), // cuDoubleComplex *out + CUFFT_FORWARD); + if (result != CUFFT_SUCCESS) { + amrex::Print() << " forward transform using cufftExecZ2Z failed ! \n"; + } + } +#else + fftw_execute(forward_plan[mfi]); +#endif + + // Copy the spectral-space field `tmpSpectralField` to the appropriate + // index of the FabArray `fields` (specified by `field_index`) + // and apply correcting shift factor if the real space data comes + // from a cell-centered grid in real space instead of a nodal grid. + amrex::Array4<const Complex> const& tmp_arr = tmpSpectralField[mfi].array(); + amrex::Array4<Complex> const& fields_arr = fields[mfi].array(); + Complex const* zshift_arr = zshift_FFTfromCell[mfi].dataPtr(); + + // Loop over indices within one box, all components. + // The fields are organized so that the fields for each mode + // are grouped together in memory. + amrex::Box const& spectralspace_bx = tmpSpectralField[mfi].box(); + int const nz = spectralspace_bx.length(1); + amrex::Real inv_nz = 1._rt/nz; + constexpr int n_fields = SpectralFieldIndex::n_fields; + + ParallelFor(spectralspace_bx, modes, + [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept { + Complex spectral_field_value = tmp_arr(i,j,k,mode); + // Apply proper shift. + if (is_nodal_z==false) spectral_field_value *= zshift_arr[j]; + // Copy field into the correct index. + int const ic = field_index + mode*n_fields; + fields_arr(i,j,k,ic) = spectral_field_value*inv_nz; + }); +} + +/* \brief Backward Z Transform the data from the fields + * (in the spectral field specified by `field_index`) + * to physical space, and return the resulting FArrayBox. + * The output, tempHTransformedSplit, is the complex, Hankel transformed + * data, which is stored wih the real and imaginary parts split. + * The output includes the imaginary component of mode 0 + * (even though it is all zeros). */ +void +SpectralFieldDataRZ::FABZBackwardTransform (amrex::MFIter const & mfi, int const field_index, + amrex::MultiFab & tempHTransformedSplit, + const bool is_nodal_z) +{ + // Copy the spectral-space field from the appropriate index of the FabArray + // `fields` (specified by `field_index`) to field `tmpSpectralField` + // and apply correcting shift factor if the real space data is on + // a cell-centered grid in real space instead of a nodal grid. + amrex::Array4<const Complex> const& fields_arr = fields[mfi].array(); + amrex::Array4<Complex> const& tmp_arr = tmpSpectralField[mfi].array(); + Complex const* zshift_arr = zshift_FFTtoCell[mfi].dataPtr(); + + // Loop over indices within one box, all components. + amrex::Box const& spectralspace_bx = tmpSpectralField[mfi].box(); + + int const modes = n_rz_azimuthal_modes; + constexpr int n_fields = SpectralFieldIndex::n_fields; + ParallelFor(spectralspace_bx, modes, + [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept { + int const ic = field_index + mode*n_fields; + Complex spectral_field_value = fields_arr(i,j,k,ic); + // Apply proper shift. + if (is_nodal_z==false) spectral_field_value *= zshift_arr[j]; + // Copy field into the right index. + tmp_arr(i,j,k,mode) = spectral_field_value; + }); + + // Perform Fourier transform from `tmpSpectralField` to `tempHTransformed`. +#ifdef AMREX_USE_GPU + // Perform Fast Fourier Transform on GPU using cuFFT. + // Make sure that this is done on the same + // GPU stream as the above copy. + cufftResult result; + cudaStream_t stream = amrex::Gpu::Device::cudaStream(); + cufftSetStream(forward_plan[mfi], stream); + for (int mode=0 ; mode < n_rz_azimuthal_modes ; mode++) { + result = cufftExecZ2Z(forward_plan[mfi], + reinterpret_cast<cuDoubleComplex*>(tmpSpectralField[mfi].dataPtr(mode)), // cuDoubleComplex *in + reinterpret_cast<cuDoubleComplex*>(tempHTransformed[mfi].dataPtr(mode)), // cuDoubleComplex *out + CUFFT_INVERSE); + if (result != CUFFT_SUCCESS) { + amrex::Print() << " backwardtransform using cufftExecZ2Z failed ! \n"; + } + } +#else + fftw_execute(backward_plan[mfi]); +#endif + + // Copy the interleaved complex to the split complex. + amrex::Box const& realspace_bx = tempHTransformed[mfi].box(); + + amrex::Array4<amrex::Real> const& split_arr = tempHTransformedSplit[mfi].array(); + amrex::Array4<const Complex> const& complex_arr = tempHTransformed[mfi].array(); + + ParallelFor(realspace_bx, modes, + [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept { + int const mode_r = 2*mode; + int const mode_i = 2*mode + 1; + split_arr(i,j,k,mode_r) = complex_arr(i,j,k,mode).real(); + split_arr(i,j,k,mode_i) = complex_arr(i,j,k,mode).imag(); + }); + +} + +/* \brief Transform the component `i_comp` of MultiFab `field_mf` + * to spectral space, and store the corresponding result internally + * (in the spectral field specified by `field_index`) */ +void +SpectralFieldDataRZ::ForwardTransform (amrex::MultiFab const & field_mf, int const field_index, + int const i_comp) +{ + // Check field index type, in order to apply proper shift in spectral space. + // Only cell centered in r is supported. + bool const is_nodal_z = field_mf.is_nodal(1); + + int const ncomp = 2*n_rz_azimuthal_modes - 1; + + // This will hold the Hankel transformed data, with the real and imaginary parts split. + // A full multifab is created so that each GPU stream has its own temp space. + amrex::MultiFab tempHTransformedSplit(tempHTransformed.boxArray(), tempHTransformed.DistributionMap(), 2*n_rz_azimuthal_modes, 0); + + // Loop over boxes. + for (amrex::MFIter mfi(field_mf); mfi.isValid(); ++mfi){ + + // Perform the Hankel transform first. + // tempHTransformedSplit includes the imaginary component of mode 0. + // field_mf does not. + amrex::Box const& realspace_bx = tempHTransformed[mfi].box(); + amrex::FArrayBox field_comp(field_mf[mfi], amrex::make_alias, i_comp*ncomp, ncomp); + multi_spectral_hankel_transformer[mfi].PhysicalToSpectral_Scalar(realspace_bx, field_comp, tempHTransformedSplit[mfi]); + + FABZForwardTransform(mfi, tempHTransformedSplit, field_index, is_nodal_z); + + } +} + +/* \brief Transform the coupled components of MultiFabs `field_mf_r` and `field_mf_t` + * to spectral space, and store the corresponding result internally + * (in the spectral fields specified by `field_index_r` and `field_index_t`) */ +void +SpectralFieldDataRZ::ForwardTransform (amrex::MultiFab const & field_mf_r, int const field_index_r, + amrex::MultiFab const & field_mf_t, int const field_index_t) +{ + // Check field index type, in order to apply proper shift in spectral space. + // Only cell centered in r is supported. + bool const is_nodal_z = field_mf_r.is_nodal(1); + + // Create copies of the input multifabs. The copies will include the imaginary part of mode 0. + // Also note that the Hankel transform will overwrite the copies. + // Full multifabs are created for the temps so that each GPU stream has its own temp space. + amrex::MultiFab field_mf_r_copy(field_mf_r.boxArray(), field_mf_r.DistributionMap(), 2*n_rz_azimuthal_modes, field_mf_r.nGrowVect()); + amrex::MultiFab field_mf_t_copy(field_mf_t.boxArray(), field_mf_t.DistributionMap(), 2*n_rz_azimuthal_modes, field_mf_t.nGrowVect()); + amrex::MultiFab::Copy(field_mf_r_copy, field_mf_r, 0, 0, 1, field_mf_r.nGrowVect()); // Real part of mode 0 + amrex::MultiFab::Copy(field_mf_t_copy, field_mf_t, 0, 0, 1, field_mf_t.nGrowVect()); // Real part of mode 0 + field_mf_r_copy.setVal(0._rt, 1, 1, field_mf_r.nGrowVect()); // Imaginary part of mode 0 + field_mf_t_copy.setVal(0._rt, 1, 1, field_mf_t.nGrowVect()); // Imaginary part of mode 0 + amrex::MultiFab::Copy(field_mf_r_copy, field_mf_r, 1, 2, 2*n_rz_azimuthal_modes-2, field_mf_r.nGrowVect()); + amrex::MultiFab::Copy(field_mf_t_copy, field_mf_t, 1, 2, 2*n_rz_azimuthal_modes-2, field_mf_t.nGrowVect()); + + amrex::MultiFab tempHTransformedSplit_p(tempHTransformed.boxArray(), tempHTransformed.DistributionMap(), 2*n_rz_azimuthal_modes, 0); + amrex::MultiFab tempHTransformedSplit_m(tempHTransformed.boxArray(), tempHTransformed.DistributionMap(), 2*n_rz_azimuthal_modes, 0); + + // Loop over boxes. + for (amrex::MFIter mfi(field_mf_r); mfi.isValid(); ++mfi){ + + // Perform the Hankel transform first. + amrex::Box const& realspace_bx = tempHTransformed[mfi].box(); + multi_spectral_hankel_transformer[mfi].PhysicalToSpectral_Vector(realspace_bx, + field_mf_r_copy[mfi], field_mf_t_copy[mfi], + tempHTransformedSplit_p[mfi], tempHTransformedSplit_m[mfi]); + + FABZForwardTransform(mfi, tempHTransformedSplit_p, field_index_r, is_nodal_z); + FABZForwardTransform(mfi, tempHTransformedSplit_m, field_index_t, is_nodal_z); + + } +} + +/* \brief Transform spectral field specified by `field_index` back to + * real space, and store it in the component `i_comp` of `field_mf` */ +void +SpectralFieldDataRZ::BackwardTransform (amrex::MultiFab& field_mf, int const field_index, + int const i_comp) +{ + // Check field index type, in order to apply proper shift in spectral space. + bool const is_nodal_z = field_mf.is_nodal(1); + + int const ncomp = 2*n_rz_azimuthal_modes - 1; + + // A full multifab is created so that each GPU stream has its own temp space. + amrex::MultiFab tempHTransformedSplit(tempHTransformed.boxArray(), tempHTransformed.DistributionMap(), 2*n_rz_azimuthal_modes, 0); + + // Loop over boxes. + for (amrex::MFIter mfi(field_mf); mfi.isValid(); ++mfi){ + + FABZBackwardTransform(mfi, field_index, tempHTransformedSplit, is_nodal_z); + + // Perform the Hankel inverse transform last. + // tempHTransformedSplit includes the imaginary component of mode 0. + // field_mf does not. + amrex::Box const& realspace_bx = tempHTransformed[mfi].box(); + amrex::FArrayBox field_comp(field_mf[mfi], amrex::make_alias, i_comp*ncomp, ncomp); + multi_spectral_hankel_transformer[mfi].SpectralToPhysical_Scalar(realspace_bx, tempHTransformedSplit[mfi], field_comp); + + } +} + +/* \brief Transform spectral fields specified by `field_index_r` and + * `field_index_t` back to real space, and store them in `field_mf_r` and `field_mf_t` */ +void +SpectralFieldDataRZ::BackwardTransform (amrex::MultiFab& field_mf_r, int const field_index_r, + amrex::MultiFab& field_mf_t, int const field_index_t) +{ + // Check field index type, in order to apply proper shift in spectral space. + bool const is_nodal_z = field_mf_r.is_nodal(1); + + // Full multifabs are created for the temps so that each GPU stream has its own temp space. + amrex::MultiFab tempHTransformedSplit_p(tempHTransformed.boxArray(), tempHTransformed.DistributionMap(), 2*n_rz_azimuthal_modes, 0); + amrex::MultiFab tempHTransformedSplit_m(tempHTransformed.boxArray(), tempHTransformed.DistributionMap(), 2*n_rz_azimuthal_modes, 0); + + // Create copies of the input multifabs. The copies will include the imaginary part of mode 0. + amrex::MultiFab field_mf_r_copy(field_mf_r.boxArray(), field_mf_r.DistributionMap(), 2*n_rz_azimuthal_modes, field_mf_r.nGrowVect()); + amrex::MultiFab field_mf_t_copy(field_mf_t.boxArray(), field_mf_t.DistributionMap(), 2*n_rz_azimuthal_modes, field_mf_t.nGrowVect()); + + // Loop over boxes. + for (amrex::MFIter mfi(field_mf_r); mfi.isValid(); ++mfi){ + + FABZBackwardTransform(mfi, field_index_r, tempHTransformedSplit_p, is_nodal_z); + FABZBackwardTransform(mfi, field_index_t, tempHTransformedSplit_m, is_nodal_z); + + // Perform the Hankel inverse transform last. + // tempHTransformedSplit includes the imaginary component of mode 0. + // field_mf_[ri] do not. + amrex::Box const& realspace_bx = tempHTransformed[mfi].box(); + multi_spectral_hankel_transformer[mfi].SpectralToPhysical_Vector(realspace_bx, + tempHTransformedSplit_p[mfi], tempHTransformedSplit_m[mfi], + field_mf_r_copy[mfi], field_mf_t_copy[mfi]); + + amrex::Array4<amrex::Real> const & field_mf_r_array = field_mf_r[mfi].array(); + amrex::Array4<amrex::Real> const & field_mf_t_array = field_mf_t[mfi].array(); + amrex::Array4<amrex::Real> const & field_mf_r_copy_array = field_mf_r_copy[mfi].array(); + amrex::Array4<amrex::Real> const & field_mf_t_copy_array = field_mf_t_copy[mfi].array(); + + ParallelFor(realspace_bx, 2*n_rz_azimuthal_modes-1, + [=] AMREX_GPU_DEVICE(int i, int j, int k, int icomp) noexcept { + if (icomp == 0) { + // mode 0 + field_mf_r_array(i,j,k,icomp) = field_mf_r_copy_array(i,j,k,icomp); + field_mf_t_array(i,j,k,icomp) = field_mf_t_copy_array(i,j,k,icomp); + } else { + field_mf_r_array(i,j,k,icomp) = field_mf_r_copy_array(i,j,k,icomp+1); + field_mf_t_array(i,j,k,icomp) = field_mf_t_copy_array(i,j,k,icomp+1); + } + }); + + } + +} diff --git a/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/BesselRoots.H b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/BesselRoots.H new file mode 100644 index 000000000..c7a816c61 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/BesselRoots.H @@ -0,0 +1,158 @@ +/* Copyright 2019 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +/* ------------------------------------------------------------------------- +! program to calculate the first zeroes (root abscissas) of the first +! kind bessel function of integer order n using the subroutine rootj. +! -------------------------------------------------------------------------- +! sample run: +! +! (calculate the first 10 zeroes of 1st kind bessel function of order 2). +! +! zeroes of bessel function of order: 2 +! +! number of calculated zeroes: 10 +! +! table of root abcissas (5 items per line) +! 5.135622 8.417244 11.619841 14.795952 17.959819 + 21.116997 24.270112 27.420574 30.569204 33.716520 +! +! table of error codes (5 items per line) +! 0 0 0 0 0 +! 0 0 0 0 0 +! +! -------------------------------------------------------------------------- +! reference: from numath library by tuan dang trong in fortran 77 +! [bibli 18]. +! +! c++ release 1.0 by j-p moreau, paris +! (www.jpmoreau.fr) +! ------------------------------------------------------------------------ */ + +using amrex::operator""_rt; + +void SecantRootFinder(int n, int nitmx, amrex::Real tol, amrex::Real *zeroj, int *ier); + +/*---------------------------------------------------------------------- + * calculate the first nk zeroes of bessel function j(n, x) + * including the trivial root (when n > 0) + * + * inputs: + * n order of function j (integer >= 0) i*4 + * nk number of first zeroes (integer > 0) i*4 + * outputs: + * roots(nk) table of first zeroes (abcissas) r*8 + * ier(nk) table of error codes (must be zeroes) i*4 + * + * reference : + * abramowitz m. & stegun irene a. + * handbook of mathematical functions + */ +void GetBesselRoots(int n, int nk, HankelTransform::RealVector& roots, amrex::Vector<int> &ier) { + amrex::Real zeroj; + int ierror, ik, k; + + const amrex::Real tol = 1e-14_rt; + const amrex::Real nitmx = 10; + + const amrex::Real c1 = 1.8557571_rt; + const amrex::Real c2 = 1.033150_rt; + const amrex::Real c3 = 0.00397_rt; + const amrex::Real c4 = 0.0908_rt; + const amrex::Real c5 = 0.043_rt; + + const amrex::Real t0 = 4.0_rt*n*n; + const amrex::Real t1 = t0 - 1.0_rt; + const amrex::Real t3 = 4.0_rt*t1*(7.0_rt*t0 - 31.0_rt); + const amrex::Real t5 = 32.0_rt*t1*((83.0_rt*t0 - 982.0_rt)*t0 + 3779.0_rt); + const amrex::Real t7 = 64.0_rt*t1*(((6949.0_rt*t0 - 153855.0_rt)*t0 + 1585743.0_rt)*t0 - 6277237.0_rt); + + roots.resize(nk); + ier.resize(nk); + + // first zero + if (n == 0) { + zeroj = c1 + c2 - c3 - c4 + c5; + SecantRootFinder(n, nitmx, tol, &zeroj, &ierror); + ier[0] = ierror; + roots[0] = zeroj; + ik = 1; + } + else { + // Include the trivial root + ier[0] = 0; + roots[0] = 0.; + const amrex::Real f1 = std::pow(n, (1.0_rt/3.0_rt)); + const amrex::Real f2 = f1*f1*n; + const amrex::Real f3 = f1*n*n; + zeroj = n + c1*f1 + (c2/f1) - (c3/n) - (c4/f2) + (c5/f3); + SecantRootFinder(n, nitmx, tol, &zeroj, &ierror); + ier[1] = ierror; + roots[1] = zeroj; + ik = 2; + } + + // other zeroes + // k counts the nontrivial roots + // ik counts all roots + k = 2; + while (ik < nk) { + + // mac mahon's series for k >> n + const amrex::Real b0 = (k + 0.5_rt*n - 0.25_rt)*MathConst::pi; + const amrex::Real b1 = 8.0_rt*b0; + const amrex::Real b2 = b1*b1; + const amrex::Real b3 = 3.0_rt*b1*b2; + const amrex::Real b5 = 5.0_rt*b3*b2; + const amrex::Real b7 = 7.0_rt*b5*b2; + + zeroj = b0 - (t1/b1) - (t3/b3) - (t5/b5) - (t7/b7); + + const amrex::Real errj = std::abs(jn(n, zeroj)); + + // improve solution using procedure SecantRootFinder + if (errj > tol) SecantRootFinder(n, nitmx, tol, &zeroj, &ierror); + + roots[ik] = zeroj; + ier[ik] = ierror; + + k++; + ik++; + } +} + +void SecantRootFinder(int n, int nitmx, amrex::Real tol, amrex::Real *zeroj, int *ier) { + + amrex::Real p0, p1, q0, q1, dp, p; + amrex::Real c[2]; + + c[0] = 0.95_rt; + c[1] = 0.999_rt; + *ier = 0; + + for (int ntry=0 ; ntry <= 1 ; ntry++) { + p0 = c[ntry]*(*zeroj); + + p1 = *zeroj; + q0 = jn(n, p0); + q1 = jn(n, p1); + for (int it=1; it <= nitmx; it++) { + if (q1 == q0) break; + p = p1 - q1*(p1 - p0)/(q1 - q0); + dp = p - p1; + if (it > 1 && std::abs(dp) < tol) { + *zeroj = p; + return; + } + p0 = p1; + q0 = q1; + p1 = p; + q1 = jn(n, p1); + } + } + *ier = 3; + *zeroj = p; +} diff --git a/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/HankelTransform.H b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/HankelTransform.H new file mode 100644 index 000000000..0640b181a --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/HankelTransform.H @@ -0,0 +1,50 @@ +/* Copyright 2019-2020 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#ifndef WARPX_HANKEL_TRANSFORM_H_ +#define WARPX_HANKEL_TRANSFORM_H_ + +#include <AMReX_FArrayBox.H> + +/* \brief This defines the class that performs the Hankel transform. + * Original authors: Remi Lehe, Manuel Kirchen + * + * + * Definition of the Hankel forward and backward transform of order p: + * g(kr) = \int_0^\infty f(r) J_p(kr r) r dr + * f(r ) = \int_0^\infty g(kr) J_p(kr r) kr dkr +*/ +class HankelTransform +{ + public: + + using RealVector = amrex::Gpu::ManagedDeviceVector<amrex::Real>; + + // Constructor + HankelTransform(const int hankel_order, + const int azimuthal_mode, + const int nr, + const amrex::Real rmax); + + const RealVector & getSpectralWavenumbers() {return m_kr;} + + void HankelForwardTransform(amrex::FArrayBox const& F, int const F_icomp, + amrex::FArrayBox & G, int const G_icomp); + + void HankelInverseTransform(amrex::FArrayBox const& G, int const G_icomp, + amrex::FArrayBox & F, int const F_icomp); + + private: + // Even though nk == nr always, use a seperate variable for clarity. + int m_nr, m_nk; + + RealVector m_kr; + + RealVector invM; + RealVector M; +}; + +#endif diff --git a/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/HankelTransform.cpp b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/HankelTransform.cpp new file mode 100644 index 000000000..c5249d54f --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/HankelTransform.cpp @@ -0,0 +1,272 @@ +/* Copyright 2019 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#include "WarpX.H" + +#include "HankelTransform.H" +#include "BesselRoots.H" + +#include <blas.hh> +#include <lapack.hh> + +using amrex::operator""_rt; + +HankelTransform::HankelTransform (int const hankel_order, + int const azimuthal_mode, + int const nr, + const amrex::Real rmax) +: m_nr(nr), m_nk(nr) +{ + + // Check that azimuthal_mode has a valid value + AMREX_ALWAYS_ASSERT_WITH_MESSAGE(hankel_order-1 <= azimuthal_mode && azimuthal_mode <= hankel_order+1, + "azimuthal_mode must be either hankel_order-1, hankel_order or hankel_order+1"); + + RealVector alphas; + amrex::Vector<int> alpha_errors; + + GetBesselRoots(azimuthal_mode, m_nk, alphas, alpha_errors); + // Add of check of alpha_errors, should be all zeros + AMREX_ALWAYS_ASSERT(std::all_of(alpha_errors.begin(), alpha_errors.end(), [](int i) { return i == 0; })); + + // Calculate the spectral grid + m_kr.resize(m_nk); + for (int ik=0 ; ik < m_nk ; ik++) { + m_kr[ik] = alphas[ik]/rmax; + } + + // Calculate the spatial grid (Uniform grid with a half-cell offset) + RealVector rmesh(m_nr); + amrex::Real dr = rmax/m_nr; + for (int ir=0 ; ir < m_nr ; ir++) { + rmesh[ir] = dr*(ir + 0.5_rt); + } + + // Calculate and store the inverse matrix invM + // (imposed by the constraints on the DHT of Bessel modes) + // NB: When compared with the FBPIC article, all the matrices here + // are calculated in transposed form. This is done so as to use the + // `dot` and `gemm` functions, in the `transform` method. + int p_denom; + if (hankel_order == azimuthal_mode) { + p_denom = hankel_order + 1; + } else { + p_denom = hankel_order; + } + + RealVector denom(m_nk); + for (int ik=0 ; ik < m_nk ; ik++) { + const amrex::Real jna = jn(p_denom, alphas[ik]); + denom[ik] = MathConst::pi*rmax*rmax*jna*jna; + } + + RealVector num(m_nk*m_nr); + for (int ir=0 ; ir < m_nr ; ir++) { + for (int ik=0 ; ik < m_nk ; ik++) { + int const ii = ik + ir*m_nk; + num[ii] = jn(hankel_order, rmesh[ir]*m_kr[ik]); + } + } + + // Get the inverse matrix + invM.resize(m_nk*m_nr); + if (azimuthal_mode > 0) { + for (int ir=0 ; ir < m_nr ; ir++) { + for (int ik=1 ; ik < m_nk ; ik++) { + int const ii = ik + ir*m_nk; + invM[ii] = num[ii]/denom[ik]; + } + } + // ik = 0 + // In this case, the functions are represented by Bessel functions + // *and* an additional mode (below) which satisfies the same + // algebric relations for curl/div/grad as the regular Bessel modes, + // with the value kperp=0. + // The normalization of this mode is arbitrary, and is chosen + // so that the condition number of invM is close to 1 + if (hankel_order == azimuthal_mode-1) { + for (int ir=0 ; ir < m_nr ; ir++) { + int const ii = ir*m_nk; + invM[ii] = std::pow(rmesh[ir], (azimuthal_mode-1))/(MathConst::pi*std::pow(rmax, (azimuthal_mode+1))); + } + } else { + for (int ir=0 ; ir < m_nr ; ir++) { + int const ii = ir*m_nk; + invM[ii] = 0.; + } + } + } else { + for (int ir=0 ; ir < m_nr ; ir++) { + for (int ik=0 ; ik < m_nk ; ik++) { + int const ii = ik + ir*m_nk; + invM[ii] = num[ii]/denom[ik]; + } + } + } + + // Calculate the matrix M by inverting invM + if (azimuthal_mode !=0 && hankel_order != azimuthal_mode-1) { + // In this case, invM is singular, thus we calculate the pseudo-inverse. + // The Moore-Penrose psuedo-inverse is calculated using the SVD method. + + M.resize(m_nk*m_nr, 0.); + RealVector invMcopy(invM); + RealVector sdiag(m_nk-1, 0.); + RealVector u((m_nk-1)*(m_nk-1), 0.); + RealVector vt((m_nr)*(m_nr), 0.); + RealVector sp((m_nr)*(m_nk-1), 0.); + RealVector temp((m_nr)*(m_nk-1), 0.); + + // Calculate the singlular-value-decomposition of invM (leaving out the first row). + // invM = u*sdiag*vt + // Note that invMcopy.dataPtr()+1 is passed in so that the first ik row is skipped + // A copy is passed in since the matrix is destroyed + lapack::gesvd(lapack::Job::AllVec, lapack::Job::AllVec, + m_nk-1, m_nr, invMcopy.dataPtr()+1, m_nk, + sdiag.dataPtr(), u.dataPtr(), m_nk-1, + vt.dataPtr(), m_nr); + + // Calculate the pseudo-inverse of sdiag, which is trivial since it + // only has values of the diagonal. + for (int i=0 ; i < m_nk-1 ; i++) { + if (sdiag[i] != 0.) { + int const j = i + i*m_nk; + sp[j] = 1._rt/sdiag[i]; + } + } + + // M can be calculated from the SVD + // M = v*sp*u_transpose + + // Do the second multiply, temp = sp*u_transpose + // temp(1:n,1:m) = matmul(transpose(vt(1:n,1:n)), matmul(sp(1:n,1:m), transpose(u(1:m,1:m)))) + blas::gemm(blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::Trans, + m_nr, m_nk-1, m_nk-1, 1._rt, + sp.dataPtr(), m_nr, + u.dataPtr(), m_nk-1, 0._rt, + temp.dataPtr(), m_nr); + + // Do the first mutiply, M = vt*temp + // Note that M.dataPtr()+m_nr is passed in so that the first ir column is skipped + blas::gemm(blas::Layout::ColMajor, blas::Op::Trans, blas::Op::NoTrans, + m_nr, m_nk-1, m_nr, 1., + vt.dataPtr(), m_nr, + temp.dataPtr(), m_nr, 0._rt, + M.dataPtr()+m_nr, m_nr); + + } else { + // In this case, invM is invertible; calculate the inverse. + // getrf calculates the LU decomposition + // getri calculates the inverse from the LU decomposition + + M = invM; + amrex::Vector<int64_t> ipiv(m_nr); + lapack::getrf(m_nk, m_nr, M.dataPtr(), m_nk, ipiv.dataPtr()); + lapack::getri(m_nr, M.dataPtr(), m_nr, ipiv.dataPtr()); + + } + +} + +void +HankelTransform::HankelForwardTransform (amrex::FArrayBox const& F, int const F_icomp, + amrex::FArrayBox & G, int const G_icomp) +{ + amrex::Box const& F_box = F.box(); + amrex::Box const& G_box = G.box(); + + int const nrF = F_box.length(0); + int const nz = F_box.length(1); + int const ngr = G_box.smallEnd(0) - F_box.smallEnd(0); + + AMREX_ALWAYS_ASSERT(m_nr == G_box.length(0)); + AMREX_ALWAYS_ASSERT(nz == G_box.length(1)); + AMREX_ALWAYS_ASSERT(ngr >= 0); + AMREX_ALWAYS_ASSERT(F_box.bigEnd(0)+1 >= m_nr); + +#ifndef AMREX_USE_GPU + // On CPU, the blas::gemm is significantly faster + + // Note that M is flagged to be transposed since it has dimensions (m_nr, m_nk) + blas::gemm(blas::Layout::ColMajor, blas::Op::Trans, blas::Op::NoTrans, + m_nk, nz, m_nr, 1._rt, + M.dataPtr(), m_nk, + F.dataPtr(F_icomp)+ngr, nrF, 0._rt, + G.dataPtr(G_icomp), m_nk); + +#else + // On GPU, the explicit loop is significantly faster + // It is not clear if the GPU gemm wasn't build properly, it is cycling data out and back + // in to the device, or if it is because gemm is launching its own threads. + + amrex::Real const * M_arr = M.dataPtr(); + amrex::Array4<const amrex::Real> const & F_arr = F.array(); + amrex::Array4< amrex::Real> const & G_arr = G.array(); + + int const nr = m_nr; + + ParallelFor(G_box, + [=] AMREX_GPU_DEVICE(int ik, int iz, int inotused) noexcept { + G_arr(ik,iz,G_icomp) = 0.; + for (int ir=0 ; ir < nr ; ir++) { + int const ii = ir + ik*nr; + G_arr(ik,iz,G_icomp) += M_arr[ii]*F_arr(ir,iz,F_icomp); + } + }); + +#endif + +} + +void +HankelTransform::HankelInverseTransform (amrex::FArrayBox const& G, int const G_icomp, + amrex::FArrayBox & F, int const F_icomp) +{ + amrex::Box const& G_box = G.box(); + amrex::Box const& F_box = F.box(); + + int const nrF = F_box.length(0); + int const nz = F_box.length(1); + int const ngr = G_box.smallEnd(0) - F_box.smallEnd(0); + + AMREX_ALWAYS_ASSERT(m_nr == G_box.length(0)); + AMREX_ALWAYS_ASSERT(nz == G_box.length(1)); + AMREX_ALWAYS_ASSERT(ngr >= 0); + AMREX_ALWAYS_ASSERT(F_box.bigEnd(0)+1 >= m_nr); + +#ifndef AMREX_USE_GPU + // On CPU, the blas::gemm is significantly faster + + // Note that invM is flagged to be transposed since it has dimensions (m_nk, m_nr) + blas::gemm(blas::Layout::ColMajor, blas::Op::Trans, blas::Op::NoTrans, + m_nr, nz, m_nk, 1._rt, + invM.dataPtr(), m_nr, + G.dataPtr(G_icomp), m_nk, 0._rt, + F.dataPtr(F_icomp)+ngr, nrF); + +#else + // On GPU, the explicit loop is significantly faster + // It is not clear if the GPU gemm wasn't build properly, it is cycling data out and back + // in to the device, or if it is because gemm is launching its own threads. + + amrex::Real const * invM_arr = invM.dataPtr(); + amrex::Array4<const amrex::Real> const & G_arr = G.array(); + amrex::Array4< amrex::Real> const & F_arr = F.array(); + + int const nk = m_nk; + + ParallelFor(G_box, + [=] AMREX_GPU_DEVICE(int ir, int iz, int inotused) noexcept { + F_arr(ir,iz,F_icomp) = 0.; + for (int ik=0 ; ik < nk ; ik++) { + int const ii = ik + ir*nk; + F_arr(ir,iz,F_icomp) += invM_arr[ii]*G_arr(ik,iz,G_icomp); + } + }); + +#endif + +} diff --git a/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/Make.package b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/Make.package new file mode 100644 index 000000000..a3c22d64a --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/Make.package @@ -0,0 +1,6 @@ +ifeq ($(USE_RZ),TRUE) + CEXE_sources += SpectralHankelTransformer.cpp + CEXE_sources += HankelTransform.cpp + + VPATH_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/SpectralSolver/SpectralHankelTransform +endif diff --git a/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/SpectralHankelTransformer.H b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/SpectralHankelTransformer.H new file mode 100644 index 000000000..d82fb9b84 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/SpectralHankelTransformer.H @@ -0,0 +1,78 @@ +/* Copyright 2019-2020 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#ifndef WARPX_SPECTRALHANKELTRANSFORMER_H_ +#define WARPX_SPECTRALHANKELTRANSFORMER_H_ + +#include <AMReX_FArrayBox.H> +#include "HankelTransform.H" + +/* \brief Object that allows to transform the fields back and forth between the + * spectral and interpolation grid. + * + * Attributes : + * - dht0, dhtm, dhtp : the discrete Hankel transform objects for the modes, + * operating along r +*/ + +class SpectralHankelTransformer +{ + + public: + + SpectralHankelTransformer() {}; + + SpectralHankelTransformer(const int nr, + const int n_rz_azimuthal_modes, + const amrex::Real rmax); + + void + ExtractKrArray(); + + // Returns an array that holds the kr for all of the modes + HankelTransform::RealVector const & getKrArray() const {return m_kr;} + + // Converts a scalar field from the physical to the spectral space + void + PhysicalToSpectral_Scalar(amrex::Box const & box, + amrex::FArrayBox const & F_physical, + amrex::FArrayBox & G_spectral); + + // Converts a vector field from the physical to the spectral space + void + PhysicalToSpectral_Vector(amrex::Box const & box, + amrex::FArrayBox & F_r_physical, + amrex::FArrayBox & F_t_physical, + amrex::FArrayBox & G_p_spectral, + amrex::FArrayBox & G_m_spectral); + + // Converts a scalar field from the spectral to the physical space + void + SpectralToPhysical_Scalar(amrex::Box const & box, + amrex::FArrayBox const & G_spectral, + amrex::FArrayBox & F_physical); + + // Converts a vector field from the spectral to the physical space + void + SpectralToPhysical_Vector(amrex::Box const & box, + amrex::FArrayBox const & G_p_spectral, + amrex::FArrayBox const & G_m_spectral, + amrex::FArrayBox & F_r_physical, + amrex::FArrayBox & F_t_physical); + + private: + + int m_nr; + int m_n_rz_azimuthal_modes; + HankelTransform::RealVector m_kr; + + amrex::Vector< std::unique_ptr<HankelTransform> > dht0; + amrex::Vector< std::unique_ptr<HankelTransform> > dhtm; + amrex::Vector< std::unique_ptr<HankelTransform> > dhtp; + +}; + +#endif diff --git a/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/SpectralHankelTransformer.cpp b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/SpectralHankelTransformer.cpp new file mode 100644 index 000000000..7dd297418 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/SpectralHankelTransformer.cpp @@ -0,0 +1,199 @@ +/* Copyright 2019-2020 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#include "Utils/WarpXConst.H" +#include "SpectralHankelTransformer.H" + +SpectralHankelTransformer::SpectralHankelTransformer (int const nr, + int const n_rz_azimuthal_modes, + amrex::Real const rmax) +: m_nr(nr), m_n_rz_azimuthal_modes(n_rz_azimuthal_modes) +{ + + dht0.resize(m_n_rz_azimuthal_modes); + dhtp.resize(m_n_rz_azimuthal_modes); + dhtm.resize(m_n_rz_azimuthal_modes); + + for (int mode=0 ; mode < m_n_rz_azimuthal_modes ; mode++) { + dht0[mode].reset( new HankelTransform(mode , mode, m_nr, rmax) ); + dhtp[mode].reset( new HankelTransform(mode+1, mode, m_nr, rmax) ); + dhtm[mode].reset( new HankelTransform(mode-1, mode, m_nr, rmax) ); + } + + ExtractKrArray(); + +} + +/* \brief Extracts the kr for all of the modes + * This needs to be separate since the ParallelFor cannot be in the constructor. */ +void +SpectralHankelTransformer::ExtractKrArray () +{ + m_kr.resize(m_nr*m_n_rz_azimuthal_modes); + + for (int mode=0 ; mode < m_n_rz_azimuthal_modes ; mode++) { + + // Save a copy of all of the kr's in one place to allow easy access later. + // They are stored with the kr's of each mode grouped together. + amrex::Real *kr_array = m_kr.dataPtr(); + auto const & kr_mode = dht0[mode]->getSpectralWavenumbers(); + auto const & kr_m_array = kr_mode.dataPtr(); + int const nr_temp = m_nr; + amrex::ParallelFor(m_nr, + [=] AMREX_GPU_DEVICE (int ir) + { + int const ii = ir + mode*nr_temp; + kr_array[ii] = kr_m_array[ir]; + }); + } +} + +/* \brief Converts a scalar field from the physical to the spectral space for all modes */ +void +SpectralHankelTransformer::PhysicalToSpectral_Scalar (amrex::Box const & box, + amrex::FArrayBox const & F_physical, + amrex::FArrayBox & G_spectral) +{ + // The Hankel transform is purely real, so the real and imaginary parts of + // F can be transformed separately, so a simple loop over components + // can be done. + // Note that F_physical does not include the imaginary part of mode 0, + // but G_spectral does. + for (int mode=0 ; mode < m_n_rz_azimuthal_modes ; mode++) { + int const mode_r = 2*mode; + int const mode_i = 2*mode + 1; + if (mode == 0) { + int const icomp = 0; + dht0[mode]->HankelForwardTransform(F_physical, icomp, G_spectral, mode_r); + G_spectral.setVal<amrex::RunOn::Device>(0., mode_i); + } else { + int const icomp = 2*mode - 1; + dht0[mode]->HankelForwardTransform(F_physical, icomp , G_spectral, mode_r); + dht0[mode]->HankelForwardTransform(F_physical, icomp+1, G_spectral, mode_i); + } + } +} + +/* \brief Converts a vector field from the physical to the spectral space for all modes */ +void +SpectralHankelTransformer::PhysicalToSpectral_Vector (amrex::Box const & box, + amrex::FArrayBox & F_r_physical, + amrex::FArrayBox & F_t_physical, + amrex::FArrayBox & G_p_spectral, + amrex::FArrayBox & G_m_spectral) +{ + // Note that F and G include the imaginary part of mode 0. + // F will be overwritten (by the + and - data). + + using amrex::operator""_rt; + + amrex::Array4<amrex::Real> const & F_r_physical_array = F_r_physical.array(); + amrex::Array4<amrex::Real> const & F_t_physical_array = F_t_physical.array(); + + for (int mode=0 ; mode < m_n_rz_azimuthal_modes ; mode++) { + + int const mode_r = 2*mode; + int const mode_i = 2*mode + 1; + + amrex::ParallelFor(box, + [=] AMREX_GPU_DEVICE (int i, int j, int k) + { + amrex::Real const r_real = F_r_physical_array(i,j,k,mode_r); + amrex::Real const r_imag = F_r_physical_array(i,j,k,mode_i); + amrex::Real const t_real = F_t_physical_array(i,j,k,mode_r); + amrex::Real const t_imag = F_t_physical_array(i,j,k,mode_i); + // Combine the values + // temp_p = (F_r - I*F_t)/2 + // temp_m = (F_r + I*F_t)/2 + F_r_physical_array(i,j,k,mode_r) = 0.5_rt*(r_real + t_imag); + F_r_physical_array(i,j,k,mode_i) = 0.5_rt*(r_imag - t_real); + F_t_physical_array(i,j,k,mode_r) = 0.5_rt*(r_real - t_imag); + F_t_physical_array(i,j,k,mode_i) = 0.5_rt*(r_imag + t_real); + }); + + amrex::Gpu::streamSynchronize(); + + dhtp[mode]->HankelForwardTransform(F_r_physical, mode_r, G_p_spectral, mode_r); + dhtp[mode]->HankelForwardTransform(F_r_physical, mode_i, G_p_spectral, mode_i); + dhtm[mode]->HankelForwardTransform(F_t_physical, mode_r, G_m_spectral, mode_r); + dhtm[mode]->HankelForwardTransform(F_t_physical, mode_i, G_m_spectral, mode_i); + + } +} + +/* \brief Converts a scalar field from the spectral to the physical space for all modes */ +void +SpectralHankelTransformer::SpectralToPhysical_Scalar (amrex::Box const & box, + amrex::FArrayBox const & G_spectral, + amrex::FArrayBox & F_physical) +{ + // The Hankel inverse transform is purely real, so the real and imaginary parts of + // F can be transformed separately, so a simple loop over components + // can be done. + // Note that F_physical does not include the imaginary part of mode 0, + // but G_spectral does. + + amrex::Gpu::streamSynchronize(); + + for (int mode=0 ; mode < m_n_rz_azimuthal_modes ; mode++) { + int const mode_r = 2*mode; + int const mode_i = 2*mode + 1; + if (mode == 0) { + int const icomp = 0; + dht0[mode]->HankelInverseTransform(G_spectral, mode_r, F_physical, icomp); + } else { + int const icomp = 2*mode - 1; + dht0[mode]->HankelInverseTransform(G_spectral, mode_r, F_physical, icomp); + dht0[mode]->HankelInverseTransform(G_spectral, mode_i, F_physical, icomp+1); + } + } +} + +/* \brief Converts a vector field from the spectral to the physical space for all modes */ +void +SpectralHankelTransformer::SpectralToPhysical_Vector (amrex::Box const & box, + amrex::FArrayBox const& G_p_spectral, + amrex::FArrayBox const& G_m_spectral, + amrex::FArrayBox & F_r_physical, + amrex::FArrayBox & F_t_physical) +{ + // Note that F and G include the imaginary part of mode 0. + + amrex::Array4<amrex::Real> const & F_r_physical_array = F_r_physical.array(); + amrex::Array4<amrex::Real> const & F_t_physical_array = F_t_physical.array(); + + for (int mode=0 ; mode < m_n_rz_azimuthal_modes ; mode++) { + + int const mode_r = 2*mode; + int const mode_i = 2*mode + 1; + + amrex::Gpu::streamSynchronize(); + + dhtp[mode]->HankelInverseTransform(G_p_spectral, mode_r, F_r_physical, mode_r); + dhtp[mode]->HankelInverseTransform(G_p_spectral, mode_i, F_r_physical, mode_i); + dhtm[mode]->HankelInverseTransform(G_m_spectral, mode_r, F_t_physical, mode_r); + dhtm[mode]->HankelInverseTransform(G_m_spectral, mode_i, F_t_physical, mode_i); + + amrex::Gpu::streamSynchronize(); + + amrex::ParallelFor(box, + [=] AMREX_GPU_DEVICE (int i, int j, int k) + { + amrex::Real const p_real = F_r_physical_array(i,j,k,mode_r); + amrex::Real const p_imag = F_r_physical_array(i,j,k,mode_i); + amrex::Real const m_real = F_t_physical_array(i,j,k,mode_r); + amrex::Real const m_imag = F_t_physical_array(i,j,k,mode_i); + // Combine the values + // F_r = G_p + G_m + // F_t = I*(G_p - G_m) + F_r_physical_array(i,j,k,mode_r) = p_real + m_real; + F_r_physical_array(i,j,k,mode_i) = p_imag + m_imag; + F_t_physical_array(i,j,k,mode_r) = -p_imag + m_imag; + F_t_physical_array(i,j,k,mode_i) = p_real - m_real; + }); + + } +} diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H index 2cd3e093a..f26677700 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H +++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H @@ -39,6 +39,7 @@ class SpectralKSpace { public: amrex::BoxArray spectralspace_ba; + SpectralKSpace() : dx(amrex::RealVect::Zero) {}; SpectralKSpace( const amrex::BoxArray& realspace_ba, const amrex::DistributionMapping& dm, const amrex::RealVect realspace_dx ); @@ -53,7 +54,7 @@ class SpectralKSpace const amrex::DistributionMapping& dm, const int i_dim, const int shift_type ) const; - private: + protected: amrex::Array<KVectorComponent, AMREX_SPACEDIM> k_vec; // 3D: k_vec is an Array of 3 components, corresponding to kx, ky, kz // 2D: k_vec is an Array of 2 components, corresponding to kx, kz diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpaceRZ.H b/Source/FieldSolver/SpectralSolver/SpectralKSpaceRZ.H new file mode 100644 index 000000000..594b06764 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralKSpaceRZ.H @@ -0,0 +1,32 @@ +/* Copyright 2019-2020 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#ifndef WARPX_SPECTRAL_K_SPACE_RZ_H_ +#define WARPX_SPECTRAL_K_SPACE_RZ_H_ + +#include "SpectralKSpace.H" + +/* \brief Class that represents the spectral, Hankel/FFT, space. + * + * (Contains info about the size of the spectral space corresponding + * to each box in `realspace_ba`, as well as the value of the + * corresponding kz coordinates) + */ +class SpectralKSpaceRZ + : + public SpectralKSpace +{ + public: + SpectralKSpaceRZ(const amrex::BoxArray& realspace_ba, + const amrex::DistributionMapping& dm, + const amrex::RealVect realspace_dx); + + KVectorComponent const & getKzArray () const {return k_vec[1];} + amrex::RealVect const & getCellSize () const {return dx;} + +}; + +#endif // WARPX_SPECTRAL_K_SPACE_RZ_H_ diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpaceRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralKSpaceRZ.cpp new file mode 100644 index 000000000..133b2ef65 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralKSpaceRZ.cpp @@ -0,0 +1,52 @@ +/* Copyright 2019 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#include "Utils/WarpXConst.H" +#include "SpectralKSpaceRZ.H" + +#include <cmath> + +/* \brief Initialize k space object. + * + * \param realspace_ba Box array that corresponds to the decomposition + * of the fields in real space (cell-centered ; includes guard cells) + * \param dm Indicates which MPI proc owns which box, in realspace_ba. + * \param realspace_dx Cell size of the grid in real space + */ +SpectralKSpaceRZ::SpectralKSpaceRZ (const amrex::BoxArray& realspace_ba, + const amrex::DistributionMapping& dm, + const amrex::RealVect realspace_dx) +{ + dx = realspace_dx; // Store the cell size as member `dx` + + AMREX_ALWAYS_ASSERT_WITH_MESSAGE( + realspace_ba.ixType() == amrex::IndexType::TheCellType(), + "SpectralKSpaceRZ expects a cell-centered box."); + + // Create the box array that corresponds to spectral space + amrex::BoxList spectral_bl; // Create empty box list + + // Loop over boxes and fill the box list + for (int i=0; i < realspace_ba.size(); i++ ) { + // 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 + amrex::Box realspace_bx = realspace_ba[i]; + amrex::IntVect fft_size = realspace_bx.length(); + amrex::IntVect spectral_bx_size = fft_size; + // Define the corresponding box + amrex::Box spectral_bx = amrex::Box(amrex::IntVect::TheZeroVector(), + spectral_bx_size - amrex::IntVect::TheUnitVector() ); + spectral_bl.push_back(spectral_bx); + } + spectralspace_ba.define(spectral_bl); + + // Allocate the components of the kz vector + const int i_dim = 1; + const bool only_positive_k = false; + k_vec[i_dim] = getKComponent(dm, realspace_ba, i_dim, only_positive_k); + +} diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H new file mode 100644 index 000000000..b0894df91 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H @@ -0,0 +1,104 @@ +/* Copyright 2019-2020 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#ifndef WARPX_SPECTRAL_SOLVER_RZ_H_ +#define WARPX_SPECTRAL_SOLVER_RZ_H_ + +#include "SpectralAlgorithms/SpectralBaseAlgorithmRZ.H" +#include "SpectralFieldDataRZ.H" + +/* \brief Top-level class for the electromagnetic spectral solver + * + * Stores the field in spectral space, and has member functions + * to Fourier-transform the fields between real space and spectral space + * and to update fields in spectral space over one time step. + */ +class SpectralSolverRZ +{ + public: + // Inline definition of the member functions of `SpectralSolverRZ`, + // except the constructor (see `SpectralSolverRZ.cpp`) + // The body of these functions is short, since the work is done in the + // underlying classes `SpectralFieldData` and `PsatdAlgorithm` + + // Constructor + SpectralSolverRZ(amrex::BoxArray const & realspace_ba, + amrex::DistributionMapping const & dm, + int const n_rz_azimuthal_modes, + int const norder_z, bool const nodal, + amrex::RealVect const dx, amrex::Real const dt, + int const lev, + bool const pml=false); + + /* \brief Transform the component `i_comp` of MultiFab `field_mf` + * to spectral space, and store the corresponding result internally + * (in the spectral field specified by `field_index`) */ + void ForwardTransform (amrex::MultiFab const & field_mf, + int const field_index, + int const i_comp=0 ) { + BL_PROFILE("SpectralSolverRZ::ForwardTransform"); + field_data.ForwardTransform(field_mf, field_index, i_comp); + }; + + /* \brief Transform the two MultiFabs `field_mf1` and `field_mf2` + * to spectral space, and store the corresponding results internally + * (in the spectral field specified by `field_index1` and `field_index2`) */ + void ForwardTransform (amrex::MultiFab const & field_mf1, int const field_index1, + amrex::MultiFab const & field_mf2, int const field_index2) { + BL_PROFILE("SpectralSolverRZ::ForwardTransform"); + field_data.ForwardTransform(field_mf1, field_index1, + field_mf2, field_index2); + }; + + /* \brief Transform spectral field specified by `field_index` back to + * real space, and store it in the component `i_comp` of `field_mf` */ + void BackwardTransform (amrex::MultiFab& field_mf, + int const field_index, + int const i_comp=0) { + BL_PROFILE("SpectralSolverRZ::BackwardTransform"); + field_data.BackwardTransform(field_mf, field_index, i_comp); + }; + + /* \brief Transform spectral fields specified by `field_index1` and `field_index2` + * back to real space, and store it in `field_mf1` and `field_mf2`*/ + void BackwardTransform (amrex::MultiFab& field_mf1, int const field_index1, + amrex::MultiFab& field_mf2, int const field_index2) { + BL_PROFILE("SpectralSolverRZ::BackwardTransform"); + field_data.BackwardTransform(field_mf1, field_index1, + field_mf2, field_index2); + }; + + /* \brief Update the fields in spectral space, over one timestep */ + void pushSpectralFields () { + BL_PROFILE("SpectralSolverRZ::pushSpectralFields"); + // Virtual function: the actual function used here depends + // on the sub-class of `SpectralBaseAlgorithm` that was + // initialized in the constructor of `SpectralSolverRZ` + algorithm->pushSpectralFields(field_data); + }; + + /** + * \brief Public interface to call the member function ComputeSpectralDivE + * of the base class SpectralBaseAlgorithmRZ from objects of class SpectralSolverRZ + */ + void ComputeSpectralDivE ( const std::array<std::unique_ptr<amrex::MultiFab>,3>& Efield, + amrex::MultiFab& divE ) { + algorithm->ComputeSpectralDivE( field_data, Efield, divE ); + }; + + private: + + SpectralFieldDataRZ field_data; // Store field in spectral space + // and perform the Fourier transforms + std::unique_ptr<SpectralBaseAlgorithmRZ> algorithm; + // Defines field update equation in spectral space, + // and the associated coefficients. + // SpectralBaseAlgorithmRZ is a base class ; this pointer is meant + // to point an instance of a *sub-class* defining a specific algorithm + +}; + +#endif // WARPX_SPECTRAL_SOLVER_RZ_H_ diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp new file mode 100644 index 000000000..3d7aa77b4 --- /dev/null +++ b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp @@ -0,0 +1,52 @@ +/* Copyright 2019-2020 David Grote + * + * This file is part of WarpX. + * + * License: BSD-3-Clause-LBNL + */ +#include "SpectralKSpaceRZ.H" +#include "SpectralSolverRZ.H" +#include "SpectralAlgorithms/PsatdAlgorithmRZ.H" + +/* \brief Initialize the spectral Maxwell solver + * + * This function selects the spectral algorithm to be used, allocates the + * corresponding coefficients for the discretized field update equation, + * and prepares the structures that store the fields in spectral space. + * + * \param n_rz_azimuthal_modes Number of azimuthal modes + * \param norder_z Order of accuracy of the spatial derivatives along z + * \param nodal Whether the solver is applied to a nodal or staggered grid + * \param dx Cell size along each dimension + * \param dt Time step + * \param pml Whether the boxes in which the solver is applied are PML boxes + * PML is not supported. + */ +SpectralSolverRZ::SpectralSolverRZ(amrex::BoxArray const & realspace_ba, + amrex::DistributionMapping const & dm, + int const n_rz_azimuthal_modes, + int const norder_z, bool const nodal, + amrex::RealVect const dx, amrex::Real const dt, + int const lev, + bool const pml ) +{ + + // Initialize all structures using the same distribution mapping dm + + // - The k space object contains info about the size of + // the spectral space corresponding to each box in `realspace_ba`, + // as well as the value of the corresponding k coordinates. + SpectralKSpaceRZ k_space(realspace_ba, dm, dx); + + // - Select the algorithm depending on the input parameters + // Initialize the corresponding coefficients over k space + // PML is not supported. + algorithm = std::unique_ptr<PsatdAlgorithmRZ>( + new PsatdAlgorithmRZ(k_space, dm, n_rz_azimuthal_modes, norder_z, nodal, dt)); + + // - Initialize arrays for fields in spectral space + FFT plans + field_data = SpectralFieldDataRZ(realspace_ba, k_space, dm, + algorithm->getRequiredNumberOfFields(), + n_rz_azimuthal_modes, lev); + +}; |