From 6723fec929d24828acb80c6b0045558d7cb2db0b Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Mon, 18 Nov 2024 17:30:46 -0800 Subject: [PATCH] Add benchmark test. --- examples/CMakeLists.txt | 16 ++++ examples/aperture/README.rst | 54 ++++++++++++ examples/aperture/analysis_absorber.py | 110 +++++++++++++++++++++++++ examples/aperture/input_absorber.in | 50 +++++++++++ examples/aperture/run_absorber.py | 69 ++++++++++++++++ 5 files changed, 299 insertions(+) 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/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..5ca200de8 100644 --- a/examples/aperture/README.rst +++ b/examples/aperture/README.rst @@ -101,3 +101,57 @@ 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..9e33df183 --- /dev/null +++ b/examples/aperture/run_absorber.py @@ -0,0 +1,69 @@ +#!/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()