QuadEdge.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_QUADEDGE_H
#define IMPACTX_QUADEDGE_H
 
#include "particles/ImpactXParticleContainer.H"
#include "mixin/alignment.H"
#include "mixin/beamoptic.H"
#include "mixin/thin.H"
#include "mixin/lineartransport.H"
#include "mixin/named.H"
#include "mixin/nofinalize.H"
 
#include <AMReX_Enum.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 QuadEdge
    : public mixin::Named,
      public mixin::BeamOptic<QuadEdge>,
      public mixin::LinearTransport<QuadEdge>,
      public mixin::Thin,
      public mixin::Alignment,
      public mixin::NoFinalize,
      public amrex::simd::Vectorized<amrex::simd::native_simd_size_particlereal>
    {
        static constexpr auto type = "QuadEdge";
        using PType = ImpactXParticleContainer::ParticleType;
 
        AMREX_ENUM_IN_CLASS(Location,
            entry = 1,  
            exit = -1,  
        );
 
        static constexpr int DEFAULT_unit = 0;
        static constexpr Location DEFAULT_flag = Location::entry;

        QuadEdge (
            amrex::ParticleReal k,
            int unit = DEFAULT_unit,
            Location flag = DEFAULT_flag,
            amrex::ParticleReal dx = DEFAULT_dx,
            amrex::ParticleReal dy = DEFAULT_dy,
            amrex::ParticleReal rotation_degree = DEFAULT_rotation_degree,
            std::optional<std::string> name = DEFAULT_name
        )
        : Named(std::move(name)),
          Alignment(dx, dy, rotation_degree),
          m_k(k), m_unit(unit), m_flag(flag)
        {
        }

        void reverse ()
        {
            m_flag = (m_flag == Location::entry)
                ? Location::exit
                : Location::entry;
        }

        using BeamOptic::operator();

        void compute_constants ([[maybe_unused]] RefPart const & refpart)
        {
            using namespace amrex::literals; // for _rt and _prt
 
            Alignment::compute_constants(refpart);
 
            // access reference particle values to find beta
            m_beta = refpart.beta();
 
            // normalize quad units to MAD-X convention if needed
            m_g = m_unit == 1 ? m_k / refpart.rigidity_Tm() : m_k;
 
            // the global scale constant; sign depends on entry or exit
            m_a = static_cast<amrex::ParticleReal>(m_flag) * (-1_prt * m_g / 12.0_prt);
        }

        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 const & AMREX_RESTRICT pt,
            [[maybe_unused]] T_IdCpu const & 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)
 
            // compute particle momentum deviation delta + 1
            T_Real const delta1 = sqrt(1_prt - 2_prt*pt/m_beta + powi<2>(pt));
            T_Real const a = m_a / delta1;
 
            // initialize output values
            T_Real xout = x;
            T_Real yout = y;
            T_Real tout = t;
            T_Real pxout = px;
            T_Real pyout = py;
 
            // intermediate quantities to be used below
            T_Real const x2my2 = x*x - y*y;
            T_Real const x2py2 = x*x + y*y;
            T_Real const denominator = 1_prt - 9_prt*powi<2>(a)*powi<2>(x2my2);
 
            // apply edge focusing (transverse phase space)
            xout = x - a * (powi<3>(x) + 3_prt * x * powi<2>(y));
            yout = y + a * (3_prt * powi<2>(x) * y + powi<3>(y));
 
            pxout = (px + 3_prt*a*px*x2py2 - 6_prt*a*py*x*y) / denominator;
            pyout = (py - 3_prt*a*py*x2py2 + 6_prt*a*px*x*y) / denominator;
 
            // apply edge focusing (longitudinal phase space)
            T_Real const f4 = a * ( pxout * (powi<3>(x) + 3_prt*x*powi<2>(y))
                                         - pyout * (powi<3>(y) + 3_prt*y*powi<2>(x)) );
            tout = t + f4 * (pt - 1_prt/m_beta) / powi<2>(delta1);
 
            // update output values
            x = xout;
            y = yout;
            t = tout;
            px = pxout;
            py = pyout;
        }

        using Thin::operator();

        using LinearTransport::operator();

        AMREX_GPU_HOST AMREX_FORCE_INLINE
        Map6x6
        transport_map ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const
        {
            using namespace amrex::literals; // for _rt and _prt
 
            // initialize linear map matrix elements
            Map6x6 R = Map6x6::Identity();
 
            // apply the transverse rotation (roll) alignment error
            return rotate_aligned_map(R);
        }
 
        amrex::ParticleReal m_k;  
        int m_unit; 
        Location m_flag; 
 
      private:
        // constants that are independent of the individually tracked particle,
        // see: compute_constants() to refresh
        amrex::ParticleReal m_g;
        amrex::ParticleReal m_a;
        amrex::ParticleReal m_beta;
    };
 
} // namespace impactx
 
IMPACTX_PUSH_EXTERN_TEMPLATE(impactx::elements::QuadEdge)
 
#endif // IMPACTX_QUADEDGE_H