Add Spin to Particle Container

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

cmake/dependencies/pyAMReX.cmake

@@ -71,10 +71,10 @@ set(ImpactX_pyamrex_src ""
 
 # Git fetcher
 option(ImpactX_pyamrex_internal "Download & build pyAMReX" ON)
-set(ImpactX_pyamrex_repo "https://github.com/AMReX-Codes/pyamrex.git"
+set(ImpactX_pyamrex_repo "https://github.com/ax3l/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 "topic-pc-impactx-spin"
     CACHE STRING
     "Repository branch for ImpactX_pyamrex_repo if(ImpactX_pyamrex_internal)")
 

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:
 

examples/CMakeLists.txt

@@ -1815,3 +1815,19 @@ 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
+)
+#TODO: add Python test
+#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
+#)

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

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

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

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

src/initialization/InitDistribution.cpp

@@ -14,6 +14,7 @@
 #include "particles/CovarianceMatrix.H"
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/distribution/All.H"
+#include "particles/distribution/SpinvMF.H"
 #include "particles/SplitEqually.H"
 
 #include <ablastr/constant.H>
@@ -298,7 +299,8 @@ namespace impactx
     ImpactX::add_particles (
         amrex::ParticleReal bunch_charge,
         distribution::KnownDistributions distr,
-        amrex::Long npart
+        amrex::Long npart,
+        std::optional<distribution::SpinvMF> spin_distr
     )
     {
         BL_PROFILE("ImpactX::add_particles");
@@ -340,13 +342,21 @@ namespace impactx
         // alloc data for particle attributes
         amrex::Gpu::DeviceVector<amrex::ParticleReal> x, y, t;
         amrex::Gpu::DeviceVector<amrex::ParticleReal> px, py, pt;
+        amrex::Gpu::DeviceVector<amrex::ParticleReal> sx, sy, sz;
         x.resize(npart_this_proc);
         y.resize(npart_this_proc);
         t.resize(npart_this_proc);
         px.resize(npart_this_proc);
         py.resize(npart_this_proc);
         pt.resize(npart_this_proc);
 
+        bool const has_spin = spin_distr.has_value();
+        if (has_spin) {
+            sx.resize(npart_this_proc);
+            sy.resize(npart_this_proc);
+            sz.resize(npart_this_proc);
+        }
+
         std::visit([&](auto&& distribution){
             // initialize distributions
             distribution.initialize(bunch_charge, ref);
@@ -357,6 +367,9 @@ namespace impactx
             amrex::ParticleReal * const AMREX_RESTRICT px_ptr = px.data();
             amrex::ParticleReal * const AMREX_RESTRICT py_ptr = py.data();
             amrex::ParticleReal * const AMREX_RESTRICT pt_ptr = pt.data();
+            amrex::ParticleReal * const AMREX_RESTRICT sx_ptr = sx.data();
+            amrex::ParticleReal * const AMREX_RESTRICT sy_ptr = sy.data();
+            amrex::ParticleReal * const AMREX_RESTRICT sz_ptr = sz.data();
 
             using Distribution = std::decay_t<decltype(distribution)>;
 
@@ -376,6 +389,7 @@ namespace impactx
                 npart_this_thread = thread_chunk.size;
 #endif
 
+                // phase space init
                 initialization::InitSingleParticleData<Distribution> const init_single_particle_data(
                     distribution,
                     x_ptr + my_offset,
@@ -387,16 +401,43 @@ namespace impactx
                 );
 
                 amrex::ParallelForRNG(npart_this_thread, init_single_particle_data);
+
+                // spin init
+                if (has_spin) {
+                    auto spinv = spin_distr.value();
+                    amrex::ParallelForRNG(npart_this_thread,
+                        [=] AMREX_GPU_DEVICE (int i, amrex::RandomEngine const& engine) noexcept {
+                        spinv(
+                            sx_ptr[i + my_offset],
+                            sy_ptr[i + my_offset],
+                            sz_ptr[i + my_offset],
+                            engine
+                        );
+                    });
+                }
             }
 
             // finalize distributions and deallocate temporary device global memory
             amrex::Gpu::streamSynchronize();
             distribution.finalize();
         }, distr);
 
-        amr_data->track_particles.m_particle_container->AddNParticles(x, y, t, px, py, pt,
-                                                      ref.qm_ratio_SI(),
-                                            bunch_charge * rel_part_this_proc);
+        if (has_spin) {
+            amr_data->track_particles.m_particle_container->AddNParticles(
+                x, y, t,
+                px, py, pt,
+                ref.qm_ratio_SI(),
+                bunch_charge * rel_part_this_proc, std::nullopt,
+                sx, sy, sz
+            );
+        } else {
+            amr_data->track_particles.m_particle_container->AddNParticles(
+                x, y, t,
+                px, py, pt,
+                ref.qm_ratio_SI(),
+                bunch_charge * rel_part_this_proc
+            );
+        }
 
         auto space_charge = get_space_charge_algo();
 
@@ -551,15 +592,36 @@ namespace impactx
             std::string unit_type;  // System of units
             pp_dist.get("units", unit_type);
 
+            // phase space distribution
             distribution::KnownDistributions dist = initialization::read_distribution(pp_dist);
             std::string distribution;
             pp_dist.get("distribution", distribution);
 
+            // spin distribution
+            amrex::ParticleReal polarization_x = 0.0_prt,
+                                polarization_y = 0.0_prt,
+                                polarization_z = 0.0_prt;
+            pp_dist.queryWithParser("polarization_x", polarization_x);
+            pp_dist.queryWithParser("polarization_y", polarization_y);
+            pp_dist.queryWithParser("polarization_z", polarization_z);
+            amrex::ParticleReal pmag = 0.0_prt;
+            amrex::ParticleReal kappa = 0.0_prt;
+
+            // magnitude of the polarization vector (= polarization magnitude)
+            pmag = std::sqrt(polarization_x*polarization_x+polarization_y*polarization_y+polarization_z*polarization_z);
+
+            if (pmag < 1) {
+               kappa = distribution::SpinvMF::inverse_Langevin(pmag);
+            } else {
+                throw std::runtime_error("Initial polarization magnitude must currently be < 1.");
+            }
+            distribution::SpinvMF spin_dist(polarization_x, polarization_y, polarization_z, kappa);
+
             amrex::Long npart = 0;  // Number of simulation particles
             if (distribution != "empty")
             {
                 pp_dist.getWithParser("npart", npart);
-                add_particles(bunch_charge, dist, npart);
+                add_particles(bunch_charge, dist, npart, spin_dist);
             }
 
             // print information on the initialized beam