Aperture.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_APERTURE_H
#define IMPACTX_APERTURE_H
 
#include "particles/ImpactXParticleContainer.H"
#include "mixin/alignment.H"
#include "mixin/beamoptic.H"
#include "mixin/thin.H"
#include "mixin/lineartransport.H"
#include "mixin/spintransport.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 Aperture
    : public mixin::Named,
      public mixin::BeamOptic<Aperture>,
      public mixin::LinearTransport<Aperture>,
      public mixin::Thin,
      public mixin::Alignment,
      public mixin::SpinTransport,
      public mixin::NoFinalize,
      public amrex::simd::Vectorized<amrex::simd::native_simd_size_particlereal>
    {
        static constexpr auto type = "Aperture";
        using PType = ImpactXParticleContainer::ParticleType;
 
        AMREX_ENUM_IN_CLASS(Shape,
            rectangular,  
            elliptical,   
        );
 
        AMREX_ENUM_IN_CLASS(Action,
            transmit,  
            absorb,    
        );
 
        static constexpr amrex::ParticleReal DEFAULT_repeat_x = 0;
        static constexpr amrex::ParticleReal DEFAULT_repeat_y = 0;
        static constexpr bool DEFAULT_shift_odd_x = false;
        static constexpr Shape DEFAULT_shape = Shape::rectangular;
        static constexpr Action DEFAULT_action = Action::transmit;

        Aperture (
            amrex::ParticleReal aperture_x,
            amrex::ParticleReal aperture_y,
            amrex::ParticleReal repeat_x = DEFAULT_repeat_x,
            amrex::ParticleReal repeat_y = DEFAULT_repeat_y,
            bool shift_odd_x = DEFAULT_shift_odd_x,
            Shape shape = DEFAULT_shape,
            Action action = DEFAULT_action,
            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_shape(shape), m_action(action), m_repeat_x(repeat_x), m_repeat_y(repeat_y), m_shift_odd_x(shift_odd_x)
        {
            using namespace amrex::literals; // for _rt and _prt
 
            if (aperture_x > 0_prt) {
                m_inv_aperture_x = 1_prt / aperture_x;
            } else {
                throw std::runtime_error("Aperture: aperture_x must be > 0!");
            }
 
            if (aperture_y > 0_prt) {
                m_inv_aperture_y = 1_prt / aperture_y;
            } else {
                throw std::runtime_error("Aperture: aperture_y must be > 0!");
            }
 
            if (m_repeat_y == 0_prt && m_shift_odd_x) {
                throw std::runtime_error("Aperture: shift_odd_x can only be used if repeat_y is set, too!");
            }
        }

        void reverse () { /* no-op */ }

        using BeamOptic::operator();

        void compute_constants ([[maybe_unused]] RefPart const & refpart)
        {
            Alignment::compute_constants(refpart);
        }

        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,
            [[maybe_unused]] T_Real const & AMREX_RESTRICT px,
            [[maybe_unused]] T_Real const & AMREX_RESTRICT py,
            [[maybe_unused]] T_Real const & AMREX_RESTRICT pt,
            T_IdCpu & AMREX_RESTRICT idcpu,
            [[maybe_unused]] RefPart const & AMREX_RESTRICT refpart
        ) const
        {
            using namespace amrex::literals; // for _rt and _prt
            namespace stdx = amrex::simd::stdx;
            using amrex::Math::powi;
            using std::fmod;
            using std::abs;
            using VBool = decltype(x>0_prt);
 
            // define intermediate variables
            amrex::ParticleReal const dx = m_repeat_x * 0.5_prt;
            amrex::ParticleReal const dy = m_repeat_y * 0.5_prt;
 
            // shift every 2nd row by half of m_repeat_x
            T_Real sx = x;
            T_Real const sy = y;
            VBool const shift_x = m_shift_odd_x == false ? VBool{false} : fmod(floor((y+dy)/m_repeat_y), 2) != T_Real{0};
#ifdef AMREX_USE_SIMD
            amrex::simd::stdx::where(shift_x, sx) += dx;
#else
            sx += shift_x ? dx : 0_prt;
#endif
 
            // if the aperture is periodic, shift sx,sy coordinates to the fundamental domain
            T_Real u = (m_repeat_x == 0_prt) ? sx : (fmod(abs(sx)+dx, m_repeat_x)-dx);
            T_Real v = (m_repeat_y == 0_prt) ? sy : (fmod(abs(sy)+dy, m_repeat_y)-dy);
 
            // scale horizontal and vertical coordinates
            // performance optimization: we intentionally store the inverse of
            //   the aperture, so we can do a quick multiplication (instead of a
            //   costly division) here.
            u = u * m_inv_aperture_x;
            v = v * m_inv_aperture_y;
 
            // compare against the aperture boundary
            VBool inside_aperture;
 
            switch (m_shape)
            {
                case Shape::rectangular:  // default
                    inside_aperture = (powi<2>(u) <= 1_prt && powi<2>(v) <= 1_prt);
                    break;
 
                case Shape::elliptical:
                    inside_aperture = (powi<2>(u) + powi<2>(v) <= 1_prt);
                    break;
 
                default:
                    inside_aperture = VBool{true};
                    break;
            }
 
            // For transmit (default): invalidate if outside aperture
            // For absorb: invalidate if inside aperture
            auto const invalid_mask = m_action == Action::transmit ? !inside_aperture : inside_aperture;
            amrex::ParticleIDWrapper<T_IdCpu>{idcpu}.make_invalid(invalid_mask);
        }

        using Thin::operator();

        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 const & AMREX_RESTRICT x,
            T_Real const & AMREX_RESTRICT y,
            T_Real const & AMREX_RESTRICT t,
            T_Real const & AMREX_RESTRICT px,
            T_Real const & AMREX_RESTRICT py,
            T_Real const & 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 & 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
        {
            // apply the transverse rotation (roll) alignment error
            return rotate_aligned_map(Map6x6::Identity());
        }
 
        Shape m_shape; 
        Action m_action; 
        // performance optimization: we intentionally store the inverse of
        //   the aperture, so we can do a quick multiplication instead of a costly
        //   division in the comparison operation per particle
        amrex::ParticleReal m_inv_aperture_x; 
        amrex::ParticleReal m_inv_aperture_y; 
        amrex::ParticleReal m_repeat_x; 
        amrex::ParticleReal m_repeat_y; 
        bool m_shift_odd_x; 
 
    };
 
} // namespace impactx
 
IMPACTX_PUSH_EXTERN_TEMPLATE(impactx::elements::Aperture)
 
#endif // IMPACTX_APERTURE_H