PRot.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_PROT_H
#define IMPACTX_PROT_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 <ablastr/constant.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 PRot
    : public mixin::Named,
      public mixin::BeamOptic<PRot>,
      public mixin::LinearTransport<PRot>,
      public mixin::Thin,
      public mixin::NoFinalize,
      public amrex::simd::Vectorized<amrex::simd::native_simd_size_particlereal>
    {
        static constexpr auto type = "PRot";
        using PType = ImpactXParticleContainer::ParticleType;

        PRot (
            amrex::ParticleReal phi_in,
            amrex::ParticleReal phi_out,
            std::optional<std::string> name = DEFAULT_name
        )
        : Named(std::move(name)),
          m_phi_in(phi_in * mixin::Alignment::degree2rad),
          m_phi_out(phi_out * mixin::Alignment::degree2rad)
        {
        }

        void reverse () { std::swap(m_phi_in, m_phi_out); }

        using BeamOptic::operator();

        void compute_constants (RefPart const & refpart)
        {
            using namespace amrex::literals; // for _rt and _prt
 
            // access reference particle values to find 1/beta
            m_inv_beta = 1.0_prt / refpart.beta();
 
            // store rotation angle and initial, final values of pz
            amrex::ParticleReal const theta = m_phi_out - m_phi_in;
            auto const [sin_theta, cos_theta] = amrex::Math::sincos(theta);
            m_sin_theta = sin_theta;
            m_cos_theta = cos_theta;
            auto const [sin_phi_in, cos_phi_in] = amrex::Math::sincos(m_phi_in);
            m_sin_phi_in = sin_phi_in;
            m_cos_phi_in = cos_phi_in;
        }

        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 const & 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 namespace std; // for cmath(float)
            using amrex::Math::powi;
 
            // 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;
 
            T_Real const pz = sqrt(
                1.0_prt -
                2.0_prt * pt * m_inv_beta +
                powi<2>(pt) -
                powi<2>(py) -
                powi<2>(px + m_sin_phi_in)
            );
            T_Real const pzf = pz * m_cos_theta - (px + m_sin_phi_in) * m_sin_theta;
 
            // advance position and momentum
            xout = x * pz / pzf;
            pxout = px * m_cos_theta + (pz - m_cos_phi_in) * m_sin_theta;
 
            yout = y + py * x * m_sin_theta / pzf;
            // pyout = py;
 
            tout = t - (pt - m_inv_beta) * x * m_sin_theta / pzf;
            // ptout = pt;
 
            // assign updated values
            x = xout;
            y = yout;
            t = tout;
            px = pxout;
            // py = pyout;
            // pt = ptout;
        }

        using Thin::operator();

        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
 
            amrex::ParticleReal const beta = refpart.beta();
 
            // Precompute trigonometric quantities.
            amrex::ParticleReal const theta = m_phi_out - m_phi_in;
            auto const [sin_theta, cos_theta] = amrex::Math::sincos(theta);
            (void) cos_theta;  // not needed for the linear map
            auto const [sin_phi_in, cos_phi_in] = amrex::Math::sincos(m_phi_in);
            (void) sin_phi_in;
            auto const [sin_phi_out, cos_phi_out] = amrex::Math::sincos(m_phi_out);
            (void) sin_phi_out;
 
            Map6x6 R = Map6x6::Identity();
            R(1,1) = cos_phi_in / cos_phi_out;
            R(2,2) = cos_phi_out / cos_phi_in;
            R(2,6) = -sin_theta / (beta * cos_phi_in);
            R(5,1) = +sin_theta / (beta * cos_phi_out);
 
            return R;
        }
 
        amrex::ParticleReal m_phi_in; 
        amrex::ParticleReal m_phi_out; 
 
      private:
        // constants that are independent of the individually tracked particle,
        // see: compute_constants() to refresh
        amrex::ParticleReal m_inv_beta, m_sin_theta, m_cos_theta, m_sin_phi_in, m_cos_phi_in;
    };
 
} // namespace impactx
 
IMPACTX_PUSH_EXTERN_TEMPLATE(impactx::elements::PRot)
 
#endif // IMPACTX_PROT_H