alignment.H

Go to the documentation of this file.

/* Copyright 2022-2023 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: Axel Huebl
 * License: BSD-3-Clause-LBNL
 */
#ifndef IMPACTX_ELEMENTS_MIXIN_ALIGNMENT_H
#define IMPACTX_ELEMENTS_MIXIN_ALIGNMENT_H
 
#include "particles/CovarianceMatrix.H"
#include "particles/ImpactXParticleContainer.H"
 
#include <ablastr/constant.H>
 
#include <AMReX_Math.H>
#include <AMReX_Extension.H>
#include <AMReX_REAL.H>
#include <AMReX_SmallMatrix.H>
 
 
namespace impactx::elements::mixin
{
    struct Alignment
    {
        static constexpr amrex::ParticleReal degree2rad =
            ablastr::constant::math::pi_v<amrex::ParticleReal> / amrex::ParticleReal{180};

        Alignment (
            amrex::ParticleReal dx,
            amrex::ParticleReal dy,
            amrex::ParticleReal rotation_degree
        )
        : m_dx(dx), m_dy(dy), m_rotation(rotation_degree * degree2rad)
        {
        }
 
        Alignment () = default;
        Alignment (Alignment const &) = default;
        Alignment& operator= (Alignment const &) = default;
        Alignment (Alignment&&) = default;
        Alignment& operator= (Alignment&& rhs) = default;
 
        ~Alignment () = default;

        void compute_constants ([[maybe_unused]] RefPart const & refpart)
        {
            auto const [sin_rotation, cos_rotation] = amrex::Math::sincos(m_rotation);
            m_sin_rotation = sin_rotation;
            m_cos_rotation = cos_rotation;
        }

        template<typename T_Real>
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void shift_in (
            T_Real & AMREX_RESTRICT x,
            T_Real & AMREX_RESTRICT y,
            T_Real & AMREX_RESTRICT px,
            T_Real & AMREX_RESTRICT py
        ) const
        {
            // position
            T_Real const xc = x - m_dx;
            T_Real const yc = y - m_dy;
            x =  xc * m_cos_rotation + yc * m_sin_rotation;
            y = -xc * m_sin_rotation + yc * m_cos_rotation;
 
            // momentum
            T_Real const pxc = px;
            T_Real const pyc = py;
            px =  pxc * m_cos_rotation + pyc * m_sin_rotation;
            py = -pxc * m_sin_rotation + pyc * m_cos_rotation;
        }

        template<typename T_Real>
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void rotate_in (
            T_Real & AMREX_RESTRICT sx,
            T_Real & AMREX_RESTRICT sy
        ) const
        {
            T_Real const sxc = sx;
            T_Real const syc = sy;
            sx =  sxc * m_cos_rotation + syc * m_sin_rotation;
            sy = -sxc * m_sin_rotation + syc * m_cos_rotation;
        }

        template<typename T_Real>
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void shift_out (
            T_Real & AMREX_RESTRICT x,
            T_Real & AMREX_RESTRICT y,
            T_Real & AMREX_RESTRICT px,
            T_Real & AMREX_RESTRICT py
        ) const
        {
            // position
            T_Real const xc = x;
            T_Real const yc = y;
            x = xc * m_cos_rotation - yc * m_sin_rotation;
            y = xc * m_sin_rotation + yc * m_cos_rotation;
            x += m_dx;
            y += m_dy;
 
            // momentum
            T_Real const pxc = px;
            T_Real const pyc = py;
            px = pxc * m_cos_rotation - pyc * m_sin_rotation;
            py = pxc * m_sin_rotation + pyc * m_cos_rotation;
        }

        template<typename T_Real>
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void rotate_out (
            T_Real & AMREX_RESTRICT sx,
            T_Real & AMREX_RESTRICT sy
        ) const
        {
            T_Real const sxc = sx;
            T_Real const syc = sy;
            sx = sxc * m_cos_rotation - syc * m_sin_rotation;
            sy = sxc * m_sin_rotation + syc * m_cos_rotation;
        }

        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        amrex::ParticleReal dx () const
        {
            return m_dx;
        }

        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        amrex::ParticleReal dy () const
        {
            return m_dy;
        }

        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        amrex::ParticleReal rotation () const
        {
            return m_rotation / degree2rad;
        }

        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        bool misaligned () const
        {
            return (m_dx != amrex::ParticleReal(0)) ||
                   (m_dy != amrex::ParticleReal(0)) ||
                   (m_rotation != amrex::ParticleReal(0));
        }

        AMREX_GPU_HOST AMREX_FORCE_INLINE
        Map6x6
        rotation_map () const
        {
            auto const [sin_rotation, cos_rotation] = amrex::Math::sincos(m_rotation);
 
            Map6x6 A = Map6x6::Identity();
            A(1, 1) =  cos_rotation;  A(1, 3) = -sin_rotation;
            A(3, 1) =  sin_rotation;  A(3, 3) =  cos_rotation;
            A(2, 2) =  cos_rotation;  A(2, 4) = -sin_rotation;
            A(4, 2) =  sin_rotation;  A(4, 4) =  cos_rotation;
            return A;
        }

        AMREX_GPU_HOST AMREX_FORCE_INLINE
        Map6x6
        rotate_aligned_map (Map6x6 const & R) const
        {
            if (m_rotation == amrex::ParticleReal(0)) { return R; }
            Map6x6 const A = rotation_map();
            return A * R * A.transpose();
        }
 
        amrex::ParticleReal m_dx = 0; 
        amrex::ParticleReal m_dy = 0; 
        amrex::ParticleReal m_rotation = 0; 
 
    private:
        // constants that are independent of the individually tracked particle,
        // see: compute_constants() to refresh
        amrex::ParticleReal m_sin_rotation;  
        amrex::ParticleReal m_cos_rotation;  
    };
 
} // namespace impactx::elements::mixin
 
#endif // IMPACTX_ELEMENTS_MIXIN_ALIGNMENT_H