ChrUniformAcc.H

Go to the documentation of this file.

/* Copyright 2022-2026 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_CHRACC_H
#define IMPACTX_CHRACC_H
 
#include "particles/ImpactXParticleContainer.H"
#include "mixin/alignment.H"
#include "mixin/pipeaperture.H"
#include "mixin/beamoptic.H"
#include "mixin/thick.H"
#include "mixin/lineartransport.H"
#include "mixin/named.H"
#include "mixin/nofinalize.H"
 
#include <AMReX_Extension.H>
#include <AMReX_Math.H>
#include <AMReX_REAL.H>
#include <AMReX_SIMD.H>
 
#include <cmath>
#include <stdexcept>
 
namespace impactx::elements
{
    struct ChrAcc
    : public mixin::Named,
      public mixin::BeamOptic<ChrAcc>,
      public mixin::LinearTransport<ChrAcc>,
      public mixin::Thick,
      public mixin::Alignment,
      public mixin::PipeAperture,
      public mixin::NoFinalize,
      public amrex::simd::Vectorized<amrex::simd::native_simd_size_particlereal>
    {
        static constexpr auto type = "ChrAcc";
        using PType = ImpactXParticleContainer::ParticleType;

        ChrAcc (
            amrex::ParticleReal ds,
            amrex::ParticleReal ez,
            amrex::ParticleReal bz,
            amrex::ParticleReal dx = DEFAULT_dx,
            amrex::ParticleReal dy = DEFAULT_dy,
            amrex::ParticleReal rotation_degree = DEFAULT_rotation_degree,
            amrex::ParticleReal aperture_x = DEFAULT_aperture_x,
            amrex::ParticleReal aperture_y = DEFAULT_aperture_y,
            int nslice = DEFAULT_nslice,
            std::optional<std::string> name = DEFAULT_name
        )
        : Named(std::move(name)),
          Thick(ds, nslice),
          Alignment(dx, dy, rotation_degree),
          PipeAperture(aperture_x, aperture_y),
          m_ez(ez), m_bz(bz)
        {
        }

        void reverse () { Thick::reverse(); }

        using BeamOptic::operator();

        void compute_constants (RefPart const & refpart)
        {
            using namespace amrex::literals; // for _rt and _prt
            using amrex::Math::powi;
 
            Alignment::compute_constants(refpart);
 
            // length of the current slice
            m_slice_ds = m_ds / nslice();
 
            // access reference particle values (final, initial):
            m_ptf_ref = refpart.pt;
            m_pti_ref = m_ptf_ref + m_ez * m_slice_ds;
            m_bgf = std::sqrt(powi<2>(m_ptf_ref) - 1.0_prt);
            m_bgi = std::sqrt(powi<2>(m_pti_ref) - 1.0_prt);
 
            // compute focusing constant (1/m) and rotation angle (in rad)
            m_alpha = m_bz * 0.5_prt;
            m_alpha_iez = m_alpha / m_ez;
        }

        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,
            [[maybe_unused]] RefPart const & AMREX_RESTRICT refpart
        ) const
        {
            using namespace amrex::literals; // for _rt and _prt
            using amrex::Math::powi;
            using namespace std; // for cmath(float)
            namespace stdx = amrex::simd::stdx;
 
            // initial conversion from static to dynamic units:
            px = px * m_bgi;
            py = py * m_bgi;
            pt = pt * m_bgi;
 
            // compute intermediate quantities related to acceleration
            T_Real const pti_tot = m_pti_ref + pt;
            T_Real const ptf_tot = m_ptf_ref + pt;
            T_Real const pti_tot2 = powi<2>(pti_tot);
            T_Real const ptf_tot2 = powi<2>(ptf_tot);
 
            // check whether particle lies within the domain of map definition
            auto const mask = pti_tot2 <= 1_prt || ptf_tot2 <= 1_prt;
            amrex::ParticleIDWrapper<T_IdCpu>{idcpu}.make_invalid(mask);
            {
                T_Real const pzi_tot = sqrt(powi<2>(pti_tot) - 1_prt);
                T_Real const pzf_tot = sqrt(powi<2>(ptf_tot) - 1_prt);
                T_Real const pzi_ref = sqrt(powi<2>(m_pti_ref) - 1_prt);
                T_Real const pzf_ref = sqrt(powi<2>(m_ptf_ref) - 1_prt);
 
                T_Real const numer = -ptf_tot + pzf_tot;
                T_Real const denom = -pti_tot + pzi_tot;
 
                // compute focusing constant (1/m) and rotation angle (in rad)
                T_Real const theta = m_alpha_iez * log(numer / denom);
                auto const [sin_theta, cos_theta] = amrex::Math::sincos(theta);
 
                // initialize output values
                T_Real xout = x;
                T_Real yout = y;
                T_Real tout = t;
                T_Real pxout = px;
                T_Real pyout = py;
                T_Real ptout = pt;
 
                // advance positions and momenta using map for focusing
                xout = cos_theta * x + sin_theta / m_alpha * px;
                pxout = -m_alpha * sin_theta * x + cos_theta * px;
 
                yout = cos_theta * y + sin_theta / m_alpha * py;
                pyout = -m_alpha * sin_theta * y + cos_theta * py;
 
                // the correct symplectic update for t
                tout = t + (pzf_tot - pzf_ref - pzi_tot + pzi_ref) / m_ez;
                tout += (1_prt/pzi_tot - 1_prt/pzf_tot)
                      * (powi<2>(py - m_alpha * x) +
                         powi<2>(px + m_alpha * y))
                      / (2_prt * m_ez);
                ptout = pt;
 
                // assign intermediate momenta
                px = pxout;
                py = pyout;
                pt = ptout;
 
                // advance positions and momenta using map for rotation
                x = cos_theta * xout + sin_theta * yout;
                pxout = cos_theta * px + sin_theta * py;
 
                y = -sin_theta * xout + cos_theta * yout;
                pyout = -sin_theta * px + cos_theta * py;
 
                t = tout;
                ptout = pt;
 
                // assign updated momenta
                px = pxout;
                py = pyout;
                pt = ptout;
 
            }
 
            // final conversion from dynamic to static units:
            px = px / m_bgf;
            py = py / m_bgf;
            pt = pt / m_bgf;
        }

        AMREX_GPU_HOST_DEVICE 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 t = refpart.t;
            amrex::ParticleReal const pt = refpart.pt;
            amrex::ParticleReal const s = refpart.s;
 
            // 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);
 
            // advance pt (uniform acceleration)
            refpart.pt = pt - m_ez*slice_ds;
 
            // compute final value of beta*gamma
            amrex::ParticleReal const ptf = refpart.pt;
            amrex::ParticleReal const bgf = std::sqrt(powi<2>(ptf) - 1.0_prt);
 
            // update t
            refpart.t = t + (bgf - bgi)/m_ez;
 
            // 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;
 
            // advance momentum (px,py,pz)
            refpart.px = px*bgf/bgi;
            refpart.py = py*bgf/bgi;
            refpart.pz = pz*bgf/bgi;
 
            // advance integrated path length
            refpart.s = s + slice_ds;
        }

        using LinearTransport::operator();

        AMREX_GPU_HOST AMREX_FORCE_INLINE
        Map6x6
        transport_map (RefPart const & AMREX_RESTRICT refpart) const
        {
            using namespace amrex::literals; // for _rt and _prt
            using amrex::Math::powi;
 
            // slice length and initial/final reference pt:
            // transport_map is called after the reference advance, so
            // refpart.pt holds the post-slice (final) pt.
            amrex::ParticleReal const slice_ds = m_ds / nslice();
            amrex::ParticleReal const ptf_ref = refpart.pt;
            amrex::ParticleReal const pti_ref = ptf_ref + m_ez * slice_ds;
 
            amrex::ParticleReal const pzf_ref = std::sqrt(powi<2>(ptf_ref) - 1.0_prt);
            amrex::ParticleReal const pzi_ref = std::sqrt(powi<2>(pti_ref) - 1.0_prt);
            amrex::ParticleReal const bgf = pzf_ref;  // |beta*gamma|_final
            amrex::ParticleReal const bgi = pzi_ref;  // |beta*gamma|_initial
 
            // Solenoid focusing constant alpha = Bz / 2. The "log"
            // factor is the integrated drift length in dynamic units and
            // is well-defined even when alpha = 0 (no solenoid).
            amrex::ParticleReal const alpha = 0.5_prt * m_bz;
            amrex::ParticleReal const drift_log =
                std::log((pzf_ref - ptf_ref) / (pzi_ref - pti_ref)) / m_ez;
            amrex::ParticleReal const rbg = bgi / bgf;
 
            Map6x6 R = Map6x6::Identity();
 
            // t couples to pt only: symplectic correction to the
            // time-of-flight, linearized at pt = 0.
            R(5,6) = bgi * (ptf_ref / pzf_ref - pti_ref / pzi_ref) / m_ez;
            // pt scales by the static<->dynamic unit change bgi/bgf.
            R(6,6) = rbg;
 
            if (alpha == 0.0_prt)
            {
                // No solenoid: pure accelerating "drift" with adiabatic
                // damping on the transverse momenta.
                R(1,2) = bgi * drift_log;
                R(2,2) = rbg;
                R(3,4) = bgi * drift_log;
                R(4,4) = rbg;
            }
            else
            {
                // Solenoid + acceleration: rotation * focusing, bracketed
                // by static<->dynamic unit conversions on (px, py).
                amrex::ParticleReal const theta0 = alpha * drift_log;
                auto const [s, c] = amrex::Math::sincos(theta0);
 
                amrex::ParticleReal const c2 = c * c;
                amrex::ParticleReal const s2 = s * s;
                amrex::ParticleReal const cs = c * s;
 
                R(1,1) =  c2;                R(1,2) =  bgi * cs / alpha;
                R(1,3) =  cs;                R(1,4) =  bgi * s2 / alpha;
                R(2,1) = -alpha * cs / bgf;  R(2,2) =  rbg * c2;
                R(2,3) = -alpha * s2 / bgf;  R(2,4) =  rbg * cs;
                R(3,1) = -cs;                R(3,2) = -bgi * s2 / alpha;
                R(3,3) =  c2;                R(3,4) =  bgi * cs / alpha;
                R(4,1) =  alpha * s2 / bgf;  R(4,2) = -rbg * cs;
                R(4,3) = -alpha * cs / bgf;  R(4,4) =  rbg * c2;
            }
 
            // apply the transverse rotation (roll) alignment error
            return rotate_aligned_map(R);
        }
 
        amrex::ParticleReal m_ez; 
        amrex::ParticleReal m_bz; 
 
      private:
        // constants that are independent of the individually tracked particle,
        // see: compute_constants() to refresh
        amrex::ParticleReal m_slice_ds;  
        amrex::ParticleReal m_ptf_ref, m_pti_ref, m_bgf, m_bgi;
        amrex::ParticleReal m_alpha, m_alpha_iez;
    };
 
} // namespace impactx
 
IMPACTX_PUSH_EXTERN_TEMPLATE(impactx::elements::ChrAcc)
 
#endif // IMPACTX_CHRACC_H