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/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/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..8cf0b1fa0 --- /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) {