Element to apply a simple x-y rotation (#747)

* Add element for simple x-y rotation.

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

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

* Add documentation.

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

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

* Add ParmParse for element intialization.

---------

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

docs/source/usage/parameters.rst

@@ -378,6 +378,13 @@ Lattice Elements
             * ``<element_name>.phi_in`` (``float``, in degrees) angle of the reference particle with respect to the longitudinal (z) axis in the original frame
             * ``<element_name>.phi_out`` (``float``, in degrees) angle of the reference particle with respect to the longitudinal (z) axis in the rotated frame
 
+        * ``plane_xyrotation`` for a rotation in the x-y plane (i.e., about the reference velocity vector). This requires these additional parameters:
+
+            * ``<element_name>.angle`` (``float``, in degrees) nominal angle of rotation
+            * ``<element_name>.dx`` (``float``, in meters) horizontal translation error
+            * ``<element_name>.dy`` (``float``, in meters) vertical translation error
+            * ``<element_name>.rotation`` (``float``, in degrees) rotation error in the transverse plane
+
         * ``kicker`` for a thin transverse kicker. This requires these additional parameters:
 
             * ``<element_name>.xkick`` (``float``, dimensionless OR in T-m) the horizontal kick strength

docs/source/usage/python.rst

@@ -934,6 +934,16 @@ This module provides elements for the accelerator lattice.
    :param phi_out: angle of the reference particle with respect to the longitudinal (z) axis in the rotated frame in degrees
    :param name: an optional name for the element
 
+.. py:class:: impactx.elements.PlaneXYRot(angle,  dx=0, dy=0, rotation=0, name=None)
+
+   Map for a transverse rotation in the x-y plane (i.e., about the reference velocity vector).
+
+   :param angle: nominal angle of rotation in the x-y plane, in degrees
+   :param dx: horizontal translation error in m
+   :param dy: vertical translation error in m
+   :param rotation: rotation error in the transverse plane [degrees]
+   :param name: an optional name for the element
+
 .. py:class:: impactx.elements.Aperture(xmax, ymax, shape="rectangular", dx=0, dy=0, rotation=0, name=None)
 
    A thin collimator element, applying a transverse aperture boundary.

examples/CMakeLists.txt

@@ -1008,3 +1008,19 @@ add_impactx_test(aperture-periodic.py
     examples/aperture/analysis_aperture_periodic.py
     OFF  # no plot script yet
 )
+
+# Rotation in the transverse plane #########################################################
+#
+# w/o space charge
+add_impactx_test(rotation-xy
+    examples/rotation/input_rotation_xy.in
+    ON   # ImpactX MPI-parallel
+    examples/rotation/analysis_rotation_xy.py
+    OFF  # no plot script yet
+)
+add_impactx_test(rotation-xy.py
+    examples/rotation/run_rotation_xy.py
+    OFF  # ImpactX MPI-parallel
+    examples/rotation/analysis_rotation_xy.py
+    OFF  # no plot script yet
+)

examples/rotation/README.rst

@@ -48,3 +48,54 @@ We run the following script to analyze correctness:
    .. literalinclude:: analysis_rotation.py
       :language: python3
       :caption: You can copy this file from ``examples/rotation/analysis_rotation.py``.
+
+
+.. _examples-rotation-xy:
+
+Simple rotation in the x-y plane
+=================================
+
+A beam is rotated counter-clockwise by 90 degrees about the direction of motion.
+
+We use a 2 GeV electron beam.
+
+The second beam moments associated with the phase planes (x,px) and (y,py) are unequal, and these are exchanged under a 90 degree rotation.
+
+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.
+
+
+Run
+---
+
+This example can be run **either** as:
+
+* **Python** script: ``python3 run_rotation_xy.py`` or
+* ImpactX **executable** using an input file: ``impactx input_rotation_xy.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_rotation_xy.py
+          :language: python3
+          :caption: You can copy this file from ``examples/rotation/run_rotation_xy.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_rotation_xy.in
+          :language: ini
+          :caption: You can copy this file from ``examples/rotation/input_rotation_xy.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_rotation_xy.py``
+
+   .. literalinclude:: analysis_rotation_xy.py
+      :language: python3
+      :caption: You can copy this file from ``examples/rotation/analysis_rotation_xy.py``.

examples/rotation/analysis_rotation_xy.py

@@ -0,0 +1,98 @@
+#!/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 = 1.5e12 * 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],
+    [
+        4.0e-03,
+        2.0e-03,
+        1.0e-03,
+        8.0e-07,
+        1.0e-06,
+        2.0e-06,
+    ],
+    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 = 1.5 * 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],
+    [
+        2.0e-03,
+        4.0e-03,
+        1.0e-03,
+        1.0e-06,
+        8.0e-07,
+        2.0e-06,
+    ],
+    rtol=rtol,
+    atol=atol,
+)

examples/rotation/input_rotation_xy.in

@@ -0,0 +1,33 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.kin_energy = 2.0e3
+beam.charge = 1.0e-9
+beam.particle = electron
+beam.distribution = waterbag
+beam.lambdaX = 4.0e-3
+beam.lambdaY = 2.0e-3
+beam.lambdaT = 1.0e-3
+beam.lambdaPx = 2.0e-4
+beam.lambdaPy = 5.0e-4
+beam.lambdaPt = 2.0e-3
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor rotation1 monitor
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+rotation1.type = plane_xyrotation
+rotation1.angle = 90.0
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false

examples/rotation/run_rotation_xy.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Ryan Sandberg, Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+from impactx import ImpactX, 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 electron beam with an initial
+# unnormalized rms emittance of  nm
+kin_energy_MeV = 2.0e3  # reference energy
+bunch_charge_C = 1.0e-9  # used without 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(0.510998950).set_kin_energy_MeV(kin_energy_MeV)
+
+#   particle bunch
+distr = distribution.Waterbag(
+    lambdaX=4.0e-3,
+    lambdaY=2.0e-3,
+    lambdaT=1.0e-3,
+    lambdaPx=2.0e-4,
+    lambdaPy=5.0e-4,
+    lambdaPt=2.0e-3,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+
+# 90 degree rotation
+rotation = elements.PlaneXYRot(name="rotation1", angle=90.0)
+
+# assign a lattice segment
+sim.lattice.append(monitor)
+sim.lattice.append(rotation)
+sim.lattice.append(monitor)
+
+# run simulation
+sim.track_particles()
+
+# clean shutdown
+sim.finalize()

src/initialization/InitElement.cpp

@@ -245,6 +245,14 @@ namespace detail
             pp_element.get("phi_out", phi_out);
 
             m_lattice.emplace_back( PRot(phi_in, phi_out, element_name) );
+        } else if (element_type == "plane_xyrotation")
+        {
+            auto a = detail::query_alignment(pp_element);
+
+            amrex::ParticleReal phi;
+            pp_element.get("angle", phi);
+
+            m_lattice.emplace_back( PlaneXYRot(phi, a["dx"], a["dy"], a["rotation_degree"], element_name) );
         } else if (element_type == "solenoid_softedge")
         {
             auto const [ds, nslice] = detail::query_ds(pp_element, nslice_default);

src/particles/elements/All.H

@@ -27,6 +27,7 @@
 #include "Multipole.H"
 #include "Empty.H"
 #include "NonlinearLens.H"
+#include "PlaneXYRot.H"
 #include "Programmable.H"
 #include "Quad.H"
 #include "RFCavity.H"
@@ -64,6 +65,7 @@ namespace impactx
         Marker,
         Multipole,
         NonlinearLens,
+        PlaneXYRot,
         Programmable,
         PRot,
         Quad,