diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 3d391093c..4d4c5be6b 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -135,6 +135,8 @@ Lattice Elements * ``.dx`` (``float``, in meters) horizontal translation error * ``.dy`` (``float``, in meters) vertical translation error * ``.rotation`` (``float``, in degrees) rotation error in the transverse plane + * ``.aperture_x`` (``float``, in meters) horizontal half-aperture (elliptical) + * ``.aperture_y`` (``float``, in meters) vertical half-aperture (elliptical) * ``.nslice`` (``integer``) number of slices used for the application of space charge (default: ``1``) * ``drift_chromatic`` for a free drift, with chromatic effects included. diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 16df118cc..8bdcd1dbe 100644 --- a/docs/source/usage/python.rst +++ b/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, aperture_x=0, aperture_y=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 aperture_x: horizontal half-aperture (elliptical) in m + :param aperture_y: 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 diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 1783ad36d..7610f93cf 100644 --- a/examples/CMakeLists.txt +++ b/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 +) diff --git a/examples/aperture/analysis_aperture_thick.py b/examples/aperture/analysis_aperture_thick.py new file mode 100755 index 000000000..702cb6bc9 --- /dev/null +++ b/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) diff --git a/examples/aperture/input_aperture_thick.in b/examples/aperture/input_aperture_thick.in new file mode 100644 index 000000000..0921240b0 --- /dev/null +++ b/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.aperture_x = 1.0e-3 +drift.aperture_y = 1.5e-3 + + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true +diag.backend = h5 diff --git a/examples/aperture/run_aperture_thick.py b/examples/aperture/run_aperture_thick.py new file mode 100755 index 000000000..6184c7946 --- /dev/null +++ b/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, aperture_x=1.0e-3, aperture_y=1.5e-3), + monitor, + ] +) + +# run simulation +sim.track_particles() + +# clean shutdown +sim.finalize() diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index d106a4b86..b049c9a82 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -93,6 +93,28 @@ namespace detail return values; } + + /** Read the Aperture parameters aperture_x and aperture_y from inputs + * + * @param pp_element the element being read + * @return key-value pairs for aperture_x and aperture_y + */ + std::map + query_aperture (amrex::ParmParse& pp_element) + { + amrex::ParticleReal aperture_x = 0; + amrex::ParticleReal aperture_y = 0; + pp_element.query("aperture_x", aperture_x); + pp_element.query("aperture_y", aperture_y); + + std::map values = { + {"aperture_x", aperture_x}, + {"aperture_y", aperture_y} + }; + + 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["aperture_x"], b["aperture_y"], nslice, element_name) ); } else if (element_type == "sbend") { auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default); diff --git a/src/particles/elements/Drift.H b/src/particles/elements/Drift.H index ef1e889af..129f3412f 100644 --- a/src/particles/elements/Drift.H +++ b/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, 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 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 */ @@ -50,12 +54,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 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) { } @@ -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); } diff --git a/src/particles/elements/mixin/aperture.H b/src/particles/elements/mixin/aperture.H new file mode 100644 index 000000000..f77f6c2fe --- /dev/null +++ b/src/particles/elements/mixin/aperture.H @@ -0,0 +1,103 @@ +/* Copyright 2022-2025 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, Chad Mitchell + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_ELEMENTS_MIXIN_APERTURE_H +#define IMPACTX_ELEMENTS_MIXIN_APERTURE_H + +#include +#include +#include +#include +#include + + +namespace impactx::elements +{ + /** This is a helper class for applying a transverse aperture restriction to thick lattice elements + */ + struct Aperture + { + + /** A finite-length element + * + * @param aperture_x horizontal half-aperture size in m + * @param aperture_y vertical half-aperture size in m + */ + Aperture ( + amrex::ParticleReal aperture_x, + amrex::ParticleReal aperture_y + ) + : m_aperture_x(aperture_x), m_aperture_y(aperture_y) + { + } + + Aperture () = default; + Aperture (Aperture const &) = default; + Aperture& operator= (Aperture const &) = default; + Aperture (Aperture&&) = default; + Aperture& operator= (Aperture&& rhs) = default; + + ~Aperture () = default; + + /** Apply the transverse aperture + * + * @param[inout] x horizontal position relative to reference particle + * @param[inout] y vertical position relative to reference particle + * @param idcpu particle global index + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void apply_aperture ( + amrex::ParticleReal & AMREX_RESTRICT x, + amrex::ParticleReal & AMREX_RESTRICT y, + uint64_t & AMREX_RESTRICT idcpu + ) const { + using namespace amrex::literals; // for _rt and _prt + + // skip aperture application if aperture_x <= 0 or aperture_y <= 0 + if (m_aperture_x > 0 && m_aperture_y > 0) { + + // scale horizontal and vertical coordinates + amrex::ParticleReal const u = x / m_aperture_x; + amrex::ParticleReal const v = y / m_aperture_y; + + // compare against the aperture boundary + if (std::pow(u,2) + std::pow(v,2) > 1_prt) { + amrex::ParticleIDWrapper{idcpu}.make_invalid(); + } + + } + } + + /** Horizontal aperture size + * + * @return horizontal aperture size in m + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + amrex::ParticleReal aperture_x () const + { + return m_aperture_x; + } + + /** Vertical aperture size + * + * @return vertical aperture size in m + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + amrex::ParticleReal aperture_y () const + { + return m_aperture_y; + } + + amrex::ParticleReal m_aperture_x = 0; //! horizontal aperture size [m] + amrex::ParticleReal m_aperture_y = 0; //! vertical aperture size [m] + }; + +} // namespace impactx::elements + +#endif // IMPACTX_ELEMENTS_MIXIN_APERTURE_H diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 7d1488a55..58d1e8358 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -215,6 +215,21 @@ void init_elements(py::module& m) ) ; + py::class_(mx, "Aperture") + .def(py::init<>(), + "Mixin class for lattice elements with a transverse aperture." + ) + .def_property_readonly("aperture_x", + &elements::Aperture::aperture_x, + "horizontal aperture in m" + ) + .def_property_readonly("aperture_y", + &elements::Aperture::aperture_y, + "vertical aperture in m" + ) + ; + + // diagnostics py::class_ py_BeamMonitor(me, "BeamMonitor"); @@ -654,7 +669,7 @@ void init_elements(py::module& m) ; register_beamoptics_push(py_DipEdge); - py::class_ py_Drift(me, "Drift"); + py::class_ py_Drift(me, "Drift"); py_Drift .def("__repr__", [](Drift const & drift) { @@ -669,6 +684,8 @@ void init_elements(py::module& m) amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, int, std::optional >(), @@ -676,6 +693,8 @@ void init_elements(py::module& m) py::arg("dx") = 0, py::arg("dy") = 0, py::arg("rotation") = 0, + py::arg("aperture_x") = 0, + py::arg("aperture_y") = 0, py::arg("nslice") = 1, py::arg("name") = py::none(), "A drift."