Add aperture input to all thick elements.

ECP-WarpX · Jan 14, 2025 · 14f4a4c · 14f4a4c
1 parent 58c1a2c
commit 14f4a4c
Show file tree

Hide file tree

Showing 15 changed files with 251 additions and 56 deletions.
diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp
@@ -144,11 +144,12 @@ namespace detail
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
 
             amrex::ParticleReal k;
             pp_element.get("k", k);
 
-            m_lattice.emplace_back( Quad(ds, k, a["dx"], a["dy"], a["rotation_degree"], nslice, element_name) );
+            m_lattice.emplace_back( Quad(ds, k, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], nslice, element_name) );
         } else if (element_type == "drift")
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
@@ -160,11 +161,12 @@ namespace detail
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
 
             amrex::ParticleReal rc;
             pp_element.get("rc", rc);
 
-            m_lattice.emplace_back( Sbend(ds, rc, a["dx"], a["dy"], a["rotation_degree"], nslice, element_name) );
+            m_lattice.emplace_back( Sbend(ds, rc, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], nslice, element_name) );
         } else if (element_type == "cfbend")
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
@@ -191,13 +193,14 @@ namespace detail
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
 
             amrex::Real kx, ky, kt;
             pp_element.get("kx", kx);
             pp_element.get("ky", ky);
             pp_element.get("kt", kt);
 
-            m_lattice.emplace_back( ConstF(ds, kx, ky, kt, a["dx"], a["dy"], a["rotation_degree"], nslice, element_name) );
+            m_lattice.emplace_back( ConstF(ds, kx, ky, kt, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], nslice, element_name) );
         } else if (element_type == "buncher")
         {
             auto a = detail::query_alignment(pp_element);
@@ -242,6 +245,7 @@ namespace detail
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
 
             amrex::ParticleReal escale, freq, phase;
             int mapsteps = mapsteps_default;
@@ -255,16 +259,17 @@ namespace detail
             detail::queryAddResize(pp_element, "cos_coefficients", cos_coef);
             detail::queryAddResize(pp_element, "sin_coefficients", sin_coef);
 
-            m_lattice.emplace_back( RFCavity(ds, escale, freq, phase, cos_coef, sin_coef, a["dx"], a["dy"], a["rotation_degree"], mapsteps, nslice, element_name) );
+            m_lattice.emplace_back( RFCavity(ds, escale, freq, phase, cos_coef, sin_coef, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], mapsteps, nslice, element_name) );
         } else if (element_type == "solenoid")
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
 
             amrex::ParticleReal ks;
             pp_element.get("ks", ks);
 
-            m_lattice.emplace_back( Sol(ds, ks, a["dx"], a["dy"], a["rotation_degree"], nslice, element_name) );
+            m_lattice.emplace_back( Sol(ds, ks, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], nslice, element_name) );
         } else if (element_type == "prot")
         {
             amrex::ParticleReal phi_in, phi_out;
@@ -284,6 +289,7 @@ namespace detail
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
 
             amrex::ParticleReal bscale;
             int mapsteps = mapsteps_default;
@@ -297,11 +303,12 @@ namespace detail
             detail::queryAddResize(pp_element, "cos_coefficients", cos_coef);
             detail::queryAddResize(pp_element, "sin_coefficients", sin_coef);
 
-            m_lattice.emplace_back( SoftSolenoid(ds, bscale, cos_coef, sin_coef, units, a["dx"], a["dy"], a["rotation_degree"], mapsteps, nslice, element_name) );
+            m_lattice.emplace_back( SoftSolenoid(ds, bscale, cos_coef, sin_coef, units, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], mapsteps, nslice, element_name) );
         } else if (element_type == "quadrupole_softedge")
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
 
             amrex::ParticleReal gscale;
             int mapsteps = mapsteps_default;
