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>

1 files changed, 198 insertions, 0 deletions

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);
+            }
+        });
+     }
+}