From 54092faa878e5c10a9bfc55710ebdcb55b95e6ed Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Wed, 12 Jul 2023 12:44:32 -0700 Subject: [PATCH] Add first draft of Aperture element. --- docs/source/usage/parameters.rst | 11 +++ docs/source/usage/python.rst | 8 ++ src/initialization/InitElement.cpp | 9 ++- src/particles/elements/All.H | 6 +- src/particles/elements/Aperture.H | 114 +++++++++++++++++++++++++++++ src/python/elements.cpp | 12 +++ 6 files changed, 157 insertions(+), 3 deletions(-) create mode 100644 src/particles/elements/Aperture.H diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 123488f6d..f22f4cc4c 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -466,6 +466,17 @@ Lattice Elements * ``.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: + + * ``.shape`` (``integer``, dimensionless) shape of the aperture boundary + + (m = 0) rectangular, (m = 1) elliptical + + * ``.xmax`` (``float``, in meters) maximum value of the horizontal coordinate + + * ``.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. diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index e85fcbd72..6686ea5dd 100644 --- a/docs/source/usage/python.rst +++ b/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. diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index 873041596..873d66122 100644 --- a/src/initialization/InitElement.cpp +++ b/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"; diff --git a/src/particles/elements/All.H b/src/particles/elements/All.H index 27fe1a176..2262f69d2 100644 --- a/src/particles/elements/All.H +++ b/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, diff --git a/src/particles/elements/Aperture.H b/src/particles/elements/Aperture.H new file mode 100644 index 000000000..a9125cac5 --- /dev/null +++ b/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 +#include +#include + +#include + +namespace impactx +{ + struct Aperture + : public elements::BeamOptic, + 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 diff --git a/src/python/elements.cpp b/src/python/elements.cpp index d734b76b4..53e8e3ce3 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -358,5 +358,17 @@ void init_elements(py::module& m) "A soft-edge quadrupole." ) ; + + py::class_(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." + ) + ; + + }