Add first draft of Aperture element.

docs/source/usage/parameters.rst

@@ -466,6 +466,17 @@ Lattice Elements
 
             * ``<element_name>.phi_out`` (``float``, in degrees) angle of the reference particle with respect to the longitudinal (z) axis in the rotated frame
 
+        * ``aperture`` for a thin collimator element applying a transverse aperture boundary.
+          This requires these additional parameters:
+
+            * ``<element_name>.shape`` (``integer``, dimensionless) shape of the aperture boundary
+
+                (m = 0) rectangular, (m = 1) elliptical
+
+            * ``<element_name>.xmax`` (``float``, in meters) maximum value of the horizontal coordinate
+
+            * ``<element_name>.ymax`` (``float``, in meters) maximum value of the vertical coordinate
+
         * ``beam_monitor`` a beam monitor, writing all beam particles at fixed ``s`` to openPMD files.
           If the same element name is used multiple times, then an output series is created with multiple outputs.
 

docs/source/usage/python.rst

@@ -651,6 +651,14 @@ References:
    :param phi_in: angle of the reference particle with respect to the longitudinal (z) axis in the original frame in degrees
    :param phi_out: angle of the reference particle with respect to the longitudinal (z) axis in the rotated frame in degrees
 
+.. py:class:: impactx.elements.Aperture(shape=0, xmax, ymax)
+
+   A thin collimator element, applying a transverse aperture boundary.
+
+   :param shape: aperture boundary shape (0 rectangular, 1 elliptical)
+   :param xmax: maximum value of the horizontal coordinate (meter)
+   :param ymax: maximum value of the vertical coordinate (meter)
+
 .. py:class:: impactx.elements.SoftQuadrupole(ds, gscale, cos_coefficients, sin_coefficients, nslice=1)
 
    A soft-edge quadrupole.

src/initialization/InitElement.cpp

@@ -224,7 +224,14 @@ namespace detail
             pp_element.get("bz", bz);
             pp_element.queryAdd("nslice", nslice);
             m_lattice.emplace_back( ChrAcc(ds, ez, bz, nslice) );
