diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 3723f33cf..14def9ae2 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -1121,6 +1121,22 @@ However, a Taylor expansion is used to evaluate the dependence on the quantum pa ISR effects are only calculated for lattice elements that include bending, such as ``Sbend``, ``ExactSbend`` and ``CFbend``. +.. _running-cpp-parameters-spin: + +Spin Tracking +^^^^^^^^^^^^^ + +Spin tracking is performed by updating the particle spin 3-vector in the presence of each element's electromagnetic fields, using methods based on the Thomas-BMT equation. +By construction, all spin tracking methods rely on pushing particles using spin maps that lie in SO(3). The algorithm implemented for each element is consistent with the algorithm +used for pushing the phase space vector. + +Currently, the implementation of spin tracking is a work in progress, and this feature is not yet supported. + +* ``algo.spin`` (``boolean``, optional, default: ``false``) + + Whether to track particle spin. + + .. _running-cpp-parameters-parser: Math parser and user-defined constants diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 68d51948f..49c6f122b 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -192,6 +192,13 @@ Collective Effects & Overall Simulation Parameters mean energy of the beam particles will decrease. This option is natural if the lattice optics, magnet settings, etc. are chosen without accounting for radiative energy loss. When ``sim.isr_on_ref_part = True``, the reference particle does lose energy due to radiation, and little centroid evolution is expected in the beam particles. This option is natural if the lattice optics, magnet settings, etc. are chosen to account for radiative energy loss. + .. py:property:: spin + + Enable (``True``) or disable (``False``) particle spin tracking (default: ``False``). + + Whether to track particle spin. + Currently, the implementation of spin tracking is a work in progress, and this feature is not yet supported. + .. py:property:: diagnostics Enable (``True``) or disable (``False``) diagnostics generally (default: ``True``). diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 388104738..0a863f17f 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1830,3 +1830,18 @@ add_impactx_test(distgen-spinvmf.py examples/distgen/analysis_kurth4d_spin.py OFF # no plot script yet ) + +# Spin Depolarization in a Quadrupole ###################### +# without space charge +add_impactx_test(quad-spin + examples/spin_tracking/input_quad_spin.in + ON # ImpactX MPI-parallel + examples/spin_tracking/analysis_quad_spin.py + OFF # no plot script yet +) +add_impactx_test(quad-spin.py + examples/spin_tracking/run_quad_spin.py + ON # ImpactX MPI-parallel + examples/spin_tracking/analysis_quad_spin.py + OFF # no plot script yet +) diff --git a/examples/spin_tracking/README.rst b/examples/spin_tracking/README.rst new file mode 100644 index 000000000..79b301a32 --- /dev/null +++ b/examples/spin_tracking/README.rst @@ -0,0 +1,50 @@ +.. _examples-quad-spin: + +Spin Depolarization in a Quadrupole +=================================== + +This example illustrates the decay of the polarization vector (describing the mean of the three spin components) along the vertical y and longitudinal z directions for a beam undergoing +horizontal focusing in a quadrupole. + +We use a 250 MeV proton beam with initial unnormalized rms emittance of 1 micron in the horizontal plane, and 2 micron in the vertical plane. + +The beam propagates over one horizontal betatron period, to a location where the polarization vector is described by a simple expression. + +In this test, the initial and final values of :math:`\polarization_x`, :math:`\polarization_y`, and :math:`\polarization_z` must agree with nominal values. + + +Run +--- + +This example can be run **either** as: + +* **Python** script: ``python3 run_quad_spin.py`` or +* ImpactX **executable** using an input file: ``impactx input_quad_spin.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_quad_spin.py + :language: python3 + :caption: You can copy this file from ``examples/spin_tracking/run_quad_spin.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_quad_spin.in + :language: ini + :caption: You can copy this file from ``examples/spin_tracking/input_quad_spin.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_quad_spin.py`` + + .. literalinclude:: analysis_quad_spin.py + :language: python3 + :caption: You can copy this file from ``examples/spin_tracking/analysis_quad_spin.py``. diff --git a/examples/spin_tracking/analysis_quad_spin.py b/examples/spin_tracking/analysis_quad_spin.py new file mode 100755 index 000000000..9c3446239 --- /dev/null +++ b/examples/spin_tracking/analysis_quad_spin.py @@ -0,0 +1,74 @@ +#!/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 + + +def get_polarization(beam): + """Calculate polarization vector, given by the mean values of spin components. + + Returns + ------- + polarization_x, polarization_y, polarization_z + """ + polarization_x = np.mean(beam["spin_x"]) + polarization_y = np.mean(beam["spin_y"]) + polarization_z = np.mean(beam["spin_z"]) + + return (polarization_x, polarization_y, polarization_z) + + +# 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:") +polarization_x, polarization_y, polarization_z = get_polarization(initial) +print( + f" polarization_x={polarization_x:e} polarization_y={polarization_y:e} polarization_z={polarization_z:e}" +) + +atol = 1.3e12 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" atol={atol}") + +assert np.allclose( + [polarization_x, polarization_y, polarization_z], + [ + 0.7, + 0.0, + 0.0, + ], + atol=atol, +) + +print("") +print("Final Beam:") +polarization_x, polarization_y, polarization_z = get_polarization(final) +print( + f" polarization_x={polarization_x:e} polarization_y={polarization_y:e} polarization_z={polarization_z:e}" +) + +atol = 1.3e12 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" atol={atol}") + +assert np.allclose( + [polarization_x, polarization_y, polarization_z], + [ + 0.7, + 0.0, + 0.0, + ], + atol=atol, +) diff --git a/examples/spin_tracking/input_quad_spin.in b/examples/spin_tracking/input_quad_spin.in new file mode 100644 index 000000000..55276a5a2 --- /dev/null +++ b/examples/spin_tracking/input_quad_spin.in @@ -0,0 +1,50 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 10000 +beam.units = static +beam.kin_energy = 99.48900104 +beam.charge = 25.0e-12 +beam.particle = electron +beam.distribution = triangle_from_twiss +beam.alphaX = 0.0 +beam.alphaY = 0.0 +beam.alphaT = 0.0 +beam.betaX = 0.002 +beam.betaY = beam.betaX +beam.betaT = 0.00004 +beam.emittX = 7.626114e-9 +beam.emittY = beam.emittX +beam.emittT = 2.5e-08 +beam.polarization_x = 0.7 +beam.polarization_y = 0.0 +beam.polarization_z = 0.0 + + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = monitor quad1 monitor +lattice.nslice = 1 + +monitor.type = beam_monitor +monitor.name = "monitor" +monitor.backend = h5 + +quad1.type = quad_exact +quad1.ds = 0.02903 +quad1.k = 207.0 +quad1.units = 1 +quad1.mapsteps = 5 + +############################################################################### +# Algorithms +############################################################################### +algo.space_charge = false +algo.spin = true + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/spin_tracking/run_quad_spin.py b/examples/spin_tracking/run_quad_spin.py new file mode 100755 index 000000000..078239cd9 --- /dev/null +++ b/examples/spin_tracking/run_quad_spin.py @@ -0,0 +1,89 @@ +#!/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 +from scipy.constants import c, e, m_e + +from impactx import ImpactX, distribution, elements, twiss + +sim = ImpactX() + +# set numerical parameters and IO control +sim.space_charge = False +sim.spin = True +sim.slice_step_diagnostics = True + +# domain decomposition & space charge mesh +sim.init_grids() + +# basic beam parameters +total_energy_MeV = 100.0 # reference energy (total) +mass_MeV = 0.510998950 # particle mass +kin_energy_MeV = total_energy_MeV - mass_MeV +bunch_charge_C = 25.0e-12 # used with space charge +npart = 10000 # number of macro particles + +# set reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(-1.0).set_mass_MeV(mass_MeV).set_kin_energy_MeV(kin_energy_MeV) + +# factors converting the beam distribution to ImpactX input +gamma = total_energy_MeV / mass_MeV +bg = np.sqrt(gamma**2 - 1.0) +sigma_tau = 1e-6 # in m +sigma_p = 2.5e-2 # dimensionless +rigidity = m_e * c * bg / e + +# particle bunch +distr = distribution.Triangle( + **twiss( + beta_x=0.002, + beta_y=0.002, + beta_t=sigma_tau / sigma_p, + emitt_x=1.5e-6 / bg, + emitt_y=1.5e-6 / bg, + emitt_t=sigma_tau * sigma_p, + alpha_x=0.0, + alpha_y=0.0, + alpha_t=0.0, + ) +) +spin = distribution.SpinvMF( + 0.7, + 0.0, + 0.0, +) + +sim.add_particles(bunch_charge_C, distr, npart, spin) + +# design the accelerator lattice +ns = 1 # number of slices per ds in the element + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# element names consistent with HTU_base_lattice_matched.lte Elegant input + +drift1 = elements.ExactDrift(name="drift1", ds=0.046, nslice=ns) +quad1 = elements.ExactQuad( + name="quad1", ds=0.02903, k=207.0, unit=1, mapsteps=5, nslice=ns +) +quad2 = elements.ExactQuad( + name="quad2", ds=0.02890, k=-207.0, unit=1, mapsteps=5, nslice=ns +) + +# set the lattice +sim.lattice.append(monitor) +sim.lattice.append(quad1) +sim.lattice.append(monitor) + +# run simulation +sim.track_particles() + +# clean shutdown +sim.finalize() diff --git a/src/elements/Quad.H b/src/elements/Quad.H index 36868a935..d3a73ace5 100644 --- a/src/elements/Quad.H +++ b/src/elements/Quad.H @@ -18,6 +18,7 @@ #include "mixin/named.H" #include "mixin/nofinalize.H" #include "mixin/lineartransport.H" +#include "mixin/spintransport.H" #include #include @@ -36,6 +37,7 @@ namespace impactx::elements public mixin::Thick, public mixin::Alignment, public mixin::PipeAperture, + public mixin::SpinTransport, public mixin::NoFinalize // At least on Intel AVX512, there is a small overhead to vectorize this element for DP and a gain for SP, see // https://github.com/BLAST-ImpactX/impactx/pull/1002 @@ -289,6 +291,82 @@ namespace impactx::elements return R; } + /** This operator pushes the particle spin. + * + * @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 + * @param refpart reference particle (unused) + * @param sx particle spin x-component + * @param sy particle spin y-component + * @param sz particle spin z-component + */ + template + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void spin_push ( + [[maybe_unused]] T_Real & AMREX_RESTRICT x, + [[maybe_unused]] T_Real & AMREX_RESTRICT y, + [[maybe_unused]] T_Real & AMREX_RESTRICT t, + [[maybe_unused]] T_Real & AMREX_RESTRICT px, + [[maybe_unused]] T_Real & AMREX_RESTRICT py, + [[maybe_unused]] T_Real & AMREX_RESTRICT pt, + [[maybe_unused]] T_IdCpu & AMREX_RESTRICT idcpu, + [[maybe_unused]] RefPart const & AMREX_RESTRICT refpart, + T_Real & AMREX_RESTRICT sx, + T_Real & AMREX_RESTRICT sy, + T_Real & AMREX_RESTRICT sz + ) const + { + using namespace amrex::literals; // for _rt and _prt + + // initialize the three components of the axis-angle vector + T_Real lambdax = 0_prt; + T_Real lambday = 0_prt; + T_Real lambdaz = 0_prt; + + // access reference particle values + amrex::ParticleReal G = refpart.gyromagnetic_anomaly; + amrex::ParticleReal const gyro_const = 1_prt + G * refpart.gamma(); + + // length of the current slice + amrex::ParticleReal const slice_ds = m_ds / nslice(); + + // compute phase advance per unit length in s (in rad/m) + amrex::ParticleReal const omega = std::sqrt(std::abs(m_k)); + + // compute trigonometric quantities + auto const [sin_omega_ds, cos_omega_ds] = amrex::Math::sincos(omega*slice_ds); + amrex::ParticleReal const sinh_omega_ds = std::sinh(omega*slice_ds); + amrex::ParticleReal const cosh_omega_ds = std::cosh(omega*slice_ds); + + if (m_k > 0.0_prt) + { + + // horizontally focusing quad case + lambdax = -gyro_const * ( py*(cosh_omega_ds - 1_prt) + y*omega*sinh_omega_ds ); + lambday = -gyro_const * ( px*(1_prt - cos_omega_ds) + x*omega*sin_omega_ds ); + + } else if (m_k < 0.0_prt) + { + + // horizontally defocusing quad case + lambdax = -gyro_const * ( px*(1_prt - cos_omega_ds) + x*omega*sin_omega_ds ); + lambday = -gyro_const * ( py*(cosh_omega_ds - 1_prt) + y*omega*sinh_omega_ds ); + + } else { + // treat as a drift + } + + // push the spin vector using the generator just determined + rotate_spin(lambdax,lambday,lambdaz,sx,sy,sz); + + } + + /** This pushes the covariance matrix. */ using LinearTransport::operator(); diff --git a/src/elements/mixin/spintransport.H b/src/elements/mixin/spintransport.H new file mode 100644 index 000000000..dfc0df12e --- /dev/null +++ b/src/elements/mixin/spintransport.H @@ -0,0 +1,77 @@ +/* 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 + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_ELEMENTS_MIXIN_SPIN_TRANSPORT_H +#define IMPACTX_ELEMENTS_MIXIN_SPIN_TRANSPORT_H + +#include "particles/ImpactXParticleContainer.H" + +#include + +#include +#include +#include + +#include + + +namespace impactx::elements::mixin +{ + /** This is a helper class for applying a spin map generated by a three-component vector + * lambda, whose magnitude defines the angle of rotation and whose direction defines + * the axis of rotation. + */ + struct SpinTransport + { + /** Rotate the spin vector (sx,sy,sz). + * + * @param[in] lambdax generator x-component + * @param[in] lambday generator y-component + * @param[in] lambdaz generator z-component + * @param[inout] sx spin vector x-component + * @param[inout] sy spin vector y-component + * @param[inout] sz spin vector z-component + */ + template + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void rotate_spin ( + T_Real lambdax, + T_Real lambday, + T_Real lambdaz, + T_Real & AMREX_RESTRICT sx, + T_Real & AMREX_RESTRICT sy, + T_Real & AMREX_RESTRICT sz + ) const + { + using namespace amrex::literals; // for _prt + + // Compute quaternion coefficients + amrex::ParticleReal angle = std::sqrt(lambdax*lambdax+lambday*lambday+lambdaz*lambdaz); + auto const [sin_half_angle, cos_half_angle] = amrex::Math::sincos(angle*0.5_prt); + amrex::ParticleReal w0 = cos_half_angle; + amrex::ParticleReal w1 = (angle==0_prt) ? 0_prt : -(lambdax/angle)*sin_half_angle; + amrex::ParticleReal w2 = (angle==0_prt) ? 0_prt : -(lambday/angle)*sin_half_angle; + amrex::ParticleReal w3 = (angle==0_prt) ? 0_prt : -(lambdaz/angle)*sin_half_angle; + + // Apply rotation + amrex::ParticleReal sxout = (1_prt-2_prt*(w2*w2+w3*w3))*sx + 2_prt*(w1*w2+w0*w3)*sy + 2_prt*(w1*w3-w0*w2)*sz; + amrex::ParticleReal syout = 2_prt*(w1*w2-w0*w3)*sx + (1_prt-2_prt*(w1*w1+w3*w3))*sy + 2_prt*(w0*w1+w2*w3)*sz; + amrex::ParticleReal szout = 2_prt*(w0*w2+w1*w3)*sx + 2_prt*(w2*w3-w0*w1)*sy + (1_prt-2_prt*(w1*w1+w2*w2))*sz; + + // Update spin vector + sx = sxout; + sy = syout; + sz = szout; + } + + }; + +} // namespace impactx::elements::mixin + +#endif // IMPACTX_ELEMENTS_MIXIN_SPIN_TRANSPORT_H diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 5fee178ee..ada50354f 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -50,18 +50,23 @@ namespace impactx constexpr amrex::ParticleReal m_e = ablastr::constant::SI::m_e_v; constexpr amrex::ParticleReal m_p = ablastr::constant::SI::m_p_v; constexpr amrex::ParticleReal MeV_inv_c2 = ablastr::constant::SI::MeV_inv_c2_v; + amrex::ParticleReal gyromagnetic_anomaly; // anomalous magnetic dipole moment if (particle_type == "electron") { qe = -1.0; massE = m_e / MeV_inv_c2; + gyromagnetic_anomaly = 0.00115965218062; } else if (particle_type == "positron") { qe = 1.0; massE = m_e / MeV_inv_c2; + gyromagnetic_anomaly = 0.00115965218062; } else if (particle_type == "proton") { qe = 1.0; massE = m_p / MeV_inv_c2; + gyromagnetic_anomaly = 1.7928473446; } else if (particle_type == "Hminus") { qe = -1.0; massE = 939.294308; // value used in TraceWin + gyromagnetic_anomaly = 1.7928473446; // this value is difficult to find, and needs to be checked } else { // default to electron ablastr::warn_manager::WMRecordWarning( @@ -71,11 +76,12 @@ namespace impactx ); qe = -1.0; massE = m_e / MeV_inv_c2; + gyromagnetic_anomaly = 0.00115965218062; } // configure a new reference particle RefPart ref; - ref.set_charge_qe(qe).set_mass_MeV(massE).set_kin_energy_MeV(kin_energy); + ref.set_charge_qe(qe).set_mass_MeV(massE).set_kin_energy_MeV(kin_energy).set_gyromagnetic_anomaly(gyromagnetic_anomaly); return ref; } diff --git a/src/particles/ReferenceParticle.H b/src/particles/ReferenceParticle.H index 7a63a78af..f4e0536ed 100644 --- a/src/particles/ReferenceParticle.H +++ b/src/particles/ReferenceParticle.H @@ -40,6 +40,7 @@ namespace impactx amrex::ParticleReal pt = 0.0; ///< energy, normalized by rest energy amrex::ParticleReal mass = 0.0; ///< reference rest mass, in kg amrex::ParticleReal charge = 0.0; ///< reference charge, in C + amrex::ParticleReal gyromagnetic_anomaly = 0.0; ///< anomalous magnetic moment [unitless] amrex::ParticleReal sedge = 0.0; ///< value of s at entrance of the current beamline element amrex::SmallMatrix map{}; ///< linearized map @@ -250,6 +251,23 @@ namespace impactx { return charge / mass; } + + /** Set reference particle anomalous magnetic dipole moment + * (aka gyromagnetic anomaly). + * + * @param gyromagnetic_anomaly_ref unitless + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + RefPart & + set_gyromagnetic_anomaly (amrex::ParticleReal const gyromagnetic_anomaly_ref) + { + using namespace amrex::literals; + + this->gyromagnetic_anomaly = gyromagnetic_anomaly_ref; + + return *this; + } + }; } // namespace impactx diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index a02d69df3..7add1f6da 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -256,6 +256,16 @@ void init_ImpactX (py::module& m) }, "Flag to determine whether ISR radiation loss is applied to the reference particle (default: False)." ) + .def_property("spin", + [](ImpactX & /* ix */) { + return detail::get_or_throw("algo", "spin"); + }, + [](ImpactX & /* ix */, bool const enable) { + amrex::ParmParse pp_algo("algo"); + pp_algo.add("spin", enable); + }, + "Enable or disable particle spin tracking (default: disabled)." + ) .def_property("eigenemittances", [](ImpactX & /* ix */) { return detail::get_or_throw("diag", "eigenemittances"); diff --git a/src/tracking/particles.cpp b/src/tracking/particles.cpp index 70ec84dc6..9e1461d89 100644 --- a/src/tracking/particles.cpp +++ b/src/tracking/particles.cpp @@ -91,10 +91,13 @@ namespace impactx pp_algo.query("csr", csr); bool isr = false; pp_algo.query("isr", isr); + bool spin = false; + pp_algo.query("spin", spin); if (verbose > 0) { amrex::Print() << " CSR effects: " << csr << "\n"; amrex::Print() << " ISR effects: " << isr << "\n"; + amrex::Print() << " Spin tracking: " << spin << "\n"; } // periods through the lattice