diff options
Diffstat (limited to 'Source/FieldSolver/SpectralSolver/SpectralHankelTransform')
6 files changed, 763 insertions, 0 deletions
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; + }); + + } +} |