Transverse Alignment Support (#490)

* Transversal Alignment: Mixin Class Create a mixin class that can be used for transversal alignment offsets and tilts. Co-authored-by: Chad Mitchell <[email protected]> * Update Elements for `Alignment` Support Co-authored-by: Chad Mitchell <[email protected]> * Inputs: `ParmParse` w/ Alignment * Inputs: Python w/ Alignment * Docs: Alignment Input * Reword Rotation Error Docstring Now: `rotation error in the transverse plane` * Add misalignment example. * Alignment Example: Docs & CMake * Fix Python Example: Protons --------- Co-authored-by: Chad Mitchell <[email protected]>
BLAST-ImpactX · Jan 3, 2024 · 0cc637c · 0cc637c
1 parent 8d64edb
commit 0cc637c
Show file tree

Hide file tree

Showing 35 changed files with 1,657 additions and 396 deletions.
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -65,9 +65,9 @@ Usage
    :hidden:
 
    usage/how_to_run
+   usage/examples
    usage/python
    usage/parameters
-   usage/examples
    usage/workflows
 
 Data Analysis

diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst
@@ -13,6 +13,7 @@ This section allows you to **download input files** that correspond to different
    examples/cfchannel/README.rst
    examples/expanding_beam/README.rst
    examples/kurth/README.rst
+   examples/alignment/README.rst
    examples/rfcavity/README.rst
    examples/fodo_rf/README.rst
    examples/fodo_chromatic/README.rst

diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst
diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst
diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -780,3 +780,21 @@ add_impactx_test(apochromat.py
     examples/apochromatic/analysis_apochromatic.py
     OFF  # no plot script yet
 )
+
+# Misalignment Quad ###########################################################
+#
+# w/o space charge
+add_impactx_test(alignment
+    examples/alignment/input_alignment.in
+      ON   # ImpactX MPI-parallel
+      OFF  # ImpactX Python interface
+    examples/alignment/analysis_alignment.py
+    OFF  # no plot script yet
+)
+add_impactx_test(alignment.py
+    examples/alignment/run_alignment.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/alignment/analysis_alignment.py
+    OFF  # no plot script yet
+)
diff --git a/examples/alignment/README.rst b/examples/alignment/README.rst
@@ -0,0 +1,49 @@
+.. _examples-alignment:
+
+Quadrupole with Alignment Errors
+================================
+
+A 2 GeV proton beam propagates through a single quadrupole with 3 mm horizontal misalignment and 30 degree rotation error.
+
+The first and second moments of the particle distribution before and after the quadrupole should coincide with
+analytical predictions, to within the level expected due to noise due to statistical sampling.
+
+In this test, the initial and final values of :math:`\mu_x`, :math:`\mu_y`, :math:`\sigma_x`, :math:`\sigma_y`,
+:math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values.
+
+
+Run
+---
+
+This example can be run **either** as:
+
+* **Python** script: ``python3 run_alignment.py`` or
+* ImpactX **executable** using an input file: ``impactx input_alignment.in``
+
+For `MPI-parallel <https://www.mpi-forum.org>`__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
+
+.. tab-set::
+
+   .. tab-item:: Python: Script
+
+       .. literalinclude:: run_alignment.py
+          :language: python3
+          :caption: You can copy this file from ``examples/alignment/run_alignment.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_alignment.in
+          :language: ini
+          :caption: You can copy this file from ``examples/alignment/input_alignment.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_alignment.py``
+
+   .. literalinclude:: analysis_alignment.py
+      :language: python3
+      :caption: You can copy this file from ``examples/alignment/analysis_alignment.py``.
diff --git a/examples/alignment/analysis_alignment.py b/examples/alignment/analysis_alignment.py
@@ -0,0 +1,137 @@
+#!/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, tmean
+
+
+def get_moments(beam):
+    """Calculate standard deviations of beam position & momenta
+    and emittance values
+
+    Returns
+    -------
+    meanx, meany, sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
+    """
+    meanx = tmean(beam["position_x"])
+    meany = tmean(beam["position_y"])
+    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 (meanx, meany, 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()
+
+# compare number of particles
+num_particles = 100000
+assert num_particles == len(initial)
+assert num_particles == len(final)
+
+print("Initial Beam:")
+meanx, meany, sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(
+    initial
+)
+print(f"  meanx={meanx:e} meany={meany:e}")
+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 = 3.0 * num_particles ** (-0.5) * sigx
+print(f" atol~={atol}")
+
+assert np.allclose(
+    [meanx, meany],
+    [
+        0.0,
+        0.0,
+    ],
+    atol=atol,
+)
+
+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.160982600086e-3,
+        1.160982600086e-3,
+        1.0e-3,
+        6.73940299e-7,
+        6.73940299e-7,
+        2.0e-6,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+
+print("")
+print("Final Beam:")
+meanx, meany, sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(
+    final
+)
+print(f"  meanx={meanx:e} meany={meany:e}")
+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 = 3.0 * num_particles ** (-0.5) * sigx
+print(f" atol~={atol}")
+
+assert np.allclose(
+    [meanx, meany],
+    [
+        1.79719761842e-4,
+        3.24815908981e-4,
+    ],
+    atol=atol,
+)
+
+atol = 0.0  # ignored
+rtol = 3.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.2372883901369e-3,
+        1.3772750218080e-3,
+        1.027364e-03,
+        7.39388142e-7,
+        7.39388142e-7,
+        2.0e-6,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
diff --git a/examples/alignment/input_alignment.in b/examples/alignment/input_alignment.in
@@ -0,0 +1,47 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 100000
+beam.units = static
+beam.kin_energy = 2.0e3
+beam.charge = 1.0e-9
+beam.particle = proton
+beam.distribution = waterbag
+beam.sigmaX = 1.16098260008648811e-3
+beam.sigmaY = 1.16098260008648811e-3
+beam.sigmaT = 1.0e-3
+beam.sigmaPx = 0.580491300043e-3
+beam.sigmaPy = 0.580491300043e-3
+beam.sigmaPt = 2.0e-3
+beam.muxpx = 0.0
+beam.muypy = 0.0
+beam.mutpt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor quad_err monitor
+lattice.nslice = 1
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+quad_err.type = quad
+quad_err.ds = 1.0
+quad_err.k = 0.25
+quad_err.dx = 0.003
+quad_err.rotation = 30.0
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
+
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true
diff --git a/examples/alignment/run_alignment.py b/examples/alignment/run_alignment.py
@@ -0,0 +1,63 @@
+#!/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
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a 2 GeV proton beam
+kin_energy_MeV = 2.0e3  # 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(
+    sigmaX=1.16098260008648811e-3,
+    sigmaY=1.16098260008648811e-3,
+    sigmaT=1.0e-3,
+    sigmaPx=0.580491300043e-3,
+    sigmaPy=0.580491300043e-3,
+    sigmaPt=2.0e-3,
+    muxpx=0.0,
+    muypy=0.0,
+    mutpt=0.0,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+
+# design the accelerator lattice)
+ns = 1  # number of slices per ds in the element
+lattice = [
+    monitor,
+    elements.Quad(ds=1.0, k=0.25, dx=0.003, dy=0.0, rotation=30.0, nslice=ns),
+    monitor,
+]
+sim.lattice.extend(lattice)
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+del sim
+amr.finalize()