RFCavity.H

Go to the documentation of this file.

/* Copyright 2022-2023 The Regents of the University of California, through Lawrence
 *           Berkeley National Laboratory (subject to receipt of any required
 *           approvals from the U.S. Dept. of Energy). All rights reserved.
 *
 * This file is part of ImpactX.
 *
 * Authors: Chad Mitchell, Axel Huebl
 * License: BSD-3-Clause-LBNL
 */
#ifndef IMPACTX_RFCAVITY_H
#define IMPACTX_RFCAVITY_H
 
#include "particles/ImpactXParticleContainer.H"
#include "particles/integrators/Integrators.H"
#include "mixin/alignment.H"
#include "mixin/pipeaperture.H"
#include "mixin/beamoptic.H"
#include "mixin/lineartransport.H"
#include "mixin/named.H"
#include "mixin/thick.H"
 
#include <ablastr/constant.H>
 
#include <AMReX.H>
#include <AMReX_Extension.H>
#include <AMReX_Math.H>
#include <AMReX_REAL.H>
#include <AMReX_SIMD.H>
#include <AMReX_SmallMatrix.H>
 
#include <cmath>
#include <stdexcept>
#include <tuple>
#include <vector>
 
 
namespace impactx::elements
{
    struct RF_field_data
    {
        amrex::Vector<amrex::ParticleReal> default_cos_coef = {
            0.1644024074311037,
            -0.1324009958969339,
            4.3443060026047219e-002,
            8.5602654094946495e-002,
            -0.2433578169042885,
            0.5297150596779437,
            0.7164884680963959,
            -5.2579522442877296e-003,
            -5.5025369142193678e-002,
            4.6845673335028933e-002,
            -2.3279346335638568e-002,
            4.0800777539657775e-003,
            4.1378326533752169e-003,
            -2.5040533340490805e-003,
            -4.0654981400000964e-003,
            9.6630592067498289e-003,
            -8.5275895985990214e-003,
            -5.8078747006425020e-002,
            -2.4044337836660403e-002,
            1.0968240064697212e-002,
            -3.4461179858301418e-003,
            -8.1201564869443749e-004,
            2.1438992904959380e-003,
            -1.4997753525697276e-003,
            1.8685171825676386e-004
        };
 
        amrex::Vector<amrex::ParticleReal> default_sin_coef = {
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0
        };
    };

namespace RFCavityData
{
    inline int next_id = 0;

    inline std::map<int, std::vector<amrex::ParticleReal>> h_cos_coef = {};
    inline std::map<int, std::vector<amrex::ParticleReal>> h_sin_coef = {};

    inline std::map<int, amrex::Gpu::DeviceVector<amrex::ParticleReal>> d_cos_coef = {};
    inline std::map<int, amrex::Gpu::DeviceVector<amrex::ParticleReal>> d_sin_coef = {};
 
} // namespace RFCavityData
 
