From 79fab7cbdf95c78369078b675fa0b6fbe9464c38 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Thu, 23 Nov 2023 16:11:36 -0800 Subject: [PATCH] Add aperture element (#398) * Add first draft of Aperture element. * Update Aperture.H Correct element docs. * Add benchmark example. * Clean up switch. * Add index change. * Update python.rst Place arguments with defaults last (Python docs). * Update parameters.rst Place arguments with defaults last (C++ docs). * Update Aperture.H Place arguments with defaults last (C++ src). * Update elements.cpp Place arguments with defaults last (Python equivalent). * Apply suggestions from code review Co-authored-by: Axel Huebl * Update elements.cpp Add missing , * Aperture Update API: Enums & Strings Update the C++ API to use an enum for the aperture shape and the Python API & inputs file syntax to accept a string. That makes the parameters at the call sites self-describing. * Draft: Copy Lost Particles to do: - remove in beam species - output * Local Move of Lost Particles * Container: Reference Lost Container * Aperture Example: Kin Energy Update changed input API * Lost Particles: Backend Control * Update Aperture Example & Analysis * Output Runtime Attribute "s_lost" * Aperture Test: Add Drift This way, we can test if "s" was set to the right value. * Fix GPU Compile * Cleaning * Lost Output: Only N>0 Do not generate empty files with half-ready meta data. * openPMD: Remove Outdated Species Outdated line from an earlier design created an empty species in output. --------- Co-authored-by: Axel Huebl --- docs/source/usage/examples.rst | 1 + docs/source/usage/parameters.rst | 14 ++ docs/source/usage/python.rst | 13 ++ examples/CMakeLists.txt | 18 ++ examples/aperture/README.rst | 52 ++++++ examples/aperture/analysis_aperture.py | 116 ++++++++++++ examples/aperture/input_aperture.in | 50 ++++++ examples/aperture/run_aperture.py | 64 +++++++ src/ImpactX.H | 3 + src/ImpactX.cpp | 33 +++- src/initialization/InitElement.cpp | 12 ++ src/particles/CMakeLists.txt | 1 + src/particles/CollectLost.H | 30 ++++ src/particles/CollectLost.cpp | 166 ++++++++++++++++++ src/particles/ImpactXParticleContainer.H | 33 ++++ src/particles/ImpactXParticleContainer.cpp | 64 ++++++- src/particles/elements/All.H | 2 + src/particles/elements/Aperture.H | 115 ++++++++++++ .../elements/diagnostics/openPMD.cpp | 14 +- src/python/ImpactX.cpp | 15 +- src/python/ImpactXParticleContainer.cpp | 47 ++--- src/python/elements.cpp | 21 +++ .../impactx/ImpactXParticleContainer.py | 11 +- 23 files changed, 837 insertions(+), 58 deletions(-) create mode 100644 examples/aperture/README.rst create mode 100755 examples/aperture/analysis_aperture.py create mode 100644 examples/aperture/input_aperture.in create mode 100755 examples/aperture/run_aperture.py create mode 100644 src/particles/CollectLost.H create mode 100644 src/particles/CollectLost.cpp create mode 100644 src/particles/elements/Aperture.H diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index e19d73e47..ddf9b2011 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -29,6 +29,7 @@ This section allows you to **download input files** that correspond to different examples/compression/README.rst examples/kicker/README.rst examples/thin_dipole/README.rst + examples/aperture/README.rst Unit tests diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index d60d2ca0e..2494651b5 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -516,6 +516,15 @@ Lattice Elements * ``.rc`` (``float``, in meters) effective radius of curvature + * ``aperture`` for a thin collimator element applying a transverse aperture boundary. + This requires these additional parameters: + + * ``.xmax`` (``float``, in meters) maximum value of the horizontal coordinate + + * ``.ymax`` (``float``, in meters) maximum value of the vertical coordinate + + * ``.shape`` (``string``) shape of the aperture boundary: ``rectangular`` (default) or ``elliptical`` + * ``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. @@ -691,6 +700,11 @@ Diagnostics and output * ``diag.file_min_digits`` (``integer``, optional, default: ``6``) The minimum number of digits used for the step number appended to the diagnostic file names. +* ``diag.backend`` (``string``, default value: ``default``) + + Diagnostics for particles lost in apertures, stored as ``diags/openPMD/particles_lost.*`` at the end of the simulation. + See the ``beam_monitor`` element for backend values. + .. _running-cpp-parameters-diagnostics-reduced: Reduced Diagnostics diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index ad8057911..7117aca1f 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -116,6 +116,11 @@ General The minimum number of digits (default: ``6``) used for the step number appended to the diagnostic file names. + .. py:property:: particle_lost_diagnostics_backend + + Diagnostics for particles lost in apertures. + See the ``BeamMonitor`` element for backend values. + .. py:method:: init_grids() Initialize AMReX blocks/grids for domain decomposition & space charge mesh. @@ -712,6 +717,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(xmax, ymax, shape="rectangular") + + A thin collimator element, applying a transverse aperture boundary. + + :param xmax: maximum allowed value of the horizontal coordinate (meter) + :param ymax: maximum allowed value of the vertical coordinate (meter) + :param shape: aperture boundary shape: ``"rectangular"`` (default) or ``"elliptical"`` + .. py:class:: impactx.elements.SoftQuadrupole(ds, gscale, cos_coefficients, sin_coefficients, nslice=1) A soft-edge quadrupole. diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index f57193e53..e16b12350 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -657,3 +657,21 @@ add_impactx_test(thin_dipole.py examples/thin_dipole/analysis_thin_dipole.py OFF # no plot script yet ) + +# Aperture collimation ############################################################ +# +# w/o space charge +add_impactx_test(aperture + examples/aperture/input_aperture.in + ON # ImpactX MPI-parallel + OFF # ImpactX Python interface + examples/aperture/analysis_aperture.py + OFF # no plot script yet +) +add_impactx_test(aperture.py + examples/aperture/run_aperture.py + OFF # ImpactX MPI-parallel + ON # ImpactX Python interface + examples/aperture/analysis_aperture.py + OFF # no plot script yet +) diff --git a/examples/aperture/README.rst b/examples/aperture/README.rst new file mode 100644 index 000000000..2592ebc38 --- /dev/null +++ b/examples/aperture/README.rst @@ -0,0 +1,52 @@ +.. _examples-aperture: + +Aperture Collimation +==================== + +Proton beam undergoing collimation by a rectangular boundary aperture. + + +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 is scraped by a 1 mm x 1.5 mm rectangular aperture. + +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 nominal values. +The test fails if: + +* any of the final coordinates for the valid (not lost) particles lie outside the aperture boundary or +* any of the lost particles are inside the aperture 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_aperture.py``) or with an app with an input file (``impactx input_aperture.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_aperture.py + :language: python3 + :caption: You can copy this file from ``examples/aperture/run_aperture.py``. + + .. tab-item:: App Input File + + .. literalinclude:: input_aperture.in + :language: ini + :caption: You can copy this file from ``examples/aperture/input_aperture.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_aperture.py`` + + .. literalinclude:: analysis_aperture.py + :language: python3 + :caption: You can copy this file from ``examples/aperture/analysis_aperture.py``. diff --git a/examples/aperture/analysis_aperture.py b/examples/aperture/analysis_aperture.py new file mode 100755 index 000000000..cbbd53c7c --- /dev/null +++ b/examples/aperture/analysis_aperture.py @@ -0,0 +1,116 @@ +#!/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.in b/examples/aperture/input_aperture.in new file mode 100644 index 000000000..a6cb8cf41 --- /dev/null +++ b/examples/aperture/input_aperture.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.sigmaX = 1.559531175539e-3 +beam.sigmaY = 2.205510139392e-3 +beam.sigmaT = 1.0e-3 +beam.sigmaPx = 6.41218345413e-4 +beam.sigmaPy = 9.06819680526e-4 +beam.sigmaPt = 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 + + +############################################################################### +# 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.py b/examples/aperture/run_aperture.py new file mode 100755 index 000000000..dc89c062e --- /dev/null +++ b/examples/aperture/run_aperture.py @@ -0,0 +1,64 @@ +#!/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, RefPart, 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( + sigmaX=1.559531175539e-3, + sigmaY=2.205510139392e-3, + sigmaT=1.0e-3, + sigmaPx=6.41218345413e-4, + sigmaPy=9.06819680526e-4, + sigmaPt=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(0.123), + elements.Aperture(xmax=1.0e-3, ymax=1.5e-3, shape="rectangular"), + monitor, + ] +) + +# run simulation +sim.evolve() + +# clean shutdown +del sim +amr.finalize() diff --git a/src/ImpactX.H b/src/ImpactX.H index 15ae35ba3..b3c2dacac 100644 --- a/src/ImpactX.H +++ b/src/ImpactX.H @@ -134,6 +134,9 @@ namespace impactx /** these are the physical/beam particles of the simulation */ std::unique_ptr m_particle_container; + /** former beam particles that got lost in apertures, the wall, etc. */ + std::unique_ptr m_particles_lost; + /** charge density per level */ std::unordered_map m_rho; /** scalar potential per level */ diff --git a/src/ImpactX.cpp b/src/ImpactX.cpp index 9baf7aa85..b6b0610d7 100644 --- a/src/ImpactX.cpp +++ b/src/ImpactX.cpp @@ -9,6 +9,7 @@ */ #include "ImpactX.H" #include "initialization/InitAmrCore.H" +#include "particles/CollectLost.H" #include "particles/ImpactXParticleContainer.H" #include "particles/Push.H" #include "particles/diagnostics/DiagnosticOutput.H" @@ -33,7 +34,8 @@ namespace impactx { ImpactX::ImpactX () : AmrCore(initialization::init_amr_core()), - m_particle_container(std::make_unique(this)) + m_particle_container(std::make_unique(this)), + m_particles_lost(std::make_unique(this)) { // todo: if amr.n_cells is provided, overwrite/redefine AmrCore object @@ -81,6 +83,21 @@ namespace impactx // init blocks / grids & MultiFabs AmrCore::InitFromScratch(0.0); amrex::Print() << "boxArray(0) " << boxArray(0) << std::endl; + + // alloc particle containers + // the lost particles have an extra runtime attribute: s when it was lost + bool comm = true; + m_particles_lost->AddRealComp(comm); + + // have to resize here, not in the constructor because grids have not + // been built when constructor was called. + m_particle_container->reserveData(); + m_particle_container->resizeData(); + m_particles_lost->reserveData(); + m_particles_lost->resizeData(); + + // register shortcut + m_particle_container->SetLostParticleContainer(m_particles_lost.get()); } void ImpactX::evolve () @@ -210,6 +227,9 @@ namespace impactx // push all particles with external maps Push(*m_particle_container, element_variant, global_step); + // move "lost" particles to another particle container + collect_lost_particles(*m_particle_container); + // just prints an empty newline at the end of the slice_step amrex::Print() << "\n"; @@ -261,6 +281,17 @@ namespace impactx diagnostics::OutputType::PrintReducedBeamCharacteristics, "diags/reduced_beam_characteristics_final", global_step); + + // output particles lost in apertures + if (m_particles_lost->TotalNumberOfParticles() > 0) + { + std::string openpmd_backend = "default"; + pp_diag.queryAdd("backend", openpmd_backend); + + diagnostics::BeamMonitor output_lost("particles_lost", openpmd_backend, "g"); + output_lost(*m_particles_lost, 0); + output_lost.finalize(); + } } // loop over all beamline elements & finalize them diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index 2d6931945..b5f5de705 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -248,6 +248,18 @@ namespace detail Kicker::UnitSystem::dimensionless : Kicker::UnitSystem::Tm; m_lattice.emplace_back( Kicker(xkick, ykick, units) ); + } else if (element_type == "aperture") { + amrex::Real xmax, ymax; + std::string shape_str = "rectangular"; + pp_element.get("xmax", xmax); + pp_element.get("ymax", ymax); + pp_element.queryAdd("shape", shape_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; + m_lattice.emplace_back( Aperture(xmax, ymax, shape) ); } else if (element_type == "beam_monitor") { std::string openpmd_name = element_name; pp_element.queryAdd("name", openpmd_name); diff --git a/src/particles/CMakeLists.txt b/src/particles/CMakeLists.txt index c2b14cfa9..85e8252e4 100644 --- a/src/particles/CMakeLists.txt +++ b/src/particles/CMakeLists.txt @@ -1,6 +1,7 @@ target_sources(lib PRIVATE ChargeDeposition.cpp + CollectLost.cpp ImpactXParticleContainer.cpp Push.cpp ) diff --git a/src/particles/CollectLost.H b/src/particles/CollectLost.H new file mode 100644 index 000000000..dbc65cadb --- /dev/null +++ b/src/particles/CollectLost.H @@ -0,0 +1,30 @@ +/* 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: Axel Huebl, Chad Mitchell + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_COLLECT_LOST_H +#define IMPACTX_COLLECT_LOST_H + +#include "particles/ImpactXParticleContainer.H" + + +namespace impactx +{ + /** Move lost particles into a separate container + * + * If particles are marked as lost, by setting their id to negative, we + * will move them to another particle container, store their position when + * lost and stop pushing them in the beamline. + * + * @param source the beam particle container that might loose particles + */ + void collect_lost_particles (ImpactXParticleContainer& source); + +} // namespace impactx + +#endif // IMPACTX_COLLECT_LOST_H diff --git a/src/particles/CollectLost.cpp b/src/particles/CollectLost.cpp new file mode 100644 index 000000000..a8f8b8ede --- /dev/null +++ b/src/particles/CollectLost.cpp @@ -0,0 +1,166 @@ +/* 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: Axel Huebl, Chad Mitchell + * License: BSD-3-Clause-LBNL + */ +#include "CollectLost.H" + +#include +#include +#include +#include +#include + + +namespace impactx +{ + struct CopyAndMarkNegative + { + static constexpr int s_index = 0; //!< index of runtime attribute in destination for position s where particle got lost + amrex::ParticleReal s_lost; //!< position s in meters where particle got lost + + using SrcData = ImpactXParticleContainer::ParticleTileType::ConstParticleTileDataType; + using DstData = ImpactXParticleContainer::ParticleTileType::ParticleTileDataType; + + AMREX_GPU_HOST_DEVICE + void operator() (DstData const &dst, SrcData const &src, int src_ip, int dst_ip) const noexcept { + dst.m_aos[dst_ip] = src.m_aos[src_ip]; + + for (int j = 0; j < SrcData::NAR; ++j) + dst.m_rdata[j][dst_ip] = src.m_rdata[j][src_ip]; + for (int j = 0; j < src.m_num_runtime_real; ++j) + dst.m_runtime_rdata[j][dst_ip] = src.m_runtime_rdata[j][src_ip]; + + // unused: integer compile-time or runtime attributes + //for (int j = 0; j < SrcData::NAI; ++j) + // dst.m_idata[j][dst_ip] = src.m_idata[j][src_ip]; + //for (int j = 0; j < src.m_num_runtime_int; ++j) + // dst.m_runtime_idata[j][dst_ip] = src.m_runtime_idata[j][src_ip]; + + // flip id to positive in destination + dst.id(dst_ip) = amrex::Math::abs(dst.id(dst_ip)); + + // remember the current s of the ref particle when lost + dst.m_runtime_rdata[s_index][dst_ip] = s_lost; + } + }; + + void collect_lost_particles (ImpactXParticleContainer& source) + { + BL_PROFILE("impactX::collect_lost_particles"); + + using SrcData = ImpactXParticleContainer::ParticleTileType::ConstParticleTileDataType; + + ImpactXParticleContainer& dest = *source.GetLostParticleContainer(); + + RefPart const ref_part = source.GetRefParticle(); + auto const s_lost = ref_part.s; + + // have to resize here, not in the constructor because grids have not + // been built when constructor was called. + dest.reserveData(); + dest.resizeData(); + + // copy all particles marked with a negative ID from source to destination + int const nLevel = source.finestLevel(); + for (int lev = 0; lev <= nLevel; ++lev) { + // loop over all particle boxes + using ParIt = ImpactXParticleContainer::iterator; + auto& plevel = source.GetParticles(lev); + + // TODO: is it safe to add OpenMP parallelism here? + for (ParIt pti(source, lev); pti.isValid(); ++pti) { + auto index = std::make_pair(pti.index(), pti.LocalTileIndex()); + if (plevel.find(index) == plevel.end()) continue; + + auto& ptile_source = plevel.at(index); + auto const np = ptile_source.numParticles(); + if (np == 0) continue; // no particles in source tile + + // we will copy particles that were marked as lost, with a negative id + auto const predicate = [] AMREX_GPU_HOST_DEVICE (const SrcData& src, int ip) + /* NVCC 11.3.109 chokes in C++17 on this: noexcept */ + { + return src.id(ip) < 0; + }; + + auto& ptile_dest = dest.DefineAndReturnParticleTile( + lev, pti.index(), pti.LocalTileIndex()); + + // count how many particles we will copy + amrex::ReduceOps reduce_op; + amrex::ReduceData reduce_data(reduce_op); + { + auto const src_data = ptile_source.getConstParticleTileData(); + + reduce_op.eval(np, reduce_data, [=] AMREX_GPU_HOST_DEVICE (int ip) + { + return predicate(src_data, ip); + }); + } + int const np_to_move = amrex::get<0>(reduce_data.value()); + if (np_to_move == 0) continue; // no particles to move from source tile + + // allocate memory in destination + int const dst_index = ptile_dest.numParticles(); + ptile_dest.resize(dst_index + np_to_move); + + // copy particles + // skipped in loop below: integer compile-time or runtime attributes + AMREX_ALWAYS_ASSERT(SrcData::NAI == 0); + AMREX_ALWAYS_ASSERT(ptile_source.NumRuntimeIntComps() == 0); + + // first runtime attribute in destination is s position where particle got lost + AMREX_ALWAYS_ASSERT(dest.NumRuntimeRealComps() > 0); + + amrex::filterAndTransformParticles( + ptile_dest, + ptile_source, + predicate, + CopyAndMarkNegative{s_lost}, + 0, + dst_index + ); + + // remove particles with negative ids in source + { + int n_removed = 0; + auto ptile_src_data = ptile_source.getParticleTileData(); + for (int ip = 0; ip < np; ++ip) + { + if (ptile_source.id(ip) < 0) + n_removed++; + else + { + if (n_removed > 0) + { + // move down + int const new_index = ip - n_removed; + + ptile_src_data.m_aos[new_index] = ptile_src_data.m_aos[ip]; + + for (int j = 0; j < SrcData::NAR; ++j) + ptile_src_data.m_rdata[j][new_index] = ptile_src_data.m_rdata[j][ip]; + for (int j = 0; j < ptile_src_data.m_num_runtime_real; ++j) + ptile_src_data.m_runtime_rdata[j][new_index] = ptile_src_data.m_runtime_rdata[j][ip]; + + // unused: integer compile-time or runtime attributes + //for (int j = 0; j < SrcData::NAI; ++j) + // dst.m_idata[j][new_index] = src.m_idata[j][ip]; + //for (int j = 0; j < ptile_src_data.m_num_runtime_int; ++j) + // dst.m_runtime_idata[j][new_index] = src.m_runtime_idata[j][ip]; + } + } + } + AMREX_ALWAYS_ASSERT(np_to_move == n_removed); + ptile_source.resize(np - n_removed); + } + + } // particle tile loop + } // lev + } +} // namespace impactx diff --git a/src/particles/ImpactXParticleContainer.H b/src/particles/ImpactXParticleContainer.H index 6c0adf1d1..7ab1abed6 100644 --- a/src/particles/ImpactXParticleContainer.H +++ b/src/particles/ImpactXParticleContainer.H @@ -177,6 +177,16 @@ namespace impactx amrex::ParticleReal const & qm, amrex::ParticleReal const & bchchg); + /** Register storage for lost particles + * + * @param lost_pc particle container for lost particles + */ + void + SetLostParticleContainer (ImpactXParticleContainer * lost_pc); + + ImpactXParticleContainer * + GetLostParticleContainer (); + /** Set reference particle attributes * * @param refpart reference particle @@ -259,6 +269,14 @@ namespace impactx DepositCharge (std::unordered_map & rho, amrex::Vector const & ref_ratio); + /** Get the name of each Real AoS component */ + std::vector + RealAoS_names () const; + + /** Get the name of each Real SoA component */ + std::vector + RealSoA_names () const; + private: //! the reference particle for the beam in the particle container @@ -267,8 +285,23 @@ namespace impactx //! the particle shape std::optional m_particle_shape; + //! a non-owning reference to lost particles, i.e., due to apertures + ImpactXParticleContainer* m_particles_lost = nullptr; + }; // ImpactXParticleContainer + /** Get the name of each Real AoS component */ + std::vector + get_RealAoS_names (); + + /** Get the name of each Real SoA component + * + * @param num_real_comps number of compile-time + runtime arrays + * @return names + */ + std::vector + get_RealSoA_names (int num_real_comps); + } // namespace impactx #endif // IMPACTX_PARTICLE_CONTAINER_H diff --git a/src/particles/ImpactXParticleContainer.cpp b/src/particles/ImpactXParticleContainer.cpp index c587694d4..d333b1e75 100644 --- a/src/particles/ImpactXParticleContainer.cpp +++ b/src/particles/ImpactXParticleContainer.cpp @@ -57,6 +57,24 @@ namespace impactx SetParticleSize(); } + void + ImpactXParticleContainer::SetLostParticleContainer (ImpactXParticleContainer * lost_pc) + { + m_particles_lost = lost_pc; + } + + ImpactXParticleContainer * + ImpactXParticleContainer::GetLostParticleContainer () + { + if (m_particles_lost == nullptr) + { + throw std::logic_error( + "ImpactXParticleContainer::GetLostParticleContainer No lost particle container registered yet."); + } else { + return m_particles_lost; + } + } + void ImpactXParticleContainer::SetParticleShape (int order) { if (m_particle_shape.has_value()) { @@ -104,11 +122,6 @@ namespace impactx // number of particles to add int const np = x.size(); - // have to resize here, not in the constructor because grids have not - // been built when constructor was called. - reserveData(); - resizeData(); - auto& particle_tile = DefineAndReturnParticleTile(0, 0, 0); /* Create a temporary tile to obtain data from simulation. This data @@ -199,4 +212,45 @@ namespace impactx ImpactXParticleContainer, RealSoA::w >(*this); } + + std::vector + ImpactXParticleContainer::RealAoS_names () const + { + return get_RealAoS_names(); + } + + std::vector + ImpactXParticleContainer::RealSoA_names () const + { + return get_RealSoA_names(this->NumRealComps()); + } + + std::vector + get_RealAoS_names () + { + std::vector real_aos_names(RealAoS::names_s.size()); + + // compile-time attributes + std::copy(RealAoS::names_s.begin(), RealAoS::names_s.end(), real_aos_names.begin()); + + return real_aos_names; + } + + std::vector + get_RealSoA_names (int num_real_comps) + { + std::vector real_soa_names(num_real_comps); + + // compile-time attributes + std::copy(RealSoA::names_s.begin(), RealSoA::names_s.end(), real_soa_names.begin()); + + // runtime attributes + if (num_real_comps > int(RealSoA::names_s.size())) + { + // particles lost record their "s" position where they got lost + real_soa_names[RealSoA::nattribs] = "s_lost"; + } + + return real_soa_names; + } } // namespace impactx diff --git a/src/particles/elements/All.H b/src/particles/elements/All.H index ab1ce8870..44e6bc7e7 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 "Buncher.H" #include "CFbend.H" #include "ChrDrift.H" @@ -43,6 +44,7 @@ namespace impactx { using KnownElements = std::variant< None, /* must be first, so KnownElements creates a default constructor */ + Aperture, Buncher, CFbend, ChrAcc, diff --git a/src/particles/elements/Aperture.H b/src/particles/elements/Aperture.H new file mode 100644 index 000000000..3325b7718 --- /dev/null +++ b/src/particles/elements/Aperture.H @@ -0,0 +1,115 @@ +/* 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; + + enum Shape + { + rectangular, + elliptical + }; + + /** A thin collimator element that applies a transverse aperture boundary. + * Particles outside the boundary are considered lost. + * + * @param shape aperture shape + * @param xmax maximum value of horizontal coordinate (m) + * @param ymax maximum value of vertical coordinate (m) + */ + Aperture (amrex::ParticleReal xmax, + amrex::ParticleReal ymax, + Shape shape) + : 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, + [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT px, + [[maybe_unused]] amrex::ParticleReal & AMREX_RESTRICT py, + [[maybe_unused]] 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); + auto const id = p.id(); + + // 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) + { + case Shape::rectangular : // default + if (pow(u,2)>1 || pow(v,2) > 1_prt) { + p.id() = -id; + } + break; + + case Shape::elliptical : + if (pow(u,2)+pow(v,2) > 1_prt) { + p.id() = -id; + } + break; + } + } + + /** This pushes the reference particle. */ + using Thin::operator(); + + private: + Shape m_shape; //! aperture type (rectangular, 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/particles/elements/diagnostics/openPMD.cpp b/src/particles/elements/diagnostics/openPMD.cpp index c2fd6cbcc..b325e94e9 100644 --- a/src/particles/elements/diagnostics/openPMD.cpp +++ b/src/particles/elements/diagnostics/openPMD.cpp @@ -254,8 +254,7 @@ namespace detail // AoS: Real { - std::vector real_aos_names(RealAoS::names_s.size()); - std::copy(RealAoS::names_s.begin(), RealAoS::names_s.end(), real_aos_names.begin()); + std::vector real_aos_names = get_RealAoS_names(); for (auto real_idx=0; real_idx < RealAoS::nattribs; real_idx++) { auto const component_name = real_aos_names.at(real_idx); getComponentRecord(component_name).resetDataset(d_fl); @@ -292,9 +291,8 @@ namespace detail // SoA: Real { - std::vector real_soa_names(RealSoA::names_s.size()); - std::copy(RealSoA::names_s.begin(), RealSoA::names_s.end(), real_soa_names.begin()); - for (auto real_idx = 0; real_idx < RealSoA::nattribs; real_idx++) { + std::vector real_soa_names = get_RealSoA_names(pc.NumRealComps()); + for (auto real_idx = 0; real_idx < pc.NumRealComps(); real_idx++) { auto const component_name = real_soa_names.at(real_idx); getComponentRecord(component_name).resetDataset(d_fl); } @@ -425,10 +423,8 @@ namespace detail auto const& soa = pti.GetStructOfArrays(); // SoA floating point (ParticleReal) properties { - std::vector real_soa_names(RealSoA::names_s.size()); - std::copy(RealSoA::names_s.begin(), RealSoA::names_s.end(), real_soa_names.begin()); - - for (auto real_idx=0; real_idx < RealSoA::nattribs; real_idx++) { + std::vector real_soa_names = get_RealSoA_names(soa.NumRealComps()); + for (auto real_idx=0; real_idx < soa.NumRealComps(); real_idx++) { auto const component_name = real_soa_names.at(real_idx); getComponentRecord(component_name).storeChunkRaw( soa.GetRealData(real_idx).data(), {offset}, {numParticleOnTile64}); diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index 8d9fa61b8..9ba47bf32 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -260,7 +260,7 @@ void init_ImpactX (py::module& m) "Enable or disable diagnostics every slice step in elements (default: disabled).\n\n" "By default, diagnostics is performed at the beginning and end of the simulation.\n" "Enabling this flag will write diagnostics every step and slice step." - ) + ) .def_property("diag_file_min_digits", [](ImpactX & /* ix */) { return detail::get_or_throw("diag", "file_min_digits"); @@ -271,7 +271,18 @@ void init_ImpactX (py::module& m) }, "The minimum number of digits (default: 6) used for the step\n" "number appended to the diagnostic file names." - ) + ) + .def_property("particle_lost_diagnostics_backend", + [](ImpactX & /* ix */) { + return detail::get_or_throw("diag", "backend"); + }, + [](ImpactX & /* ix */, std::string const backend) { + amrex::ParmParse pp_diag("diag"); + pp_diag.add("backend", backend); + }, + "Diagnostics for particles lost in apertures.\n\n" + "See the ``BeamMonitor`` element for backend values." + ) .def_property("abort_on_warning_threshold", [](ImpactX & /* ix */){ return detail::get_or_throw("impactx", "abort_on_warning_threshold"); diff --git a/src/python/ImpactXParticleContainer.cpp b/src/python/ImpactXParticleContainer.cpp index 91561faed..c1a645b8a 100644 --- a/src/python/ImpactXParticleContainer.cpp +++ b/src/python/ImpactXParticleContainer.cpp @@ -22,40 +22,6 @@ using namespace impactx; void init_impactxparticlecontainer(py::module& m) { - py::class_(m, "RealAoS") - .def_property_readonly_static("names_s", - [](py::object) { - std::vector real_aos_names(RealAoS::names_s.size()); - std::copy(RealAoS::names_s.begin(), RealAoS::names_s.end(), real_aos_names.begin()); - return real_aos_names; - }, - "named labels for fixed s") - .def_property_readonly_static("names_t", - [](py::object) { - std::vector real_aos_names(RealAoS::names_t.size()); - std::copy(RealAoS::names_t.begin(), RealAoS::names_t.end(), real_aos_names.begin()); - return real_aos_names; - }, - "named labels for fixed t") - ; - - py::class_(m, "RealSoA") - .def_property_readonly_static("names_s", - [](py::object) { - std::vector real_soa_names(RealSoA::names_s.size()); - std::copy(RealSoA::names_s.begin(), RealSoA::names_s.end(), real_soa_names.begin()); - return real_soa_names; - }, - "named labels for fixed s") - .def_property_readonly_static("names_t", - [](py::object) { - std::vector real_soa_names(RealSoA::names_t.size()); - std::copy(RealSoA::names_t.begin(), RealSoA::names_t.end(), real_soa_names.begin()); - return real_soa_names; - }, - "named labels for fixed t") - ; - py::class_< ParIter, amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs> @@ -150,5 +116,18 @@ void init_impactxparticlecontainer(py::module& m) "Charge deposition" ) */ + + .def_property_readonly("RealAoS_names", &ImpactXParticleContainer::RealAoS_names, + "Get the name of each Real AoS component") + + .def_property_readonly("RealSoA_names", &ImpactXParticleContainer::RealSoA_names, + "Get the name of each Real SoA component") ; + + m.def("get_RealAoS_names", &get_RealAoS_names, + "Get the name of each Real AoS component"); + + m.def("get_RealSoA_names", &get_RealSoA_names, + py::arg("num_real_comps"), + "Get the name of each Real SoA component\n\nnum_real_comps: pass number of compile-time + runtime arrays"); } diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 0305a637d..03012e6c5 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -123,6 +123,27 @@ void init_elements(py::module& m) // beam optics + py::class_ py_Aperture(me, "Aperture"); + py_Aperture + .def(py::init([]( + amrex::ParticleReal xmax, + amrex::ParticleReal ymax, + std::string shape) + { + if (shape != "rectangular" && shape != "elliptical") + throw std::runtime_error("shape must be \"rectangular\" or \"elliptical\""); + + Aperture::Shape s = shape == "rectangular" ? + Aperture::Shape::rectangular : + Aperture::Shape::elliptical; + return new Aperture(xmax, ymax, s); + }), + py::arg("xmax"), py::arg("ymax"), py::arg("shape") = "rectangular", + "A short collimator element applying a transverse aperture boundary." + ) + ; + register_beamoptics_push(py_Aperture); + py::class_ py_ChrDrift(me, "ChrDrift"); py_ChrDrift .def(py::init< diff --git a/src/python/impactx/ImpactXParticleContainer.py b/src/python/impactx/ImpactXParticleContainer.py index 9fcb56f9d..8547348db 100644 --- a/src/python/impactx/ImpactXParticleContainer.py +++ b/src/python/impactx/ImpactXParticleContainer.py @@ -13,8 +13,8 @@ def ix_pc_to_df(self, local=True, comm=None, root_rank=0): Parameters ---------- - self : amrex.ParticleContainer_* - A ParticleContainer class in pyAMReX + self : ImpactXParticleContainer_* + The particle container class in ImpactX local : bool MPI-local particles comm : MPI Communicator @@ -36,17 +36,14 @@ def ix_pc_to_df(self, local=True, comm=None, root_rank=0): # todo: check if currently in fixed s or fixed t and pick name accordingly names = [] - for n in self.RealAoS.names_s: + for n in self.RealAoS_names: names.append(n) names.append("cpuid") - for n in self.RealSoA.names_s: + for n in self.RealSoA_names: names.append(n) df.columns.values[0 : len(names)] = names - # todo: also rename runtime attributes (e.g., "s_lost") - # https://github.com/ECP-WarpX/impactx/pull/398 - return df