Add periodic masking option to Aperture. (#739)

* Add periodic masking option to Aperture.

* Add draft files for benchmark example.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix modular arithmetic.

* Add analysis script for testing.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Reduce particle population for tests.

* Update examples/aperture/README.rst

Correct type in example README.

* Update Aperture.H

Generalize periodic aperture to allow periodicity in one direction only.

* Rename analysis_aperture_pepperpot.py to analysis_aperture_periodic.py

Rename example:  analysis script.

* Rename input_aperture_pepperpot.in to input_aperture_periodic.in

Rename example:  C++ input

* Rename run_aperture_pepperpot.py to run_aperture_periodic.py

Rename example:  Python input

* Update CMakeLists.txt

Rename example:  CMakeLists

* Update README.rst

Rename example:  README documentation

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

docs/source/usage/parameters.rst

@@ -401,6 +401,8 @@ Lattice Elements
 
             * ``<element_name>.xmax`` (``float``, in meters) maximum value of the horizontal coordinate
             * ``<element_name>.ymax`` (``float``, in meters) maximum value of the vertical coordinate
+            * ``<element_name>.repeat_x`` (``float``, in meters) horizontal period for repeated aperture masking (inactive by default)
+            * ``<element_name>.repeat_y`` (``float``, in meters) vertical period for repeated aperture masking (inactive by default)
             * ``<element_name>.shape`` (``string``) shape of the aperture boundary: ``rectangular`` (default) or ``elliptical``
             * ``<element_name>.dx`` (``float``, in meters) horizontal translation error
             * ``<element_name>.dy`` (``float``, in meters) vertical translation error

docs/source/usage/python.rst

@@ -940,6 +940,8 @@ This module provides elements for the accelerator lattice.
 
    :param xmax: maximum allowed value of the horizontal coordinate (meter)
    :param ymax: maximum allowed value of the vertical coordinate (meter)
+   :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 dx: horizontal translation error in m
    :param dy: vertical translation error in m

examples/CMakeLists.txt

@@ -992,3 +992,19 @@ add_impactx_test(coupled-optics.py
     examples/coupled_optics/analysis_coupled_optics.py
     OFF  # no plot script yet
 )
+
+# Aperture with Periodic Masking #########################################################
+#
+# w/o space charge
+add_impactx_test(aperture-periodic
+    examples/aperture/input_aperture_periodic.in
+    ON   # ImpactX MPI-parallel
+    examples/aperture/analysis_aperture_periodic.py
+    OFF  # no plot script yet
+)
+add_impactx_test(aperture-periodic.py
+    examples/aperture/run_aperture_periodic.py
+    OFF  # ImpactX MPI-parallel
+    examples/aperture/analysis_aperture_periodic.py
+    OFF  # no plot script yet
+)

examples/aperture/README.rst

@@ -50,3 +50,54 @@ We run the following script to analyze correctness:
    .. literalinclude:: analysis_aperture.py
       :language: python3
       :caption: You can copy this file from ``examples/aperture/analysis_aperture.py``.
+
+
+.. _examples-aperture-periodic:
+
+Aperture Collimation with Periodic Masking
+===========================================
+
+Proton beam undergoing collimation by a periodic array of rectangular apertures, such as those used in a pepperpot emittance measurement.
+
+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 m, the beam intercepts a periodic array of apertures of period 1 mm (in the horizontal and vertical), each of which is 0.15 mm x 0.1 mm in size.
+
+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 lie inside the aperture boundary or
+* if the sum of the numbers of lost and retained particles is not equal to the number of initial particles
+
+
+Run
+---
+
+This example can be run as a Python script (``python3 run_aperture_periodic.py``) or with an app with an input file (``impactx input_aperture_periodic.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_aperture_periodic.py
+          :language: python3
+          :caption: You can copy this file from ``examples/aperture/run_aperture_periodic.py``.
+
+   .. tab-item:: App Input File
+
+       .. literalinclude:: input_aperture_periodic.in
+          :language: ini
+          :caption: You can copy this file from ``examples/aperture/input_aperture_periodic.in``.
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_aperture_periodic.py``
+
+   .. literalinclude:: analysis_aperture_periodic.py
+      :language: python3
+      :caption: You can copy this file from ``examples/aperture/analysis_aperture_periodic.py``.

examples/aperture/analysis_aperture_periodic.py

@@ -0,0 +1,120 @@
+#!/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 = 100000
+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 = 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, 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 periodic rectangular aperture boundary
+xmax = 1.5e-4
+ymax = 1.0e-4
+repeat_x = 1.0e-3
+repeat_y = 1.0e-3
+
+
+# kept particles, shifted to the fundamental domain
+xshifted = abs(final["position_x"]) + xmax
+yshifted = abs(final["position_y"]) + ymax
+u = np.fmod(xshifted, repeat_x) - xmax
+v = np.fmod(yshifted, repeat_y) - ymax
+
+# difference from maximum aperture
+dx = abs(u) - xmax
+dy = abs(v) - ymax
+
+print()
+print(f"  fundamental x_max={u.max()}")
+print(f"  fundamental x_min={u.min()}")
+assert np.less_equal(dx.max(), 0.0)
+
+print(f"  fundamental y_max={v.max()}")
+print(f"  fundamental y_min={v.min()}")
+assert np.less_equal(dy.max(), 0.0)
+
+# lost particles, shifted to the fundamental domain
+xshifted = abs(particles_lost["position_x"]) - xmax
+yshifted = abs(particles_lost["position_y"]) - ymax
+u = np.fmod(xshifted, repeat_x) - xmax
+v = np.fmod(yshifted, repeat_y) - ymax
+
+# difference from maximum aperture
+dx = abs(u) - xmax
+dy = abs(v) - ymax
+
+print()
+print(f"  fundamental x_max={u.max()}")
+print(f"  fundamental x_min={u.min()}")
+assert np.greater_equal(dx.max(), 0.0)
+
+print(f"  fundamental y_max={v.max()}")
+print(f"  fundamental y_min={v.min()}")
+assert np.greater_equal(dy.max(), 0.0)

examples/aperture/input_aperture_periodic.in

@@ -0,0 +1,52 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 100000
+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 pepperpot monitor
+lattice.nslice = 1
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+drift.type = drift
+drift.ds = 0.123
+
+pepperpot.type = aperture
+pepperpot.shape = rectangular
+pepperpot.xmax = 1.5e-4
+pepperpot.ymax = 1.0e-4
+pepperpot.repeat_x = 1.0e-3
+pepperpot.repeat_y = 1.0e-3
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
+
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true
+diag.backend = h5

examples/aperture/run_aperture_periodic.py

@@ -0,0 +1,74 @@
+#!/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 = 100000  # 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="pepperpot",
+            xmax=1.5e-4,
+            ymax=1.0e-4,
+            repeat_x=1.0e-3,
+            repeat_y=1.0e-3,
+            shape="rectangular",
+        ),
+        monitor,
+    ]
+)
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+sim.finalize()

src/initialization/InitElement.cpp

@@ -376,17 +376,21 @@ namespace detail
             auto a = detail::query_alignment(pp_element);
 
             amrex::Real xmax, ymax;
+            amrex::ParticleReal repeat_x = 0.0;
+            amrex::ParticleReal repeat_y = 0.0;
             std::string shape_str = "rectangular";
             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);
             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, a["dx"], a["dy"], a["rotation_degree"], element_name) );
+            m_lattice.emplace_back( Aperture(xmax, ymax, repeat_x, repeat_y, shape, a["dx"], a["dy"], a["rotation_degree"], element_name) );
         } else if (element_type == "beam_monitor")
         {
             std::string openpmd_name = element_name;