    struct RFCavity
    : public mixin::Named,
      public mixin::BeamOptic<RFCavity>,
      public mixin::LinearTransport<RFCavity>,
      public mixin::Thick,
      public mixin::Alignment,
      public mixin::PipeAperture,
      public amrex::simd::Vectorized<amrex::simd::native_simd_size_particlereal>
    {
        static constexpr auto type = "RFCavity";
        using PType = ImpactXParticleContainer::ParticleType;

        RFCavity (
            amrex::ParticleReal ds,
            amrex::ParticleReal escale,
            amrex::ParticleReal freq,
            amrex::ParticleReal phase,
            std::vector<amrex::ParticleReal> cos_coef,
            std::vector<amrex::ParticleReal> sin_coef,
            amrex::ParticleReal dx = 0,
            amrex::ParticleReal dy = 0,
            amrex::ParticleReal rotation_degree = 0,
            amrex::ParticleReal aperture_x = 0,
            amrex::ParticleReal aperture_y = 0,
            int mapsteps = 1,
            int nslice = 1,
            std::optional<std::string> name = std::nullopt
        )
          : Named(std::move(name)),
            Thick(ds, nslice),
            Alignment(dx, dy, rotation_degree),
            PipeAperture(aperture_x, aperture_y),
            m_escale(escale), m_freq(freq), m_phase(phase), m_mapsteps(mapsteps), m_id(RFCavityData::next_id)
        {
            // next created RF cavity has another id for its data
            RFCavityData::next_id++;
 
            // validate sin and cos coefficients are the same length
            m_ncoef = int(cos_coef.size());
            if (m_ncoef != int(sin_coef.size()))
                throw std::runtime_error("RFCavity: cos and sin coefficients must have same length!");
 
            // host data
            RFCavityData::h_cos_coef[m_id] = cos_coef;
            RFCavityData::h_sin_coef[m_id] = sin_coef;
            m_cos_h_data = RFCavityData::h_cos_coef[m_id].data();
            m_sin_h_data = RFCavityData::h_sin_coef[m_id].data();
 
            // device data
            RFCavityData::d_cos_coef.emplace(m_id, amrex::Gpu::DeviceVector<amrex::ParticleReal>(m_ncoef));
            RFCavityData::d_sin_coef.emplace(m_id, amrex::Gpu::DeviceVector<amrex::ParticleReal>(m_ncoef));
            amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice,
                                  cos_coef.begin(), cos_coef.end(),
                                  RFCavityData::d_cos_coef[m_id].begin());
            amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice,
                                  sin_coef.begin(), sin_coef.end(),
                                  RFCavityData::d_sin_coef[m_id].begin());
            amrex::Gpu::streamSynchronize();
 
            // low-level objects we can use on device
            m_cos_d_data = RFCavityData::d_cos_coef[m_id].data();
            m_sin_d_data = RFCavityData::d_sin_coef[m_id].data();
        }

        using BeamOptic::operator();

        void compute_constants (RefPart const & refpart)
        {
            using namespace amrex::literals; // for _rt and _prt
 
            Alignment::compute_constants(refpart);
        }

        template<typename T_Real=amrex::ParticleReal, typename T_IdCpu=uint64_t>
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void operator() (
            T_Real & AMREX_RESTRICT x,
            T_Real & AMREX_RESTRICT y,
            T_Real & AMREX_RESTRICT t,
            T_Real & AMREX_RESTRICT px,
            T_Real & AMREX_RESTRICT py,
            T_Real & AMREX_RESTRICT pt,
            T_IdCpu & AMREX_RESTRICT idcpu,
            RefPart const & AMREX_RESTRICT refpart
        ) const
        {
            using namespace amrex::literals; // for _rt and _prt
 
            // shift due to alignment errors of the element
            shift_in(x, y, px, py);
 
            // get the linear map
            amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> const R = refpart.map;
 
            // symplectic linear map for the RF cavity is computed using the
            // Hamiltonian formalism as described in:
            // https://uspas.fnal.gov/materials/09UNM/ComputationalMethods.pdf.
            // R denotes the transfer matrix in the basis (x,px,y,py,t,pt),
            // so that, e.g., R(3,4) = dyf/dpyi.
            amrex::SmallVector<T_Real, 6, 1> const v{x, px, y, py, t, pt};
 
            // push particles using the linear map
            auto const out = R * v;
 
            // assign updated values
            x  = out[1];
            px = out[2];
            y  = out[3];
            py = out[4];
            t  = out[5];
            pt = out[6];
 
            // apply transverse aperture
            apply_aperture(x, y, idcpu);
 
            // undo shift due to alignment errors of the element
            shift_out(x, y, px, py);
        }

        AMREX_GPU_HOST AMREX_FORCE_INLINE
        void operator() (RefPart & AMREX_RESTRICT refpart) const
        {
            using namespace amrex::literals; // for _rt and _prt
            using amrex::Math::powi;
 
            // assign input reference particle values
            amrex::ParticleReal const x = refpart.x;
            amrex::ParticleReal const px = refpart.px;
            amrex::ParticleReal const y = refpart.y;
            amrex::ParticleReal const py = refpart.py;
            amrex::ParticleReal const z = refpart.z;
            amrex::ParticleReal const pz = refpart.pz;
            amrex::ParticleReal const pt = refpart.pt;
            amrex::ParticleReal const s = refpart.s;
            amrex::ParticleReal const sedge = refpart.sedge;
 
            // initialize linear map (deviation) values
            refpart.map = decltype(refpart.map)::Identity();
 
            // length of the current slice
            amrex::ParticleReal const slice_ds = m_ds / nslice();
 
            // compute initial value of beta*gamma
            amrex::ParticleReal const bgi = std::sqrt(powi<2>(pt) - 1.0_prt);
 
            // call integrator to advance (t,pt)
            amrex::ParticleReal const zin = s - sedge;
            amrex::ParticleReal const zout = zin + slice_ds;
            int const nsteps = m_mapsteps;
 
            integrators::symp2_integrate_split3(refpart,zin,zout,nsteps,*this);
            amrex::ParticleReal const ptf = refpart.pt;
 
            // advance position (x,y,z)
            refpart.x = x + slice_ds*px/bgi;
            refpart.y = y + slice_ds*py/bgi;
            refpart.z = z + slice_ds*pz/bgi;
 
            // compute final value of beta*gamma
            amrex::ParticleReal const bgf = std::sqrt(powi<2>(ptf) - 1.0_prt);
 
            // advance momentum (px,py,pz)
            refpart.px = px*bgf/bgi;
            refpart.py = py*bgf/bgi;
            refpart.pz = pz*bgf/bgi;
 
            // convert linear map from dynamic to static units
            amrex::ParticleReal scale_in = 1.0_prt;
            amrex::ParticleReal scale_fin = 1.0_prt;
 
            for (int i=1; i<7; i++) {
               for (int j=1; j<7; j++) {
                   if( i % 2 == 0)
                      scale_fin = bgf;
                   else
                      scale_fin = 1.0_prt;
                   if( j % 2 == 0)
                      scale_in = bgi;
                   else
                      scale_in = 1.0_prt;
                   refpart.map(i, j) = refpart.map(i, j) * scale_in / scale_fin;
               }
            }
 
            // advance integrated path length
            refpart.s = s + slice_ds;
        }

        using LinearTransport::operator();

        AMREX_GPU_HOST AMREX_FORCE_INLINE
        Map6x6
        transport_map ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const
        {
 
            Map6x6 R = Map6x6::Identity();
            R = refpart.map;
 
            return R;
        }

        std::tuple<amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal>
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        RF_Efield (amrex::ParticleReal const zeval) const
        {
            using namespace amrex::literals; // for _rt and _prt
            using amrex::Math::powi;
 
            // pick the right data depending if we are on the host side
            // (reference particle push) or device side (particles):
#if AMREX_DEVICE_COMPILE
            amrex::ParticleReal* cos_data = m_cos_d_data;
            amrex::ParticleReal* sin_data = m_sin_d_data;
#else
            amrex::ParticleReal* cos_data = m_cos_h_data;
            amrex::ParticleReal* sin_data = m_sin_h_data;
#endif
 
            // specify constants
            using ablastr::constant::math::pi;
            amrex::ParticleReal const zlen = m_ds;
            amrex::ParticleReal const zmid = zlen * 0.5_prt;
 
            // compute on-axis electric field (z is relative to cavity midpoint)
            amrex::ParticleReal efield = 0.0;
            amrex::ParticleReal efieldp = 0.0;
            amrex::ParticleReal efieldpp = 0.0;
            amrex::ParticleReal efieldint = 0.0;
            amrex::ParticleReal const z = zeval - zmid;
 
            if (std::abs(z) <= zmid)
            {
               efield = 0.5_prt*cos_data[0];
               efieldint = z*efield;
               for (int j=1; j < m_ncoef; ++j)
               {
                 efield = efield + cos_data[j] * std::cos(j*2*pi*z/zlen) +
                     sin_data[j] * std::sin(j*2*pi*z/zlen);
                 efieldp = efieldp-j*2*pi*cos_data[j] * std::sin(j*2*pi*z/zlen)/zlen +
                      j*2*pi*sin_data[j] * std::cos(j*2*pi*z/zlen)/zlen;
                 efieldpp = efieldpp- powi<2>(j*2*pi*cos_data[j]/zlen)  * std::cos(j*2*pi*z/zlen) -
                     powi<2>(j*2*pi*sin_data[j]/zlen)  * std::sin(j*2*pi*z/zlen);
                 efieldint = efieldint + zlen*cos_data[j] * std::sin(j*2*pi*z/zlen)/(j*2*pi) -
                      zlen*sin_data[j] * std::cos(j*2*pi*z/zlen)/(j*2*pi);
               }
            }
            else  // endpoint of the RF, outsize zlen
            {
               efieldint = std::copysign(z, z)*zmid*0.5_prt*cos_data[0];;
               for (int j=1; j < m_ncoef; ++j)
               {
                 efieldint = efieldint - zlen*sin_data[j] * std::cos(j*pi)/(j*2*pi);
               }
            }
            return std::make_tuple(efield, efieldp, efieldint);
        }

        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void map3 (amrex::ParticleReal const tau,
                   RefPart & refpart,
                   [[maybe_unused]] amrex::ParticleReal & zeval) const
        {
            using namespace amrex::literals; // for _rt and _prt
            using amrex::Math::powi;
 
            // push the reference particle
            amrex::ParticleReal const t = refpart.t;
            amrex::ParticleReal const pt = refpart.pt;
 
            if (pt < -1.0_prt) {
                refpart.t = t + tau/std::sqrt(1.0_prt - powi<-2>(pt));
                refpart.pt = pt;
            }
            else {
                refpart.t = t;
                refpart.pt = pt;
            }
 
            // push the linear map equations
            amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> const R = refpart.map;
            amrex::ParticleReal const betgam = refpart.beta_gamma();
 
            refpart.map(5,5) = R(5,5) + tau*R(6,5)/powi<3>(betgam);
            refpart.map(5,6) = R(5,6) + tau*R(6,6)/powi<3>(betgam);
        }

        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void map2 (amrex::ParticleReal const tau,
                   RefPart & refpart,
                   amrex::ParticleReal & zeval) const
        {
            using namespace amrex::literals; // for _rt and _prt
            using amrex::Math::powi;
 
            amrex::ParticleReal const t = refpart.t;
            amrex::ParticleReal const pt = refpart.pt;
 
            // Define parameters and intermediate constants
            using ablastr::constant::math::pi;
            using ablastr::constant::SI::c;
            amrex::ParticleReal const k = (2_prt*pi/c)*m_freq;
            amrex::ParticleReal const phi = m_phase*(pi/180_prt);
            amrex::ParticleReal const E0 = m_escale;
 
            // push the reference particle
            auto [ez, ezp, ezint] = RF_Efield(zeval);
            amrex::ignore_unused(ez, ezint);
 
            refpart.t = t;
            refpart.pt = pt;
 
            // push the linear map equations
            amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> const R = refpart.map;
            amrex::ParticleReal const s = tau/refpart.beta_gamma();
            amrex::ParticleReal const L = E0*ezp * std::sin(k*t+phi)/(2_prt*k);
 
            refpart.map(1,1) = (1_prt-s*L)*R(1,1) + s*R(2,1);
            refpart.map(1,2) = (1_prt-s*L)*R(1,2) + s*R(2,2);
            refpart.map(2,1) = -s * powi<2>(L)*R(1,1) + (1_prt+s*L)*R(2,1);
            refpart.map(2,2) = -s * powi<2>(L)*R(1,2) + (1_prt+s*L)*R(2,2);
 
            refpart.map(3,3) = (1_prt-s*L)*R(3,3) + s*R(4,3);
            refpart.map(3,4) = (1_prt-s*L)*R(3,4) + s*R(4,4);
            refpart.map(4,3) = -s * powi<2>(L)*R(3,3) + (1_prt+s*L)*R(4,3);
            refpart.map(4,4) = -s * powi<2>(L)*R(3,4) + (1_prt+s*L)*R(4,4);
        }

        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void map1 (amrex::ParticleReal const tau,
                   RefPart & refpart,
                   amrex::ParticleReal & zeval) const
        {
            using namespace amrex::literals; // for _rt and _prt
 
            amrex::ParticleReal const t = refpart.t;
            amrex::ParticleReal const pt = refpart.pt;
            amrex::ParticleReal const z = zeval;
 
            // Define parameters and intermediate constants
            using ablastr::constant::math::pi;
            using ablastr::constant::SI::c;
            amrex::ParticleReal const k = (2_prt*pi/c)*m_freq;
            amrex::ParticleReal const phi = m_phase*(pi/180_prt);
            amrex::ParticleReal const E0 = m_escale;
 
            // push the reference particle
            auto [ez, ezp, ezint] = RF_Efield(z);
            amrex::ignore_unused(ez);
            zeval = z + tau;
            auto [ezf, ezpf, ezintf] = RF_Efield(zeval);
            amrex::ignore_unused(ezf);
 
            refpart.t = t;
            refpart.pt = pt - E0*(ezintf-ezint) * std::cos(k*t+phi);
 
            // push the linear map equations
            amrex::SmallMatrix<amrex::ParticleReal, 6, 6, amrex::Order::F, 1> const R = refpart.map;
            amrex::ParticleReal const M = E0*(ezintf-ezint)*k * std::sin(k*t+phi);
            amrex::ParticleReal const L = E0*(ezpf-ezp) * std::sin(k*t+phi)/(2_prt*k)+M*0.5_prt;
 
            refpart.map(2,1) = L*R(1,1) + R(2,1);
            refpart.map(2,2) = L*R(1,2) + R(2,2);
 
            refpart.map(4,3) = L*R(3,3) + R(4,3);
            refpart.map(4,4) = L*R(3,4) + R(4,4);
 
            refpart.map(6,5) = M*R(5,5) + R(6,5);
            refpart.map(6,6) = M*R(5,6) + R(6,6);
        }

        void
        finalize ()
        {
            // remove from unique data map
            if (RFCavityData::h_cos_coef.count(m_id) != 0u)
                RFCavityData::h_cos_coef.erase(m_id);
            if (RFCavityData::h_sin_coef.count(m_id) != 0u)
                RFCavityData::h_sin_coef.erase(m_id);
 
            if (RFCavityData::d_cos_coef.count(m_id) != 0u)
                RFCavityData::d_cos_coef.erase(m_id);
            if (RFCavityData::d_sin_coef.count(m_id) != 0u)
                RFCavityData::d_sin_coef.erase(m_id);
        }
 
        amrex::ParticleReal m_escale; 
        amrex::ParticleReal m_freq; 
        amrex::ParticleReal m_phase; 
        int m_mapsteps; 
        int m_id; 
 
        int m_ncoef = 0; 
        amrex::ParticleReal* m_cos_h_data = nullptr; 
        amrex::ParticleReal* m_sin_h_data = nullptr; 
        amrex::ParticleReal* m_cos_d_data = nullptr; 
        amrex::ParticleReal* m_sin_d_data = nullptr; 
    };
 
} // namespace impactx
 
#endif // IMPACTX_RFCAVITY_H