Element to apply a simple x-y rotation

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/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/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,

src/particles/elements/PlaneXYRot.H

@@ -0,0 +1,139 @@
+/* Copyright 2022-2023 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Chad Mitchell, Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACTX_PLANE_XYROT_H
+#define IMPACTX_PLANE_XYROT_H
+
+#include "particles/ImpactXParticleContainer.H"
+#include "mixin/alignment.H"
+#include "mixin/beamoptic.H"
+#include "mixin/thin.H"
+#include "mixin/named.H"
+#include "mixin/nofinalize.H"
+
+#include <ablastr/constant.H>
+
+#include <AMReX_Extension.H>
+#include <AMReX_Math.H>
+#include <AMReX_REAL.H>
+
+#include <cmath>
+
+
+namespace impactx
+{
+    struct PlaneXYRot
+    : public elements::Named,
+      public elements::BeamOptic<PlaneXYRot>,
+      public elements::Thin,
+      public elements::Alignment,
+      public elements::NoFinalize
+    {
+        static constexpr auto type = "PlaneXYRot";
+        using PType = ImpactXParticleContainer::ParticleType;
+
+        static constexpr amrex::ParticleReal degree2rad = ablastr::constant::math::pi / 180.0;
+
+        /** A simple rotation by an angle phi in the plane transverse to the velocity
+         *  of the reference particle (x-y plane).  This is the linear symplectic map
+         *  generated by the longitudinal angular momentum J_z.
+         *
+         * @param phi angle of rotation in the x-y plane (about the direction of the ref traj) (degrees)
+         * @param dx horizontal translation error in m
+         * @param dy vertical translation error in m
+         * @param rotation_degree rotation error in the transverse plane [degrees]
+         * @param name a user defined and not necessarily unique name of the element
+         */
+        PlaneXYRot (
+            amrex::ParticleReal phi,
+            amrex::ParticleReal dx = 0,
+            amrex::ParticleReal dy = 0,
+            amrex::ParticleReal rotation_degree = 0,
+            std::optional<std::string> name = std::nullopt
+        )
+        : Named(name),
+          Alignment(dx, dy, rotation_degree),
+          m_phi(phi * degree2rad)
+        {
+        }
+
+        /** Push all particles */
+        using BeamOptic::operator();
+
+        /** This is a prot functor, so that a variable of this type can be used like a
+         *  prot function.
+         *
+         * @param x particle position in x
+         * @param y particle position in y
+         * @param t particle position in t
+         * @param px particle momentum in x
+         * @param py particle momentum in y
+         * @param pt particle momentum in t
+         * @param idcpu particle global index (unused)
+         * @param refpart reference particle
+         */
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        void operator() (
+            amrex::ParticleReal & AMREX_RESTRICT x,
+            amrex::ParticleReal & AMREX_RESTRICT y,
+            amrex::ParticleReal & AMREX_RESTRICT t,
+            amrex::ParticleReal & AMREX_RESTRICT px,
+            amrex::ParticleReal & AMREX_RESTRICT py,
+            amrex::ParticleReal & AMREX_RESTRICT pt,
+            [[maybe_unused]] uint64_t & AMREX_RESTRICT idcpu,
+            [[maybe_unused]] RefPart const & refpart
+        ) const
+        {
+            using namespace amrex::literals; // for _rt and _prt
+
+            // shift due to alignment errors of the element
+            shift_in(x, y, px, py);
+
+            // intialize output values
+            amrex::ParticleReal xout = x;
+            amrex::ParticleReal yout = y;
+            amrex::ParticleReal tout = t;
+            amrex::ParticleReal pxout = px;
+            amrex::ParticleReal pyout = py;
+            amrex::ParticleReal ptout = pt;
+
+            // store sin(phi) and cos(phi)
+            auto const [sin_phi, cos_phi] = amrex::Math::sincos(m_phi);
+
+            // advance position and momentum
+            xout = x*cos_phi - y*sin_phi;
+            pxout = px*cos_phi - py*sin_phi;
+
+            yout = x*sin_phi + y*cos_phi;
+            pyout = px*sin_phi + py*cos_phi;
+
+            // tout = t;
+            // ptout = pt;
+
+            // assign updated values
+            x = xout;
+            y = yout;
+            t = tout;
+            px = pxout;
+            py = pyout;
+            pt = ptout;
+
+            // undo shift due to alignment errors of the element
+            shift_out(x, y, px, py);
+        }
+
+        /** This pushes the reference particle. */
+        using Thin::operator();
+
+        amrex::ParticleReal m_phi; //! angle of rotation (rad)
+    };
+
+} // namespace impactx
+
+#endif // IMPACTX_PLANE_XYROT_H