Solvable example of an expanding beam scraping an aperture (#813)

* Add scraping beam example.
* Add negligible p-spread in app input
* Add negligible p-spread in Python input
* Delete unnecessary emittance comparison
* Restore exact zero p-spread in app input
* Restore exact zero p-spread in Python input

docs/source/usage/examples.rst

@@ -36,6 +36,8 @@ Single Particle Dynamics
    examples/dogleg/README.rst
    examples/coupled_optics/README.rst
    examples/linear_map/README.rst
+   examples/scraping_beam/README.rst
+
 
 Collective Effects
 ------------------

examples/CMakeLists.txt

@@ -1185,3 +1185,19 @@ add_impactx_test(IOTA_nll_aperture.py
     examples/iota_lens/analysis_iotalens_sdep_aperture.py
     OFF  # no plot script yet
 )
+
+# Expanding beam scraping against a vacuum pipe ##############################
+#
+# w/o space charge
+add_impactx_test(examples-scraping
+    examples/scraping_beam/input_scraping.in
+    ON   # ImpactX MPI-parallel
+    examples/scraping_beam/analysis_scraping.py
+    OFF  # no plot script yet
+)
+add_impactx_test(examples-scraping.py
+    examples/scraping_beam/run_scraping.py
+    OFF   # ImpactX MPI-parallel
+    examples/scraping_beam/analysis_scraping.py
+    OFF  # no plot script yet
+)

examples/scraping_beam/README.rst

@@ -0,0 +1,74 @@
+.. _examples-scraping:
+
+Expanding Beam Scraping Against a Vacuum Pipe
+=============================================
+
+This example describes a coasting bunch, expanding transversely and encountering the aperture defined by the vacuum pipe.
+Space charge is neglected, making the problem analytically soluble.
+
+We use a cold (zero emittance) 250 MeV proton bunch whose
+initial distribution is a uniformly-populated cylinder of transverse radius :math:`r_b = 2 \mathrm{mm}` with zero momentum spread.
+
+The beam propagates in a drift with a vacuum chamber radius of :math:`R = 3.5 \mathrm{mm}`.
+
+To generate an expanding beam, a linear map is first applied.  This map applies a radial kick to each particle that is proportional to the particle's initial distance from the axis.
+This induces a phase space correlation within the beam, such that :math:`p_x = k \cdot x` and :math:`p_y = k \cdot y`, similar to what would be induced by a space charge kick.
+
+The beam remains cylindrical with zero emittance during its evolution in a 6 m drift.
+In the absence of an aperture, the beam radius evolves as :math:`r_b(s) = r_b(1 + k\cdot s)`.
+In the presence of an aperture, particles are lost during the transverse expansion.  The fraction of charge remaining after a distance s is given by:
+
+.. math::
+
+   \frac{Q_s}{Q_0} = \min\left[1,R^2/(r_b^2(1+s\cdot k)^2)\right].
+
+In this test, the initial and final 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.
+
+In addition, the initial and final values of :math:`\sigma_{p_x}`, :math:`\sigma_{p_y}`, and :math:`\sigma_{p_t}` must agree with nominal values.
+
+Finally, the fraction of charge lost against the aperture at the exit of the drift must agree with nominal values.
+
+The physical problem is defined by four relevant parameters, defined within ``run_scraping.py``, that can be modified by the user:
+
+.. code-block:: python
+
+   # problem parameters
+   beam_radius = r_b
+   aperture_radius = R
+   correlation_k = k
+   drift_distance = s
+
+These parameters should also be modified inside ``analysis_scraping.py`` for testing.
+
+
+Run
+---
+
+This example can be run as a Python script (``python3 run_scraping.py``) or with an app with an input file (``impactx input_scraping.in``).
+Each can also be prefixed with an `MPI executor <https://www.mpi-forum.org>`__, such as ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
+
+.. tab-set::
+
+   .. tab-item:: Python Script
+
+       .. literalinclude:: run_scraping.py
+          :language: python3
+          :caption: You can copy this file from ``examples/scraping_beam/run_scraping.py``.
+
+   .. tab-item:: App Input File
+
+       .. literalinclude:: input_scraping.in
+          :language: ini
+          :caption: You can copy this file from ``examples/scraping_beam/input_scraping.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_scraping.py``
+
+   .. literalinclude:: analysis_scraping.py
+      :language: python3
+      :caption: You can copy this file from ``examples/scraping_beam/analysis_scraping.py``.

examples/scraping_beam/analysis_scraping.py

@@ -0,0 +1,163 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+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, sigpx, sigpy, sigpt, 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,
+        sigpx,
+        sigpy,
+        sigpt,
+        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_beam = series.iterations[1].particles["beam"]
+final_beam = series.iterations[last_step].particles["beam"]
+initial = initial_beam.to_df()
+final = final_beam.to_df()
+
+# compare number of particles
+num_particles = 100000
+assert num_particles == len(initial)
+
+# problem parameters
+beam_radius = 2.0e-3
+aperture_radius = 3.5e-3
+correlation_k = 0.5
+drift_distance = 6.0
+
+print("Initial Beam:")
+sigx, sigy, sigt, sigpx, sigpy, sigpt, emittance_x, emittance_y, emittance_t = (
+    get_moments(initial)
+)
+print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
+print(f"  sigpx={sigpx:e} sigpy={sigpy:e} sigpt={sigpt:e}")
+print(
+    f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
+)
+
+atol = 0.0  # ignored
+rtol = 2.0 * num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sigx, sigy, sigt],
+    [
+        beam_radius / 2.0,
+        beam_radius / 2.0,
+        beam_radius / 2.0,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+atol = 1.0e-11  # ignored
+print(f"  atol={atol}")
+
+assert np.allclose(
+    [sigpx, sigpy, sigpt, emittance_x, emittance_y, emittance_t],
+    [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+    ],
+    atol=atol,
+)
+
+
+# calculation of predicted final beam parameters
+beam_radius_no_aperture = beam_radius * (1.0 + correlation_k * drift_distance)
+beam_radius_with_aperture = min(beam_radius_no_aperture, aperture_radius)
+
+fractional_loss = 1.0 - min(1.0, (aperture_radius / beam_radius_no_aperture) ** 2)
+sigma_x_final = beam_radius_with_aperture / 2.0
+sigma_px_final = correlation_k / (1.0 + correlation_k * drift_distance) * sigma_x_final
+
+print("")
+print("Predicted Final Beam:")
+print(f"  sigx={sigma_x_final:e} sigy={sigma_x_final:e} sigt={beam_radius / 2.0:e}")
+print(f"  sigpx={sigma_px_final:e} sigpy={sigma_px_final:e} sigpt=0.0")
+print(f"  fractional_loss={fractional_loss:e}")
+
+
+print("")
+print("Final Beam:")
+sigx, sigy, sigt, sigpx, sigpy, sigpt, emittance_x, emittance_y, emittance_t = (
+    get_moments(final)
+)
+print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
+print(f"  sigpx={sigpx:e} sigpy={sigpy:e} sigpt={sigpt:e}")
+print(
+    f"  emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
+)
+
+atol = 0.0  # ignored
+rtol = 2.0 * num_particles**-0.5  # from random sampling of a smooth distribution
+print(f"  rtol={rtol} (ignored: atol~={atol})")
+
+assert np.allclose(
+    [sigx, sigy, sigt, sigpx, sigpy],
+    [
+        sigma_x_final,
+        sigma_x_final,
+        beam_radius / 2.0,
+        sigma_px_final,
+        sigma_px_final,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+charge_i = initial_beam.get_attribute("charge_C")
+charge_f = final_beam.get_attribute("charge_C")
+
+loss_pct = 100.0 * (charge_i - charge_f) / charge_i
+
+print(f" fractional loss (%) = {loss_pct}")
+
+atol = 0.2  # tolerance 0.2%
+print(f"  atol={atol}")
+assert np.allclose(
+    [loss_pct],
+    [100 * fractional_loss],
+    atol=atol,
+)

examples/scraping_beam/input_scraping.in

@@ -0,0 +1,50 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 100000
+beam.units = static
+beam.kin_energy = 250.0
+beam.charge = 1.0e-9
+beam.particle = proton
+beam.distribution = kurth4d
+beam.lambdaX = 1.0e-3
+beam.lambdaY = beam.lambdaX
+beam.lambdaT = 1.0e-3
+beam.lambdaPx = 0.0
+beam.lambdaPy = 0.0
+beam.lambdaPt = 0.0
+beam.muxpx = 0.0
+beam.muypy = 0.0
+beam.mutpt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor map1 drift1 monitor
+lattice.nslice = 40
+
+map1.type = linear_map
+map1.R21 = 0.5
+map1.R43 = map1.R21
+
+drift1.type = drift
+drift1.ds = 6.0
+drift1.aperture_x = 3.5e-3
+drift1.aperture_y = 3.5e-3
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
+
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true