@@ -313,35 +320,38 @@ namespace detail
             detail::queryAddResize(pp_element, "cos_coefficients", cos_coef);
             detail::queryAddResize(pp_element, "sin_coefficients", sin_coef);
 
-            m_lattice.emplace_back( SoftQuadrupole(ds, gscale, cos_coef, sin_coef, a["dx"], a["dy"], a["rotation_degree"], mapsteps, nslice, element_name) );
+            m_lattice.emplace_back( SoftQuadrupole(ds, gscale, cos_coef, sin_coef, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], mapsteps, nslice, element_name) );
         } else if (element_type == "drift_chromatic")
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
 
-            m_lattice.emplace_back( ChrDrift(ds, a["dx"], a["dy"], a["rotation_degree"], nslice, element_name) );
+            m_lattice.emplace_back( ChrDrift(ds, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], nslice, element_name) );
         } else if (element_type == "quad_chromatic")
         {
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
 
             amrex::ParticleReal k;
             int units = 0;
             pp_element.get("k", k);
             pp_element.queryAdd("units", units);
 
-            m_lattice.emplace_back( ChrQuad(ds, k, units, a["dx"], a["dy"], a["rotation_degree"], nslice, element_name) );
+            m_lattice.emplace_back( ChrQuad(ds, k, units, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], nslice, element_name) );
         } else if (element_type == "plasma_lens_chromatic")
         {
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
 
             amrex::ParticleReal k;
             int units = 0;
             pp_element.get("k", k);
             pp_element.queryAdd("units", units);
 
-            m_lattice.emplace_back( ChrPlasmaLens(ds, k, units, a["dx"], a["dy"], a["rotation_degree"], nslice, element_name) );
+            m_lattice.emplace_back( ChrPlasmaLens(ds, k, units, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], nslice, element_name) );
         } else if (element_type == "tapered_plasma_lens")
         {
             auto a = detail::query_alignment(pp_element);
@@ -358,29 +368,32 @@ namespace detail
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
 
-            m_lattice.emplace_back( ExactDrift(ds, a["dx"], a["dy"], a["rotation_degree"], nslice, element_name) );
+            m_lattice.emplace_back( ExactDrift(ds, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], nslice, element_name) );
         } else if (element_type == "sbend_exact")
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
 
             amrex::ParticleReal phi;
             amrex::ParticleReal B = 0.0;
             pp_element.get("phi", phi);
             pp_element.queryAdd("B", B);
 
-            m_lattice.emplace_back( ExactSbend(ds, phi, B, a["dx"], a["dy"], a["rotation_degree"], nslice, element_name) );
+            m_lattice.emplace_back( ExactSbend(ds, phi, B, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], nslice, element_name) );
         } else if (element_type == "uniform_acc_chromatic")
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);
             auto a = detail::query_alignment(pp_element);
+            auto b = detail::query_aperture(pp_element);
 
             amrex::ParticleReal ez, bz;
             pp_element.get("ez", ez);
             pp_element.get("bz", bz);
 
