Apply transverse aperture to thick elements.

docs/source/usage/parameters.rst

@@ -135,6 +135,8 @@ Lattice Elements
             * ``<element_name>.dx`` (``float``, in meters) horizontal translation error
             * ``<element_name>.dy`` (``float``, in meters) vertical translation error
             * ``<element_name>.rotation`` (``float``, in degrees) rotation error in the transverse plane
+            * ``<element_name>.x_aperture`` (``float``, in meters) horizontal half-aperture (elliptical)
+            * ``<element_name>.y_aperture`` (``float``, in meters) vertical half-aperture (elliptical)
             * ``<element_name>.nslice`` (``integer``) number of slices used for the application of space charge (default: ``1``)
 
         * ``drift_chromatic`` for a free drift, with chromatic effects included.

docs/source/usage/python.rst

@@ -570,11 +570,16 @@ This module provides elements for the accelerator lattice.
    :param rotation: rotation error in the transverse plane [degrees]
    :param name: an optional name for the element
 
-.. py:class:: impactx.elements.Drift(ds, dx=0, dy=0, rotation=0, nslice=1, name=None)
+.. py:class:: impactx.elements.Drift(ds, dx=0, dy=0, rotation=0, x_aperture=0, y_aperture=0, nslice=1, name=None)
 
    A drift.
 
    :param ds: Segment length in m
+   :param dx: horizontal translation error in m
+   :param dy: vertical translation error in m
+   :param rotation: rotation error in the transverse plane [degrees]
+   :param x_aperture: horizontal half-aperture (elliptical) in m
+   :param y_aperture: vertical half-aperture (elliptical) in m
    :param nslice: number of slices used for the application of space charge
    :param name: an optional name for the element
 

examples/CMakeLists.txt

@@ -1056,3 +1056,19 @@ add_impactx_test(linac-segment.py
     examples/linac_segment/analysis_linac_segment.py
     OFF  # no plot script yet
 )
+
+# Drift with transverse aperture ################################################
+#
+# w/o space charge
+add_impactx_test(aperture-thick
+    examples/aperture/input_aperture_thick.in
+    ON   # ImpactX MPI-parallel
+    examples/aperture/analysis_aperture_thick.py
+    OFF  # no plot script yet
+)
+add_impactx_test(aperture-thick.py
+    examples/aperture/run_aperture_thick.py
+    OFF  # ImpactX MPI-parallel
+    examples/aperture/analysis_aperture_thick.py
+    OFF  # no plot script yet
+)

