diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index 547f1a1b3..e19d73e47 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -28,6 +28,7 @@ This section allows you to **download input files** that correspond to different examples/cfbend/README.rst examples/compression/README.rst examples/kicker/README.rst + examples/thin_dipole/README.rst Unit tests diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 77d438d71..d60d2ca0e 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -509,6 +509,13 @@ Lattice Elements * ``.units`` (``string``) specification of units: ``dimensionless`` (default, in units of the magnetic rigidity of the reference particle) or ``T-m`` + * ``thin_dipole`` for a thin dipole element. + This requires these additional parameters: + + * ``.theta`` (``float``, in degrees) dipole bend angle + + * ``.rc`` (``float``, in meters) effective radius of curvature + * ``beam_monitor`` a beam monitor, writing all beam particles at fixed ``s`` to openPMD files. If the same element name is used multiple times, then an output series is created with multiple outputs. diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 94d250ad2..ad8057911 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -725,6 +725,15 @@ References: :param mapsteps: number of integration steps per slice used for map and reference particle push in applied fields :param nslice: number of slices used for the application of space charge +.. py:class:: impactx.elements.ThinDipole(theta, rc) + + A general thin dipole element. + + :param theta: Bend angle (degrees) + :param rc: Effective curvature radius (meters) + +Reference: + * G. Ripken and F. Schmidt, Thin-Lens Formalism for Tracking, CERN/SL/95-12 (AP), 1995. Coordinate Transformation ------------------------- diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 042b0c1ea..f57193e53 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -639,3 +639,21 @@ add_impactx_test(hvkicker_madx.py examples/kicker/analysis_kicker.py OFF # no plot script yet ) + +# Thin dipole ######################################################## +# +# w/o space charge +add_impactx_test(thin_dipole + examples/thin_dipole/input_thin_dipole.in + ON # ImpactX MPI-parallel + OFF # ImpactX Python interface + examples/thin_dipole/analysis_thin_dipole.py + OFF # no plot script yet +) +add_impactx_test(thin_dipole.py + examples/thin_dipole/run_thin_dipole.py + OFF # ImpactX MPI-parallel + ON # ImpactX Python interface + examples/thin_dipole/analysis_thin_dipole.py + OFF # no plot script yet +) diff --git a/examples/thin_dipole/README.rst b/examples/thin_dipole/README.rst new file mode 100644 index 000000000..258768709 --- /dev/null +++ b/examples/thin_dipole/README.rst @@ -0,0 +1,61 @@ +.. _examples-thin-dipole: + +Thin Dipole +=========== + +This test involves tracking a 5 MeV proton beam through a 90 degree sector bend, using two different methods: + +#. Using a sequence of drifts and thin kicks as described by: +G. Ripken and F. Schmidt, "A Symplectic Six-Dimensional Thin-Lens Formalism for Tracking," CERN/SL/95-12 (1995). + +#. Using the exact nonlinear transfer map for a sector bend as described by: +D. Bruhwiler et al, "Symplectic Propagation of the Map, Tangent Map and Tangent Map Derivative through +Quadrupole and Combined-Function Dipole Magnets without Truncation," THP41C, EPAC98, pp. 1171-1173 (1998). + +To compare the two models, the beam is tracked using 1), followed by the inverse of 2). +The beam moments and emittances should coincide with those of the initial distribution. + + +Run +--- + +This example can be run **either** as: + +* **Python** script: ``python3 run_thin_dipole.py`` or +* ImpactX **executable** using an input file: ``impactx input_thin_dipole.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_thin_dipole.py + :language: python3 + :caption: You can copy this file from ``examples/thin_dipole/run_thin_dipole.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_thin_dipole.in + :language: ini + :caption: You can copy this file from ``examples/thin_dipole/input_thin_dipole.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_thin_dipole.py`` + + .. literalinclude:: analysis_thin_dipole.py + :language: python3 + :caption: You can copy this file from ``examples/thin_dipole/analysis_thin_dipole.py``. + + +Visualize +--------- + +.. note:: + + TODO :) diff --git a/examples/thin_dipole/analysis_thin_dipole.py b/examples/thin_dipole/analysis_thin_dipole.py new file mode 100755 index 000000000..6559e2552 --- /dev/null +++ b/examples/thin_dipole/analysis_thin_dipole.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +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.8 * 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], + [ + 1.0e-3, + 1.0e-3, + 3.0e-1, + 2.0e-7, + 2.0e-7, + 6.0e-5, + ], + 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.8 * 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], + [ + 1.0e-3, + 1.0e-3, + 3.0e-1, + 2.0e-7, + 2.0e-7, + 6.0e-5, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/thin_dipole/input_thin_dipole.in b/examples/thin_dipole/input_thin_dipole.in new file mode 100644 index 000000000..6bbec0a48 --- /dev/null +++ b/examples/thin_dipole/input_thin_dipole.in @@ -0,0 +1,57 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 10000 +beam.units = static +beam.kin_energy = 5.0 +beam.charge = 1.0e-9 +beam.particle = proton +beam.distribution = waterbag +beam.sigmaX = 1.0e-3 +beam.sigmaY = 1.0e-3 +beam.sigmaT = 0.3 +beam.sigmaPx = 2.0e-4 +beam.sigmaPy = 2.0e-4 +beam.sigmaPt = 2.0e-4 +beam.muxpx = 0.0 +beam.muypy = 0.0 +beam.mutpt = 0.0 + + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = monitor bend inverse_bend monitor +lattice.nslice = 1 + +monitor.type = beam_monitor +monitor.backend = h5 + +#90 degree sbend using drift-kick-drift: +bend.type = line +bend.elements = dr kick dr +bend.repeat = 200 + +dr.type = drift +dr.ds = 0.003926990816987 + +kick.type = thin_dipole +kick.theta = 0.45 +kick.rc = 1.0 + +#inverse of a 90 degree sbend using exact nonlinear map: +inverse_bend.type = sbend_exact +inverse_bend.ds = -1.570796326794897 +inverse_bend.phi = -90.0 + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/thin_dipole/run_thin_dipole.py b/examples/thin_dipole/run_thin_dipole.py new file mode 100755 index 000000000..ebde9e7eb --- /dev/null +++ b/examples/thin_dipole/run_thin_dipole.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import amrex.space3d as amr +from impactx import ImpactX, RefPart, 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 initial beam +kin_energy_MeV = 5.0 # reference energy +bunch_charge_C = 1.0e-9 # used with 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV) + +# particle bunch +distr = distribution.Waterbag( + sigmaX=1.0e-3, + sigmaY=1.0e-3, + sigmaT=0.3, + sigmaPx=2.0e-4, + sigmaPy=2.0e-4, + sigmaPt=2.0e-4, + muxpx=-0.0, + muypy=0.0, + mutpt=0.0, +) +sim.add_particles(bunch_charge_C, distr, npart) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# design the accelerator lattice) +ns = 1 # number of slices per ds in the element +segment = [ + elements.Drift(ds=0.003926990816987, nslice=ns), + elements.ThinDipole(theta=0.45, rc=1.0), + elements.Drift(ds=0.003926990816987, nslice=ns), +] +bend = 200 * segment + +inverse_bend = elements.ExactSbend(ds=-1.570796326794897, phi=-90.0) + +sim.lattice.append(monitor) +sim.lattice.extend(bend) +sim.lattice.append(inverse_bend) +sim.lattice.append(monitor) + +# run simulation +sim.evolve() + +# clean shutdown +del sim +amr.finalize() diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index 78630acba..2d6931945 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -231,6 +231,11 @@ namespace detail pp_element.get("bz", bz); pp_element.queryAdd("nslice", nslice); m_lattice.emplace_back( ChrAcc(ds, ez, bz, nslice) ); + } else if (element_type == "thin_dipole") { + amrex::Real theta, rc; + pp_element.get("theta", theta); + pp_element.get("rc", rc); + m_lattice.emplace_back( ThinDipole(theta, rc) ); } else if (element_type == "kicker") { amrex::Real xkick, ykick; std::string units_str = "dimensionless"; diff --git a/src/particles/elements/All.H b/src/particles/elements/All.H index de2544ef7..ab1ce8870 100644 --- a/src/particles/elements/All.H +++ b/src/particles/elements/All.H @@ -33,6 +33,7 @@ #include "PRot.H" #include "SoftSol.H" #include "SoftQuad.H" +#include "ThinDipole.H" #include "diagnostics/openPMD.H" #include @@ -64,7 +65,8 @@ namespace impactx ShortRF, SoftSolenoid, SoftQuadrupole, - Sol + Sol, + ThinDipole >; } // namespace impactx diff --git a/src/particles/elements/ThinDipole.H b/src/particles/elements/ThinDipole.H new file mode 100644 index 000000000..5489d4b57 --- /dev/null +++ b/src/particles/elements/ThinDipole.H @@ -0,0 +1,120 @@ +/* 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_THINDIPOLE_H +#define IMPACTX_THINDIPOLE_H + +#include "particles/ImpactXParticleContainer.H" +#include "mixin/beamoptic.H" +#include "mixin/thin.H" +#include "mixin/nofinalize.H" + +#include +#include + +namespace impactx +{ + struct ThinDipole + : public elements::BeamOptic, + public elements::Thin, + public elements::NoFinalize + { + static constexpr auto name = "ThinDipole"; + using PType = ImpactXParticleContainer::ParticleType; + + static constexpr amrex::ParticleReal degree2rad = ablastr::constant::math::pi / 180.0; + + /** A general thin-kick dipole element with chromatic effects + * + * This model is equivalent to the model described in Section 3.1 of + * G. Ripken, F. Schmidt, "A Symplectic Six-Dimensional Thin-Lens Formalism + * for Tracking," CERN/SL/95-12 (AP), 1995. It is intended to replicate + * the thin-lens dipole model in MAD-X. + * + * @param theta - the total bending angle (degrees) + * @param rc - the curvature radius (m) + */ + ThinDipole( amrex::ParticleReal const theta, + amrex::ParticleReal const rc) + : m_theta(theta * degree2rad), m_rc(rc) + { + } + + /** Push all particles */ + using BeamOptic::operator(); + + /** This is a multipole functor, so that a variable of this type can be used like a + * multipole function. + * + * @param p Particle AoS data for positions and cpu/id + * @param px particle momentum in x + * @param py particle momentum in y + * @param pt particle momentum in t + * @param refpart reference particle + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void operator() ( + PType& AMREX_RESTRICT p, + amrex::ParticleReal & AMREX_RESTRICT px, + amrex::ParticleReal & AMREX_RESTRICT py, + amrex::ParticleReal & AMREX_RESTRICT pt, + RefPart const & refpart) const { + + using namespace amrex::literals; // for _rt and _prt + + // access AoS data such as positions and cpu/id + amrex::ParticleReal const x = p.pos(RealAoS::x); + amrex::ParticleReal const y = p.pos(RealAoS::y); + amrex::ParticleReal const t = p.pos(RealAoS::t); + + // access reference particle to find relativistic beta + amrex::ParticleReal const beta_ref = refpart.beta(); + + // intialize output values of momenta + amrex::ParticleReal pxout = px; + amrex::ParticleReal pyout = py; + amrex::ParticleReal ptout = pt; + + // compute the function expressing dp/p in terms of pt (labeled f in Ripken etc.) + amrex::ParticleReal f = -1.0_prt + sqrt(1.0_prt - 2.0_prt*pt/beta_ref + pow(pt,2)); + amrex::ParticleReal fprime = (1.0_prt - beta_ref*pt)/(beta_ref*(1.0_prt + f)); + + // compute the effective (equivalent) arc length and curvature + amrex::ParticleReal ds = m_theta*m_rc; + amrex::ParticleReal kx = 1.0_prt/m_rc; + + // advance position and momentum + p.pos(RealAoS::x) = x; + pxout = px - pow(kx,2)*ds*x + kx*ds*f; //eq (3.2b) + + p.pos(RealAoS::y) = y; + pyout = py; + + p.pos(RealAoS::t) = t + kx*x*ds*fprime; //eq (3.2e) + ptout = pt; + + // assign updated momenta + px = pxout; + py = pyout; + pt = ptout; + + } + + /** This pushes the reference particle. */ + using Thin::operator(); + + private: + amrex::ParticleReal m_theta; //! dipole bending angle (rad) + amrex::ParticleReal m_rc; //! curvature radius (m) + + }; + +} // namespace impactx + +#endif // IMPACTX_THINDIPOLE_H diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 328eb9fca..0305a637d 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -471,6 +471,18 @@ void init_elements(py::module& m) ; register_beamoptics_push(py_SoftQuadrupole); + py::class_ py_ThinDipole(me, "ThinDipole"); + py_ThinDipole + .def(py::init< + amrex::ParticleReal const, + amrex::ParticleReal const>(), + py::arg("theta"), py::arg("rc"), + "A thin kick model of a dipole bend." + ) + ; + register_beamoptics_push(py_ThinDipole); + + // free-standing push function m.def("push", &Push, py::arg("pc"), py::arg("element"), py::arg("step")=0,