-        } else if (element_type == "beam_monitor") {
+        } else if (element_type == "aperture") {
+            amrex::Real xmax, ymax;
+            int shape = 0;
+            pp_element.get("xmax", xmax);
+	    pp_element.get("ymax", ymax);
+            pp_element.queryAdd("shape", shape);
+            m_lattice.emplace_back( Aperture(shape, xmax, ymax) );
+    } else if (element_type == "beam_monitor") {
             std::string openpmd_name = element_name;
             pp_element.queryAdd("name", openpmd_name);
             std::string openpmd_backend = "default";

src/particles/elements/All.H

@@ -10,6 +10,7 @@
 #ifndef IMPACTX_ELEMENTS_ALL_H
 #define IMPACTX_ELEMENTS_ALL_H
 
+#include "Aperture.H"
 #include "CFbend.H"
 #include "ChrDrift.H"
 #include "ChrQuad.H"
@@ -40,7 +41,8 @@ namespace impactx
 {
     using KnownElements = std::variant<
         None, /* must be first, so KnownElements creates a default constructor */
-        CFbend,
+	Aperture,
+   	CFbend,
         ChrAcc,
         ChrDrift,
         ChrQuad,
@@ -55,7 +57,7 @@ namespace impactx
         Programmable,
         PRot,
         Quad,
-        RFCavity,
+	RFCavity,
         Sbend,
         ShortRF,
         SoftSolenoid,

src/particles/elements/Aperture.H

@@ -0,0 +1,114 @@
+/* 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: Chad Mitchell, Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_APERTURE_H
+#define IMPACTX_APERTURE_H
+
+
+#include "particles/ImpactXParticleContainer.H"
+#include "mixin/beamoptic.H"
+#include "mixin/thin.H"
+#include "mixin/nofinalize.H"
+
+#include <ablastr/particles/IndexHandling.H>
+#include <AMReX_Extension.H>
+#include <AMReX_REAL.H>
+
+#include <cmath>
+
+namespace impactx
+{
+    struct Aperture
+    : public elements::BeamOptic<Aperture>,
+      public elements::Thin,
+      public elements::NoFinalize
+    {
+        static constexpr auto name = "Aperture";
+        using PType = ImpactXParticleContainer::ParticleType;
+
+        /** A thin collimator element that applies a transverse aperture boundary.
+         *  Particles outside the boundary are considered lost.
+         *
+         * @param type aperture type (0 rectangular, 1 elliptical)
+         * @param xmax maximum value of horizontal coordinate (m)
+         * @param ymax maximum value of vertical coordinate (m)
+         */
+        Aperture( int const shape,
+                   amrex::ParticleReal const xmax,
+                   amrex::ParticleReal const ymax )
+        : m_shape(shape), m_xmax(xmax), m_ymax(ymax)
+        {
+        }
+
+        /** Push all particles */
+        using BeamOptic::operator();
+
+        /** This is an aperture functor, so that a variable of this type can be used like an
+         *  aperture function.
+         *
+         * @param p Particle AoS data for positions and cpu/id
+         * @param px particle momentum in x
+         * @param py particle momentum in y
+         * @param pt particle momentum in t
+         * @param refpart reference particle (unused)
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+                PType& AMREX_RESTRICT p,
+		amrex::ParticleReal & AMREX_RESTRICT px,
+                amrex::ParticleReal & AMREX_RESTRICT py,
+                amrex::ParticleReal & AMREX_RESTRICT pt,
+                [[maybe_unused]] RefPart const & refpart) const {
+
+            using namespace amrex::literals; // for _rt and _prt
+
+            // access AoS data such as positions and cpu/id
+            amrex::ParticleReal const x = p.pos(RealAoS::x);
+            amrex::ParticleReal const y = p.pos(RealAoS::y);
+            uint64_t const global_id = ablastr::particles::localIDtoGlobal(p.id(), p.cpu());
+
+            // scale horizontal and vertical coordinates
+            amrex::ParticleReal const u = x/m_xmax;
+            amrex::ParticleReal const v = y/m_ymax;
+
+            // compare against the aperture boundary
+
+            switch(m_shape) {
+
+	       // rectangular aperture (default)
+               case(0) :                      
+                  if (pow(u,2)>1 || pow(v,2)>1){
+                     p.pos(RealAoS::x) = 0.0_prt; //replace with id change of sign
+		     p.pos(RealAoS::y) = 0.0_prt; //replace with id change of sign
+                  }
+
+	       // elliptical aperture
+               case(1) :                
+                  if (pow(u,2)+pow(v,2)>1){
+                     p.pos(RealAoS::x) = 0.0_prt; //replace with id change of sign
+                     p.pos(RealAoS::y) = 0.0_prt; //replace with id change of sign
+                  }
+
+            }
+
+        }
+
+        /** This pushes the reference particle. */
+        using Thin::operator();
+
+    private:
+        int m_shape; //! aperture type (1 rectangular, 2 elliptical)
+        amrex::ParticleReal m_xmax; //! maximum horizontal coordinate
+        amrex::ParticleReal m_ymax; //! maximum vertical coordinate
+
+    };
+
+} // namespace impactx
+
+#endif // IMPACTX_APERTURE_H

src/python/elements.cpp

@@ -358,5 +358,17 @@ void init_elements(py::module& m)
              "A soft-edge quadrupole."
         )
     ;
+
+    py::class_<Aperture, elements::Thin>(me, "Aperture")
+        .def(py::init<
+		int,
+                amrex::ParticleReal const,
+                amrex::ParticleReal const>(),
+             py::arg("shape") = 0, py::arg("xmax"), py::arg("ymax"),
+             "A short collimator element applying a transverse aperture boundary."
+        )    
+    ;        
+
 
+
 }