User-Defined Linear Map

docs/source/usage/examples.rst

@@ -35,7 +35,7 @@ Single Particle Dynamics
    examples/fodo_programmable/README.rst
    examples/dogleg/README.rst
    examples/coupled_optics/README.rst
-
+   examples/linear_map/README.rst
 
 Collective Effects
 ------------------

examples/CMakeLists.txt

@@ -1071,3 +1071,19 @@ add_impactx_test(linac-segment.py
     examples/linac_segment/analysis_linac_segment.py
     OFF  # no plot script yet
 )
+
+# Iteration of a linear one-turn map #########################################
+#
+# w/o space charge
+add_impactx_test(linear-map
+    examples/linear_map/input_map.in
+    ON   # ImpactX MPI-parallel
+    examples/linear_map/analysis_map.py
+    OFF  # no plot script yet
+)
+add_impactx_test(linear-map.py
+    examples/linear_map/run_map.py
+    ON  # ImpactX MPI-parallel
+    examples/linear_map/analysis_map.py
+    OFF  # no plot script yet
+)

examples/linear_map/README.rst

@@ -0,0 +1,58 @@
+.. _examples-linear-map:
+
+Iteration of a user-defined linear map
+=======================================
+
+This example illustrates the application of a user-defined linear map via a matrix.
+
+Here, the linear map represents an abstract symplectic transformation of the beam in 6D phase space.
+If desired, the user may interpret the matrix as the one-turn map of a storage ring or circular collider.
+
+The (fractional) tunes (Qx, Qy, Qt) of the map are given by (0.139, 0.219, 0.0250).
+We use a 45.6 GeV electron beam that is invariant under the action of the linear map (matched).
+The horizontal and vertical unnormalized emittances are 0.27 nm and 1.0 pm, respectively.
+
+These parameters are based on the single-beam parameters of FCC-ee (Z-mode) appearing here:
+https://twiki.cern.ch/twiki/bin/view/FCC/FCCeeParameters_CDRBaseline-1_0
+
+The second moments of the phase space variables should be unchanged under application of the map.
+
+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 tunes associated with a single particle orbit are extracted, and must agree with the values given above.
+
+Run
+---
+
+This example can be run **either** as:
+
+* **Python** script: ``python3 run_multipole.py`` or
+* ImpactX **executable** using an input file: ``impactx input_multipole.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_multipole.py
+          :language: python3
+          :caption: You can copy this file from ``examples/multipole/run_multipole.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_multipole.in
+          :language: ini
+          :caption: You can copy this file from ``examples/multipole/input_multipole.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_multipole.py``
+
+   .. literalinclude:: analysis_multipole.py
+      :language: python3
+      :caption: You can copy this file from ``examples/multipole/analysis_multipole.py``.

examples/linear_map/analysis_map.py

@@ -0,0 +1,154 @@
+#!/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()
+
+# compare number of particles
+num_particles = 10000
+assert num_particles == len(initial)
+assert num_particles == 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.2 * 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],
+    [
+        6.363961030678928e-6,
+        28.284271247461902e-9,
+        0.0035,
+        0.27e-9,
+        1.0e-12,
+        1.33e-6,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+print("")
+print("Final Beam:")
+sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
+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.2 * 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],
+    [
+        6.363961030678928e-6,
+        28.284271247461902e-9,
+        0.0035,
+        0.27e-9,
+        1.0e-12,
+        1.33e-6,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+# Specify time series for particle j
+j = 5
+print(f"output for particle index = {j}")
+
+# Create array of TBT data values
+x = []
+px = []
+y = []
+py = []
+t = []
+pt = []
+n = 0
+for k_i, i in series.iterations.items():
+    beam = i.particles["beam"]
+    turn = beam.to_df()
+    x.append(turn["position_x"][j])
+    px.append(turn["momentum_x"][j])
+    y.append(turn["position_y"][j])
+    py.append(turn["momentum_y"][j])
+    t.append(turn["position_t"][j])
+    pt.append(turn["momentum_t"][j])
+    n = n + 1
+
+# Output number of periods in data series
+nturns = len(x)
+print(f"number of periods = {nturns}")
+print()
+
+# Approximate the tune and closed orbit using the 4-turn formula:
+
+# from x data only
+argument = (x[0] - x[1] + x[2] - x[3]) / (2.0 * (x[1] - x[2]))
+tunex = np.arccos(argument) / (2.0 * np.pi)
+print(f"tune output from 4-turn formula, using x data = {tunex}")
+
+# from y data only
+argument = (y[0] - y[1] + y[2] - y[3]) / (2.0 * (y[1] - y[2]))
+tuney = np.arccos(argument) / (2.0 * np.pi)
+print(f"tune output from 4-turn formula, using y data = {tuney}")
+
+# from t data only
+argument = (t[0] - t[1] + t[2] - t[3]) / (2.0 * (t[1] - t[2]))
+tunet = np.arccos(argument) / (2.0 * np.pi)
+print(f"tune output from 4-turn formula, using t data = {tunet}")
+
+rtol = 1.0e-3
+print(f"  rtol={rtol}")
+
+assert np.allclose(
+    [tunex, tuney, tunet],
+    [
+        0.139,
+        0.219,
+        0.0250,
+    ],
+    rtol=rtol,
+)

examples/linear_map/input_map.in

@@ -0,0 +1,57 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.kin_energy = 45.6e3
+beam.charge = 2.72370027e-8  #population 1.7e11
+beam.particle = electron
+beam.distribution = waterbag_from_twiss
+beam.alphaX = 0.0
+beam.alphaY = 0.0
+beam.alphaT = 0.0
+beam.betaX = 0.15
+beam.betaY = 0.8e-3
+beam.betaT = 9.210526315789473
+beam.emittX = 0.27e-09
+beam.emittY = 1e-12
+beam.emittT = 1.33e-6
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.periods = 5
+lattice.elements = monitor map1
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+map1.type = linear_map
+# horizontal plane
+map1.R11 = 0.642252653176584
+map1.R12 = 0.114973951021402
+map1.R21 = -5.109953378728999
+map1.R22 = 0.642252653176584
+# vertical plane
+map1.R33 = 0.193549468050860
+map1.R34 = 0.0007848724139547
+map1.R43 = -1226.363146804167548
+map1.R44 = 0.193549468050860
+# longitudinal plane
+map1.R55 = 0.987688340595138
+map1.R56 = 1.440843756949495
+map1.R65 = -0.016984313347225
+map1.R66 = 0.987688340595138
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
+
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = false