Quad spin push

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

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``).

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
+)

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 <https://www.mpi-forum.org>`__ 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``.

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,
+)

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

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()