diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 9d4c63620..7a654075d 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -378,6 +378,13 @@ Lattice Elements * ``.phi_in`` (``float``, in degrees) angle of the reference particle with respect to the longitudinal (z) axis in the original frame * ``.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: + + * ``.angle`` (``float``, in degrees) nominal angle of rotation + * ``.dx`` (``float``, in meters) horizontal translation error + * ``.dy`` (``float``, in meters) vertical translation error + * ``.rotation`` (``float``, in degrees) rotation error in the transverse plane + * ``kicker`` for a thin transverse kicker. This requires these additional parameters: * ``.xkick`` (``float``, dimensionless OR in T-m) the horizontal kick strength diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 95e774ab1..2e24c86fb 100644 --- a/docs/source/usage/python.rst +++ b/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. diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 8da70838a..d9ade3e5a 100644 --- a/examples/CMakeLists.txt +++ b/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 +) diff --git a/examples/rotation/README.rst b/examples/rotation/README.rst index 621357f67..2f8369818 100644 --- a/examples/rotation/README.rst +++ b/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 `__ 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``. diff --git a/examples/rotation/analysis_rotation_xy.py b/examples/rotation/analysis_rotation_xy.py new file mode 100755 index 000000000..b40a3ed5d --- /dev/null +++ b/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, +) diff --git a/examples/rotation/input_rotation_xy.in b/examples/rotation/input_rotation_xy.in new file mode 100644 index 000000000..2b129487b --- /dev/null +++ b/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 diff --git a/examples/rotation/run_rotation_xy.py b/examples/rotation/run_rotation_xy.py new file mode 100644 index 000000000..586832a6d --- /dev/null +++ b/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() diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index d666322b5..148e719d7 100644 --- a/src/initialization/InitElement.cpp +++ b/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); diff --git a/src/particles/elements/All.H b/src/particles/elements/All.H index ba108c9fe..fbd0edc10 100644 --- a/src/particles/elements/All.H +++ b/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, diff --git a/src/particles/elements/PlaneXYRot.H b/src/particles/elements/PlaneXYRot.H new file mode 100644 index 000000000..9070afb81 --- /dev/null +++ b/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 + +#include +#include +#include + +#include + + +namespace impactx +{ + struct PlaneXYRot + : public elements::Named, + public elements::BeamOptic, + 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 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 diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 4d162573a..fc60ad89a 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -894,6 +894,38 @@ void init_elements(py::module& m) ; register_beamoptics_push(py_NonlinearLens); + py::class_ py_PlaneXYRot(me, "PlaneXYRot"); + py_PlaneXYRot + .def("__repr__", + [](PlaneXYRot const & plane_xyrot) { + return element_name( + plane_xyrot, + std::make_pair("angle", plane_xyrot.m_phi) + ); + } + ) + .def(py::init< + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + std::optional + >(), + py::arg("angle"), + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("name") = py::none(), + "A rotation in the x-y plane." + ) + .def_property("angle", + [](PlaneXYRot & plane_xyrot) { return plane_xyrot.m_phi; }, + [](PlaneXYRot & plane_xyrot, amrex::ParticleReal phi) { plane_xyrot.m_phi = phi; }, + "Rotation angle (rad)." + ) + ; + register_beamoptics_push(py_PlaneXYRot); + py::class_(me, "Programmable", py::dynamic_attr()) .def("__repr__", [](Programmable const & prg) {