BLAST-ImpactX · ax3l · Dec 16, 2025 · Nov 21, 2025 · Nov 24, 2025 · Nov 24, 2025
diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml
@@ -64,4 +64,4 @@ jobs:
           -DImpactX_PRECISION=SINGLE   \
           -DAMReX_CUDA_ERROR_CROSS_EXECUTION_SPACE_CALL=ON \
           -DAMReX_CUDA_ERROR_CAPTURE_THIS=ON
-        cmake --build build -j 4
+        cmake --build build -j 3
diff --git a/cmake/dependencies/ABLASTR.cmake b/cmake/dependencies/ABLASTR.cmake
@@ -186,7 +186,7 @@ set(ImpactX_ablastr_branch "25.12"
 set(ImpactX_amrex_repo "https://github.com/AMReX-Codes/amrex.git"
     CACHE STRING
     "Repository URI to pull and build AMReX from if(ImpactX_amrex_internal)")
-set(ImpactX_amrex_branch ""
+set(ImpactX_amrex_branch "4dad1664d7467ae00c133c1425c1a300003fa885"
     CACHE STRING
     "Repository branch for ImpactX_amrex_repo if(ImpactX_amrex_internal)")
 

diff --git a/cmake/dependencies/pyAMReX.cmake b/cmake/dependencies/pyAMReX.cmake
@@ -74,7 +74,7 @@ option(ImpactX_pyamrex_internal "Download & build pyAMReX" ON)
 set(ImpactX_pyamrex_repo "https://github.com/AMReX-Codes/pyamrex.git"
     CACHE STRING
     "Repository URI to pull and build pyamrex from if(ImpactX_pyamrex_internal)")
-set(ImpactX_pyamrex_branch "25.12"
+set(ImpactX_pyamrex_branch "74e213e3186371d5a7412af3d2a5d6040fab51ce"
     CACHE STRING
     "Repository branch for ImpactX_pyamrex_repo if(ImpactX_pyamrex_internal)")
 

diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst
@@ -138,6 +138,16 @@ Initial Beam Distributions
     * ``beam.normalize_halo`` (``float``, dimensionless) normalizing constant for halo population
     * ``beam.halo`` (``float``, dimensionless) fraction of charge in halo
 
+Initial Spin Distributions
+--------------------------
+
+  The specification of an initial particle spin distribution is optional, and is required only if spin tracking is used.
+  The default distribution type is the von Mises-Fisher distribution, uniquely determined by the input polarization vector.
+  The polarization vector provided by the user must lie within the unit ball.
+
+* ``beam.polarization_x`` (``float``, dimensionless) mean value of the spin vector x-component
+* ``beam.polarization_y`` (``float``, dimensionless) mean value of the spin vector y-component
+* ``beam.polarization_z`` (``float``, dimensionless) mean value of the spin vector z-component
 
 .. _running-cpp-parameters-lattice:
 

diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst
@@ -227,7 +227,7 @@ Collective Effects & Overall Simulation Parameters
 
       This must come first, before particle beams and lattice elements are initialized.
 
-   .. py:method:: add_particles(charge_C, distr, npart)
+   .. py:method:: add_particles(charge_C, distr, npart, spinv=None)
 
       Particle tracking mode: Generate and add n particles to the particle container.
       Note: Set the reference particle properties (charge, mass, energy) first.
@@ -237,6 +237,7 @@ Collective Effects & Overall Simulation Parameters
       :param float charge_C: bunch charge (C)
       :param distr: distribution function to draw from (object from :py:mod:`impactx.distribution`)
       :param int npart: number of particles to draw
+      :param SpinvMF spinv: optional spin distribution
 
    .. py:method:: init_envelope(ref, distr, intensity=None)
 
@@ -578,8 +579,8 @@ Particles
       :param madx_file: file name to MAD-X file with a ``BEAM`` entry
 
 
-Initial Beam Distributions
---------------------------
+Initial Beam Phase Space Distributions
+--------------------------------------
 
 This module provides particle beam distributions that can be used to initialize particle beams in an :py:class:`impactx.ParticleContainer`.
 
@@ -659,6 +660,25 @@ For the input from Twiss parameters in Python, please use the helper function ``
 
    A 6D stationary thermal or bithermal distribution.
 
+Initial Beam Spin Distribution
+------------------------------
+
+.. py:class:: impactx.distribution.SpinvMF(mux, muy, muz)
+
+   A von Mises-Fisher (vMF) distribution on the unit 2-sphere.
+
+   This is used for initializing particle spin. There is a natural bijective correspondence between vMF distributions and mean (polarization) vectors.
+
+   The algorithm used here is a simplification of the algorithm described in:
+   C. Pinzon and K. Jung, "Fast Python sampler of the von Mises Fisher distribution", in the special case of the 2-sphere. Additional references used include:
+
+   - K. V. Mardia and P. E. Jupp, Directional Statistics, Wiley, 1999;
+   - S. Kang and H-S. Oh, "Novel sampling method for the von Mises-Fisher distribution", Stat. and Comput. 34, 106 (2024), `DOI:10.1007/s11222-024-10419-3 <https://doi.org/10.1007/s11222-024-10419-3>`__
+
+   :param mux: x component of the unit vector specifying the mean direction
+   :param muy: y component of the unit vector specifying the mean direction
+   :param muz: z component of the unit vector specifying the mean direction
+
 
 Lattice Elements
 ----------------

diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -1815,3 +1815,18 @@ add_impactx_test(dipedge-nonlinear.py
     examples/edge_effects/analysis_dipedge.py
     OFF  # no plot script yet
 )
+
+# Spin Sampling from a vMF Distribution ######################
+# without space charge
+add_impactx_test(distgen-spinvmf
+    examples/distgen/input_kurth4d_spin.in
+    ON   # ImpactX MPI-parallel
+    examples/distgen/analysis_kurth4d_spin.py
+    OFF  # no plot script yet
+)
+add_impactx_test(distgen-spinvmf.py
+    examples/distgen/run_kurth4d_spin.py
+    ON    # ImpactX MPI-parallel
+    examples/distgen/analysis_kurth4d_spin.py
+    OFF  # no plot script yet
+)
diff --git a/examples/distgen/README.rst b/examples/distgen/README.rst
@@ -219,3 +219,50 @@ We run the following script to analyze correctness:
    .. literalinclude:: analysis_semigaussian.py
       :language: python3
       :caption: You can copy this file from ``examples/distgen/analysis_semigaussian.py``.
+
+
+.. _examples-distgen-spinvmf:
+
+Spin Sampling from a vMF Distribution
+=====================================
+
+This tests the sampling of initial particle spin from a von Mises-Fisher distribution, given an initial input polarization.
+The phase space distribution coincides with the 4D Kurth distribution used in examples-distgen-kurth4d.
+
+In this test, the initial and final values of of the mean spin 3-vector (i.e., the polarization vector) must agree with nominal values.
+
+Run
+---
+
+This example can be run **either** as:
+
+* **Python** script: ``python3 run_kurth4d_spin.py`` or
+* ImpactX **executable** using an input file: ``impactx input_kurth4d_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_kurth4d_spin.py
+          :language: python3
+          :caption: You can copy this file from ``examples/distgen/run_kurth4d_spin.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_kurth4d_spin.in
+          :language: ini
+          :caption: You can copy this file from ``examples/distgen/input_kurth4d_spin.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_kurth4d_spin.py``
+
+   .. literalinclude:: analysis_kurth4d_spin.py
+      :language: python3
+      :caption: You can copy this file from ``examples/distgen/analysis_kurth4d_spin.py``.
diff --git a/examples/distgen/analysis_kurth4d_spin.py b/examples/distgen/analysis_kurth4d_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.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],
+    [
+        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.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],
+    [
+        0.7,
+        0.0,
+        0.0,
+    ],
+    atol=atol,
+)
diff --git a/examples/distgen/input_kurth4d_spin.in b/examples/distgen/input_kurth4d_spin.in
@@ -0,0 +1,41 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.kin_energy = 2.0e3
+beam.charge = 1.0e-9
+beam.particle = proton
+beam.distribution = kurth4d
+beam.lambdaX = 1.0e-3
+beam.lambdaY = beam.lambdaX
+beam.lambdaT = 1.0e-3
+beam.lambdaPx = 1.0e-3
+beam.lambdaPy = beam.lambdaPx
+beam.lambdaPt = 2.0e-3
+beam.muxpx = 0.0
+beam.muypy = 0.0
+beam.mutpt = 0.0
+beam.polarization_x = 0.7
+beam.polarization_y = 0.0
+beam.polarization_z = 0.0
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor constf1 monitor
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+constf1.type = constf
+constf1.ds = 2.0
+constf1.kx = 1.0
+constf1.ky = 1.0
+constf1.kt = 1.0e-4
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.space_charge = false
diff --git a/examples/distgen/run_kurth4d_spin.py b/examples/distgen/run_kurth4d_spin.py
@@ -0,0 +1,68 @@
+#!/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.slice_step_diagnostics = True
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a 2 GeV proton 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.particle_container().ref_particle()
+ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
+
+#   particle bunch
+distr = distribution.Kurth4D(
+    lambdaX=1.0e-3,
+    lambdaY=1.0e-3,
+    lambdaT=1.0e-3,
+    lambdaPx=1.0e-3,
+    lambdaPy=1.0e-3,
+    lambdaPt=2.0e-3,
+    muxpx=-0.0,
+    muypy=0.0,
+    mutpt=0.0,
+)
+spinv = distribution.SpinvMF(
+    0.7,
+    0.0,
+    0.0,
+)
+
+sim.add_particles(bunch_charge_C, distr, npart, spinv)
+
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+
+# design the accelerator lattice)
+constf = [
+    monitor,
+    elements.ConstF(name="constf1", ds=2.0, kx=1.0, ky=1.0, kt=1.0e-4),
+    monitor,
+]
+
+# assign a constf segment
+sim.lattice.extend(constf)
+
+# run simulation
+sim.track_particles()
+
+# clean shutdown
+sim.finalize()
diff --git a/examples/solenoid_restart/analysis_solenoid.py b/examples/solenoid_restart/analysis_solenoid.py
@@ -79,7 +79,7 @@ def get_moments(beam):
 )
 
 atol = 0.0  # ignored
-rtol = 2.3 * num_particles**-0.5  # from random sampling of a smooth distribution
+rtol = 1.4 * num_particles**-0.5  # from random sampling of a smooth distribution
 print(f"  rtol={rtol} (ignored: atol~={atol})")
 
 assert np.allclose(

diff --git a/src/ImpactX.H b/src/ImpactX.H
@@ -85,12 +85,14 @@ namespace impactx
          * @param bunch_charge bunch charge (C)
          * @param distr distribution function to draw from (object)
          * @param npart number of particles to draw
+         * @param spin_distr optional spin distribution
          */
         void
         add_particles (
             amrex::ParticleReal bunch_charge,
             distribution::KnownDistributions distr,
-            amrex::Long npart
+            amrex::Long npart,
+            std::optional<distribution::SpinvMF> spin_distr = std::nullopt
         );
 
         /** Validate the simulation is ready to run the particle tracking loop via @see track_particles