-            m_lattice.emplace_back( ChrAcc(ds, ez, bz, a["dx"], a["dy"], a["rotation_degree"], nslice, element_name) );
+            m_lattice.emplace_back( ChrAcc(ds, ez, bz, a["dx"], a["dy"], a["rotation_degree"], b["aperture_x"], b["aperture_y"], nslice, element_name) );
         } else if (element_type == "thin_dipole")
         {
             auto a = detail::query_alignment(pp_element);

diff --git a/src/particles/elements/ChrDrift.H b/src/particles/elements/ChrDrift.H
@@ -12,6 +12,7 @@
 
 #include "particles/ImpactXParticleContainer.H"
 #include "mixin/alignment.H"
+#include "mixin/aperture.H"
 #include "mixin/beamoptic.H"
 #include "mixin/thick.H"
 #include "mixin/named.H"
@@ -30,6 +31,7 @@ namespace impactx
       public elements::BeamOptic<ChrDrift>,
       public elements::Thick,
       public elements::Alignment,
+      public elements::Aperture,
       public elements::NoFinalize
     {
         static constexpr auto type = "ChrDrift";
@@ -44,6 +46,8 @@ namespace impactx
          * @param dx horizontal translation error in m
          * @param dy vertical translation error in m
          * @param rotation_degree rotation error in the transverse plane [degrees]
+         * @param aperture_x horizontal half-aperture in m
+         * @param aperture_y vertical half-aperture in m
          * @param nslice number of slices used for the application of space charge
          * @param name a user defined and not necessarily unique name of the element
          */
@@ -52,12 +56,15 @@ namespace impactx
             amrex::ParticleReal dx = 0,
             amrex::ParticleReal dy = 0,
             amrex::ParticleReal rotation_degree = 0,
+            amrex::ParticleReal aperture_x = 0,
+            amrex::ParticleReal aperture_y = 0,
             int nslice = 1,
             std::optional<std::string> name = std::nullopt
         )
         : Named(std::move(name)),
           Thick(ds, nslice),
-          Alignment(dx, dy, rotation_degree)
+          Alignment(dx, dy, rotation_degree),
+          Aperture(aperture_x, aperture_y)
         {
         }
 
@@ -72,7 +79,7 @@ namespace impactx
          * @param px particle momentum in x
          * @param py particle momentum in y
          * @param pt particle momentum in t
-         * @param idcpu particle global index (unused)
+         * @param idcpu particle global index
          * @param refpart reference particle
          */
         AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
@@ -83,7 +90,7 @@ namespace impactx
             amrex::ParticleReal & AMREX_RESTRICT px,
             amrex::ParticleReal & AMREX_RESTRICT py,
             amrex::ParticleReal & AMREX_RESTRICT pt,
-            [[maybe_unused]] uint64_t & AMREX_RESTRICT idcpu,
+            uint64_t & AMREX_RESTRICT idcpu,
             RefPart const & refpart
         ) const
         {
@@ -133,6 +140,9 @@ namespace impactx
             py = pyout;
             pt = ptout;
 
+            // apply transverse aperture
+            apply_aperture(x, y, idcpu);
+
             // undo shift due to alignment errors of the element
             shift_out(x, y, px, py);
         }

diff --git a/src/particles/elements/ChrPlasmaLens.H b/src/particles/elements/ChrPlasmaLens.H
@@ -12,6 +12,7 @@
 
 #include "particles/ImpactXParticleContainer.H"
 #include "mixin/alignment.H"
+#include "mixin/aperture.H"
 #include "mixin/beamoptic.H"
 #include "mixin/thick.H"
 #include "mixin/named.H"
@@ -30,6 +31,7 @@ namespace impactx
       public elements::BeamOptic<ChrPlasmaLens>,
       public elements::Thick,
       public elements::Alignment,
+      public elements::Aperture,
       public elements::NoFinalize
     {
         static constexpr auto type = "ChrPlasmaLens";
@@ -50,6 +52,8 @@ namespace impactx
          * @param dx horizontal translation error in m
          * @param dy vertical translation error in m
          * @param rotation_degree rotation error in the transverse plane [degrees]
+         * @param aperture_x horizontal half-aperture in m
+         * @param aperture_y vertical half-aperture in m
          * @param nslice number of slices used for the application of space charge
          * @param name a user defined and not necessarily unique name of the element
          */
@@ -60,12 +64,15 @@ namespace impactx
             amrex::ParticleReal dx = 0,
             amrex::ParticleReal dy = 0,
             amrex::ParticleReal rotation_degree = 0,
+            amrex::ParticleReal aperture_x = 0,
+            amrex::ParticleReal aperture_y = 0,
             int nslice = 1,
             std::optional<std::string> name = std::nullopt
         )
         : Named(std::move(name)),
           Thick(ds, nslice),
           Alignment(dx, dy, rotation_degree),