examples/aperture/analysis_aperture_thick.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+
+import numpy as np
+import openpmd_api as io
+from scipy.stats import moment
+
+
+def get_moments(beam):
+    """Calculate standard deviations of beam position & momenta
+    and emittance values
+
+    Returns
+    -------
+    sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
+    """
+    sigx = moment(beam["position_x"], moment=2) ** 0.5  # variance -> std dev.
+    sigpx = moment(beam["momentum_x"], moment=2) ** 0.5
+    sigy = moment(beam["position_y"], moment=2) ** 0.5
+    sigpy = moment(beam["momentum_y"], moment=2) ** 0.5
+    sigt = moment(beam["position_t"], moment=2) ** 0.5
+    sigpt = moment(beam["momentum_t"], moment=2) ** 0.5
+
+    epstrms = beam.cov(ddof=0)
+    emittance_x = (sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5
+    emittance_y = (sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5
+    emittance_t = (sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2) ** 0.5
+
+    return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t)
+
+
+# initial/final beam
+series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
+last_step = list(series.iterations)[-1]
+initial = series.iterations[1].particles["beam"].to_df()
+final = series.iterations[last_step].particles["beam"].to_df()
+
+series_lost = io.Series("diags/openPMD/particles_lost.h5", io.Access.read_only)
+particles_lost = series_lost.iterations[0].particles["beam"].to_df()
+
+# compare number of particles
+num_particles = 10000
+assert num_particles == len(initial)
+# we lost particles in apertures
+assert num_particles > len(final)
+assert num_particles == len(particles_lost) + len(final)
+
+print("Initial Beam:")
+sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial)
+print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
+print(
+    f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
+)
+
+atol = 0.0  # ignored
+rtol = 1.8 * num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
+    [
+        1.559531175539e-3,
+        2.205510139392e-3,
+        1.0e-3,
+        1.0e-6,
+        2.0e-6,
+        1.0e-6,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+# particle-wise comparison against the rectangular aperture boundary
+xmax = 1.0e-3
+ymax = 1.5e-3
+
+# kept particles
+dx = abs(final["position_x"]) - xmax
+dy = abs(final["position_y"]) - ymax
+
+print()
+print(f"  x_max={final['position_x'].max()}")
+print(f"  x_min={final['position_x'].min()}")
+assert np.less_equal(dx.max(), 0.0)
+
+print(f"  y_max={final['position_y'].max()}")
+print(f"  y_min={final['position_y'].min()}")
+assert np.less_equal(dy.max(), 0.0)
+
+# lost particles
+dx = abs(particles_lost["position_x"]) - xmax
+dy = abs(particles_lost["position_y"]) - ymax
+
+print()
+print(f"  x_max={particles_lost['position_x'].max()}")
+print(f"  x_min={particles_lost['position_x'].min()}")
+assert np.greater_equal(dx.max(), 0.0)
+
+print(f"  y_max={particles_lost['position_y'].max()}")
+print(f"  y_min={particles_lost['position_y'].min()}")
+assert np.greater_equal(dy.max(), 0.0)
+
+# check that s is set correctly
+lost_at_s = particles_lost["s_lost"]
+drift_s = np.ones_like(lost_at_s) * 0.123
+assert np.allclose(lost_at_s, drift_s)

examples/aperture/input_aperture_thick.in

@@ -0,0 +1,47 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.kin_energy = 250.0
+beam.charge = 1.0e-9
+beam.particle = proton
+beam.distribution = waterbag
+beam.lambdaX = 1.559531175539e-3
+beam.lambdaY = 2.205510139392e-3
+beam.lambdaT = 1.0e-3
+beam.lambdaPx = 6.41218345413e-4
+beam.lambdaPy = 9.06819680526e-4
+beam.lambdaPt = 1.0e-3
+beam.muxpx = 0.0
+beam.muypy = 0.0
+beam.mutpt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor drift monitor
+lattice.nslice = 1
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+drift.type = drift
+drift.ds = 0.123
+drift.x_aperture = 1.0e-3
+drift.y_aperture = 1.5e-3
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
+
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true
+diag.backend = h5

examples/aperture/run_aperture_thick.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+from impactx import ImpactX, distribution, elements
+
+sim = ImpactX()
+
+# set numerical parameters and IO control
+sim.particle_shape = 2  # B-spline order
+sim.space_charge = False
+# sim.diagnostics = False  # benchmarking
+sim.slice_step_diagnostics = True
+sim.particle_lost_diagnostics_backend = "h5"
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a 250 MeV proton beam with an initial
+# horizontal rms emittance of 1 um and an
+# initial vertical rms emittance of 2 um
+kin_energy_MeV = 250.0  # reference energy
+bunch_charge_C = 1.0e-9  # used with space charge
+npart = 10000  # number of macro particles
+
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
+
+#   particle bunch
+distr = distribution.Waterbag(
+    lambdaX=1.559531175539e-3,
+    lambdaY=2.205510139392e-3,
+    lambdaT=1.0e-3,
+    lambdaPx=6.41218345413e-4,
+    lambdaPy=9.06819680526e-4,
+    lambdaPt=1.0e-3,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+
+# design the accelerator lattice
+sim.lattice.extend(
+    [
+        monitor,
+        elements.Drift(name="drift", ds=0.123, x_aperture=1.0e-3, y_aperture=1.5e-3),
+        monitor,
+    ]
+)
+
+# run simulation
+sim.track_particles()
+
+# clean shutdown
+sim.finalize()

src/initialization/InitElement.cpp

@@ -93,6 +93,28 @@ namespace detail
 
         return values;
     }
+
+    /** Read the Aperture parameters x_aperture and y_aperture from inputs
+     *
+     * @param pp_element the element being read
+     * @return key-value pairs for x_aperture and y_aperture
+     */
+    std::map<std::string, amrex::ParticleReal>
+    query_aperture (amrex::ParmParse& pp_element)
+    {
+        amrex::ParticleReal x_aperture = 0;
+        amrex::ParticleReal y_aperture = 0;
+        pp_element.query("x_aperture", x_aperture);
+        pp_element.query("y_aperture", y_aperture);
+
+        std::map<std::string, amrex::ParticleReal> values = {
+                {"x_aperture", x_aperture},
+                {"y_aperture", y_aperture}
+        };
+
+        return values;
+    }
+
 } // namespace detail
 
     /** Read a lattice element
@@ -128,8 +150,9 @@ 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( Drift(ds, a["dx"], a["dy"], a["rotation_degree"], nslice, element_name) );
+            m_lattice.emplace_back( Drift(ds, a["dx"], a["dy"], a["rotation_degree"], b["x_aperture"], b["y_aperture"], nslice, element_name) );
         } else if (element_type == "sbend")
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);

src/particles/elements/Drift.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<Drift>,
       public elements::Thick,
       public elements::Alignment,
+      public elements::Aperture,
       public elements::NoFinalize
     {
         static constexpr auto type = "Drift";
@@ -42,6 +44,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 x_aperture horizontal half-aperture in m
+         * @param y_aperture 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
          */
@@ -50,12 +54,15 @@ namespace impactx
             amrex::ParticleReal dx = 0,
             amrex::ParticleReal dy = 0,
             amrex::ParticleReal rotation_degree = 0,
+            amrex::ParticleReal x_aperture = 0,
+            amrex::ParticleReal y_aperture = 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(x_aperture, y_aperture)
         {
         }
 
@@ -70,7 +77,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
@@ -81,7 +88,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
         {
@@ -121,6 +128,9 @@ namespace impactx
             py = pyout;
             pt = ptout;
 
+            // apply transverse aperture
+            apply_aperture(x, y, idcpu);  //value of idcpu?
+
             // undo shift due to alignment errors of the element
             shift_out(x, y, px, py);
         }