LinearMap.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_ELEMENT_LINEAR_MAP_H
#define IMPACTX_ELEMENT_LINEAR_MAP_H
 
#include "particles/ImpactXParticleContainer.H"
#include "mixin/alignment.H"
#include "mixin/beamoptic.H"
#include "mixin/lineartransport.H"
#include "mixin/spintransport.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>
#include <type_traits>
 
namespace impactx::elements
{
    struct LinearMap
    : public mixin::Named,
      public mixin::BeamOptic<LinearMap>,
      public mixin::LinearTransport<LinearMap>,
      public mixin::Alignment,
      public mixin::SpinTransport,
      public mixin::NoFinalize,
      public amrex::simd::Vectorized<amrex::simd::native_simd_size_particlereal>
    {
        static constexpr auto type = "LinearMap";
        using PType = ImpactXParticleContainer::ParticleType;
 
        static constexpr amrex::ParticleReal DEFAULT_ds = 0;

        LinearMap (
            Map6x6 const & R,
            amrex::ParticleReal ds = DEFAULT_ds,
            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_transport_map = R;
            m_ds = ds;
        }

        bool symplectic () const
        {
            using namespace amrex::literals;
 
            Map6x6 const J = impactx::symplectic_form();
            Map6x6 const RT = m_transport_map.transpose();
            Map6x6 const check = RT * J * m_transport_map;
            amrex::ParticleReal max_err = 0_prt;
            for (int i = 1; i <= 6; ++i) {
                for (int j = 1; j <= 6; ++j) {
                    amrex::ParticleReal const err = std::abs(check(i,j) - J(i,j));
                    max_err = std::max(max_err, err);
                }
            }
 
            amrex::ParticleReal constexpr tol = std::is_same_v<amrex::ParticleReal, float>
                ? 1.0e-4_prt : 1.0e-10_prt;
            return max_err <= tol;
        }

        void reverse ()
        {
            if (!symplectic()) {
                throw std::runtime_error(
                    "LinearMap::reverse: transport map is not symplectic.");
            }
 
            Map6x6 const Omega = impactx::symplectic_form();
 
            // R_inv = Omega^T * R^T * Omega = -Omega * R^T * Omega
            Map6x6 const RT = m_transport_map.transpose();
            Map6x6 const R_inv = (-Omega) * RT * Omega;
            m_transport_map = R_inv;
            m_ds = -m_ds;
        }

        using BeamOptic::operator();

        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,
            [[maybe_unused]] T_IdCpu const & AMREX_RESTRICT idcpu,
            [[maybe_unused]] RefPart const & AMREX_RESTRICT refpart
        ) const
        {
            using namespace amrex::literals; // for _rt and _prt
 
            // input and output phase space vectors
            amrex::SmallVector<T_Real, 6, 1> const vectorin{
                x, px, y, py, t, pt
            };
 
            // TODO amrex::SmallMatrix::operator* overloads on T for SIMD*scalar / scalar*SIMD
            amrex::SmallVector<T_Real, 6, 1> const vectorout = m_transport_map * vectorin;
 
            // assign updated values
            x = vectorout(1);
            px = vectorout(2);
            y = vectorout(3);
            py = vectorout(4);
            t = vectorout(5);
            pt = vectorout(6);
        }

        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void operator() (RefPart & AMREX_RESTRICT refpart) const
        {
            if (m_ds != 0)  // Drift forward or backward
            {
                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();
 
                // assign intermediate parameter
                amrex::ParticleReal const step = slice_ds / std::sqrt(powi<2>(pt)-1.0_prt);
 
                // advance position and momentum (drift)
                refpart.x = x + step*px;
                refpart.y = y + step*py;
                refpart.z = z + step*pz;
                refpart.t = t - step*pt;
 
                // advance integrated path length
                refpart.s = s + slice_ds;
            }
            // else nothing to do for a zero-length element
        }

        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        int nslice () const
        {
            return 1;
        }

        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        amrex::ParticleReal ds () const
        {
            return m_ds;
        }

        template<typename T_Real=amrex::ParticleReal, typename T_IdCpu=uint64_t>
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void spin_and_phasespace_push (
            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,
            [[maybe_unused]] T_Real const & AMREX_RESTRICT sx,
            [[maybe_unused]] T_Real const & AMREX_RESTRICT sy,
            [[maybe_unused]] T_Real const & AMREX_RESTRICT sz,
            T_IdCpu const & AMREX_RESTRICT idcpu,
            RefPart const & AMREX_RESTRICT refpart
        ) const
        {
             // spin is unaffected
 
            // phase space push
            (*this)(x, y, t, px, py, pt, idcpu, refpart);
 
        }

        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 = m_transport_map;
 
            // apply the transverse rotation (roll) alignment error
            return rotate_aligned_map(R);
        }
 
        Map6x6 m_transport_map;  // 6x6 transport map
        amrex::ParticleReal m_ds;  // finite ds allowed for bookkeeping, but we do not allow slicing
    };
 
} // namespace impactx
 
IMPACTX_PUSH_EXTERN_TEMPLATE(impactx::elements::LinearMap)
 
#endif // IMPACTX_ELEMENT_LINEAR_MAP_H