+          Aperture(aperture_x, aperture_y),
           m_k(k), m_unit(unit)
         {
         }
@@ -81,7 +88,7 @@ namespace impactx
          * @param px particle momentum in x
          * @param py particle momentum in y
          * @param pt particle momentum in t
-         * @param idcpu particle global index (unused)
+         * @param idcpu particle global index
          * @param refpart reference particle
          */
         AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
@@ -92,7 +99,7 @@ namespace impactx
             amrex::ParticleReal & AMREX_RESTRICT px,
             amrex::ParticleReal & AMREX_RESTRICT py,
             amrex::ParticleReal & AMREX_RESTRICT pt,
-            [[maybe_unused]] uint64_t & AMREX_RESTRICT idcpu,
+            uint64_t & AMREX_RESTRICT idcpu,
             RefPart const & refpart
         ) const
         {
@@ -171,6 +178,9 @@ namespace impactx
             py = pyout;
             pt = ptout;
 
+            // apply transverse aperture
+            apply_aperture(x, y, idcpu);
+
             // undo shift due to alignment errors of the element
             shift_out(x, y, px, py);
         }

diff --git a/src/particles/elements/ChrQuad.H b/src/particles/elements/ChrQuad.H
@@ -12,6 +12,7 @@
 
 #include "particles/ImpactXParticleContainer.H"
 #include "mixin/alignment.H"
+#include "mixin/aperture.H"
 #include "mixin/beamoptic.H"
 #include "mixin/thick.H"
 #include "mixin/named.H"
@@ -31,6 +32,7 @@ namespace impactx
       public elements::BeamOptic<ChrQuad>,
       public elements::Thick,
       public elements::Alignment,
+      public elements::Aperture,
       public elements::NoFinalize
     {
         static constexpr auto type = "ChrQuad";
@@ -53,6 +55,8 @@ namespace impactx
          * @param dx horizontal translation error in m
          * @param dy vertical translation error in m
          * @param rotation_degree rotation error in the transverse plane [degrees]
+         * @param aperture_x horizontal half-aperture in m
+         * @param aperture_y vertical half-aperture in m
          * @param nslice number of slices used for the application of space charge
          * @param name a user defined and not necessarily unique name of the element
          */
@@ -63,12 +67,15 @@ namespace impactx
             amrex::ParticleReal dx = 0,
             amrex::ParticleReal dy = 0,
             amrex::ParticleReal rotation_degree = 0,
+            amrex::ParticleReal aperture_x = 0,
+            amrex::ParticleReal aperture_y = 0,
             int nslice = 1,
             std::optional<std::string> name = std::nullopt
         )
         : Named(std::move(name)),
           Thick(ds, nslice),
           Alignment(dx, dy, rotation_degree),
+          Aperture(aperture_x, aperture_y),
           m_k(k), m_unit(unit)
         {
         }
@@ -84,7 +91,7 @@ namespace impactx
          * @param px particle momentum in x
          * @param py particle momentum in y
          * @param pt particle momentum in t
-         * @param idcpu particle global index (unused)
+         * @param idcpu particle global index
          * @param refpart reference particle
          */
         AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
@@ -95,7 +102,7 @@ namespace impactx
             amrex::ParticleReal & AMREX_RESTRICT px,
             amrex::ParticleReal & AMREX_RESTRICT py,
             amrex::ParticleReal & AMREX_RESTRICT pt,
-            [[maybe_unused]] uint64_t & AMREX_RESTRICT idcpu,
+            uint64_t & AMREX_RESTRICT idcpu,
             RefPart const & refpart
         ) const
         {
@@ -210,6 +217,9 @@ namespace impactx
             py = pyout;
             pt = ptout;
 
+            // apply transverse aperture
+            apply_aperture(x, y, idcpu);
+
             // undo shift due to alignment errors of the element
             shift_out(x, y, px, py);
         }