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 \(\langle{s_x\rangle}\), \(\langle{s_y\rangle}\), \(\langle{s_z\rangle}\) 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.

Listing 176 You can copy this file from examples/spin_tracking/run_quad_spin.py.
#!/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 impactx import ImpactX, distribution, elements

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
kin_energy_MeV = 100.0  # reference energy (kinetic)
bunch_charge_C = 25.0e-12  # used with space charge
npart = 100000  # number of macro particles

# set reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

#   particle bunch
distr = distribution.Gaussian(
    lambdaX=0.0003,
    lambdaY=0.0003,
    lambdaT=1.0e-6,
    lambdaPx=0.0002,
    lambdaPy=0.0002,
    lambdaPt=1.0e-6,
    muxpx=0.0,
    muypy=0.0,
    mutpt=0.0,
)
spin = distribution.SpinvMF(
    0.4,
    0.9,
    0.1,
)

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

# quad focusing strength (> 0)
k_value = 100.0
print("k_value")
print(k_value)

# length for this test should be one period
ds_value = 2.0 * np.pi / np.sqrt(k_value)

quad1 = elements.Quad(name="quad1", ds=ds_value, k=k_value, 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()
Listing 177 You can copy this file from examples/spin_tracking/input_quad_spin.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 100000
beam.units = static
beam.kin_energy = 100.0
beam.charge = 25.0e-12
beam.particle = electron
beam.distribution = gaussian
beam.lambdaX = 0.0003
beam.lambdaY = 0.0003
beam.lambdaT = 1.0e-6
beam.lambdaPx = 0.0002
beam.lambdaPy = 0.0002
beam.lambdaPt = 1.0e-6
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
#beam.polarization_x = 0.0
#beam.polarization_y = 0.9
#beam.polarization_z = 0.0
beam.polarization_x = 0.4
beam.polarization_y = 0.9
beam.polarization_z = 0.1

###############################################################################
# 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
quad1.ds = 0.628318530717959
quad1.k = 100.0

###############################################################################
# Algorithms
###############################################################################
algo.space_charge = false
algo.spin = true

###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true

Analyze

The analysis can be run using either of the following scripts:

Script analysis_quad_spin.py
Listing 178 You can copy this file from examples/spin_tracking/analysis_quad_spin.py.
#!/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 = 100000
assert num_particles == len(initial)
assert num_particles == len(final)

# numerical parameters based on input file
beam_initial = series.iterations[1].particles["beam"]
gryo_anomaly = 0.001159652181644  # for electrons
rel_gamma = beam_initial.get_attribute("gamma_ref")  # for 100 MeV
quad_gradient = 100  # value in 1/m^2 from input
sigma_y = 0.0003  # value in m = lambdaY from input
sigma_py = 0.0002  # value in rad = lambdaPy from input
Pxi = 0.4  # polarization_x from input
Pyi = 0.9  # polarization_y from input
Pzi = 0.1  # polarization_z from input

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.3 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  atol={atol}")

assert np.allclose(
    [polarization_x, polarization_y, polarization_z],
    [
        Pxi,
        Pyi,
        Pzi,
    ],
    atol=atol,
)

# predicted final polarization
damping_eigenvalue = (1 + gryo_anomaly * rel_gamma) * np.sqrt(
    sigma_py**2 * (np.cosh(2 * np.pi) - 1) ** 2
    + sigma_y**2 * quad_gradient * np.sinh(2 * np.pi) ** 2
)
damping_factor = np.exp(-(damping_eigenvalue**2) / 2.0)
Pxf = Pxi
Pyf = damping_factor * Pyi
Pzf = damping_factor * Pzi

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.3 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  atol={atol}")

assert np.allclose(
    [polarization_x, polarization_y, polarization_z],
    [
        Pxf,
        Pyf,
        Pzf,
    ],
    atol=atol,
)
Script analysis_quad_spin_rbc.py
Listing 179 You can copy this file from examples/spin_tracking/analysis_quad_spin_rbc.py.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#

import numpy as np


# Load data from envelope simulation
def read_time_series(file_pattern):
    """Read in all CSV files from each MPI rank (and potentially OpenMP
    thread). Concatenate into one Pandas dataframe.

    Returns
    -------
    pandas.DataFrame
    """

    import glob
    import re

    import pandas as pd

    def read_file(file_pattern):
        for filename in glob.glob(file_pattern):
            df = pd.read_csv(filename, delimiter=r"\s+")
            if "step" not in df.columns:
                step = int(re.findall(r"[0-9]+", filename)[0])
                df["step"] = step
            yield df

    return pd.concat(
        read_file(file_pattern),
        axis=0,
        ignore_index=True,
    )  # .set_index('id')


rbc = read_time_series("diags/reduced_beam_characteristics.*")

# numerical parameters based on input file
gryo_anomaly = 0.001159652181644  # for electrons
rel_gamma = 196.69511809100055  # for 100 MeV
quad_gradient = 100  # value in 1/m^2 from input
sigma_y = 0.0003  # value in m = lambdaY from input
sigma_py = 0.0002  # value in rad = lambdaPy from input
Pxi = 0.4  # polarization_x from input
Pyi = 0.9  # polarization_y from input
Pzi = 0.1  # polarization_z from input

print("Initial Beam:")
polarization_x = rbc["mean_sx"].iloc[0]
polarization_y = rbc["mean_sy"].iloc[0]
polarization_z = rbc["mean_sz"].iloc[0]
print(
    f"  polarization_x={polarization_x:e} polarization_y={polarization_y:e} polarization_z={polarization_z:e}"
)

# Number of particles (from input file)
num_particles = 100000

atol = 1.3 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  atol={atol}")

assert np.allclose(
    [polarization_x, polarization_y, polarization_z],
    [
        Pxi,
        Pyi,
        Pzi,
    ],
    atol=atol,
)

# predicted final polarization
damping_eigenvalue = (1 + gryo_anomaly * rel_gamma) * np.sqrt(
    sigma_py**2 * (np.cosh(2 * np.pi) - 1) ** 2
    + sigma_y**2 * quad_gradient * np.sinh(2 * np.pi) ** 2
)
damping_factor = np.exp(-(damping_eigenvalue**2) / 2.0)
Pxf = Pxi
Pyf = damping_factor * Pyi
Pzf = damping_factor * Pzi

print("")
print("Final Beam:")
polarization_x = rbc["mean_sx"].iloc[-1]
polarization_y = rbc["mean_sy"].iloc[-1]
polarization_z = rbc["mean_sz"].iloc[-1]
print(
    f"  polarization_x={polarization_x:e} polarization_y={polarization_y:e} polarization_z={polarization_z:e}"
)

atol = 1.3 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  atol={atol}")

assert np.allclose(
    [polarization_x, polarization_y, polarization_z],
    [
        Pxf,
        Pyf,
        Pzf,
    ],
    atol=atol,
)

# numerical tests of spin moment conditions
sigma_sx = rbc["sigma_sx"].iloc[-1]
sigma_sy = rbc["sigma_sy"].iloc[-1]
sigma_sz = rbc["sigma_sz"].iloc[-1]
polarization = np.sqrt(polarization_x**2 + polarization_y**2 + polarization_z**2)
condition = sigma_sx**2 + sigma_sy**2 + sigma_sz**2 + polarization**2

print("")
print(f"Spin moment consistency condition = {condition:e}")

atol = 1.0e-12  # machine precision
print(f"  atol={atol}")

assert np.allclose(
    [condition],
    [
        1.0,
    ],
    atol=atol,
)

Spin Depolarization in a FODO Channel

This example illustrates the decay of the polarization vector (describing the mean of the three spin components) for a matched beam in a FODO channel.

We use a 10 GeV electron beam, with an initially 6D Gaussian distribution in the phase space. The FODO channel and Twiss parameters are otherwise identical to those appearing in examples/fodo_channel.

The beam spin undergoes rapid mixing and depolarization.

In this test, the initial and final values of \(\langle{s_x\rangle}\), \(\langle{s_y\rangle}\), \(\langle{s_z\rangle}\) must agree with nominal values.

Run

This example can be run either as:

  • Python script: python3 run_fodo_channel_spin.py or

  • ImpactX executable using an input file: impactx input_fodo_channel_spin.in

For MPI-parallel runs, prefix these lines with mpiexec -n 4 ... or srun -n 4 ..., depending on the system.

Listing 180 You can copy this file from examples/spin_tracking/run_fodo_channel_spin.py.
#!/usr/bin/env python3
#
# Copyright 2022-2024 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell, Marco Garten
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

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

# load a 2 GeV electron beam with an initial
# unnormalized rms emittance of 2 nm
kin_energy_MeV = 10.0e3  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

#   reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

#   particle bunch
distr = distribution.Gaussian(
    **twiss(
        beta_x=2.8216194100262637,
        beta_y=2.8216194100262637,
        beta_t=0.5,
        emitt_x=2e-04,
        emitt_y=2e-04,
        emitt_t=2e-05,
        alpha_x=-1.5905003499999992,
        alpha_y=1.5905003499999992,
        alpha_t=0.0,
    )
)
spin = distribution.SpinvMF(
    0.4,
    0.9,
    0.1,
)

sim.add_particles(bunch_charge_C, distr, npart, spin)

# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5", period_sample_intervals=10)

# design the accelerator lattice)
ns = 1  # number of slices per ds in the element
fodo = [
    monitor,
    elements.ChrDrift(ds=0.25, nslice=ns),
    elements.ChrQuad(ds=1.0, k=1.0, nslice=ns),
    elements.ChrDrift(ds=0.5, nslice=ns),
    elements.ChrQuad(ds=1.0, k=-1.0, nslice=ns),
    elements.ChrDrift(ds=0.25, nslice=ns),
    monitor,
]
# assign a fodo segment
sim.lattice.extend(fodo)

# FODO channel of 101 periods
sim.periods = 101

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 181 You can copy this file from examples/spin_tracking/input_fodo_channel_spin.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 10.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = gaussian_from_twiss
beam.alphaX = -1.5905003499999992
beam.alphaY = -beam.alphaX
beam.alphaT = 0.0
beam.betaX = 2.8216194100262637
beam.betaY = beam.betaX
beam.betaT = 0.5
beam.emittX = 2e-04
beam.emittY = beam.emittX
beam.emittT = 2e-05
beam.polarization_x = 0.4
beam.polarization_y = 0.9
beam.polarization_z = 0.1

###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift1 quad1 drift2 quad2 drift3 monitor
lattice.nslice = 1
lattice.periods = 101  # FODO channel of 101 periods

monitor.type = beam_monitor
monitor.backend = h5

drift1.type = drift_chromatic
drift1.ds = 0.25

quad1.type = quad_chromatic
quad1.ds = 1.0
quad1.k = 1.0

drift2.type = drift_chromatic
drift2.ds = 0.5

quad2.type = quad_chromatic
quad2.ds = 1.0
quad2.k = -1.0

drift3.type = drift_chromatic
drift3.ds = 0.25


###############################################################################
# Algorithms
###############################################################################
algo.space_charge = false
algo.spin = true

###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true

Spin Depolarization in a Dipole

This example illustrates the decay of the polarization vector (describing the mean of the three spin components) along the horizontal x and longitudinal z directions for a beam undergoing bending in the x-z plane in a sector dipole.

We use a 2 GeV electron beam. The beam parameters (in particular, the momentum and energy spread) are artificially large in order to enhance the effect.

The beam propagates over one period, as set by the design spin tune. By increasing the number of slices, and turning on diagnostics, one can view precession of the polarization vector about the vertical direction. A clean precession over 1 period becomes visible when the beam size, momentum spread, and energy spread are set to small values.

In this test, the initial and final values of \(\langle{s_x\rangle}\), \(\langle{s_y\rangle}\), \(\langle{s_z\rangle}\) must agree with nominal values.

Run

This example can be run either as:

  • Python script: python3 run_sbend_spin.py or

  • ImpactX executable using an input file: impactx input_sbend_spin.in

For MPI-parallel runs, prefix these lines with mpiexec -n 4 ... or srun -n 4 ..., depending on the system.

Listing 182 You can copy this file from examples/spin_tracking/run_sbend_spin.py.
#!/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 impactx import ImpactX, distribution, elements

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
kin_energy_MeV = 2000.0  # reference energy (kinetic)
mass_MeV = 0.510998950  # particle mass
bunch_charge_C = 25.0e-12  # used with space charge
gyromagnetic_anomaly = 0.00115965218062  # value for an electron
npart = 100000  # number of macro particles

# set reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

#   particle bunch
distr = distribution.Gaussian(
    lambdaX=0.003,
    lambdaY=0.003,
    lambdaT=1.0e-6,
    lambdaPx=0.2,
    lambdaPy=0.2,
    lambdaPt=0.2,
    muxpx=0.0,
    muypy=0.0,
    mutpt=0.0,
)
spin = distribution.SpinvMF(
    0.4,
    0.9,
    0.1,
)

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

# bend radius (> 0)
rc_value = 1.0
print("rc_value")
print(rc_value)

gamma = kin_energy_MeV / mass_MeV + 1
print("relativistic gamma")
print(gamma)

# length for this test should be one period
ds_value = 2.0 * np.pi / (gyromagnetic_anomaly * gamma)

bend1 = elements.Sbend(name="bend1", ds=ds_value, rc=rc_value, nslice=ns)

# set the lattice
sim.lattice.append(monitor)
sim.lattice.append(bend1)
sim.lattice.append(monitor)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 183 You can copy this file from examples/spin_tracking/input_sbend_spin.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 100000
beam.units = static
beam.kin_energy = 2000.0
beam.charge = 25.0e-12
beam.particle = electron
beam.distribution = gaussian
beam.lambdaX = 0.003
beam.lambdaY = 0.003
beam.lambdaT = 1.0e-6
beam.lambdaPx = 0.2
beam.lambdaPy = 0.2
beam.lambdaPt = 0.2
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
beam.polarization_x = 0.4
beam.polarization_y = 0.9
beam.polarization_z = 0.1

###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor bend1 monitor
lattice.nslice = 1

monitor.type = beam_monitor
monitor.name = "monitor"
monitor.backend = h5

bend1.type = sbend
bend1.ds = 1.3839843641758893
bend1.rc = 1.0

###############################################################################
# Algorithms
###############################################################################
algo.space_charge = false
algo.spin = true

###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true

Analyze

The analysis can be run using:

Script analysis_sbend_spin.py
Listing 184 You can copy this file from examples/spin_tracking/analysis_sbend_spin.py.
#!/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 = 100000
assert num_particles == len(initial)
assert num_particles == len(final)

# numerical parameters based on input file

gyro_anomaly = 0.001159652181644  # for electrons
quad_gradient = 100  # value in 1/m^2 from input
sigma_x = 0.003  # value in m = lambdaX from input
sigma_px = 0.2  # value in rad = lambdaPx from input
sigma_pt = 0.2  # value = lambdaPt from input
Pxi = 0.4  # polarization_x from input
Pyi = 0.9  # polarization_y from input
Pzi = 0.1  # polarization_z from input

beam_initial = series.iterations[1].particles["beam"]
rel_gamma = beam_initial.get_attribute("gamma_ref")
rel_beta = beam_initial.get_attribute("beta_ref")
h = 1.0 / 10.0  # 1/radius of bend curvature from input

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.3 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  atol={atol}")

assert np.allclose(
    [polarization_x, polarization_y, polarization_z],
    [
        Pxi,
        Pyi,
        Pzi,
    ],
    atol=atol,
)

# predicted final polarization
gyro_const = 1.0 + gyro_anomaly * rel_gamma
sin1 = np.sin(np.pi / (gyro_anomaly * rel_gamma))
sin2 = np.sin(2 * np.pi / (gyro_anomaly * rel_gamma))
damping_eigenvalue2 = (
    4 * sigma_px**2 * gyro_const**2 * sin1**4
    + h**2 * gyro_const**2 * sigma_x**2 * sin2**2
    + sigma_pt**2 * (-2.0 * np.pi * rel_beta**2 + gyro_const * sin2) ** 2 / rel_beta**2
)
damping_eigenvalue = np.sqrt(damping_eigenvalue2)
damping_factor = np.exp(-(damping_eigenvalue**2) / 2.0)

Pxf = damping_factor * Pxi
Pyf = Pyi
Pzf = damping_factor * Pzi

print("")
print("Predicted Final Polarization:")
print(f"  polarization_x={Pxf:e} polarization_y={Pyf:e} polarization_z={Pzf:e}")

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.3 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  atol={atol}")

assert np.allclose(
    [polarization_x, polarization_y, polarization_z],
    [
        Pxf,
        Pyf,
        Pzf,
    ],
    atol=atol,
)

Spin Depolarization in a Solenoid

This example illustrates the decay of the polarization vector (describing the mean of the three spin components) along the horizontal x and vertical y directions for a 2 GeV proton beam undergoing transverse focusing and rotation in a solenoid.

The beam propagates over a distance: \(\Delta s = 2\pi/(Gk_s)\),

where \(G\) is the value of the gyromagnetic anomaly, \(k_s\) denotes the solenoid focusing strength. At this location, the polarization vector is described by a simple expression.

In this test, the initial and final values of \(\langle{s_x\rangle}\), \(\langle{s_y\rangle}\), \(\langle{s_z\rangle}\) must agree with nominal values.

Run

This example can be run either as:

  • Python script: python3 run_sol_spin.py or

  • ImpactX executable using an input file: impactx input_sol_spin.in

For MPI-parallel runs, prefix these lines with mpiexec -n 4 ... or srun -n 4 ..., depending on the system.

Listing 185 You can copy this file from examples/spin_tracking/run_sol_spin.py.
#!/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 impactx import ImpactX, distribution, elements

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
kin_energy_MeV = 2000.0  # reference energy (kinetic)
bunch_charge_C = 25.0e-12  # used with space charge
npart = 100000  # number of macro particles

# set reference particle
ref = sim.beam.ref
ref.set_species("proton").set_kin_energy_MeV(kin_energy_MeV)

#   particle bunch
distr = distribution.Gaussian(
    lambdaX=0.003,
    lambdaY=0.003,
    lambdaT=1.0e-6,
    lambdaPx=0.2,
    lambdaPy=0.2,
    lambdaPt=0.2,
    muxpx=0.0,
    muypy=0.0,
    mutpt=0.0,
)
spin = distribution.SpinvMF(
    0.4,
    0.9,
    0.1,
)

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

# bend radius (> 0)
ks_value = 1.0
print(f"ks_value={ks_value}")

# length for this test should be one period
ds_value = 2.0 * np.pi / (ref.gyromagnetic_anomaly * ks_value)

sol1 = elements.Sol(name="sol1", ds=ds_value, ks=ks_value, nslice=ns)

# set the lattice
sim.lattice.append(monitor)
sim.lattice.append(sol1)
sim.lattice.append(monitor)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 186 You can copy this file from examples/spin_tracking/input_sol_spin.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 100000
beam.units = static
beam.kin_energy = 2000.0
beam.charge = 25.0e-12
beam.particle = proton
beam.distribution = gaussian
beam.lambdaX = 0.003
beam.lambdaY = 0.003
beam.lambdaT = 1.0e-6
beam.lambdaPx = 0.2
beam.lambdaPy = 0.2
beam.lambdaPt = 0.2
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
beam.polarization_x = 0.4
beam.polarization_y = 0.9
beam.polarization_z = 0.1

###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor sol1 monitor
lattice.nslice = 1

monitor.type = beam_monitor
monitor.name = "monitor"
monitor.backend = h5

sol1.type = solenoid
sol1.ds = 3.504584663108292
sol1.ks = 1.0

###############################################################################
# Algorithms
###############################################################################
algo.space_charge = false
algo.spin = true

###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true

Analyze

The analysis can be run using:

Script analysis_sol_spin.py
Listing 187 You can copy this file from examples/spin_tracking/analysis_sol_spin.py.
#!/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 = 100000
assert num_particles == len(initial)
assert num_particles == len(final)

# numerical parameters based on input file
beam_initial = series.iterations[1].particles["beam"]
gyro_anomaly = 1.7928473446  # for protons
rel_beta = beam_initial.get_attribute("beta_ref")
sol_strength_ks = 1  # value in 1/m^2 from input
sol_length_ds = 3.504584663108292
sigma_pt = 0.2  # value in rad = lambdaPt from input
Pxi = 0.4  # polarization_x from input
Pyi = 0.9  # polarization_y from input
Pzi = 0.1  # polarization_z from input

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.3 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  atol={atol}")

assert np.allclose(
    [polarization_x, polarization_y, polarization_z],
    [
        Pxi,
        Pyi,
        Pzi,
    ],
    atol=atol,
)

# predicted final polarization
damping_eigenvalue = (
    (1 + gyro_anomaly) * sol_strength_ks * sol_length_ds * sigma_pt / rel_beta
)
damping_factor = np.exp(-(damping_eigenvalue**2) / 2.0)
cosG = np.cos(2 * (1 + gyro_anomaly) * np.pi / gyro_anomaly)
sinG = np.sin(2 * (1 + gyro_anomaly) * np.pi / gyro_anomaly)
Pxf = damping_factor * (Pxi * cosG + Pyi * sinG)
Pyf = damping_factor * (-Pxi * sinG + Pyi * cosG)
Pzf = Pzi

print("")
print("Predicted Final Polarization:")
print(f"  polarization_x={Pxf:e} polarization_y={Pyf:e} polarization_z={Pzf:e}")

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 = 2.0 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  atol={atol}")

assert np.allclose(
    [polarization_x, polarization_y, polarization_z],
    [
        Pxf,
        Pyf,
        Pzf,
    ],
    atol=atol,
)

Element Reversibility with Spin

In the case of linear elements, including spin, the joint spin-orbit map has the following exact reversibility property.

The effect of setting ds -> -ds is equivalent to replacing the map by its inverse.

In this test, a beam is propagated forward through an element of length ds, following by the corresponding element with length -ds. As a result, the composite map is the identity. This provides a non-trivial test for consistency between the spin map and the orbit map.

In this test, the initial and final spin components \(s_x\), \(s_y\), and \(s_z\) are compared. The norm of the change in the spin vector must lie within a very small tolerance.

Run

This example can be run either as:

  • Python script: python3 run_reversibility_spin.py or

  • ImpactX executable using an input file: impactx input_reversibility_spin.in

For MPI-parallel runs, prefix these lines with mpiexec -n 4 ... or srun -n 4 ..., depending on the system.

Listing 188 You can copy this file from examples/spin_tracking/run_reversibility_spin.py.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: 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.space_charge = False
sim.spin = True
sim.slice_step_diagnostics = False

# domain decomposition & space charge mesh
sim.init_grids()

# basic beam parameters
kin_energy_MeV = 2.0e3  # reference energy (kinetic)
bunch_charge_C = 1.0e-9  # used with space charge
npart = 100000  # number of macro particles

# set reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

#   particle bunch
distr = distribution.Gaussian(
    lambdaX=5.0e-6,
    lambdaY=8.0e-6,
    lambdaT=0.0599584916,
    lambdaPx=2.5543422003e-9,
    lambdaPy=1.5964638752e-9,
    lambdaPt=9.0e-4,
    muxpx=0.0,
    muypy=0.0,
    mutpt=0.0,
)
spin = distribution.SpinvMF(
    0.4,
    0.9,
    0.1,
)

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

# parameters
L = 1.5
kquad = 60.0
ksol = 1.0

quad1 = elements.Quad(name="quad1", ds=L, k=kquad, nslice=ns)
iquad1 = elements.Quad(name="iquad1", ds=-L, k=kquad, nslice=ns)
quad2 = elements.ChrQuad(name="quad2", ds=L, k=-kquad, nslice=ns)
iquad2 = elements.ChrQuad(name="iquad2", ds=-L, k=-kquad, nslice=ns)
sol1 = elements.Sol(name="sol1", ds=L, ks=ksol, nslice=ns)
isol1 = elements.Sol(name="isol1", ds=-L, ks=ksol, nslice=ns)

# set the lattice
sim.lattice.append(monitor)
sim.lattice.append(quad1)
sim.lattice.append(iquad1)
sim.lattice.append(sol1)
sim.lattice.append(isol1)
sim.lattice.append(quad2)
sim.lattice.append(iquad2)
sim.lattice.append(monitor)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 189 You can copy this file from examples/spin_tracking/input_reversibility_spin.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 100000
beam.units = static
beam.kin_energy = 2.0e3  #2 GeV
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 5.0e-6            # 5 um
beam.lambdaY = 8.0e-6            # 8 um
beam.lambdaT = 0.0599584916      # 200 ps
beam.lambdaPx = 2.5543422003e-9  # exn = 50 pm-rad
beam.lambdaPy = 1.5964638752e-9  # eyn = 50 pm-rad
beam.lambdaPt = 9.0e-4           # approximately dE/E
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
beam.polarization_x = 0.4
beam.polarization_y = 0.9
beam.polarization_z = 0.1

###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor quad1 iquad1 sol1 isol1 quad2 iquad2 monitor

lattice.nslice = 1

iquad1.type = quad  #inverse quad
iquad1.ds = -1.5
iquad1.k = 60.0

quad1.type = quad
quad1.ds = 1.5
quad1.k = 60.0

isol1.type = solenoid  #inverse solenoid
isol1.ds = -1.5
isol1.ks = 1.0

sol1.type = solenoid
sol1.ds = 1.5
sol1.ks = 1.0

iquad2.type = quad_chromatic  #inverse quad
iquad2.ds = -1.5
iquad2.k = -60.0

quad2.type = quad_chromatic
quad2.ds = 1.5
quad2.k = -60.0

monitor.type = beam_monitor
monitor.backend = h5


###############################################################################
# Algorithms
###############################################################################
algo.space_charge = false
algo.spin = true

###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false

Analyze

The analysis can be run using:

Script analysis_reversibility_spin.py
Listing 190 You can copy this file from examples/spin_tracking/analysis_reversibility_spin.py.
#!/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

# 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()
beam_final = series.iterations[last_step].particles["beam"]
final = beam_final.to_df()

sxi = initial["spin_x"]
syi = initial["spin_y"]
szi = initial["spin_z"]

sxf = final["spin_x"]
syf = final["spin_y"]
szf = final["spin_z"]

dspin2 = (sxf - sxi) ** 2 + (syf - syi) ** 2 + (szf - szi) ** 2
dspin = np.sqrt(dspin2)
dspinmax = dspin.max()

print("Change in the spin:")
print("||delta s||_max", dspinmax)

atol = 1.5e-9
print(f"  atol={atol}")

assert np.allclose(
    [dspinmax],
    [
        0.0,
    ],
    atol=atol,
)

Spin Limiting Cases of a Combined-Function Bend

This example tests the spin dynamics of a combined-function bend in the quad-like and bend-like limiting cases when the curvature or the quadrupole gradient are small, respectively. The beam parameters are identical to those used in examples-cfbend, with the addition of spin.

In this test, the beam is tracked through a combined-function dipole, and then tracked backward through an element with equivalent parameters (up to a small tolerance). Due to the s-reversibility of the maps applied, the final phase space coordinates and spin variables should coincide with their initial values.

This test computes the norm of the change in the spin polarization vector, which must be zero to within a small tolerance.

Run

This example can be run either as:

  • Python script: python3 run_cfbend_spin.py or

  • ImpactX executable using an input file: impactx input_cfbend_spin.in

For MPI-parallel runs, prefix these lines with mpiexec -n 4 ... or srun -n 4 ..., depending on the system.

Listing 191 You can copy this file from examples/spin_tracking/run_cfbend_spin.py.
#!/usr/bin/env python3
#
# Copyright 2022-2026 ImpactX contributors
# Authors: Chad Mitchell, Axel Huebl
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

from impactx import ImpactX, distribution, elements

sim = ImpactX()

# set numerical parameters and IO control
sim.space_charge = False
sim.spin = True
sim.slice_step_diagnostics = False

# domain decomposition & space charge mesh
sim.init_grids()

# load a 5 GeV electron beam with an initial
# normalized transverse rms emittance of 1 um
kin_energy_MeV = 2.0e3  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

#   reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

#   particle bunch
distr = distribution.Waterbag(
    lambdaX=5.0e-6,  # 5 um
    lambdaY=8.0e-6,  # 8 um
    lambdaT=0.0599584916,  # 200 ps
    lambdaPx=2.5543422003e-9,  # exn = 50 pm-rad
    lambdaPy=1.5964638752e-9,  # eyn = 50 pm-rad
    lambdaPt=9.0e-4,  # approximately dE/E
    muxpx=0.0,
    muypy=0.0,
    mutpt=0.0,
)
spin = distribution.SpinvMF(
    0.4,
    0.9,
    0.1,
)

sim.add_particles(bunch_charge_C, distr, npart, spin)

# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")

# design the accelerator lattice
ns = 1  # number of slices per ds in the element

# element length (m)
L = 1.5

# large curvature radius value for the quad test
rclarge = 1.0e8

# small gradient value for the sbend test
ksmall = 1.0e-8

lattice = [
    monitor,
    elements.CFbend(name="cfbend1", ds=L, rc=rclarge, k=60.0, nslice=ns),
    elements.Quad(name="iquad1", ds=-L, k=60.0, nslice=ns),
    elements.CFbend(name="cfbend2", ds=L, rc=10.0, k=ksmall, nslice=ns),
    elements.Sbend(name="isbend1", ds=-L, rc=10.0, nslice=ns),
    monitor,
]

# assign a lattice segment
sim.lattice.extend(lattice)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 192 You can copy this file from examples/spin_tracking/input_cfbend_spin.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3  #2 GeV
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 5.0e-6            # 5 um
beam.lambdaY = 8.0e-6            # 8 um
beam.lambdaT = 0.0599584916      # 200 ps
beam.lambdaPx = 2.5543422003e-9  # exn = 50 pm-rad
beam.lambdaPy = 1.5964638752e-9  # eyn = 50 pm-rad
beam.lambdaPt = 9.0e-4           # approximately dE/E
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
beam.polarization_x = 0.4
beam.polarization_y = 0.9
beam.polarization_z = 0.1

###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor cfbend1 iquad1 cfbend2 isbend1 monitor

lattice.nslice = 1

#quad-like cfbend
cfbend1.type = cfbend
cfbend1.ds = 1.5
cfbend1.rc = 1.0e8   # bending radius [m]
cfbend1.k = 60.0   # (upright) quadrupole component [m^(-2)]

iquad1.type = quad  #inverse quad
iquad1.ds = -1.5
iquad1.k = 60.0

#sbend-like cfbend
cfbend2.type = cfbend
cfbend2.ds = 1.5
cfbend2.rc = 10.0   # bending radius [m]
cfbend2.k = 1.0e-8   # (upright) quadrupole component [m^(-2)]

isbend1.type = sbend  #inverse sbend
isbend1.ds = -1.5
isbend1.rc = 10.0

monitor.type = beam_monitor
monitor.backend = h5


###############################################################################
# Algorithms
###############################################################################
algo.space_charge = false
algo.spin = true

###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false

Analyze

The analysis can be run using:

Script analysis_cfbend_spin.py
Listing 193 You can copy this file from examples/spin_tracking/analysis_cfbend_spin.py.
#!/usr/bin/env python3
#
# Copyright 2022-2026 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

import numpy as np
import openpmd_api as io

# 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()
beam_final = series.iterations[last_step].particles["beam"]
final = beam_final.to_df()

sxi = initial["spin_x"]
syi = initial["spin_y"]
szi = initial["spin_z"]

sxf = final["spin_x"]
syf = final["spin_y"]
szf = final["spin_z"]

dspin2 = (sxf - sxi) ** 2 + (syf - syi) ** 2 + (szf - szi) ** 2
dspin = np.sqrt(dspin2)
dspinmax = dspin.max()

print("Change in the spin:")
print("||delta s||_max", dspinmax)

atol = 6.0e-8
print(f"  atol={atol}")

assert np.allclose(
    [dspinmax],
    [
        0.0,
    ],
    atol=atol,
)

Illustration of a User-Defined Spin Map

This example illustrates the application of a user-defined map that affects only the spin variables \((s_x,s_y,s_z)\). Spin maps are specified in the Lie-algebraic form:

\(\vec{s}_f = M(\zeta)\vec{s}_i\) where \(M(\zeta)=e^{v\cdot L}e^{A\Delta\zeta\cdot L}\).

Here \(v\) is a 3-vector that defines the axis and angle of rotation at the phase space design point, and \(A\) is a 3x6 matrix that defines the spin-orbit coupling for particles not on the design orbit. Also, \(\Delta\zeta=(x,p_x,y,p_y,t,p_t)\) denotes the 6-vector of phase space variables as deviations from the design orbit. The quantities \(L_x\), \(L_y\), and \(L_z\) are standard 3x3 matrices that define a basis for the Lie algebra \(so(3)\).

In this test, the inverses of the spin maps for a quadrupole and a sector bend are provided by the user as input. The inverse spin map for a quadrupole is composed with spin tracking in a quadrupole element, which should return all spin variables to their initial values. The same procedure is repeated for a sector bend.

In this test, the vector norm of the difference between the initial and final spins for all particles tracked must vanish to machine precision.

Run

This example can be run either as:

  • Python script: python3 run_spin_map.py or

  • ImpactX executable using an input file: impactx input_spin_map.in

For MPI-parallel runs, prefix these lines with mpiexec -n 4 ... or srun -n 4 ..., depending on the system.

Listing 194 You can copy this file from examples/spin_tracking/run_spin_map.py.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

# from elements import LinearTransport

from impactx import ImpactX, Map3x6, Vector3, distribution, elements, twiss

sim = ImpactX()

# set numerical parameters and IO control
sim.space_charge = False
sim.spin = True
sim.slice_step_diagnostics = False

# domain decomposition & space charge mesh
sim.init_grids()

# set beam reference values
kin_energy_MeV = 10.0e3  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

#   reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

#   particle bunch
distr = distribution.Gaussian(
    **twiss(
        beta_x=2.8216194100262637,
        beta_y=2.8216194100262637,
        beta_t=0.5,
        emitt_x=2e-04,
        emitt_y=2e-04,
        emitt_t=2e-12,
        alpha_x=-1.5905003499999992,
        alpha_y=1.5905003499999992,
        alpha_t=0.0,
    )
)
spin = distribution.SpinvMF(
    0.4,
    0.9,
    0.1,
)
sim.add_particles(bunch_charge_C, distr, npart, spin)

# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")

# design the accelerator lattice

# Elements for forward tracking
Q1 = elements.Quad(name="Q1", ds=1.0, k=1.0)
B1 = elements.Sbend(name="B1", ds=1.0, rc=10.0)

# Inverse spin map for Quad1
vmatQ1 = Vector3()
AmatQ1 = Map3x6()
AmatQ1[1, 3] = 27.846376647658047
AmatQ1[1, 4] = 12.868288416407257
AmatQ1[2, 1] = 19.9386437894808211
AmatQ1[2, 2] = 10.8925307463018033
Map1 = elements.SpinMap(v=vmatQ1, A=AmatQ1)

# Inverse spin map for Bend1
vmatB1 = Vector3()
AmatB1 = Map3x6()
vmatB1[2] = 2.269498669527234
AmatB1[1, 4] = 1.643140677773437
AmatB1[2, 1] = 0.236555147919017
AmatB1[2, 2] = 0.118376237268958
AmatB1[2, 6] = 0.096052815713841
AmatB1[3, 4] = -0.765638392807848
Map2 = elements.SpinMap(v=vmatB1, A=AmatB1)

line = [monitor, Map1, Q1, Map2, B1, monitor]

sim.lattice.extend(line)

# number of periods through the lattice
sim.periods = 1

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 195 You can copy this file from examples/spin_tracking/input_spin_map.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 10.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = gaussian_from_twiss
beam.alphaX = -1.5905003499999992
beam.alphaY = -beam.alphaX
beam.alphaT = 0.0
beam.betaX = 2.8216194100262637
beam.betaY = beam.betaX
beam.betaT = 0.5
beam.emittX = 2e-04
beam.emittY = 2e-04
beam.emittT = 2e-12
beam.polarization_x = 0.4
beam.polarization_y = 0.9
beam.polarization_z = 0.1

###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor map1 quad1 map2 bend1 monitor
lattice.nslice = 1

monitor.type = beam_monitor
monitor.backend = h5

quad1.type = quad
quad1.ds = 1.0
quad1.k = 1.0

map1.type = spin_map  # Inverse spin map for quad1
map1.A13 = 27.846376647658047
map1.A14 = 12.868288416407257
map1.A21 = 19.9386437894808211
map1.A22 = 10.8925307463018033

bend1.type = sbend
bend1.ds = 1.0
bend1.rc = 10.0

map2.type = spin_map  # Inverse spin map for bend1
map2.v2 = 2.269498669527234
map2.A14 = 1.643140677773437
map2.A21 = 0.236555147919017
map2.A22 = 0.118376237268958
map2.A26 = 0.096052815713841
map2.A34 = -0.765638392807848

###############################################################################
# Algorithms
###############################################################################
algo.space_charge = false
algo.spin = true

###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false

Analyze

The analysis can be run using:

Script analysis_spin_map.py
Listing 196 You can copy this file from examples/spin_tracking/analysis_spin_map.py.
#!/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

# 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()
beam_final = series.iterations[last_step].particles["beam"]
final = beam_final.to_df()

sxi = initial["spin_x"]
syi = initial["spin_y"]
szi = initial["spin_z"]

sxf = final["spin_x"]
syf = final["spin_y"]
szf = final["spin_z"]

dspin2 = (sxf - sxi) ** 2 + (syf - syi) ** 2 + (szf - szi) ** 2
dspin = np.sqrt(dspin2)
dspinmax = dspin.max()

print("Change in the spin:")
print("||delta s||_max", dspinmax)

atol = 5.0e-9
print(f"  atol={atol}")

assert np.allclose(
    [dspinmax],
    [
        0.0,
    ],
    atol=atol,
)

Test of Thomas-BMT Spin Integration in a FODO Channel

This example tests the propagation of spin in a single period of a FODO channel that uses symplectic integration for tracking via the ExactQuad element.

The ExactQuad element incorportates the full nonlinear phase space dependence appearing in both the Hamiltonian and the Thomas-BMT equations for tracking.

The physical parameters of this FODO channel are identical with those used in examples-fodo-spin.

A small number of initial conditions are tracked, with increasing distance from the phase space design point.

The final spin variables are compared against those obtained by tracking using the Quad element, which treats both orbit and spin variables to first order in the phase space variables.

In this test, a scaling exponent is extracted when comparing the vector-norm of difference in final spin with the vector-norm of the initial phase space vector. This exponent should be 2 (in the limit of small deviation).

Run

This example can be run as:

  • Python script: python3 run_exact_quad_spin_scaling.py

For MPI-parallel runs, prefix these lines with mpiexec -n 4 ... or srun -n 4 ..., depending on the system.

Listing 197 You can copy this file from examples/spin_tracking/run_exact_quad_spin_scaling.py.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

import pandas as pd

import amrex.space3d as amr
from impactx import Config, ImpactX, elements

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

# beam parameters
kin_energy_MeV = 10.0e3  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

#   reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

qm_eev = 1.0 / 0.510998950 / 1e6  # electron charge/mass in e / eV

beam = sim.beam

if amr.ParallelDescriptor.IOProcessor():
    df_initial = pd.read_csv("./initial_coords.csv", sep=" ")
    dx = df_initial["x"].to_numpy()
    dpx = df_initial["px"].to_numpy()
    dy = df_initial["y"].to_numpy()
    dpy = df_initial["py"].to_numpy()
    dt = df_initial["t"].to_numpy()
    dpt = df_initial["pt"].to_numpy()
    dw = df_initial["w"].to_numpy()
    dsx = df_initial["sx"].to_numpy()
    dsy = df_initial["sy"].to_numpy()
    dsz = df_initial["sz"].to_numpy()
    if not Config.have_gpu:  # initialize using cpu-based PODVectors
        dx_podv = amr.PODVector_real_std()
        dy_podv = amr.PODVector_real_std()
        dt_podv = amr.PODVector_real_std()
        dpx_podv = amr.PODVector_real_std()
        dpy_podv = amr.PODVector_real_std()
        dpt_podv = amr.PODVector_real_std()
        dw_podv = amr.PODVector_real_std()
        dsx_podv = amr.PODVector_real_std()
        dsy_podv = amr.PODVector_real_std()
        dsz_podv = amr.PODVector_real_std()

    else:  # initialize on device using arena/gpu-based PODVectors
        dx_podv = amr.PODVector_real_arena()
        dy_podv = amr.PODVector_real_arena()
        dt_podv = amr.PODVector_real_arena()
        dpx_podv = amr.PODVector_real_arena()
        dpy_podv = amr.PODVector_real_arena()
        dpt_podv = amr.PODVector_real_arena()
        dw_podv = amr.PODVector_real_arena()
        dsx_podv = amr.PODVector_real_arena()
        dsy_podv = amr.PODVector_real_arena()
        dsz_podv = amr.PODVector_real_arena()

    for p_dx in dx:
        dx_podv.push_back(p_dx)
    for p_dy in dy:
        dy_podv.push_back(p_dy)
    for p_dt in dt:
        dt_podv.push_back(p_dt)
    for p_dpx in dpx:
        dpx_podv.push_back(p_dpx)
    for p_dpy in dpy:
        dpy_podv.push_back(p_dpy)
    for p_dpt in dpt:
        dpt_podv.push_back(p_dpt)
    for p_dw in dw:
        dw_podv.push_back(p_dw)
    for p_dsx in dsx:
        dsx_podv.push_back(p_dsx)
    for p_dsy in dsy:
        dsy_podv.push_back(p_dsy)
    for p_dsz in dsz:
        dsz_podv.push_back(p_dsz)

    beam.add_n_particles(
        dx_podv,
        dy_podv,
        dt_podv,
        dpx_podv,
        dpy_podv,
        dpt_podv,
        qm_eev,
        bunch_charge_C,
        None,
        dsx_podv,
        dsy_podv,
        dsz_podv,
    )

# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5", period_sample_intervals=10)

# design the accelerator lattice)
ns = 1  # number of slices per ds in the element
order = 4  # order of symplectic integration
nmap = 10  # number of steps for symplectic integration
fodo = [
    monitor,
    elements.ExactDrift(ds=0.25, nslice=ns),
    elements.ExactQuad(ds=1.0, k=1.0, int_order=order, mapsteps=nmap, nslice=ns),
    elements.ExactDrift(ds=0.5, nslice=ns),
    elements.ExactQuad(ds=1.0, k=-1.0, int_order=order, mapsteps=nmap, nslice=ns),
    elements.ExactDrift(ds=0.25, nslice=ns),
    monitor,
]
# assign a fodo segment
sim.lattice.extend(fodo)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()

Analyze

The analysis can be run using the following script:

Script analysis_exact_quad_spin_scaling.py
Listing 198 You can copy this file from examples/spin_tracking/analysis_exact_quad_spin_scaling.py.
#!/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
import pandas as pd

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

# load particle data
df_initial = pd.read_csv("./initial_coords.csv", sep=" ")
df_final = pd.read_csv("./final_coords.csv", sep=" ")

# initial coordinates
xi = df_initial["x"]
pxi = df_initial["px"]
yi = df_initial["y"]
pyi = df_initial["py"]
ti = df_initial["t"]
pti = df_initial["pt"]
sxi = df_initial["sx"]
syi = df_initial["sy"]
szi = df_initial["sz"]
norm_phase_space_vector = np.sqrt(xi**2 + pxi**2 + yi**2 + pyi**2 + ti**2 + pti**2)

# difference of final coordinates from linear values
diff_xf = df_final["x"] - final["position_x"]
diff_pxf = df_final["px"] - final["momentum_x"]
diff_yf = df_final["y"] - final["position_y"]
diff_pyf = df_final["py"] - final["momentum_y"]
diff_tf = df_final["t"] - final["position_t"]
diff_ptf = df_final["pt"] - final["momentum_t"]
diff_sxf = df_final["sx"] - final["spin_x"]
diff_syf = df_final["sy"] - final["spin_y"]
diff_szf = df_final["sz"] - final["spin_z"]
norm_spin_diff = np.sqrt(diff_sxf**2 + diff_syf**2 + diff_szf**2)

print("")
print(":")
print("Phase space vector norm:")
print(norm_phase_space_vector)

print(":")
print("Spin difference norm:")
print(norm_spin_diff)

m, b = np.polyfit(np.log(norm_phase_space_vector), np.log(norm_spin_diff), deg=1)
print(f"  slope={m}")

# Test for  quadratic scaling with initial phase space vector:
rtol = 2.0e-2
print(f"  rtol={rtol}")
assert np.allclose(
    [m],
    [
        2.0,
    ],
    rtol=rtol,
)

Spin Tracking in a Multipole

The example illustrates particle tracking with spin in both a thin multipole (Multipole) and a thick multipole (ExactMultipole).

This example tests general multipole spin tracking in two ways. First, tracking through an ExactQuad is compared against tracking through an ExactMultipole, with the same quadrupole strength. To achieve this, forward tracking through the ExactMultipole is followed by reverse tracking through an ExactQuad of the same length.

Second, tracking through an ExactMultipole with a short value of length is compared against tracking through a thin Multipole element, using the same method.

In this test, the initial and final values of the beam moments \(\sigma_x\), \(\sigma_y\), \(\sigma_t\), \(\epsilon_x\), \(\epsilon_y\), and \(\epsilon_t\) must coincide.

At the same time, the initial and final spin components \(s_x\), \(s_y\), and \(s_z\) are compared for all particles. The norm of the change in the spin vector must be zero to, within a small tolerance.

Run

This example can be run either as:

  • Python script: python3 run_multipole_spin.py or

  • ImpactX executable using an input file: impactx input_multipole_spin.in

For MPI-parallel runs, prefix these lines with mpiexec -n 4 ... or srun -n 4 ..., depending on the system.

Listing 199 You can copy this file from examples/spin_tracking/run_multipole_spin.py.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: 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.space_charge = False
# sim.diagnostics = False  # benchmarking
sim.spin = True
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 2 nm
kin_energy_MeV = 2.0e3  # 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_species("electron").set_kin_energy_MeV(kin_energy_MeV)

#   particle bunch
distr = distribution.Waterbag(
    lambdaX=3.9984884770e-4,
    lambdaY=3.9984884770e-4,
    lambdaT=1.0e-3,
    lambdaPx=2.6623538760e-4,
    lambdaPy=2.6623538760e-4,
    lambdaPt=2.0e-3,
    muxpx=-0.846574929020762,
    muypy=0.846574929020762,
    mutpt=0.0,
)
spin = distribution.SpinvMF(
    0.4,
    0.9,
    0.1,
)
sim.add_particles(bunch_charge_C, distr, npart, spin)

# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")

# design the accelerator lattice)
ns = 5  # number of slices per ds in the element
Kn_quad = 2.0
L = 1.0
Kn_quad_integrated = -0.2
Ks_quad_integrated = 0.1
Kn_sext_integrated = 1.0
Ks_sext_integrated = -0.5
L2 = 1.0e-5

lattice = [
    monitor,
    elements.ExactMultipole(
        name="multipole",
        ds=L,
        k_normal=[0.0, Kn_quad],
        k_skew=[0.0, 0.0],
        mapsteps=100,
        nslice=ns,
    ),
    elements.ExactQuad(
        name="quad",
        ds=-L,
        k=Kn_quad,
        mapsteps=100,
        nslice=ns,
    ),
    #  Test 2:  thick multipole vs thin multipole
    elements.Multipole(
        name="thin_multipole",
        multipole=2,
        K_normal=-Kn_quad_integrated,
        K_skew=-Ks_quad_integrated,
    ),
    elements.Multipole(
        name="thin_multipole",
        multipole=3,
        K_normal=-Kn_sext_integrated,
        K_skew=-Ks_sext_integrated,
    ),
    elements.ExactMultipole(
        name="thick_multipole",
        ds=L2,
        k_normal=[0.0, Kn_quad_integrated / L2, Kn_sext_integrated / L2],
        k_skew=[0.0, Ks_quad_integrated / L2, Ks_sext_integrated / L2],
        mapsteps=100,
        nslice=ns,
    ),
    monitor,
]
sim.lattice.extend(lattice)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 200 You can copy this file from examples/spin_tracking/input_multipole_spin.in.
###############################################################################
# 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 = 3.9984884770e-4
beam.lambdaY = beam.lambdaX
beam.lambdaT = 1.0e-3
beam.lambdaPx = 2.6623538760e-4
beam.lambdaPy = beam.lambdaPx
beam.lambdaPt = 2.0e-3
beam.muxpx = -0.846574929020762
beam.muypy = -beam.muxpx
beam.mutpt = 0.0
beam.polarization_x = 0.4
beam.polarization_y = 0.9
beam.polarization_z = 0.1

###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor multipole quad thin_quad thin_sext thick_multipole monitor
lattice.nslice = 5

monitor.type = beam_monitor
monitor.backend = h5

multipole.type = multipole_exact
multipole.ds = 1.0
multipole.k_normal = 0 2.0
multipole.k_skew = 0 0
multipole.units = 0
multipole.mapsteps = 100

quad.type = quad_exact
quad.ds = -1.0
quad.k = 2.0
quad.mapsteps = 100

thin_quad.type = multipole
thin_quad.multipole = 2       # Thin quadrupole
thin_quad.k_normal = 0.2
thin_quad.k_skew = -0.1

thin_sext.type = multipole
thin_sext.multipole = 3       # Thin sextupole
thin_sext.k_normal = -1.0
thin_sext.k_skew = 0.5

thick_multipole.type = multipole_exact
thick_multipole.ds = 1.0e-5
thick_multipole.k_normal = 0 -0.2e5 1.0e5
thick_multipole.k_skew = 0 0.1e5 -0.5e5
thick_multipole.mapsteps = 100


###############################################################################
# Algorithms
###############################################################################
algo.space_charge = false
algo.spin = true


###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true

Analyze

The analysis can be run using the following script:

Script analysis_multipole_spin.py
Listing 201 You can copy this file from examples/spin_tracking/analysis_multipole_spin.py.
#!/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()
beam_final = series.iterations[last_step].particles["beam"]
final = beam_final.to_df()

# compare number of particles
num_particles = 10000
assert num_particles == len(initial)
assert num_particles == len(final)

print("Compare initial and final beam moments:")

sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial)
sigxf, sigyf, sigtf, emittance_xf, emittance_yf, emittance_tf = get_moments(final)

atol = 0.0  # ignored
rtol = 3.0 * 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],
    [sigxf, sigyf, sigtf, emittance_xf, emittance_yf, emittance_tf],
    rtol=rtol,
    atol=atol,
)

print("Compare initial and final spins:")

sxi = initial["spin_x"]
syi = initial["spin_y"]
szi = initial["spin_z"]

sxf = final["spin_x"]
syf = final["spin_y"]
szf = final["spin_z"]

dspin2 = (sxf - sxi) ** 2 + (syf - syi) ** 2 + (szf - szi) ** 2
dspin = np.sqrt(dspin2)
dspinmax = dspin.max()

print("Change in the spin:")
print("||delta s||_max", dspinmax)

is_double = np.dtype(sxi.dtype) == np.dtype(np.float64)
atol = 5.0e-8 if is_double else 2.2e-5
print(f"  dtype={sxi.dtype}, atol={atol}")

assert np.allclose(
    [dspinmax],
    [
        0.0,
    ],
    atol=atol,
)