From 37d09bae2aac65d3014899f08da4ffdb41e4e1a3 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Fri, 21 Nov 2025 09:16:54 -0800 Subject: [PATCH 01/28] Add spin to particle container. --- src/particles/ImpactXParticleContainer.H | 13 +++++++++++-- src/particles/ImpactXParticleContainer.cpp | 16 ++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/src/particles/ImpactXParticleContainer.H b/src/particles/ImpactXParticleContainer.H index d518b5bd6..708bff273 100644 --- a/src/particles/ImpactXParticleContainer.H +++ b/src/particles/ImpactXParticleContainer.H @@ -51,6 +51,9 @@ namespace impactx px, ///< momentum in x, scaled by the magnitude of the reference momentum [unitless] (at fixed s or t) py, ///< momentum in y, scaled by the magnitude of the reference momentum [unitless] (at fixed s or t) pt, ///< energy deviation, scaled by speed of light * the magnitude of the reference momentum [unitless] (at fixed s) + sx, ///< spin vector x-component [unitless] (at fixed s or t) + sy, ///< spin vector y-component [unitless] (at fixed s or t) + sz, ///< spin vector z-component [unitless] (at fixed s or t) qm, ///< charge to mass ratio, in q_e/m_e [q_e/eV] w, ///< particle weight, number of real particles represented by this macroparticle [unitless] nattribs ///< the number of attributes above (always last) @@ -63,9 +66,9 @@ namespace impactx }; //! named labels for fixed s - static constexpr auto names_s = { "position_x", "position_y", "position_t", "momentum_x", "momentum_y", "momentum_t", "qm", "weighting" }; + static constexpr auto names_s = { "position_x", "position_y", "position_t", "momentum_x", "momentum_y", "momentum_t", "spin_x", "spin_y", "spin_z", "qm", "weighting" }; //! named labels for fixed t - static constexpr auto names_t = { "position_x", "position_y", "position_z", "momentum_x", "momentum_y", "momentum_z", "qm", "weighting" }; + static constexpr auto names_t = { "position_x", "position_y", "position_z", "momentum_x", "momentum_y", "momentum_z", "spin_x", "spin_y", "spin_z", "qm", "weighting" }; static_assert(names_s.size() == nattribs); static_assert(names_t.size() == nattribs); }; @@ -176,6 +179,9 @@ namespace impactx * @param px momentum in x * @param py momentum in y * @param pt momentum in t + * @param sx spin component in x + * @param sy spin component in y + * @param sz spin component in z * @param qm charge over mass in 1/eV * @param bunch_charge total charge within a bunch in C * @param w weight of each particle: how many real particles to represent @@ -188,6 +194,9 @@ namespace impactx amrex::Gpu::DeviceVector const & px, amrex::Gpu::DeviceVector const & py, amrex::Gpu::DeviceVector const & pt, + amrex::Gpu::DeviceVector const & sx, + amrex::Gpu::DeviceVector const & sy, + amrex::Gpu::DeviceVector const & sz, amrex::ParticleReal qm, std::optional bunch_charge = std::nullopt, std::optional> w = std::nullopt diff --git a/src/particles/ImpactXParticleContainer.cpp b/src/particles/ImpactXParticleContainer.cpp index bb6ce3e4a..dea7fa7c1 100644 --- a/src/particles/ImpactXParticleContainer.cpp +++ b/src/particles/ImpactXParticleContainer.cpp @@ -240,6 +240,9 @@ namespace impactx amrex::Gpu::DeviceVector const & px, amrex::Gpu::DeviceVector const & py, amrex::Gpu::DeviceVector const & pt, + amrex::Gpu::DeviceVector const & sx, + amrex::Gpu::DeviceVector const & sy, + amrex::Gpu::DeviceVector const & sz, amrex::ParticleReal qm, std::optional bunch_charge, std::optional> w @@ -260,6 +263,9 @@ namespace impactx AMREX_ALWAYS_ASSERT(x.size() == px.size()); AMREX_ALWAYS_ASSERT(x.size() == py.size()); AMREX_ALWAYS_ASSERT(x.size() == pt.size()); + AMREX_ALWAYS_ASSERT(x.size() == sx.size()); + AMREX_ALWAYS_ASSERT(x.size() == sy.size()); + AMREX_ALWAYS_ASSERT(x.size() == sz.size()); if (has_w) { AMREX_ALWAYS_ASSERT(x.size() == w->size()); } // number of particles to add @@ -331,6 +337,9 @@ namespace impactx amrex::ParticleReal * const AMREX_RESTRICT px_arr = soa[RealSoA::px].dataPtr(); amrex::ParticleReal * const AMREX_RESTRICT py_arr = soa[RealSoA::py].dataPtr(); amrex::ParticleReal * const AMREX_RESTRICT pt_arr = soa[RealSoA::pt].dataPtr(); + amrex::ParticleReal * const AMREX_RESTRICT sx_arr = soa[RealSoA::sx].dataPtr(); + amrex::ParticleReal * const AMREX_RESTRICT sy_arr = soa[RealSoA::sy].dataPtr(); + amrex::ParticleReal * const AMREX_RESTRICT sz_arr = soa[RealSoA::sz].dataPtr(); amrex::ParticleReal * const AMREX_RESTRICT qm_arr = soa[RealSoA::qm].dataPtr(); amrex::ParticleReal * const AMREX_RESTRICT w_arr = soa[RealSoA::w ].dataPtr(); @@ -342,6 +351,9 @@ namespace impactx amrex::ParticleReal const * const AMREX_RESTRICT px_ptr = px.data(); amrex::ParticleReal const * const AMREX_RESTRICT py_ptr = py.data(); amrex::ParticleReal const * const AMREX_RESTRICT pt_ptr = pt.data(); + amrex::ParticleReal const * const AMREX_RESTRICT sx_ptr = sx.data(); + amrex::ParticleReal const * const AMREX_RESTRICT sy_ptr = sy.data(); + amrex::ParticleReal const * const AMREX_RESTRICT sz_ptr = sz.data(); amrex::ParticleReal const * const AMREX_RESTRICT w_ptr = has_w ? w->data() : nullptr; amrex::ParticleReal const bunch_charge_value = has_w ? 0_prt : bunch_charge.value(); @@ -358,6 +370,10 @@ namespace impactx py_arr[old_np+i] = py_ptr[my_offset+i]; pt_arr[old_np+i] = pt_ptr[my_offset+i]; + sx_arr[old_np+i] = sx_ptr[my_offset+i]; + sy_arr[old_np+i] = sy_ptr[my_offset+i]; + sz_arr[old_np+i] = sz_ptr[my_offset+i]; + qm_arr[old_np+i] = qm; w_arr[old_np+i] = has_w ? w_ptr[my_offset+i] : bunch_charge_value / ablastr::constant::SI::q_e/np; }); From 8602239289084947ae9705af0daa5842893da983 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Mon, 24 Nov 2025 15:11:29 -0800 Subject: [PATCH 02/28] Add spin sampling algorithm. --- src/particles/distribution/All.H | 4 +- src/particles/distribution/SpinvMF.H | 123 +++++++++++++++++++++++++++ 2 files changed, 126 insertions(+), 1 deletion(-) create mode 100644 src/particles/distribution/SpinvMF.H diff --git a/src/particles/distribution/All.H b/src/particles/distribution/All.H index 09606cec6..e73c984e2 100644 --- a/src/particles/distribution/All.H +++ b/src/particles/distribution/All.H @@ -19,6 +19,7 @@ #include "Thermal.H" #include "Triangle.H" #include "Waterbag.H" +#include "SpinvMF.H" #include @@ -34,7 +35,8 @@ namespace impactx::distribution Thermal, Triangle, Semigaussian, - Waterbag + Waterbag, + SpinvMF >; } // namespace impactx::distribution diff --git a/src/particles/distribution/SpinvMF.H b/src/particles/distribution/SpinvMF.H new file mode 100644 index 000000000..633bbfeb7 --- /dev/null +++ b/src/particles/distribution/SpinvMF.H @@ -0,0 +1,123 @@ +/* Copyright 2022-2023 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Chad Mitchell, Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_DISTRIBUTION_SPINVMF +#define IMPACTX_DISTRIBUTION_SPINVMF + +#include + +#include +#include +#include + +#include + + +namespace impactx::distribution +{ + struct SpinvMF + { + /** 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. + * + * Return sampling from a vMF distribution. + * + * @param mux,muy,muz components of a unit vector specifying the mean direction + * @param kappa concentration parameter + */ + SpinvMF ( + amrex::ParticleReal mux, + amrex::ParticleReal muy, + amrex::ParticleReal muz, + amrex::ParticleReal kappa + ) + : m_muX(mux), m_muY(muy), m_muZ(muz), m_kappa(kappa) + { + } + + /** Initialize the distribution. + * + * Nothing to do here. + * + * @param bunch_charge charge of the beam in C + * @param ref the reference particle + */ + void initialize ([[maybe_unused]] amrex::ParticleReal bunch_charge, [[maybe_unused]] RefPart const & ref) + { + } + + /** Close and deallocate all data and handles. + * + * Nothing to do here. + */ + void + finalize () + { + } + + /** Return 1 3D spin (unit) vector + * + * @param sx particle spin component in x + * @param sy particle spin component in y + * @param sz particle spin component in z + * @param engine a random number engine (with associated state) + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void operator() ( + amrex::ParticleReal & AMREX_RESTRICT sx, + amrex::ParticleReal & AMREX_RESTRICT sy, + amrex::ParticleReal & AMREX_RESTRICT sz, + amrex::RandomEngine const & engine + ) const + { + using namespace amrex::literals; + using amrex::Math::powi; + using ablastr::constant::math::pi; + + // Normalize the mean direction vector just in case: + amrex::ParticleReal mu_norm = std::sqrt(powi<2>(m_muX)+powi<2>(m_muY)+powi<2>(m_muZ)); + amrex::ParticleReal mux = m_muX/mu_norm; + amrex::ParticleReal muy = m_muY/mu_norm; + amrex::ParticleReal muz = m_muZ/mu_norm; + + // Evaluate constants independent of particle: + amrex::ParticleReal const muperp = std::sqrt(powi<2>(mux)+powi<2>(muy)); + amrex::ParticleReal const a = (muperp==0) ? 0_prt : mux/muperp; + amrex::ParticleReal const b = (muperp==0) ? 1_prt : muy/muperp; + + // Sample two iid random variables: + amrex::ParticleReal x1 = amrex::Random(engine); + amrex::ParticleReal x2 = amrex::Random(engine); + + // Basic transformation of the random variables: + amrex::ParticleReal c = std::cos(2_prt*pi*x1); + amrex::ParticleReal s = std::sin(2_prt*pi*x1); + amrex::ParticleReal t; + if(m_kappa > 0) { + t = 1_prt + std::log1p(x2 * std::expm1(-2*m_kappa))/m_kappa; + } else { + t = 1_prt - 2_prt * x2; + } + + // Evaluation of the spin components + amrex::ParticleReal u=std::sqrt(1_prt-powi<2>(t)); + sx = u * (b*c + a*muz*s) + t*mux; + sy = u * (-a*c + b*muz*s) + t*muy; + sz = u * (-muperp*s) + t*muz; + + } + + amrex::ParticleReal m_muX, m_muY, m_muZ, m_kappa; //! parameters specifying the vMF distribution + }; + +} // namespace impactx::distribution + +#endif // IMPACTX_DISTRIBUTION_SPINVMF From 10bd32e4b3ac499b06dd73fca64f7ddc93c09299 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Mon, 24 Nov 2025 15:57:23 -0800 Subject: [PATCH 03/28] Add citations for documentation --- src/particles/distribution/SpinvMF.H | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/particles/distribution/SpinvMF.H b/src/particles/distribution/SpinvMF.H index 633bbfeb7..386909c9e 100644 --- a/src/particles/distribution/SpinvMF.H +++ b/src/particles/distribution/SpinvMF.H @@ -28,6 +28,12 @@ namespace impactx::distribution * 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), https://doi.org/10.1007/s11222-024-10419-3 * Return sampling from a vMF distribution. * * @param mux,muy,muz components of a unit vector specifying the mean direction @@ -101,6 +107,9 @@ namespace impactx::distribution amrex::ParticleReal c = std::cos(2_prt*pi*x1); amrex::ParticleReal s = std::sin(2_prt*pi*x1); amrex::ParticleReal t; + // The form of the function below is obtained from eq (30) of: + // D. Frisch and U. D. Hanebeck, "Deterministic Von Mises-Fisher Sampling on the Sphere Using + // Fibonacci Lattices", IEEE, 2023, DOI: 10.1109/SDF-MFI59545.2023.10361396 if(m_kappa > 0) { t = 1_prt + std::log1p(x2 * std::expm1(-2*m_kappa))/m_kappa; } else { From 21424ce75f73ed8d2111c953f1bf34e62dd20937 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 24 Nov 2025 23:57:31 +0000 Subject: [PATCH 04/28] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/particles/distribution/SpinvMF.H | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/particles/distribution/SpinvMF.H b/src/particles/distribution/SpinvMF.H index 386909c9e..91b784da8 100644 --- a/src/particles/distribution/SpinvMF.H +++ b/src/particles/distribution/SpinvMF.H @@ -32,7 +32,7 @@ namespace impactx::distribution * 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. + * S. Kang and H-S. Oh, "Novel sampling method for the von Mises-Fisher distribution", Stat. * and Comput. 34, 106 (2024), https://doi.org/10.1007/s11222-024-10419-3 * Return sampling from a vMF distribution. * From b1781be9fb83774734e269a996743ff9ef41f506 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Mon, 24 Nov 2025 15:59:56 -0800 Subject: [PATCH 05/28] Update citations --- src/particles/distribution/SpinvMF.H | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/particles/distribution/SpinvMF.H b/src/particles/distribution/SpinvMF.H index 91b784da8..9ead94518 100644 --- a/src/particles/distribution/SpinvMF.H +++ b/src/particles/distribution/SpinvMF.H @@ -31,9 +31,11 @@ namespace impactx::distribution * 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 + * + * 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), https://doi.org/10.1007/s11222-024-10419-3 + * and Comput. 34, 106 (2024), https://doi.org/10.1007/s11222-024-10419-3. + * * Return sampling from a vMF distribution. * * @param mux,muy,muz components of a unit vector specifying the mean direction From 90767f52fdae7b830fe07d0d12dc78344760ab8c Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 25 Nov 2025 16:07:36 -0800 Subject: [PATCH 06/28] Spin Distribution Call --- src/ImpactX.H | 4 +- src/initialization/InitDistribution.cpp | 45 ++++++++++++++++- src/particles/ImpactXParticleContainer.H | 14 ++--- src/particles/ImpactXParticleContainer.cpp | 59 +++++++++++++--------- src/particles/distribution/All.H | 4 +- src/particles/distribution/SpinvMF.H | 24 +-------- src/python/ImpactX.cpp | 1 + src/python/ImpactXParticleContainer.cpp | 4 ++ 8 files changed, 98 insertions(+), 57 deletions(-) diff --git a/src/ImpactX.H b/src/ImpactX.H index 531dfc95f..90c257ccc 100644 --- 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 spin_distr = std::nullopt ); /** Validate the simulation is ready to run the particle tracking loop via @see track_particles diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index b5c2f40e0..4b1c38778 100644 --- a/src/initialization/InitDistribution.cpp +++ b/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 @@ -295,7 +296,8 @@ namespace impactx ImpactX::add_particles ( amrex::ParticleReal bunch_charge, distribution::KnownDistributions distr, - amrex::Long npart + amrex::Long npart, + std::optional spin_distr ) { BL_PROFILE("ImpactX::add_particles"); @@ -337,6 +339,7 @@ namespace impactx // alloc data for particle attributes amrex::Gpu::DeviceVector x, y, t; amrex::Gpu::DeviceVector px, py, pt; + amrex::Gpu::DeviceVector sx, sy, sz; x.resize(npart_this_proc); y.resize(npart_this_proc); t.resize(npart_this_proc); @@ -344,6 +347,13 @@ namespace impactx 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); @@ -354,6 +364,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; @@ -373,6 +386,7 @@ namespace impactx npart_this_thread = thread_chunk.size; #endif + // phase space init initialization::InitSingleParticleData const init_single_particle_data( distribution, x_ptr + my_offset, @@ -384,6 +398,20 @@ 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 @@ -548,15 +576,28 @@ 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); + // TODO: check magnitude of x,y,z is in range [0:1] + // TODO: calculate kappa from mag(x,y,z) + amrex::ParticleReal kappa = 1.0; // FIXME: root finding + 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 diff --git a/src/particles/ImpactXParticleContainer.H b/src/particles/ImpactXParticleContainer.H index 708bff273..950e931a7 100644 --- a/src/particles/ImpactXParticleContainer.H +++ b/src/particles/ImpactXParticleContainer.H @@ -179,12 +179,12 @@ namespace impactx * @param px momentum in x * @param py momentum in y * @param pt momentum in t - * @param sx spin component in x - * @param sy spin component in y - * @param sz spin component in z * @param qm charge over mass in 1/eV * @param bunch_charge total charge within a bunch in C * @param w weight of each particle: how many real particles to represent + * @param sx spin component in x + * @param sy spin component in y + * @param sz spin component in z */ void AddNParticles ( @@ -194,12 +194,12 @@ namespace impactx amrex::Gpu::DeviceVector const & px, amrex::Gpu::DeviceVector const & py, amrex::Gpu::DeviceVector const & pt, - amrex::Gpu::DeviceVector const & sx, - amrex::Gpu::DeviceVector const & sy, - amrex::Gpu::DeviceVector const & sz, amrex::ParticleReal qm, std::optional bunch_charge = std::nullopt, - std::optional> w = std::nullopt + std::optional> w = std::nullopt, + std::optional> sx = std::nullopt, + std::optional> sy = std::nullopt, + std::optional> sz = std::nullopt ); /** Register storage for lost particles diff --git a/src/particles/ImpactXParticleContainer.cpp b/src/particles/ImpactXParticleContainer.cpp index dea7fa7c1..9a6cbe955 100644 --- a/src/particles/ImpactXParticleContainer.cpp +++ b/src/particles/ImpactXParticleContainer.cpp @@ -240,36 +240,42 @@ namespace impactx amrex::Gpu::DeviceVector const & px, amrex::Gpu::DeviceVector const & py, amrex::Gpu::DeviceVector const & pt, - amrex::Gpu::DeviceVector const & sx, - amrex::Gpu::DeviceVector const & sy, - amrex::Gpu::DeviceVector const & sz, amrex::ParticleReal qm, std::optional bunch_charge, - std::optional> w + std::optional> w, + std::optional> sx, + std::optional> sy, + std::optional> sz ) { BL_PROFILE("ImpactX::AddNParticles"); using namespace amrex::literals; // for _rt and _prt + // number of particles to add + std::size_t const np_s = x.size(); + + // input validation bool const has_w = w.has_value(); if (!(bunch_charge.has_value() ^ has_w)) { throw std::runtime_error("AddNParticles: Exactly one of bunch_charge or w must be provided!"); } - AMREX_ALWAYS_ASSERT(x.size() == y.size()); - AMREX_ALWAYS_ASSERT(x.size() == t.size()); - AMREX_ALWAYS_ASSERT(x.size() == px.size()); - AMREX_ALWAYS_ASSERT(x.size() == py.size()); - AMREX_ALWAYS_ASSERT(x.size() == pt.size()); - AMREX_ALWAYS_ASSERT(x.size() == sx.size()); - AMREX_ALWAYS_ASSERT(x.size() == sy.size()); - AMREX_ALWAYS_ASSERT(x.size() == sz.size()); - if (has_w) { AMREX_ALWAYS_ASSERT(x.size() == w->size()); } - - // number of particles to add - amrex::Long const np = x.size(); + bool const has_spin = sx.has_value(); + + AMREX_ALWAYS_ASSERT(np_s == y.size()); + AMREX_ALWAYS_ASSERT(np_s == t.size()); + AMREX_ALWAYS_ASSERT(np_s == px.size()); + AMREX_ALWAYS_ASSERT(np_s == py.size()); + AMREX_ALWAYS_ASSERT(np_s == pt.size()); + if (has_spin) { + AMREX_ALWAYS_ASSERT(sy.has_value()); + AMREX_ALWAYS_ASSERT(sz.has_value()); + AMREX_ALWAYS_ASSERT(np_s == sy.value().size()); + AMREX_ALWAYS_ASSERT(np_s == sz.value().size()); + } + if (has_w) { AMREX_ALWAYS_ASSERT(np_s == w->size()); } // we add particles to lev 0, grid 0 int lid = 0, gid = 0; @@ -294,7 +300,8 @@ namespace impactx DefineAndReturnParticleTile(lid, gid, ithr); } - amrex::Long pid = ParticleType::NextID(); + amrex::Long const pid = ParticleType::NextID(); + amrex::Long const np = np_s; ParticleType::NextID(pid + np); AMREX_ALWAYS_ASSERT_WITH_MESSAGE( pid + np < amrex::LongParticleIds::LastParticleID, @@ -351,9 +358,9 @@ namespace impactx amrex::ParticleReal const * const AMREX_RESTRICT px_ptr = px.data(); amrex::ParticleReal const * const AMREX_RESTRICT py_ptr = py.data(); amrex::ParticleReal const * const AMREX_RESTRICT pt_ptr = pt.data(); - amrex::ParticleReal const * const AMREX_RESTRICT sx_ptr = sx.data(); - amrex::ParticleReal const * const AMREX_RESTRICT sy_ptr = sy.data(); - amrex::ParticleReal const * const AMREX_RESTRICT sz_ptr = sz.data(); + amrex::ParticleReal const * const AMREX_RESTRICT sx_ptr = has_spin ? sx.value().data() : nullptr; + amrex::ParticleReal const * const AMREX_RESTRICT sy_ptr = has_spin ? sy.value().data() : nullptr; + amrex::ParticleReal const * const AMREX_RESTRICT sz_ptr = has_spin ? sz.value().data() : nullptr; amrex::ParticleReal const * const AMREX_RESTRICT w_ptr = has_w ? w->data() : nullptr; amrex::ParticleReal const bunch_charge_value = has_w ? 0_prt : bunch_charge.value(); @@ -370,9 +377,15 @@ namespace impactx py_arr[old_np+i] = py_ptr[my_offset+i]; pt_arr[old_np+i] = pt_ptr[my_offset+i]; - sx_arr[old_np+i] = sx_ptr[my_offset+i]; - sy_arr[old_np+i] = sy_ptr[my_offset+i]; - sz_arr[old_np+i] = sz_ptr[my_offset+i]; + if (has_spin) { + sx_arr[old_np+i] = sx_ptr[my_offset+i]; + sy_arr[old_np+i] = sy_ptr[my_offset+i]; + sz_arr[old_np+i] = sz_ptr[my_offset+i]; + } else { + sx_arr[old_np+i] = 0_prt; + sy_arr[old_np+i] = 0_prt; + sz_arr[old_np+i] = 0_prt; + } qm_arr[old_np+i] = qm; w_arr[old_np+i] = has_w ? w_ptr[my_offset+i] : bunch_charge_value / ablastr::constant::SI::q_e/np; diff --git a/src/particles/distribution/All.H b/src/particles/distribution/All.H index e73c984e2..5e676fef6 100644 --- a/src/particles/distribution/All.H +++ b/src/particles/distribution/All.H @@ -26,6 +26,7 @@ namespace impactx::distribution { + // Phase space distributions using KnownDistributions = std::variant< Empty, /* must be first, so KnownDistributions creates a default constructor */ Gaussian, @@ -35,8 +36,7 @@ namespace impactx::distribution Thermal, Triangle, Semigaussian, - Waterbag, - SpinvMF + Waterbag >; } // namespace impactx::distribution diff --git a/src/particles/distribution/SpinvMF.H b/src/particles/distribution/SpinvMF.H index 9ead94518..0b4cb1b4a 100644 --- a/src/particles/distribution/SpinvMF.H +++ b/src/particles/distribution/SpinvMF.H @@ -51,26 +51,6 @@ namespace impactx::distribution { } - /** Initialize the distribution. - * - * Nothing to do here. - * - * @param bunch_charge charge of the beam in C - * @param ref the reference particle - */ - void initialize ([[maybe_unused]] amrex::ParticleReal bunch_charge, [[maybe_unused]] RefPart const & ref) - { - } - - /** Close and deallocate all data and handles. - * - * Nothing to do here. - */ - void - finalize () - { - } - /** Return 1 3D spin (unit) vector * * @param sx particle spin component in x @@ -113,13 +93,13 @@ namespace impactx::distribution // D. Frisch and U. D. Hanebeck, "Deterministic Von Mises-Fisher Sampling on the Sphere Using // Fibonacci Lattices", IEEE, 2023, DOI: 10.1109/SDF-MFI59545.2023.10361396 if(m_kappa > 0) { - t = 1_prt + std::log1p(x2 * std::expm1(-2*m_kappa))/m_kappa; + t = 1_prt + std::log1p(x2 * std::expm1(-2_prt * m_kappa)) / m_kappa; } else { t = 1_prt - 2_prt * x2; } // Evaluation of the spin components - amrex::ParticleReal u=std::sqrt(1_prt-powi<2>(t)); + amrex::ParticleReal u = std::sqrt(1_prt-powi<2>(t)); sx = u * (b*c + a*muz*s) + t*mux; sy = u * (-a*c + b*muz*s) + t*muy; sz = u * (-muperp*s) + t*muz; diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index b50876dd0..72ba86c15 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -607,6 +607,7 @@ void init_ImpactX (py::module& m) .def("add_particles", &ImpactX::add_particles, py::arg("bunch_charge"), py::arg("distr"), py::arg("npart"), + py::arg("spin_distr") = py::none(), // TODO: bind the distribution::SpinvMF object to Python "Particle tracking mode:" "Generate and add n particles to the particle container.\n\n" "Will also resize the geometry based on the updated particle\n" diff --git a/src/python/ImpactXParticleContainer.cpp b/src/python/ImpactXParticleContainer.cpp index 93f097886..f89caad6f 100644 --- a/src/python/ImpactXParticleContainer.cpp +++ b/src/python/ImpactXParticleContainer.cpp @@ -84,6 +84,7 @@ void init_impactxparticlecontainer(py::module& m) py::arg("x"), py::arg("y"), py::arg("t"), py::arg("px"), py::arg("py"), py::arg("pt"), py::arg("qm"), py::arg("bunch_charge")=py::none(), py::arg("w")=py::none(), + py::arg("sx")=py::none(), py::arg("sy")=py::none(), py::arg("sz")=py::none(), "Add new particles to the container for fixed s.\n\n" "Either the total charge (bunch_charge) or the weight of each\n" "particle (w) must be provided.\n\n" @@ -99,6 +100,9 @@ void init_impactxparticlecontainer(py::module& m) ":param qm: charge over mass in 1/eV\n" ":param bunch_charge: total charge within a bunch in C" ":param w: weight of each particle: how many real particles to represent" + ":param sx: spin component in x\n" + ":param sy: spin component in y\n" + ":param sz: spin component in z\n" ) .def("ref_particle", py::overload_cast<>(&ImpactXParticleContainer::GetRefParticle), From b0d9d1e508b6acab3c0689ef9187706412f787b2 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Wed, 26 Nov 2025 12:32:38 -0800 Subject: [PATCH 07/28] Add conditional form of call to AddNParticles. --- src/initialization/InitDistribution.cpp | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 4b1c38778..877569366 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -419,9 +419,15 @@ namespace impactx 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, 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(); From 336187411cbb11a2bed677460d81e3a3b2649daa Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Wed, 26 Nov 2025 13:14:55 -0800 Subject: [PATCH 08/28] Add input for a simple distribution test. --- examples/distgen/input_kurth4d_spin.in | 41 ++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 examples/distgen/input_kurth4d_spin.in diff --git a/examples/distgen/input_kurth4d_spin.in b/examples/distgen/input_kurth4d_spin.in new file mode 100644 index 000000000..f2794e960 --- /dev/null +++ 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 From 23e30740dd2fa62e6632080b571f817c3d74a5b3 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Wed, 26 Nov 2025 14:50:47 -0800 Subject: [PATCH 09/28] [Dependencies] pyAMReX: ax3l branch Update after 25.12 releases to development branch, once PR is merged https://github.com/AMReX-Codes/pyamrex/pull/518 --- cmake/dependencies/pyAMReX.cmake | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cmake/dependencies/pyAMReX.cmake b/cmake/dependencies/pyAMReX.cmake index 0e8f5e170..b158a36f0 100644 --- a/cmake/dependencies/pyAMReX.cmake +++ b/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.11" +set(ImpactX_pyamrex_branch "topic-pc-impactx-spin" CACHE STRING "Repository branch for ImpactX_pyamrex_repo if(ImpactX_pyamrex_internal)") From c9a417129c47681985cbdade743a86c74cbd3fcc Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Wed, 26 Nov 2025 15:26:06 -0800 Subject: [PATCH 10/28] Fix use of optional arguments in AddNParticles call. --- src/initialization/InitDistribution.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 877569366..42c9134ba 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -422,7 +422,7 @@ namespace impactx 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, sx, sy, sz); + 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(), From 88a56954225847db144f2579e56b343377839bff Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Wed, 26 Nov 2025 15:58:32 -0800 Subject: [PATCH 11/28] Formatting Just some readability, already fixed in prior commit --- src/initialization/InitDistribution.cpp | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 42c9134ba..acd250fcd 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -420,13 +420,20 @@ namespace impactx }, distr); 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); + 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); + 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(); From 6201372b6482d9aaab18ee59e7187b0dd93ba77b Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Wed, 26 Nov 2025 16:36:47 -0800 Subject: [PATCH 12/28] [Dependencies] AMReX: development branch Update after 25.12 releases to development branch, once PR is merged https://github.com/AMReX-Codes/pyamrex/pull/518 --- cmake/dependencies/ABLASTR.cmake | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmake/dependencies/ABLASTR.cmake b/cmake/dependencies/ABLASTR.cmake index 4a65f26f4..a77d47ba2 100644 --- a/cmake/dependencies/ABLASTR.cmake +++ b/cmake/dependencies/ABLASTR.cmake @@ -186,7 +186,7 @@ set(ImpactX_ablastr_branch "25.11" 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)") From c27096193bd64f05bf0a58ea92cddedd413ed2a9 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Thu, 4 Dec 2025 09:55:38 -0800 Subject: [PATCH 13/28] Add P to kappa. --- src/initialization/InitDistribution.cpp | 14 ++++++++--- src/particles/distribution/SpinvMF.H | 32 +++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index acd250fcd..d2cfc570f 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -601,9 +601,17 @@ namespace impactx pp_dist.queryWithParser("polarization_x", polarization_x); pp_dist.queryWithParser("polarization_y", polarization_y); pp_dist.queryWithParser("polarization_z", polarization_z); - // TODO: check magnitude of x,y,z is in range [0:1] - // TODO: calculate kappa from mag(x,y,z) - amrex::ParticleReal kappa = 1.0; // FIXME: root finding + 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 diff --git a/src/particles/distribution/SpinvMF.H b/src/particles/distribution/SpinvMF.H index 0b4cb1b4a..aeb581400 100644 --- a/src/particles/distribution/SpinvMF.H +++ b/src/particles/distribution/SpinvMF.H @@ -106,6 +106,38 @@ namespace impactx::distribution } + + /** This function evaluatates the inverse Langevin function, in order to return + * the value of concentration (kappa) required to produce a given polarization magnitude. + * It is used to prepare for spin sampling using the vMF distribution. + * + * Here, we use an approximate expression due to R. Petrosyan. See: + * R. Petrosyan, Rheologica Acta 56, 21-26, 2016; doi:10.1007/s00397-016-0977-9 + * R. Jedynak, Mathematics and Mechanics of Solids 24, 1-25, 2018; doi:10.1177/1081286518811395 + * + * @param pmag polarization magnitude + * @returns required value of concentration kappa + */ + AMREX_GPU_HOST AMREX_FORCE_INLINE + static amrex::ParticleReal + inverse_Langevin (amrex::ParticleReal const & AMREX_RESTRICT pmag) + { + using namespace amrex::literals; // for _rt and _prt + using amrex::Math::powi; + + amrex::ParticleReal concentration; + + // Evaluation of the inverse Langevin function. Here, we use the approximate expression + // due to R. Petrosyan. See: + // R. Petrosyan, Rheologica Acta 56, 21-26, 2016; doi:10.1007/s00397-016-0977-9 + // R. Jedynak, Mathematics and Mechanics of Solids 24, 1-25, 2018; doi:10.1177/1081286518811395 + + concentration = 3_prt*pmag + powi<2>(pmag)/5_prt*std::sin(7_prt*pmag/2_prt) + powi<3>(pmag)/(1_prt-pmag); + + return concentration; + } + + amrex::ParticleReal m_muX, m_muY, m_muZ, m_kappa; //! parameters specifying the vMF distribution }; From 57a29586ca03b8245e1b382a46767d3af224d862 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 4 Dec 2025 17:55:41 +0000 Subject: [PATCH 14/28] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/particles/distribution/SpinvMF.H | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/particles/distribution/SpinvMF.H b/src/particles/distribution/SpinvMF.H index aeb581400..2f17f9c06 100644 --- a/src/particles/distribution/SpinvMF.H +++ b/src/particles/distribution/SpinvMF.H @@ -106,7 +106,7 @@ namespace impactx::distribution } - + /** This function evaluatates the inverse Langevin function, in order to return * the value of concentration (kappa) required to produce a given polarization magnitude. * It is used to prepare for spin sampling using the vMF distribution. @@ -126,7 +126,7 @@ namespace impactx::distribution using amrex::Math::powi; amrex::ParticleReal concentration; - + // Evaluation of the inverse Langevin function. Here, we use the approximate expression // due to R. Petrosyan. See: // R. Petrosyan, Rheologica Acta 56, 21-26, 2016; doi:10.1007/s00397-016-0977-9 @@ -136,7 +136,7 @@ namespace impactx::distribution return concentration; } - + amrex::ParticleReal m_muX, m_muY, m_muZ, m_kappa; //! parameters specifying the vMF distribution }; From 7c399a1f9d45d5b8efd0a219ced8a8e7552e0333 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Thu, 4 Dec 2025 10:25:00 -0800 Subject: [PATCH 15/28] Add analysis and README for distribution test. --- examples/CMakeLists.txt | 16 +++++ examples/distgen/README.rst | 49 +++++++++++++++ examples/distgen/analysis_kurth4d_spin.py | 75 +++++++++++++++++++++++ 3 files changed, 140 insertions(+) create mode 100755 examples/distgen/analysis_kurth4d_spin.py diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index f1cd0f7e8..8871552d1 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1756,3 +1756,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 +#) diff --git a/examples/distgen/README.rst b/examples/distgen/README.rst index c81081456..02cab3331 100644 --- a/examples/distgen/README.rst +++ b/examples/distgen/README.rst @@ -219,3 +219,52 @@ 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 `__ 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 new file mode 100755 index 000000000..32986a6c9 --- /dev/null +++ b/examples/distgen/analysis_kurth4d_spin.py @@ -0,0 +1,75 @@ +#!/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_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, +) From b6d6722ed886f04dd08ead86e7e2b6f523a60aed Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 4 Dec 2025 18:24:57 +0000 Subject: [PATCH 16/28] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/distgen/README.rst | 6 ++---- examples/distgen/analysis_kurth4d_spin.py | 3 +-- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/examples/distgen/README.rst b/examples/distgen/README.rst index 02cab3331..6846a05ff 100644 --- a/examples/distgen/README.rst +++ b/examples/distgen/README.rst @@ -222,7 +222,7 @@ We run the following script to analyze correctness: .. _examples-distgen-spinvmf: - + Spin Sampling from a vMF Distribution =============================================== @@ -255,7 +255,7 @@ For `MPI-parallel `__ runs, prefix these lines with ` :language: ini :caption: You can copy this file from ``examples/distgen/input_kurth4d_spin.in``. - + Analyze ------- @@ -266,5 +266,3 @@ We run the following script to analyze correctness: .. 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 index 32986a6c9..87f49e8e8 100755 --- a/examples/distgen/analysis_kurth4d_spin.py +++ b/examples/distgen/analysis_kurth4d_spin.py @@ -7,7 +7,6 @@ import numpy as np import openpmd_api as io -from scipy.stats import moment def get_polarization(beam): @@ -60,7 +59,7 @@ def get_polarization(beam): 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}") From 172865a8f008faffab7556602fdbb4d46d8c15f2 Mon Sep 17 00:00:00 2001 From: Chad Mitchell Date: Thu, 4 Dec 2025 10:52:57 -0800 Subject: [PATCH 17/28] Add app/C++ documentation of distribution. --- docs/source/usage/parameters.rst | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 6e90e8951..8fd632b1d 100644 --- 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: From 7561fcbc23f96df9132b171c6a8249a2b1359861 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 4 Dec 2025 18:53:00 +0000 Subject: [PATCH 18/28] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/source/usage/parameters.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 8fd632b1d..7a4406bf9 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -144,7 +144,7 @@ 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 From 5437431662d8bbfa830d5d4a690daaa8e8352c7b Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 12 Dec 2025 18:52:42 -0800 Subject: [PATCH 19/28] Test: Update DataFrame Comparison --- tests/python/test_dataframe.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/python/test_dataframe.py b/tests/python/test_dataframe.py index aa45271a4..f776c4a25 100644 --- a/tests/python/test_dataframe.py +++ b/tests/python/test_dataframe.py @@ -107,6 +107,9 @@ def test_df_pandas(save_png=True): "momentum_x", "momentum_y", "momentum_t", + "spin_x", + "spin_y", + "spin_z", "qm", "weighting", ] From 370035ddec28d243715504b441643b7bacad7ff1 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 16 Dec 2025 02:31:36 +0100 Subject: [PATCH 20/28] Fix Init (No-Spin Case) --- src/initialization/InitDistribution.cpp | 46 +++++++++++-------------- src/particles/distribution/SpinvMF.H | 9 ++--- src/python/distribution.cpp | 17 +++++++++ 3 files changed, 43 insertions(+), 29 deletions(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index e0a073244..3ad2c90c6 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -342,7 +342,7 @@ namespace impactx // alloc data for particle attributes amrex::Gpu::DeviceVector x, y, t; amrex::Gpu::DeviceVector px, py, pt; - amrex::Gpu::DeviceVector sx, sy, sz; + std::optional> sx, sy, sz; x.resize(npart_this_proc); y.resize(npart_this_proc); t.resize(npart_this_proc); @@ -352,9 +352,9 @@ namespace impactx 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); + sx = amrex::Gpu::DeviceVector(npart_this_proc); + sy = amrex::Gpu::DeviceVector(npart_this_proc); + sz = amrex::Gpu::DeviceVector(npart_this_proc); } std::visit([&](auto&& distribution){ @@ -367,9 +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(); + amrex::ParticleReal * const AMREX_RESTRICT sx_ptr = has_spin ? sx->data() : nullptr; + amrex::ParticleReal * const AMREX_RESTRICT sy_ptr = has_spin ? sy->data() : nullptr; + amrex::ParticleReal * const AMREX_RESTRICT sz_ptr = has_spin ? sz->data() : nullptr; using Distribution = std::decay_t; @@ -422,22 +422,14 @@ namespace impactx distribution.finalize(); }, distr); - 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 - ); - } + 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 + ); auto space_charge = get_space_charge_algo(); @@ -610,12 +602,16 @@ namespace impactx // 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) { + std::optional spin_dist; + if (pmag == 0.0_prt) { + // no spin init + } + else if (pmag < 1_prt) { kappa = distribution::SpinvMF::inverse_Langevin(pmag); + spin_dist = distribution::SpinvMF(polarization_x, polarization_y, polarization_z, kappa); } 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") diff --git a/src/particles/distribution/SpinvMF.H b/src/particles/distribution/SpinvMF.H index 2f17f9c06..d16513847 100644 --- a/src/particles/distribution/SpinvMF.H +++ b/src/particles/distribution/SpinvMF.H @@ -24,9 +24,10 @@ namespace impactx::distribution struct SpinvMF { /** 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. + * + * 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 @@ -107,7 +108,7 @@ namespace impactx::distribution } - /** This function evaluatates the inverse Langevin function, in order to return + /** This function evaluates the inverse Langevin function, in order to return * the value of concentration (kappa) required to produce a given polarization magnitude. * It is used to prepare for spin sampling using the vMF distribution. * diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp index a074f6805..cb8f44f63 100644 --- a/src/python/distribution.cpp +++ b/src/python/distribution.cpp @@ -194,5 +194,22 @@ void init_distribution(py::module& m) .def_property("beam_intensity", &Envelope::beam_intensity, &Envelope::set_beam_intensity) ; + py::class_(m, "SpinvMF") + .def(py::init< + amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal, + amrex::ParticleReal + >(), + py::arg("mux"), py::arg("muy"), py::arg("muz"), + py::arg("kappa"), + "A von Mises-Fisher (vMF) distribution on the unit 2-sphere, for particle spin." + ) + .def_static("inverse_Langevin", + &distribution::SpinvMF::inverse_Langevin, + py::arg("pmag"), + "This function evaluates the inverse Langevin function, in order to return " + "the value of concentration (kappa) required to produce a given polarization magnitude." + ) + ; + m.def("create_envelope", &initialization::create_envelope); } From 55056325e36283441d40703cf1d50dbfe7e5bd77 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 15 Dec 2025 17:33:14 -0800 Subject: [PATCH 21/28] rst titles --- docs/source/usage/parameters.rst | 2 +- examples/distgen/README.rst | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 5e5f5a0cd..3723f33cf 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -139,7 +139,7 @@ Initial Beam Distributions * ``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. diff --git a/examples/distgen/README.rst b/examples/distgen/README.rst index 6846a05ff..7bf4e659b 100644 --- a/examples/distgen/README.rst +++ b/examples/distgen/README.rst @@ -224,7 +224,7 @@ We run the following script to analyze correctness: .. _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. From 562fd6df1d4fc1a86789684fe202499989cada20 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 15 Dec 2025 17:36:11 -0800 Subject: [PATCH 22/28] Remove todo comment --- src/python/ImpactX.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index d4eec2a99..a02d69df3 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -611,7 +611,7 @@ void init_ImpactX (py::module& m) .def("add_particles", &ImpactX::add_particles, py::arg("bunch_charge"), py::arg("distr"), py::arg("npart"), - py::arg("spin_distr") = py::none(), // TODO: bind the distribution::SpinvMF object to Python + py::arg("spin_distr") = py::none(), "Particle tracking mode:" "Generate and add n particles to the particle container.\n\n" "Will also resize the geometry based on the updated particle\n" From dd6a94fa2caca2bde9f10c2953477a5f5800ee81 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 16 Dec 2025 02:50:53 +0100 Subject: [PATCH 23/28] Improve `SpinvMF` --- src/initialization/InitDistribution.cpp | 21 +++++-------------- src/particles/distribution/SpinvMF.H | 28 +++++++++++++++++-------- src/python/distribution.cpp | 4 +--- 3 files changed, 25 insertions(+), 28 deletions(-) diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 3ad2c90c6..5cc1d81f8 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -593,24 +593,13 @@ namespace impactx 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); + bool const has_sx = pp_dist.queryWithParser("polarization_x", polarization_x); + bool const has_sy = pp_dist.queryWithParser("polarization_y", polarization_y); + bool const has_sz = pp_dist.queryWithParser("polarization_z", polarization_z); std::optional spin_dist; - if (pmag == 0.0_prt) { - // no spin init - } - else if (pmag < 1_prt) { - kappa = distribution::SpinvMF::inverse_Langevin(pmag); - spin_dist = distribution::SpinvMF(polarization_x, polarization_y, polarization_z, kappa); - } else { - throw std::runtime_error("Initial polarization magnitude must currently be < 1."); + if (has_sx || has_sy || has_sz) { + spin_dist = distribution::SpinvMF(polarization_x, polarization_y, polarization_z); } amrex::Long npart = 0; // Number of simulation particles diff --git a/src/particles/distribution/SpinvMF.H b/src/particles/distribution/SpinvMF.H index d16513847..da074941d 100644 --- a/src/particles/distribution/SpinvMF.H +++ b/src/particles/distribution/SpinvMF.H @@ -40,16 +40,22 @@ namespace impactx::distribution * Return sampling from a vMF distribution. * * @param mux,muy,muz components of a unit vector specifying the mean direction - * @param kappa concentration parameter */ SpinvMF ( amrex::ParticleReal mux, amrex::ParticleReal muy, - amrex::ParticleReal muz, - amrex::ParticleReal kappa + amrex::ParticleReal muz ) - : m_muX(mux), m_muY(muy), m_muZ(muz), m_kappa(kappa) + : m_muX(mux), m_muY(muy), m_muZ(muz) { + // magnitude of the polarization vector (= polarization magnitude) + amrex::ParticleReal const pmag = std::sqrt( + mux * mux + + muy * muy + + muz * muz + ); + + m_kappa = inverse_Langevin(pmag); } /** Return 1 3D spin (unit) vector @@ -121,25 +127,29 @@ namespace impactx::distribution */ AMREX_GPU_HOST AMREX_FORCE_INLINE static amrex::ParticleReal - inverse_Langevin (amrex::ParticleReal const & AMREX_RESTRICT pmag) + inverse_Langevin (amrex::ParticleReal pmag) { using namespace amrex::literals; // for _rt and _prt using amrex::Math::powi; - amrex::ParticleReal concentration; + AMREX_ASSERT_WITH_MESSAGE(pmag >= 0_prt, "Polarization magnitude must be greater or equal 0."); + AMREX_ASSERT_WITH_MESSAGE(pmag < 1_prt, "Polarization magnitude must be less than 1."); // Evaluation of the inverse Langevin function. Here, we use the approximate expression // due to R. Petrosyan. See: // R. Petrosyan, Rheologica Acta 56, 21-26, 2016; doi:10.1007/s00397-016-0977-9 // R. Jedynak, Mathematics and Mechanics of Solids 24, 1-25, 2018; doi:10.1177/1081286518811395 - concentration = 3_prt*pmag + powi<2>(pmag)/5_prt*std::sin(7_prt*pmag/2_prt) + powi<3>(pmag)/(1_prt-pmag); + amrex::ParticleReal concentration = + 3_prt * pmag + + powi<2>(pmag) / 5_prt * std::sin(7_prt * pmag * 0.5_prt) + + powi<3>(pmag) / (1_prt - pmag); return concentration; } - - amrex::ParticleReal m_muX, m_muY, m_muZ, m_kappa; //! parameters specifying the vMF distribution + amrex::ParticleReal m_muX, m_muY, m_muZ; //! components of a unit vector specifying the mean direction + amrex::ParticleReal m_kappa; //! concentration parameter }; } // namespace impactx::distribution diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp index cb8f44f63..1a5c5ceac 100644 --- a/src/python/distribution.cpp +++ b/src/python/distribution.cpp @@ -196,11 +196,9 @@ void init_distribution(py::module& m) py::class_(m, "SpinvMF") .def(py::init< - amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal, - amrex::ParticleReal + amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal >(), py::arg("mux"), py::arg("muy"), py::arg("muz"), - py::arg("kappa"), "A von Mises-Fisher (vMF) distribution on the unit 2-sphere, for particle spin." ) .def_static("inverse_Langevin", From b6af5c889672b34a3189bc571418511c2c205395 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 16 Dec 2025 03:07:56 +0100 Subject: [PATCH 24/28] NVCC: Work-Around Lambda Same work-around for NVCC as for phase space init. --- src/initialization/InitDistribution.H | 63 +++++++++++++++++++++++++ src/initialization/InitDistribution.cpp | 18 +++---- 2 files changed, 70 insertions(+), 11 deletions(-) diff --git a/src/initialization/InitDistribution.H b/src/initialization/InitDistribution.H index 4f1aa5874..63425ea6a 100644 --- a/src/initialization/InitDistribution.H +++ b/src/initialization/InitDistribution.H @@ -129,6 +129,69 @@ namespace impactx::initialization amrex::ParticleReal* const AMREX_RESTRICT m_part_pt; }; + /** Initialize a single particle's spin using the SpinvMF distribution + * + * Note: we usually would just write a C++ lambda below in ParallelForRNG. But, due to restrictions + * in NVCC as of 11.5, we cannot write a lambda in a lambda as we also std::visit the element + * types of our lattice elements list. + * error #3206-D: An extended __device__ lambda cannot be defined inside a generic lambda expression("operator()"). + * Thus, we fall back to writing a C++ functor here, instead of nesting two lambdas. + * + * Nvidia bug report: 3458976 + * Minimal demonstrator: https://cuda.godbolt.org/z/39e4q53Ye + */ + struct InitSingleParticleSpin + { + /** Constructor taking in pointers to particle data + * + * @param spinv spin distribution function to call + * @param part_sx the array to the particle spin in x (sx) + * @param part_sy the array to the particle spin in y (sy) + * @param part_sz the array to the particle spin in z (sz) + */ + InitSingleParticleSpin ( + distribution::SpinvMF spinv, + amrex::ParticleReal* AMREX_RESTRICT part_sx, + amrex::ParticleReal* AMREX_RESTRICT part_sy, + amrex::ParticleReal* AMREX_RESTRICT part_sz + ) + : m_spinv(std::move(spinv)), + m_part_sx(part_sx), m_part_sy(part_sy), m_part_sz(part_sz) + { + } + + InitSingleParticleSpin () = delete; + InitSingleParticleSpin (InitSingleParticleSpin const &) = default; + InitSingleParticleSpin (InitSingleParticleSpin &&) = default; + ~InitSingleParticleSpin () = default; + + /** Initialize the spin for a single particle + * + * @param i particle index in the current box + * @param engine a random number engine (with associated state) + */ + AMREX_GPU_DEVICE AMREX_FORCE_INLINE + void + operator() ( + amrex::Long i, + amrex::RandomEngine const & engine + ) const + { + m_spinv( + m_part_sx[i], + m_part_sy[i], + m_part_sz[i], + engine + ); + } + + private: + distribution::SpinvMF const m_spinv; + amrex::ParticleReal* const AMREX_RESTRICT m_part_sx; + amrex::ParticleReal* const AMREX_RESTRICT m_part_sy; + amrex::ParticleReal* const AMREX_RESTRICT m_part_sz; + }; + /** Initialize the input parameters for all distributions that read phase space ellipse parameters from the input * diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 5cc1d81f8..5fee178ee 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -399,21 +399,17 @@ namespace impactx py_ptr + my_offset, pt_ptr + my_offset ); - 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 - ); - }); + initialization::InitSingleParticleSpin const init_single_particle_spin( + spin_distr.value(), + sx_ptr + my_offset, + sy_ptr + my_offset, + sz_ptr + my_offset + ); + amrex::ParallelForRNG(npart_this_thread, init_single_particle_spin); } } From ab8a5e5e7617ac0ef4a48fb7e07ec86703b35fc2 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 16 Dec 2025 03:16:02 +0100 Subject: [PATCH 25/28] pyAMReX: `development` --- cmake/dependencies/pyAMReX.cmake | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cmake/dependencies/pyAMReX.cmake b/cmake/dependencies/pyAMReX.cmake index 180982849..cd4b3a099 100644 --- a/cmake/dependencies/pyAMReX.cmake +++ b/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/ax3l/pyamrex.git" +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 "topic-pc-impactx-spin" +set(ImpactX_pyamrex_branch "74e213e3186371d5a7412af3d2a5d6040fab51ce" CACHE STRING "Repository branch for ImpactX_pyamrex_repo if(ImpactX_pyamrex_internal)") From f2448dc6b3be8b0ac8b95b65cfcbb5c3eb8e54e8 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 16 Dec 2025 03:21:26 +0100 Subject: [PATCH 26/28] Python Test --- docs/source/usage/python.rst | 26 +++++++++-- examples/CMakeLists.txt | 13 +++--- examples/distgen/run_kurth4d_spin.py | 68 ++++++++++++++++++++++++++++ src/python/distribution.cpp | 16 +++---- 4 files changed, 105 insertions(+), 18 deletions(-) create mode 100755 examples/distgen/run_kurth4d_spin.py diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index d26ab81c9..68d51948f 100644 --- 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 `__ + + :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 index 40c59136f..388104738 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1824,10 +1824,9 @@ add_impactx_test(distgen-spinvmf 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 -#) +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/run_kurth4d_spin.py b/examples/distgen/run_kurth4d_spin.py new file mode 100755 index 000000000..093d80cd9 --- /dev/null +++ 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/src/python/distribution.cpp b/src/python/distribution.cpp index 1a5c5ceac..7f3ac79f1 100644 --- a/src/python/distribution.cpp +++ b/src/python/distribution.cpp @@ -187,14 +187,7 @@ void init_distribution(py::module& m) "A 6D Waterbag distribution" ); - py::class_(m, "Envelope") - .def(py::init<>()) - .def(py::init()) - .def_property("envelope", &Envelope::covariance_matrix, &Envelope::set_covariance_matrix) - .def_property("beam_intensity", &Envelope::beam_intensity, &Envelope::set_beam_intensity) - ; - - py::class_(m, "SpinvMF") + py::class_(md, "SpinvMF") .def(py::init< amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal >(), @@ -209,5 +202,12 @@ void init_distribution(py::module& m) ) ; + py::class_(m, "Envelope") + .def(py::init<>()) + .def(py::init()) + .def_property("envelope", &Envelope::covariance_matrix, &Envelope::set_covariance_matrix) + .def_property("beam_intensity", &Envelope::beam_intensity, &Envelope::set_beam_intensity) + ; + m.def("create_envelope", &initialization::create_envelope); } From 15cbcf93e218b3b9a2dcf567734abcdf8e72191d Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 16 Dec 2025 03:29:09 +0100 Subject: [PATCH 27/28] `Source` Element: Support Spin in Restart --- examples/solenoid_restart/analysis_solenoid.py | 2 +- src/elements/Source.cpp | 11 ++++++----- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/examples/solenoid_restart/analysis_solenoid.py b/examples/solenoid_restart/analysis_solenoid.py index 7d1d10f82..eb9f5fc93 100755 --- 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/elements/Source.cpp b/src/elements/Source.cpp index 256099e45..ca4a548fb 100644 --- a/src/elements/Source.cpp +++ b/src/elements/Source.cpp @@ -50,7 +50,7 @@ namespace impactx::elements std::string const species_name = "beam"; io::ParticleSpecies beam = iteration.particles[species_name]; // TODO: later we can make the bunch charge an option (i.e., allow rescaling a distribution) - amrex::ParticleReal bunch_charge = beam.getAttribute("charge_C").get(); + // amrex::ParticleReal bunch_charge = beam.getAttribute("charge_C").get(); auto const scalar = openPMD::RecordComponent::SCALAR; auto const getComponentRecord = [&beam](std::string comp_name) { @@ -73,8 +73,6 @@ namespace impactx::elements int const nleft = npart - navg * nprocs; std::uint64_t const npart_this_proc = (myproc < nleft) ? navg+1 : navg; // add 1 to each proc until distributed std::uint64_t npart_before_this_proc = (myproc < nleft) ? (navg+1) * myproc : navg * myproc + nleft; - auto const rel_part_this_proc = - amrex::ParticleReal(npart_this_proc) / amrex::ParticleReal(npart); // alloc data for particle attributes std::map> pinned_SoA; @@ -118,11 +116,14 @@ namespace impactx::elements // finalize distributions and deallocate temporary device global memory amrex::Gpu::streamSynchronize(); - // TODO: at this point, we ignore the "id", "qm" and "weighting" in the file. Could be improved + // TODO: at this point, we ignore the "id" and "qm" in the file. Could be improved + pc.AddNParticles(d_SoA["position_x"], d_SoA["position_y"], d_SoA["position_t"], d_SoA["momentum_x"], d_SoA["momentum_y"], d_SoA["momentum_t"], pc.GetRefParticle().qm_ratio_SI(), - bunch_charge * rel_part_this_proc); + std::nullopt, + d_SoA["weighting"], + d_SoA["spin_x"], d_SoA["spin_y"], d_SoA["spin_z"]); #else // ImpactX_USE_OPENPMD amrex::ignore_unused(pc); From a79aca03ebac5ba2ac3061cd4b760145428fb55a Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 16 Dec 2025 19:03:08 +0100 Subject: [PATCH 28/28] CUDA CI: Lower Build Parallelism Aborts close to the end of the build process. Probably a resource limit hit on the runners. --- .github/workflows/cuda.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 3f5a61d5f..a209014ab 100644 --- 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