NonlinearLens.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_NONLINEARLENS_H
#define IMPACTX_NONLINEARLENS_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_Extension.H>
#include <AMReX_GpuComplex.H>
#include <AMReX_Math.H>
#include <AMReX_REAL.H>
// TODO: #include <AMReX_SIMD.H>
//   SIMD will require better complex math support in vir-simd / C++26
//   https://github.com/mattkretz/vir-simd/issues/42
 
#include <cmath>
#include <stdexcept>
 
namespace impactx::elements
{
    struct NonlinearLens
    : public mixin::Named,
      public mixin::BeamOptic<NonlinearLens>,
      public mixin::LinearTransport<NonlinearLens>,
      public mixin::Thin,
      public mixin::Alignment,
      public mixin::NoFinalize
      // TODO: public amrex::simd::Vectorized<amrex::simd::native_simd_size_particlereal>
    {
        static constexpr auto type = "NonlinearLens";
        using PType = ImpactXParticleContainer::ParticleType;

        NonlinearLens (
            amrex::ParticleReal knll,
            amrex::ParticleReal cnll,
            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_knll(knll), m_cnll(cnll)
        {
        }

        void reverse () { m_knll = -m_knll; }

        using BeamOptic::operator();

        void compute_constants (RefPart const & refpart)
        {
            using namespace amrex::literals; // for _rt and _prt
 
            Alignment::compute_constants(refpart);
 
            m_inv_cnll = 1.0_prt / m_cnll;
            m_kick = -m_knll * m_inv_cnll;
        }

        template<typename T_Real=amrex::ParticleReal, typename T_IdCpu=uint64_t>
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void operator() (
            T_Real const & AMREX_RESTRICT x,
            T_Real const & AMREX_RESTRICT y,
            [[maybe_unused]] T_Real const & AMREX_RESTRICT t,
            T_Real & AMREX_RESTRICT px,
            T_Real & AMREX_RESTRICT py,
            [[maybe_unused]] 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;
 
            // a complex type with two T_Real
            using Complex = amrex::GpuComplex<T_Real>;
 
            // access reference particle values to find (beta*gamma)^2
            //amrex::ParticleReal const pt_ref = refpart.pt;
            //amrex::ParticleReal const betgam2 = powi<2>(pt_ref) - 1.0_prt;
 
            // 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;
 
            // assign complex position zeta = (x + iy)/cnll
            Complex zeta(x, y);
            zeta = zeta * m_inv_cnll;
            Complex const re1(1.0_prt, 0.0_prt);
            Complex const im1(0.0_prt, 1.0_prt);
 
            // compute croot = sqrt(1-zeta**2)
            Complex croot = powi<2>(zeta);
            croot = re1 - croot;
            croot = amrex::sqrt(croot);
 
            // compute carcsin = arcsin(zeta)
            Complex carcsin = im1 * zeta + croot;
            carcsin = -im1 * amrex::log(carcsin);
 
            // compute complex function F'(zeta)
            Complex dF = zeta / powi<2>(croot);
            dF = dF + carcsin / powi<3>(croot);
 
            // compute momentum kick
            T_Real const dpx = +m_kick * dF.m_real;
            T_Real const dpy = -m_kick * dF.m_imag;
 
            // advance position and momentum
            // xout = x;
            pxout = px + dpx;
 
            // yout = y;
            pyout = py + dpy;
 
            // tout = t;
            // 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 ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const
        {
            using namespace amrex::literals; // for _rt and _prt
            using amrex::Math::powi;
 
            Map6x6 R = Map6x6::Identity();
            // Leading 2*zeta term of F'(zeta); higher orders vanish at zero.
            amrex::ParticleReal const k_lin = 2.0_prt * m_knll / powi<2>(m_cnll);
            R(2,1) = -k_lin;
            R(4,3) = +k_lin;
 
            // apply the transverse rotation (roll) alignment error
            return rotate_aligned_map(R);
        }
 
        amrex::ParticleReal m_knll; 
        amrex::ParticleReal m_cnll; 
 
      private:
        // constants that are independent of the individually tracked particle,
        // see: compute_constants() to refresh
        amrex::ParticleReal m_inv_cnll;  
        amrex::ParticleReal m_kick;
    };
 
} // namespace impactx
 
IMPACTX_PUSH_EXTERN_TEMPLATE(impactx::elements::NonlinearLens)
 
#endif // IMPACTX_NONLINEARLENS_H