From f4bd0f87f83886891715baac09ae76c5402c8e60 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Tue, 19 Nov 2024 16:57:45 -0800 Subject: [PATCH] Add absorber option to Aperture element. (#758) * Add absorber capability to Aperture element. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update InitElement.cpp Restore deleted PlaneXY. * Update elements.cpp Resture deleted PlaneXYRot to elements.cpp. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update src/python/elements.cpp * Add benchmark test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update documentation. * Update Aperture.H Modify modulo step in Aperture.H to fix the case of overlapping elliptical apertures. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/source/usage/parameters.rst | 1 + docs/source/usage/python.rst | 5 ++ examples/CMakeLists.txt | 16 ++++ examples/aperture/README.rst | 53 ++++++++++++ examples/aperture/analysis_absorber.py | 110 +++++++++++++++++++++++++ examples/aperture/input_absorber.in | 50 +++++++++++ examples/aperture/run_absorber.py | 73 ++++++++++++++++ src/initialization/InitElement.cpp | 9 +- src/particles/elements/Aperture.H | 56 ++++++++++--- src/particles/elements/Aperture.cpp | 14 ++++ src/python/elements.cpp | 29 ++++++- 11 files changed, 403 insertions(+), 13 deletions(-) create mode 100755 examples/aperture/analysis_absorber.py create mode 100644 examples/aperture/input_absorber.in create mode 100755 examples/aperture/run_absorber.py diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 7a654075d..3d391093c 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -411,6 +411,7 @@ Lattice Elements * ``.repeat_x`` (``float``, in meters) horizontal period for repeated aperture masking (inactive by default) * ``.repeat_y`` (``float``, in meters) vertical period for repeated aperture masking (inactive by default) * ``.shape`` (``string``) shape of the aperture boundary: ``rectangular`` (default) or ``elliptical`` + * ``.action`` (``string``) action of the aperture domain: ``transmit`` (default) or ``absorb`` * ``.dx`` (``float``, in meters) horizontal translation error * ``.dy`` (``float``, in meters) vertical translation error * ``.rotation`` (``float``, in degrees) rotation error in the transverse plane diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 2e24c86fb..ddda0a0d0 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -953,6 +953,7 @@ This module provides elements for the accelerator lattice. :param repeat_x: horizontal period for repeated aperture masking (inactive by default) (meter) :param repeat_y: vertical period for repeated aperture masking (inactive by default) (meter) :param shape: aperture boundary shape: ``"rectangular"`` (default) or ``"elliptical"`` + :param action: aperture domain action: ``"transmit"`` (default) or ``"absorb"`` :param dx: horizontal translation error in m :param dy: vertical translation error in m :param rotation: rotation error in the transverse plane [degrees] @@ -962,6 +963,10 @@ This module provides elements for the accelerator lattice. aperture type (rectangular, elliptical) + .. py:property:: action + + aperture type (transmit, absorb) + .. py:property:: xmax maximum horizontal coordinate diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index d9ade3e5a..02e083fd8 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -843,6 +843,22 @@ add_impactx_test(aperture.py OFF # no plot script yet ) +# Absorber collimation ######################################################## +# +# w/o space charge +add_impactx_test(absorber + examples/aperture/input_absorber.in + ON # ImpactX MPI-parallel + examples/aperture/analysis_absorber.py + OFF # no plot script yet +) +add_impactx_test(absorber.py + examples/aperture/run_absorber.py + OFF # ImpactX MPI-parallel + examples/aperture/analysis_absorber.py + OFF # no plot script yet +) + # Apochromat drift-quad example ########################################################## # # w/o space charge diff --git a/examples/aperture/README.rst b/examples/aperture/README.rst index 4fac877e7..edb4247c5 100644 --- a/examples/aperture/README.rst +++ b/examples/aperture/README.rst @@ -101,3 +101,56 @@ We run the following script to analyze correctness: .. literalinclude:: analysis_aperture_periodic.py :language: python3 :caption: You can copy this file from ``examples/aperture/analysis_aperture_periodic.py``. + + +.. _examples-absorber: + +Collimation Using an Absorber +================================ + +Proton beam undergoing collimation through partial absorption by a rectangular domain. +This test is the exact negative of the previous test, and illustrates the ``absorb`` option of the ``Aperture`` element. + +We use a 250 MeV proton beam with a horizontal rms beam size of 1.56 mm and a vertical rms beam size of 2.21 mm. + +After a short drift of 0.123, the beam hits a 1 mm x 1.5 mm rectangular structure, resulting in particle loss. + +In this test, the initial values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nomin> +The test fails if: + +* any of the final coordinates for the valid (not lost) particles lie inside the absorber boundary or +* any of the lost particles are outside the absorber boundary or +* if the sum of lost and kept particles is not equal to the initial particles or +* if the recorded position :math:`s` for the lost particles does not coincide with the drift distance. + + +Run +--- + +This example can be run as a Python script (``python3 run_absorber.py``) or with an app with an input file (``impactx input_absorber.in``). +Each can also be prefixed with an `MPI executor `__, such as ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. + +.. tab-set:: + + .. tab-item:: Python Script + + .. literalinclude:: run_absorber.py + :language: python3 + :caption: You can copy this file from ``examples/aperture/run_absorber.py``. + + .. tab-item:: App Input File + + .. literalinclude:: input_absorber.in + :language: ini + :caption: You can copy this file from ``examples/aperture/input_absorber.in``. + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_absorber.py`` + + .. literalinclude:: analysis_absorber.py + :language: python3 + :caption: You can copy this file from ``examples/aperture/analysis_absorber.py``. diff --git a/examples/aperture/analysis_absorber.py b/examples/aperture/analysis_absorber.py new file mode 100755 index 000000000..e0cff8f59 --- /dev/null +++ b/examples/aperture/analysis_absorber.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.greater_equal(dx.max(), 0.0) + +print(f" y_max={final['position_y'].max()}") +print(f" y_min={final['position_y'].min()}") +assert np.greater_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.less_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.less_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_absorber.in b/examples/aperture/input_absorber.in new file mode 100644 index 000000000..129229983 --- /dev/null +++ b/examples/aperture/input_absorber.in @@ -0,0 +1,50 @@ +############################################################################### +# 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 collimator monitor +lattice.nslice = 1 + +monitor.type = beam_monitor +monitor.backend = h5 + +drift.type = drift +drift.ds = 0.123 + +collimator.type = aperture +collimator.shape = rectangular +collimator.xmax = 1.0e-3 +collimator.ymax = 1.5e-3 +collimator.action = absorb + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true +diag.backend = h5 diff --git a/examples/aperture/run_absorber.py b/examples/aperture/run_absorber.py new file mode 100755 index 000000000..4c75592ae --- /dev/null +++ b/examples/aperture/run_absorber.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import amrex.space3d as amr +from impactx import ImpactX, distribution, elements + +# work-around for https://github.com/ECP-WarpX/impactx/issues/499 +pp_amrex = amr.ParmParse("amrex") +pp_amrex.add("the_arena_is_managed", 1) + +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), + elements.Aperture( + name="collimator", + xmax=1.0e-3, + ymax=1.5e-3, + shape="rectangular", + action="absorb", + ), + monitor, + ] +) + +# run simulation +sim.track_particles() + +# clean shutdown +sim.finalize() diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index 148e719d7..d106a4b86 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -387,18 +387,25 @@ namespace detail amrex::ParticleReal repeat_x = 0.0; amrex::ParticleReal repeat_y = 0.0; std::string shape_str = "rectangular"; + std::string action_str = "transmit"; pp_element.get("xmax", xmax); pp_element.get("ymax", ymax); pp_element.queryAdd("repeat_x", repeat_x); pp_element.queryAdd("repeat_y", repeat_y); pp_element.queryAdd("shape", shape_str); + pp_element.queryAdd("action", action_str); AMREX_ALWAYS_ASSERT_WITH_MESSAGE(shape_str == "rectangular" || shape_str == "elliptical", element_name + ".shape must be \"rectangular\" or \"elliptical\""); Aperture::Shape shape = shape_str == "rectangular" ? Aperture::Shape::rectangular : Aperture::Shape::elliptical; + AMREX_ALWAYS_ASSERT_WITH_MESSAGE(action_str == "transmit" || action_str == "absorb", + element_name + ".action must be \"transmit\" or \"absorb\""); + Aperture::Action action = action_str == "transmit" ? + Aperture::Action::transmit : + Aperture::Action::absorb; - m_lattice.emplace_back( Aperture(xmax, ymax, repeat_x, repeat_y, shape, a["dx"], a["dy"], a["rotation_degree"], element_name) ); + m_lattice.emplace_back( Aperture(xmax, ymax, repeat_x, repeat_y, shape, action, a["dx"], a["dy"], a["rotation_degree"], element_name) ); } else if (element_type == "beam_monitor") { std::string openpmd_name = element_name; diff --git a/src/particles/elements/Aperture.H b/src/particles/elements/Aperture.H index d60809aa2..c7917ec27 100644 --- a/src/particles/elements/Aperture.H +++ b/src/particles/elements/Aperture.H @@ -41,14 +41,23 @@ namespace impactx rectangular, elliptical }; + enum Action + { + transmit, + absorb + }; static std::string shape_name (Shape const & shape); + static std::string + action_name (Action const & action); + /** A thin collimator element that applies a transverse aperture boundary. * Particles outside the boundary are considered lost. * * @param shape aperture shape + * @param action specify action of domain (transmit/absorb) * @param xmax maximum value of horizontal coordinate (m) * @param ymax maximum value of vertical coordinate (m) * @param repeat_x horizontal period for repeated masking, optional (m) @@ -64,6 +73,7 @@ namespace impactx amrex::ParticleReal repeat_x, amrex::ParticleReal repeat_y, Shape shape, + Action action, amrex::ParticleReal dx = 0, amrex::ParticleReal dy = 0, amrex::ParticleReal rotation_degree = 0, @@ -71,7 +81,7 @@ namespace impactx ) : Named(name), Alignment(dx, dy, rotation_degree), - m_shape(shape), m_xmax(xmax), m_ymax(ymax), m_repeat_x(repeat_x), m_repeat_y(repeat_y) + m_shape(shape), m_action(action), m_xmax(xmax), m_ymax(ymax), m_repeat_x(repeat_x), m_repeat_y(repeat_y) { } @@ -110,8 +120,8 @@ namespace impactx // define intermediate variables amrex::ParticleReal u; amrex::ParticleReal v; - amrex::ParticleReal dx = m_xmax; - amrex::ParticleReal dy = m_ymax; + amrex::ParticleReal dx = m_repeat_x/2.0; + amrex::ParticleReal dy = m_repeat_y/2.0; // if the aperture is periodic, shift x,y coordinates to the fundamental domain u = (m_repeat_x==0.0) ? x : (std::fmod(std::abs(x)+dx,m_repeat_x)-dx); @@ -122,21 +132,46 @@ namespace impactx v = v / m_ymax; // compare against the aperture boundary - switch (m_shape) + switch (m_action) { - case Shape::rectangular : // default - if (std::pow(u,2)>1 || std::pow(v,2) > 1_prt) { - amrex::ParticleIDWrapper{idcpu}.make_invalid(); + case Action::transmit : // default action + + switch (m_shape) + { + case Shape::rectangular : // default shape + if (std::pow(u,2)>1 || std::pow(v,2) > 1_prt) { + amrex::ParticleIDWrapper{idcpu}.make_invalid(); + } + break; + + case Shape::elliptical : + if (std::pow(u,2) + std::pow(v,2) > 1_prt) { + amrex::ParticleIDWrapper{idcpu}.make_invalid(); + } + break; } break; - case Shape::elliptical : - if (std::pow(u,2) + std::pow(v,2) > 1_prt) { - amrex::ParticleIDWrapper{idcpu}.make_invalid(); + case Action::absorb : + + switch (m_shape) + { + case Shape::rectangular : // default + if (std::pow(u,2) < 1 && std::pow(v,2) < 1_prt) { + amrex::ParticleIDWrapper{idcpu}.make_invalid(); + } + break; + + case Shape::elliptical : + if (std::pow(u,2) + std::pow(v,2) < 1_prt) { + amrex::ParticleIDWrapper{idcpu}.make_invalid(); + } + break; } break; } + // undo shift due to alignment errors of the element shift_out(x, y, px, py); } @@ -145,6 +180,7 @@ namespace impactx using Thin::operator(); Shape m_shape; //! aperture type (rectangular, elliptical) + Action m_action; //! action type (transmit, absorb) amrex::ParticleReal m_xmax; //! maximum horizontal coordinate amrex::ParticleReal m_ymax; //! maximum vertical coordinate amrex::ParticleReal m_repeat_x; //! horizontal period for repeated masking diff --git a/src/particles/elements/Aperture.cpp b/src/particles/elements/Aperture.cpp index c647bca7c..348a81be9 100644 --- a/src/particles/elements/Aperture.cpp +++ b/src/particles/elements/Aperture.cpp @@ -26,3 +26,17 @@ impactx::Aperture::shape_name (Shape const & shape) throw std::runtime_error("Unknown shape"); } } + +std::string +impactx::Aperture::action_name (Action const & action) +{ + switch (action) + { + case Aperture::Action::transmit : // default + return "transmit"; + case Aperture::Action::absorb : + return "absorb"; + default: + throw std::runtime_error("Unknown action"); + } +} diff --git a/src/python/elements.cpp b/src/python/elements.cpp index fc60ad89a..b9c9429d6 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -279,7 +279,8 @@ void init_elements(py::module& m) [](Aperture const & ap) { return element_name( ap, - std::make_pair("shape", ap.shape_name(ap.m_shape)) + std::make_pair("shape", ap.shape_name(ap.m_shape)), + std::make_pair("action", ap.action_name(ap.m_action)) ); } ) @@ -289,6 +290,7 @@ void init_elements(py::module& m) amrex::ParticleReal repeat_x, amrex::ParticleReal repeat_y, std::string const & shape, + std::string const & action, amrex::ParticleReal dx, amrex::ParticleReal dy, amrex::ParticleReal rotation_degree, @@ -298,16 +300,23 @@ void init_elements(py::module& m) if (shape != "rectangular" && shape != "elliptical") throw std::runtime_error(R"(shape must be "rectangular" or "elliptical")"); + if (action != "transmit" && action != "absorb") + throw std::runtime_error(R"(action must be "transmit" or "absorb")"); + Aperture::Shape const s = shape == "rectangular" ? Aperture::Shape::rectangular : Aperture::Shape::elliptical; - return new Aperture(xmax, ymax, repeat_x, repeat_y, s, dx, dy, rotation_degree, name); + Aperture::Action const a = action == "transmit" ? + Aperture::Action::transmit : + Aperture::Action::absorb; + return new Aperture(xmax, ymax, repeat_x, repeat_y, s, a, dx, dy, rotation_degree, name); }), py::arg("xmax"), py::arg("ymax"), py::arg("repeat_x") = 0, py::arg("repeat_y") = 0, py::arg("shape") = "rectangular", + py::arg("action") = "transmit", py::arg("dx") = 0, py::arg("dy") = 0, py::arg("rotation") = 0, @@ -330,6 +339,22 @@ void init_elements(py::module& m) }, "aperture type (rectangular, elliptical)" ) + .def_property("action", + [](Aperture & ap) + { + return ap.action_name(ap.m_action); + }, + [](Aperture & ap, std::string const & action) + { + if (action != "transmit" && action != "absorb") + throw std::runtime_error(R"(action must be "transmit" or "absorb")"); + + ap.m_action = action == "transmit" ? + Aperture::Action::transmit : + Aperture::Action::absorb; + }, + "action type (transmit, absorb)" + ) .def_property("xmax", [](Aperture & ap) { return ap.m_xmax; }, [](Aperture & ap, amrex::ParticleReal xmax) { ap.m_xmax = xmax; },