From c55c5788125015bfdb621c5e2a822c037b3f4f41 Mon Sep 17 00:00:00 2001 From: Ethan Peterson Date: Thu, 16 Jul 2026 10:54:56 -0400 Subject: [PATCH] Native parametric tokamak source (#3999) Co-authored-by: Paul Romano --- docs/source/io_formats/settings.rst | 80 +++- docs/source/pythonapi/base.rst | 1 + docs/source/usersguide/settings.rst | 47 ++ include/openmc/math_functions.h | 14 + include/openmc/source.h | 136 ++++++ openmc/source.py | 430 ++++++++++++++++++- src/distribution.cpp | 11 +- src/math_functions.cpp | 42 ++ src/source.cpp | 474 ++++++++++++++++++++- tests/cpp_unit_tests/test_distribution.cpp | 25 ++ tests/cpp_unit_tests/test_math.cpp | 15 + tests/unit_tests/test_source_tokamak.py | 263 ++++++++++++ 12 files changed, 1524 insertions(+), 14 deletions(-) create mode 100644 tests/unit_tests/test_source_tokamak.py diff --git a/docs/source/io_formats/settings.rst b/docs/source/io_formats/settings.rst index fb0215916..7091757f2 100644 --- a/docs/source/io_formats/settings.rst +++ b/docs/source/io_formats/settings.rst @@ -777,9 +777,9 @@ attributes/sub-elements: *Default*: 1.0 :type: - Indicator of source type. One of ``independent``, ``file``, ``compiled``, or - ``mesh``. The type of the source will be determined by this attribute if it - is present. + Indicator of source type. One of ``independent``, ``file``, ``compiled``, + ``mesh``, or ``tokamak``. The type of the source will be determined by this + attribute if it is present. :particle: The source particle type, specified as a PDG number or a string alias (e.g., @@ -1015,6 +1015,80 @@ attributes/sub-elements: mesh element and follows the format for :ref:`source_element`. The number of ```` sub-elements should correspond to the number of mesh elements. + For a source with ``type="tokamak"``, the spatial distribution is described by + a Miller-style flux-surface parameterization and the following sub-elements + are used instead of the ``space`` element: + + :major_radius: + The major radius :math:`R_0` of the plasma in [cm]. + + :minor_radius: + The minor radius :math:`a` of the plasma in [cm]. Must be smaller than + ``major_radius``. + + :elongation: + The plasma elongation :math:`\kappa` (must be > 0). + + :triangularity: + The plasma triangularity :math:`\delta` (must be in [-1, 1]). Negative + values describe negative-triangularity plasmas. + + :shafranov_shift: + The Shafranov shift :math:`\Delta` in [cm] (must be >= 0 and less than + ``minor_radius``/2). + + :r_over_a: + A list of normalized minor-radius grid points :math:`r/a`. Must be strictly + increasing, start at 0, and end at 1. + + :emission_density: + A list of neutron emission densities :math:`S(r)` evaluated at each + ``r_over_a`` grid point (arbitrary units, must be non-negative). Only the + shape matters, since the profile is normalized internally. Values are + interpolated linearly between grid points and the profile is refined on an + internal grid for radial sampling. Must have the same length as + ``r_over_a`` and contain at least one positive value. + + :phi_start: + The starting toroidal angle in [rad]. + + *Default*: 0.0 + + :phi_extent: + The toroidal angle extent in [rad]. The source is sampled uniformly in + :math:`[\phi_\text{start},\ \phi_\text{start} + \phi_\text{extent}]`. + + *Default*: :math:`2\pi` + + :n_alpha: + The number of poloidal-angle grid points used to build the sampling CDFs + (must be > 2). Larger values reduce discretization bias; values below 51 + produce a warning. + + *Default*: 101 + + :vertical_shift: + A vertical shift of the plasma center in [cm]. + + *Default*: 0.0 + + :energy: + For a tokamak source, one or more ``energy`` sub-elements specify the + neutron energy distribution(s). Either a single distribution is given (used + at all radii) or exactly one distribution per ``r_over_a`` grid point is + given, in which case the energy is sampled from one of the two + distributions bracketing the sampled radius, selected stochastically with + probability proportional to the proximity of the radius to each grid point + (stochastic interpolation). Each follows the format of a univariate + probability distribution (see :ref:`univariate`). + + :time: + An optional ``time`` sub-element specifying the time distribution of source + particles, following the format of a univariate probability distribution + (see :ref:`univariate`). + + *Default*: particles are born at :math:`t=0` + .. note:: Biased sampling can be applied to the spatial and energy distributions of a source by using the ```` sub-element (see :ref:`univariate` for details on how to specify bias distributions). diff --git a/docs/source/pythonapi/base.rst b/docs/source/pythonapi/base.rst index b3911c8b3..f06a1eba4 100644 --- a/docs/source/pythonapi/base.rst +++ b/docs/source/pythonapi/base.rst @@ -26,6 +26,7 @@ Simulation Settings openmc.FileSource openmc.CompiledSource openmc.MeshSource + openmc.TokamakSource openmc.SourceParticle openmc.VolumeCalculation openmc.Settings diff --git a/docs/source/usersguide/settings.rst b/docs/source/usersguide/settings.rst index e8514b856..6faeb59e1 100644 --- a/docs/source/usersguide/settings.rst +++ b/docs/source/usersguide/settings.rst @@ -291,6 +291,53 @@ example, the following would generate a photon source:: For a full list of all classes related to statistical distributions, see :ref:`pythonapi_stats`. +Tokamak Plasma Sources +---------------------- + +For fusion applications, the :class:`openmc.TokamakSource` class provides a +native parametric neutron source for tokamak plasmas. Rather than specifying +spatial, angular, and energy distributions separately, the source is defined by +the plasma geometry (using a `Miller-style flux-surface parameterization +`_) and a radial emission profile. Source +sites are sampled directly from the plasma volume without rejection. + +The plasma shape is described by the major radius :math:`R_0`, minor radius +:math:`a`, elongation :math:`\kappa`, triangularity :math:`\delta`, and +Shafranov shift :math:`\Delta`. The neutron birth profile is given as an +emission density :math:`S(r/a)` tabulated on a normalized minor-radius grid that +runs from 0 (magnetic axis) to 1 (last closed flux surface); only the shape of +the profile matters, since it is normalized internally. The emission density is +linearly interpolated between the supplied points and refined internally for +radial sampling. For example:: + + import numpy as np + + r_over_a = np.linspace(0.0, 1.0, 50) + emission = (1.0 - r_over_a**2)**2 # peaked on-axis profile + + source = openmc.TokamakSource( + major_radius=620.0, # cm + minor_radius=200.0, # cm + elongation=1.8, + triangularity=0.45, + shafranov_shift=10.0, # cm + r_over_a=r_over_a, + emission_density=emission, + energy=openmc.stats.muir(e0=14.08e6, m_rat=5.0, kt=20000.0), + ) + + settings.source = source + +The ``energy`` argument accepts either a single +:class:`~openmc.stats.Univariate` distribution applied at all radii, or a +sequence with one distribution per ``r_over_a`` grid point to model a +radially-varying neutron spectrum (energies are then sampled by stochastic +interpolation between the two distributions bracketing the sampled radius). A +time distribution can be given with the ``time`` argument; by default, particles +are born at :math:`t=0`. The toroidal extent can be restricted with +``phi_start`` and ``phi_extent`` to model a sector of the plasma, and +``vertical_shift`` translates the plasma center along the z-axis. + File-based Sources ------------------ diff --git a/include/openmc/math_functions.h b/include/openmc/math_functions.h index f322ad28a..d3bccca7b 100644 --- a/include/openmc/math_functions.h +++ b/include/openmc/math_functions.h @@ -213,6 +213,20 @@ double exprel(double x); //! \return log(1+x)/x without loss of precision near 0 double log1prel(double x); +//! Evaluate the cylindrical Bessel function of the first kind J_n(x) +//! +//! Uses std::cyl_bessel_j where available (e.g., libstdc++). On standard +//! library implementations lacking the C++17 special math functions (e.g., +//! libc++ on Apple Clang/LLVM), falls back to the ascending power series, +//! which converges to machine precision for the small arguments (|x| <= 2) +//! used in OpenMC. Unlike std::cyl_bessel_j, negative arguments are handled +//! via the parity relation J_n(-x) = (-1)^n J_n(x). +//! +//! \param n Non-negative integer order of the Bessel function +//! \param x Real argument +//! \return J_n(x) +double cyl_bessel_j(int n, double x); + //! Helper function to get index and interpolation function on an incident //! energy grid //! diff --git a/include/openmc/source.h b/include/openmc/source.h index e307b1ed2..51b54a1d1 100644 --- a/include/openmc/source.h +++ b/include/openmc/source.h @@ -7,9 +7,12 @@ #include #include #include +#include // for pair #include "pugixml.hpp" +#include "openmc/array.h" +#include "openmc/distribution.h" #include "openmc/distribution_multi.h" #include "openmc/distribution_spatial.h" #include "openmc/memory.h" @@ -260,6 +263,139 @@ private: vector> sources_; //!< Source distributions }; +//============================================================================== +//! Parametric tokamak plasma neutron source +//! +//! This source samples neutron positions from a tokamak plasma geometry using +//! Miller-style flux surface parameterization with user-specified emission +//! profiles and energy distributions. +//! +//! Flux surface parameterization: +//! R = R0 + r*cos(alpha + delta*sin(alpha)) + Delta*(1 - (r/a)^2) +//! Z = kappa * r * sin(alpha) +//! +//! The sampling algorithm: +//! 1. Sample minor radius r from precomputed CDF of S(r) * Jacobian +//! 2. Sample poloidal angle alpha from conditional P(alpha|r) using mixture +//! of precomputed CDFs weighted by functions of r +//! 3. Sample energy and time from user-provided distribution(s) +//! 4. Sample isotropic direction +//! 5. Sample toroidal angle phi uniformly in [phi_start, phi_start + +//! phi_extent] +//! 6. Transform (r, alpha, phi) to Cartesian (x, y, z), applying the optional +//! vertical shift of the plasma center +//! +//! The user provides the emission density S(r) directly (e.g., from transport +//! codes like TRANSP, ASTRA, etc.), allowing full flexibility in reaction +//! physics calculations. S(r) is a profile in arbitrary units sampled on the +//! r_over_a grid; only its shape matters, since it is normalized internally. +//! Energy distributions can be specified as either a single distribution for +//! all r, or one distribution per radial point. +//============================================================================== + +class TokamakSource : public Source { +public: + // Constructors + explicit TokamakSource(pugi::xml_node node); + + //! Sample from the tokamak source distribution + //! \param[inout] seed Pseudorandom seed pointer + //! \return Sampled site + SourceSite sample(uint64_t* seed) const override; + +private: + //========================================================================== + // Private methods + + //! Precompute data structures for efficient sampling + void precompute_sampling_distributions(); + + //! Sample minor radius from marginal CDF + //! \param seed Pseudorandom seed pointer + //! \return Sampled r/a value + double sample_r_over_a(uint64_t* seed) const; + + //! Sample poloidal angle given r using mixture of precomputed CDFs + //! \param r_norm Normalized minor radius r/a + //! \param seed Pseudorandom seed pointer + //! \return Sampled poloidal angle alpha [rad] + double sample_poloidal_angle(double r_norm, uint64_t* seed) const; + + //! Compute the k-th mixture weight w_k(r) * I_hat_k for poloidal sampling + //! \param k Basis function index (0-5) + //! \param r Normalized minor radius r/a + //! \return Mixture weight for component k + double mixture_weight(int k, double r) const; + + //! Sample energy from the distribution(s) + //! \param r_norm Normalized minor radius r/a (for distribution selection) + //! \param seed Pseudorandom seed pointer + //! \return (Sampled energy [eV], importance weight) + std::pair sample_energy(double r_norm, uint64_t* seed) const; + + //! Transform from flux coordinates (r, alpha, phi) to Cartesian (x, y, z) + //! \param r Minor radius [cm] + //! \param alpha Poloidal angle [rad] + //! \param phi Toroidal angle [rad] + //! \return Position in Cartesian coordinates [cm] + Position flux_to_cartesian(double r, double alpha, double phi) const; + + //========================================================================== + // Data members + + // Emission profile (input) + vector r_over_a_; //!< Normalized minor radius grid points + vector emission_density_; //!< Emission density S(r) at grid points + + // Energy distribution(s): either 1 for all r, or one per r point + vector> energy_dists_; + + // Time distribution (defaults to a delta distribution at t=0) + UPtrDist time_; + + // Angular distribution (isotropic) + UPtrAngle angle_; + + // Tokamak geometry parameters + double major_radius_; //!< Major radius R0 [cm] + double minor_radius_; //!< Minor radius a [cm] + double elongation_; //!< Elongation kappa + double triangularity_; //!< Triangularity delta + double shafranov_shift_; //!< Shafranov shift Delta [cm] + double vertical_shift_; //!< Vertical shift of plasma center [cm] + + // Normalized geometry parameters (precomputed for efficiency) + double epsilon_; //!< Inverse aspect ratio a/R0 + double delta_tilde_; //!< Normalized Shafranov shift Delta/a + + // Toroidal angle bounds + double phi_start_; //!< Starting toroidal angle [rad] + double phi_extent_; //!< Toroidal angle extent [rad] + + // Precomputed distribution for radial sampling + unique_ptr radial_dist_; + + // Coefficients of the radial geometric polynomial: A*r - B*r^2 - C*r^3 + // Also used as the analytical normalization for poloidal mixture weights + double radial_poly_a_; //!< 1 + ε·Δ̃ + double radial_poly_b_; //!< (3/8)·c₁·ε + double radial_poly_c_; //!< 2·ε·Δ̃ + + // Precomputed Bernstein basis functions for poloidal angle sampling. + // Using the factorization f(r_tilde, alpha) = R_tilde x J_tilde where: + // R_tilde = b0*(1-r)^2 + 2*b1*r*(1-r) + b2*r^2 (Bernstein quadratic) + // J_tilde = b3*(1-r) + b4*r (Bernstein linear) + // The product gives 6 non-negative basis functions g_k(alpha) with + // weights w_k(r_tilde) that are products of Bernstein polynomials. + // Distributions are tabulated on [0, pi] exploiting up-down symmetry. + static constexpr int N_POLOIDAL_BASIS = 6; //!< Number of basis functions + int n_alpha_; //!< Number of poloidal angle grid points + array, N_POLOIDAL_BASIS> + poloidal_dists_; //!< Distributions for each basis function g_k(alpha) + array + poloidal_integrals_; //!< Integrals of g_k(alpha) over [0, pi] +}; + //============================================================================== // Functions //============================================================================== diff --git a/openmc/source.py b/openmc/source.py index a11cfd6c8..5947540d6 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import ABC, abstractmethod from collections.abc import Iterable, Sequence -from numbers import Real +from numbers import Integral, Real from pathlib import Path import warnings from typing import Any @@ -46,7 +46,7 @@ class SourceBase(ABC): Attributes ---------- - type : {'independent', 'file', 'compiled', 'mesh'} + type : {'independent', 'file', 'compiled', 'mesh', 'tokamak'} Indicator of source type. strength : float Strength of the source @@ -203,6 +203,8 @@ class SourceBase(ABC): return FileSource.from_xml_element(elem) elif source_type == 'mesh': return MeshSource.from_xml_element(elem, meshes) + elif source_type == 'tokamak': + return TokamakSource.from_xml_element(elem) else: raise ValueError( f'Source type {source_type} is not recognized') @@ -896,6 +898,430 @@ class FileSource(SourceBase): return cls(**kwargs) +class TokamakSource(SourceBase): + r"""A source representing neutron emission from a tokamak plasma. + + This source samples neutron positions from a tokamak plasma geometry using + Miller-style flux surface parameterization. The user provides an emission + profile S(r/a) as a function of normalized minor radius, along with one or + more energy distributions. + + The flux surface parameterization is + + .. math:: + + \begin{aligned} + R &= R_0 + r \cos\left(\alpha + \delta \sin\alpha\right) + + \Delta \left[1 - \left(\frac{r}{a}\right)^2\right] \\ + Z &= Z_\mathrm{shift} + \kappa r \sin\alpha + \end{aligned} + + where :math:`R_0` is major radius, :math:`a` is minor radius, + :math:`\kappa` is elongation, :math:`\delta` is triangularity, + :math:`\Delta` is the Shafranov shift, and :math:`Z_\mathrm{shift}` is + the vertical shift. + + .. versionadded:: 0.15.4 + + Parameters + ---------- + major_radius : float + Major radius R0 in [cm] + minor_radius : float + Minor radius a in [cm] + elongation : float + Plasma elongation κ (must be > 0) + triangularity : float + Plasma triangularity δ (must be in [-1, 1]) + shafranov_shift : float + Shafranov shift Δ in [cm] (must be >= 0 and < a/2) + r_over_a : numpy.ndarray + Normalized minor radius grid points, must start at 0 and end at 1 + emission_density : numpy.ndarray + Emission density S(r) at each r/a point (arbitrary units, must be >= 0). + Values are linearly interpolated between grid points and refined on an + internal grid for radial sampling. Must have the same length as + ``r_over_a`` and contain at least one positive value. + energy : openmc.stats.Univariate or Sequence[openmc.stats.Univariate] + Energy distribution(s). Either a single distribution used at all radii, + or one distribution per ``r_over_a`` grid point. When one distribution + per grid point is given, the energy of a sampled particle is drawn from + one of the two distributions bracketing its sampled radius, selected + stochastically with probability proportional to the proximity of the + radius to each grid point (stochastic interpolation). + time : openmc.stats.Univariate, optional + Time distribution of the source. If None, particles are born at + :math:`t=0`, matching the default behavior of + :class:`openmc.IndependentSource`. + phi_start : float + Starting toroidal angle in [rad] (default: 0) + phi_extent : float + Toroidal angle extent in [rad] (default: 2π) + n_alpha : int + Number of poloidal angle grid points for CDF sampling (default: 101) + vertical_shift : float + Vertical shift of the plasma center in [cm] (default: 0) + strength : float + Strength of the source (default: 1.0) + constraints : dict + Constraints on sampled source particles. See :class:`SourceBase` for + valid keys and values. + + Attributes + ---------- + major_radius : float + Major radius R0 in [cm] + minor_radius : float + Minor radius a in [cm] + elongation : float + Plasma elongation κ + triangularity : float + Plasma triangularity δ + shafranov_shift : float + Shafranov shift Δ in [cm] + r_over_a : numpy.ndarray + Normalized minor radius grid points + emission_density : numpy.ndarray + Emission density S(r) at each r/a point + energy : list of openmc.stats.Univariate + Energy distribution(s) + time : openmc.stats.Univariate or None + Time distribution of the source + phi_start : float + Starting toroidal angle in [rad] + phi_extent : float + Toroidal angle extent in [rad] + n_alpha : int + Number of poloidal angle grid points + vertical_shift : float + Vertical shift of the plasma center in [cm] + strength : float + Strength of the source + type : str + Indicator of source type: 'tokamak' + constraints : dict + Constraints on sampled source particles + + """ + + def __init__( + self, + major_radius: float, + minor_radius: float, + elongation: float, + triangularity: float, + shafranov_shift: float, + r_over_a: Sequence[float], + emission_density: Sequence[float], + energy: Univariate | Sequence[Univariate], + time: Univariate | None = None, + phi_start: float = 0.0, + phi_extent: float = 2.0 * np.pi, + n_alpha: int = 101, + vertical_shift: float = 0.0, + strength: float = 1.0, + constraints: dict[str, Any] | None = None + ): + super().__init__(strength=strength, constraints=constraints) + self.major_radius = major_radius + self.minor_radius = minor_radius + self.elongation = elongation + self.triangularity = triangularity + self.shafranov_shift = shafranov_shift + self.r_over_a = r_over_a + self.emission_density = emission_density + self.phi_start = phi_start + self.phi_extent = phi_extent + self.n_alpha = n_alpha + self.vertical_shift = vertical_shift + self.energy = energy + self.time = time + + self._validate() + + def _validate(self): + """Validate relationships between tokamak source parameters.""" + if self.minor_radius >= self.major_radius: + raise ValueError( + f"minor_radius ({self.minor_radius}) must be smaller than " + f"major_radius ({self.major_radius})") + if self.shafranov_shift >= 0.5 * self.minor_radius: + raise ValueError( + f"shafranov_shift ({self.shafranov_shift}) must be smaller " + f"than half the minor_radius ({0.5 * self.minor_radius})") + if len(self.emission_density) != len(self.r_over_a): + raise ValueError( + f"emission_density (length {len(self.emission_density)}) must " + f"have the same length as r_over_a (length {len(self.r_over_a)})") + if not np.any(self.emission_density > 0.0): + raise ValueError("emission_density must contain a positive value") + if len(self.energy) not in (1, len(self.r_over_a)): + raise ValueError( + f"Number of energy distributions ({len(self.energy)}) must be " + f"either 1 or equal to the number of r_over_a grid points " + f"({len(self.r_over_a)})") + + @property + def type(self) -> str: + return "tokamak" + + @property + def major_radius(self) -> float: + return self._major_radius + + @major_radius.setter + def major_radius(self, value: float): + cv.check_type('major radius', value, Real) + cv.check_greater_than('major radius', value, 0.0) + self._major_radius = value + + @property + def minor_radius(self) -> float: + return self._minor_radius + + @minor_radius.setter + def minor_radius(self, value: float): + cv.check_type('minor radius', value, Real) + cv.check_greater_than('minor radius', value, 0.0) + self._minor_radius = value + + @property + def elongation(self) -> float: + return self._elongation + + @elongation.setter + def elongation(self, value: float): + cv.check_type('elongation', value, Real) + cv.check_greater_than('elongation', value, 0.0) + self._elongation = value + + @property + def triangularity(self) -> float: + return self._triangularity + + @triangularity.setter + def triangularity(self, value: float): + cv.check_type('triangularity', value, Real) + cv.check_greater_than('triangularity', value, -1.0, equality=True) + cv.check_less_than('triangularity', value, 1.0, equality=True) + self._triangularity = value + + @property + def shafranov_shift(self) -> float: + return self._shafranov_shift + + @shafranov_shift.setter + def shafranov_shift(self, value: float): + cv.check_type('Shafranov shift', value, Real) + cv.check_greater_than('Shafranov shift', value, 0.0, equality=True) + self._shafranov_shift = value + + @property + def r_over_a(self) -> np.ndarray: + return self._r_over_a + + @r_over_a.setter + def r_over_a(self, value: Sequence[float]): + value = np.asarray(value, dtype=float) + if value.ndim != 1 or len(value) < 2: + raise ValueError("r_over_a must be a 1-D array with at least 2 points") + if value[0] != 0.0: + raise ValueError("r_over_a must start at 0") + if value[-1] != 1.0: + raise ValueError("r_over_a must end at 1") + if not np.all(np.diff(value) > 0): + raise ValueError("r_over_a must be strictly increasing") + self._r_over_a = value + + @property + def emission_density(self) -> np.ndarray: + return self._emission_density + + @emission_density.setter + def emission_density(self, value: Sequence[float]): + value = np.asarray(value, dtype=float) + if value.ndim != 1: + raise ValueError("emission_density must be a 1-D array") + if np.any(value < 0): + raise ValueError("emission_density values cannot be negative") + self._emission_density = value + + @property + def energy(self) -> list[Univariate]: + return self._energy + + @energy.setter + def energy(self, value: Univariate | Sequence[Univariate]): + if isinstance(value, Univariate): + self._energy = [value] + else: + cv.check_iterable_type('energy distributions', value, Univariate) + self._energy = list(value) + + @property + def time(self) -> Univariate | None: + return self._time + + @time.setter + def time(self, value: Univariate | None): + if value is not None: + cv.check_type('time distribution', value, Univariate) + self._time = value + + @property + def phi_start(self) -> float: + return self._phi_start + + @phi_start.setter + def phi_start(self, value: float): + cv.check_type('phi_start', value, Real) + self._phi_start = value + + @property + def phi_extent(self) -> float: + return self._phi_extent + + @phi_extent.setter + def phi_extent(self, value: float): + cv.check_type('phi_extent', value, Real) + cv.check_greater_than('phi_extent', value, 0.0) + cv.check_less_than('phi_extent', value, 2.0 * np.pi, equality=True) + self._phi_extent = value + + @property + def n_alpha(self) -> int: + return self._n_alpha + + @n_alpha.setter + def n_alpha(self, value: int): + cv.check_type('n_alpha', value, Integral) + cv.check_greater_than('n_alpha', value, 2) + if value < 51: + warnings.warn( + "n_alpha values below 51 may introduce noticeable " + "discretization bias in tokamak source sampling", stacklevel=2) + self._n_alpha = value + + @property + def vertical_shift(self) -> float: + return self._vertical_shift + + @vertical_shift.setter + def vertical_shift(self, value: float): + cv.check_type('vertical shift', value, Real) + self._vertical_shift = value + + def populate_xml_element(self, element): + """Add necessary tokamak source information to an XML element + + Returns + ------- + element : lxml.etree._Element + XML element containing source data + + """ + self._validate() + + # Geometry parameters + ET.SubElement(element, "major_radius").text = str(self.major_radius) + ET.SubElement(element, "minor_radius").text = str(self.minor_radius) + ET.SubElement(element, "elongation").text = str(self.elongation) + ET.SubElement(element, "triangularity").text = str(self.triangularity) + ET.SubElement(element, "shafranov_shift").text = str(self.shafranov_shift) + + # Toroidal angle bounds + ET.SubElement(element, "phi_start").text = str(self.phi_start) + ET.SubElement(element, "phi_extent").text = str(self.phi_extent) + + # Poloidal sampling resolution + ET.SubElement(element, "n_alpha").text = str(self.n_alpha) + + # Vertical shift + if self.vertical_shift != 0.0: + ET.SubElement(element, "vertical_shift").text = str(self.vertical_shift) + + # Emission profile + ET.SubElement(element, "r_over_a").text = ' '.join(str(r) for r in self.r_over_a) + ET.SubElement(element, "emission_density").text = ' '.join(str(s) for s in self.emission_density) + + # Energy distribution(s) + for dist in self.energy: + element.append(dist.to_xml_element('energy')) + + # Time distribution + if self.time is not None: + element.append(self.time.to_xml_element('time')) + + @classmethod + def from_xml_element(cls, elem: ET.Element) -> TokamakSource: + """Generate tokamak source from an XML element + + Parameters + ---------- + elem : lxml.etree._Element + XML element + + Returns + ------- + openmc.TokamakSource + Source generated from XML element + + """ + # Read geometry parameters + major_radius = float(get_text(elem, 'major_radius')) + minor_radius = float(get_text(elem, 'minor_radius')) + elongation = float(get_text(elem, 'elongation')) + triangularity = float(get_text(elem, 'triangularity')) + shafranov_shift = float(get_text(elem, 'shafranov_shift')) + + # Read optional parameters + phi_start_text = get_text(elem, 'phi_start') + phi_start = float(phi_start_text) if phi_start_text else 0.0 + + phi_extent_text = get_text(elem, 'phi_extent') + phi_extent = float(phi_extent_text) if phi_extent_text else 2.0 * np.pi + + n_alpha_text = get_text(elem, 'n_alpha') + n_alpha = int(n_alpha_text) if n_alpha_text else 101 + + vertical_shift_text = get_text(elem, 'vertical_shift') + vertical_shift = float(vertical_shift_text) if vertical_shift_text else 0.0 + + # Read emission profile + r_over_a = np.array([float(x) for x in get_text(elem, 'r_over_a').split()]) + emission_density = np.array([float(x) for x in get_text(elem, 'emission_density').split()]) + + # Read energy distributions + energy = [Univariate.from_xml_element(e) for e in elem.findall('energy')] + if len(energy) == 1: + energy = energy[0] + + # Read time distribution + time_elem = elem.find('time') + time = Univariate.from_xml_element(time_elem) if time_elem is not None else None + + # Read constraints and strength + constraints = cls._get_constraints(elem) + strength_text = get_text(elem, 'strength') + strength = float(strength_text) if strength_text else 1.0 + + return cls( + major_radius=major_radius, + minor_radius=minor_radius, + elongation=elongation, + triangularity=triangularity, + shafranov_shift=shafranov_shift, + r_over_a=r_over_a, + emission_density=emission_density, + energy=energy, + time=time, + phi_start=phi_start, + phi_extent=phi_extent, + n_alpha=n_alpha, + vertical_shift=vertical_shift, + strength=strength, + constraints=constraints + ) class SourceParticle: diff --git a/src/distribution.cpp b/src/distribution.cpp index 7f5b498ad..cbdd13127 100644 --- a/src/distribution.cpp +++ b/src/distribution.cpp @@ -552,14 +552,9 @@ double Tabular::sample_unbiased(uint64_t* seed) const double c = prn(seed); // Find first CDF bin which is above the sampled value - double c_i = c_[0]; - int i; - std::size_t n = c_.size(); - for (i = 0; i < n - 1; ++i) { - if (c <= c_[i + 1]) - break; - c_i = c_[i + 1]; - } + auto c_iter = std::lower_bound(c_.begin() + 1, c_.end(), c); + int i = std::distance(c_.begin(), c_iter) - 1; + double c_i = c_[i]; // Determine bounding PDF values double x_i = x_[i]; diff --git a/src/math_functions.cpp b/src/math_functions.cpp index 3befa8bc2..ddacc2bd9 100644 --- a/src/math_functions.cpp +++ b/src/math_functions.cpp @@ -1,5 +1,7 @@ #include "openmc/math_functions.h" +#include // for numeric_limits + #include "openmc/external/Faddeeva.hh" #include "openmc/constants.h" @@ -944,6 +946,46 @@ double log1prel(double x) } } +double cyl_bessel_j(int n, double x) +{ + // Handle negative arguments via the parity relation + // J_n(-x) = (-1)^n J_n(x); std::cyl_bessel_j has a domain error for x < 0. + double sign = 1.0; + if (x < 0.0) { + x = -x; + if (n % 2 == 1) + sign = -1.0; + } + +#if defined(__cpp_lib_math_special_functions) && \ + __cpp_lib_math_special_functions >= 201603L + return sign * std::cyl_bessel_j(static_cast(n), x); +#else + // Ascending power series (e.g., Abramowitz & Stegun eq. 9.1.10): + // J_n(x) = sum_{m=0}^inf (-1)^m / (m! (m+n)!) * (x/2)^(2m+n) + // The term ratio is -(x/2)^2 / (m*(m+n)), so for |x| <= 2 the series + // converges to machine precision within ~20 terms. + double half_x = 0.5 * x; + + // First term: (x/2)^n / n! + double term = 1.0; + for (int k = 1; k <= n; ++k) { + term *= half_x / k; + } + + double sum = term; + double neg_half_x_sq = -half_x * half_x; + for (int m = 1; m <= 50; ++m) { + term *= neg_half_x_sq / (m * (m + n)); + sum += term; + if (std::abs(term) <= + std::numeric_limits::epsilon() * std::abs(sum)) + break; + } + return sign * sum; +#endif +} + // Helper function to get index and interpolation function on an incident energy // grid void get_energy_index( diff --git a/src/source.cpp b/src/source.cpp index 1e783b623..12bbdf84e 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -4,7 +4,9 @@ #define HAS_DYNAMIC_LINKING #endif -#include // for move +#include // for max +#include // for sin, cos, abs +#include // for move #ifdef HAS_DYNAMIC_LINKING #include // for dlopen, dlsym, dlclose, dlerror @@ -16,17 +18,20 @@ #include "openmc/bank.h" #include "openmc/capi.h" #include "openmc/cell.h" +#include "openmc/constants.h" #include "openmc/container_util.h" #include "openmc/error.h" #include "openmc/file_utils.h" #include "openmc/geometry.h" #include "openmc/hdf5_interface.h" #include "openmc/material.h" +#include "openmc/math_functions.h" #include "openmc/mcpl_interface.h" #include "openmc/memory.h" #include "openmc/message_passing.h" #include "openmc/mgxs_interface.h" #include "openmc/nuclide.h" +#include "openmc/random_dist.h" #include "openmc/random_lcg.h" #include "openmc/search.h" #include "openmc/settings.h" @@ -100,6 +105,8 @@ unique_ptr Source::create(pugi::xml_node node) return make_unique(node); } else if (source_type == "mesh") { return make_unique(node); + } else if (source_type == "tokamak") { + return make_unique(node); } else { fatal_error(fmt::format("Invalid source type '{}' found.", source_type)); } @@ -673,6 +680,471 @@ SourceSite MeshSource::sample(uint64_t* seed) const return source(element)->sample_with_constraints(seed); } +//============================================================================== +// TokamakSource implementation +//============================================================================== + +TokamakSource::TokamakSource(pugi::xml_node node) : Source(node) +{ + // Read geometry parameters + major_radius_ = std::stod(get_node_value(node, "major_radius")); + minor_radius_ = std::stod(get_node_value(node, "minor_radius")); + elongation_ = std::stod(get_node_value(node, "elongation")); + triangularity_ = std::stod(get_node_value(node, "triangularity")); + shafranov_shift_ = std::stod(get_node_value(node, "shafranov_shift")); + + // Read optional vertical shift + if (check_for_node(node, "vertical_shift")) { + vertical_shift_ = std::stod(get_node_value(node, "vertical_shift")); + } else { + vertical_shift_ = 0.0; + } + + // Read optional toroidal angle bounds + if (check_for_node(node, "phi_start")) { + phi_start_ = std::stod(get_node_value(node, "phi_start")); + } else { + phi_start_ = 0.0; + } + if (check_for_node(node, "phi_extent")) { + phi_extent_ = std::stod(get_node_value(node, "phi_extent")); + } else { + phi_extent_ = 2.0 * PI; + } + if (check_for_node(node, "n_alpha")) { + n_alpha_ = std::stoi(get_node_value(node, "n_alpha")); + } else { + n_alpha_ = 101; // Default + } + + // Read emission profile + r_over_a_ = get_node_array(node, "r_over_a"); + emission_density_ = get_node_array(node, "emission_density"); + + // Read energy distribution(s) + for (auto energy_node : node.children("energy")) { + energy_dists_.push_back(distribution_from_xml(energy_node)); + } + + // Read optional time distribution; default to a delta distribution at t=0 + // for the same behavior as IndependentSource + if (check_for_node(node, "time")) { + time_ = distribution_from_xml(node.child("time")); + } else { + double T[] {0.0}; + double p[] {1.0}; + time_ = UPtrDist {new Discrete {T, p, 1}}; + } + + // Validate inputs + if (emission_density_.size() != r_over_a_.size()) { + fatal_error("TokamakSource: emission_density and r_over_a must have the " + "same length."); + } + if (r_over_a_.size() < 2) { + fatal_error( + "TokamakSource: At least 2 radial points are required for profiles."); + } + if (r_over_a_.front() != 0.0) { + fatal_error("TokamakSource: r_over_a must start at 0."); + } + if (r_over_a_.back() != 1.0) { + fatal_error("TokamakSource: r_over_a must end at 1."); + } + for (size_t i = 1; i < r_over_a_.size(); ++i) { + if (r_over_a_[i] <= r_over_a_[i - 1]) { + fatal_error("TokamakSource: r_over_a must be strictly increasing."); + } + } + for (size_t i = 0; i < emission_density_.size(); ++i) { + if (emission_density_[i] < 0.0) { + fatal_error("TokamakSource: emission_density values cannot be negative."); + } + } + if (major_radius_ <= 0.0) { + fatal_error("TokamakSource: major_radius must be > 0."); + } + if (minor_radius_ <= 0.0) { + fatal_error("TokamakSource: minor_radius must be > 0."); + } + if (minor_radius_ >= major_radius_) { + fatal_error("TokamakSource: minor_radius must be less than major_radius."); + } + if (elongation_ <= 0.0) { + fatal_error("TokamakSource: elongation must be > 0."); + } + if (triangularity_ < -1.0 || triangularity_ > 1.0) { + fatal_error("TokamakSource: triangularity must be in the range [-1, 1]."); + } + if (shafranov_shift_ < 0.0) { + fatal_error("TokamakSource: shafranov_shift must be >= 0."); + } + if (shafranov_shift_ >= 0.5 * minor_radius_) { + fatal_error("TokamakSource: shafranov_shift must be less than half the " + "minor radius."); + } + if (phi_extent_ <= 0.0 || phi_extent_ > 2.0 * PI) { + fatal_error("TokamakSource: phi_extent must be > 0 and <= 2*pi."); + } + if (n_alpha_ <= 2) { + fatal_error("TokamakSource: n_alpha must be > 2."); + } + if (n_alpha_ < 51) { + warning("TokamakSource: n_alpha values below 51 may introduce noticeable " + "discretization bias in source sampling."); + } + if (energy_dists_.empty()) { + fatal_error("TokamakSource: At least one energy distribution is required."); + } + if (energy_dists_.size() != 1 && energy_dists_.size() != r_over_a_.size()) { + fatal_error("TokamakSource: energy distributions must be either 1 (for all " + "r) or match the number of r_over_a points."); + } + + // Compute normalized geometry parameters + epsilon_ = minor_radius_ / major_radius_; + delta_tilde_ = shafranov_shift_ / minor_radius_; + + // Initialize isotropic angular distribution + angle_ = UPtrAngle {new Isotropic()}; + + precompute_sampling_distributions(); +} + +void TokamakSource::precompute_sampling_distributions() +{ + // Use precomputed normalized geometry parameters + double eps = epsilon_; // Inverse aspect ratio (a/R0) + double Dt = delta_tilde_; // Normalized Shafranov shift (Delta/a) + double delta = triangularity_; + + //========================================================================== + // RADIAL CDF (computed first since it's simpler and sampled first) + //========================================================================== + // The marginal radial PDF is obtained by analytically integrating the joint + // distribution f(r_tilde, alpha) over alpha. The result is: + // + // p(r_tilde) ~ S(r_tilde) * [(1 + eps*Dt)*r_tilde + // - (3/8)*c1*eps*r_tilde^2 + // - 2*eps*Dt*r_tilde^3] + // + // where the Bessel function coefficients are: + // c0 = J_0(delta) + J_2(delta) + // c1 = (J_1(2*delta) + J_3(2*delta)) / c0 + // + // For delta -> 0, c0 -> 1 and c1 -> 0, giving the circular cross-section + // limit. + + // Compute Bessel function coefficients. openmc::cyl_bessel_j handles + // negative arguments (negative triangularity) via the parity relation + // J_n(-x) = (-1)^n * J_n(x). + double J0_d = cyl_bessel_j(0, delta); + double J2_d = cyl_bessel_j(2, delta); + double J1_2d = cyl_bessel_j(1, 2.0 * delta); + double J3_2d = cyl_bessel_j(3, 2.0 * delta); + double c0 = J0_d + J2_d; + double c1 = (J1_2d + J3_2d) / c0; + + // Coefficients for the radial polynomial: A*r - B*r^2 - C*r^3 + radial_poly_a_ = 1.0 + eps * Dt; + radial_poly_b_ = 0.375 * c1 * eps; // 3/8 * c1 * eps + radial_poly_c_ = 2.0 * eps * Dt; + + // Build a refined radial grid that retains the user-specified grid points. + // The emission density is interpreted as linear-linear between those points. + constexpr int MIN_SUBINTERVALS = 8; + constexpr double MAX_GRID_SPACING = 1.0e-3; + vector radial_grid {r_over_a_.front()}; + vector radial_emission {emission_density_.front()}; + for (size_t i = 1; i < r_over_a_.size(); ++i) { + double r_lo = r_over_a_[i - 1]; + double r_hi = r_over_a_[i]; + double s_lo = emission_density_[i - 1]; + double s_hi = emission_density_[i]; + int n_subintervals = std::max(MIN_SUBINTERVALS, + static_cast(std::ceil((r_hi - r_lo) / MAX_GRID_SPACING))); + for (int j = 1; j <= n_subintervals; ++j) { + double t = static_cast(j) / n_subintervals; + radial_grid.push_back(r_lo + t * (r_hi - r_lo)); + radial_emission.push_back(s_lo + t * (s_hi - s_lo)); + } + } + + vector radial_pdf(radial_grid.size()); + for (size_t i = 0; i < radial_grid.size(); ++i) { + double r = radial_grid[i]; + // p(r) ~ S(r) * [A*r - B*r^2 - C*r^3] + double geometric_factor = + radial_poly_a_ * r - radial_poly_b_ * r * r - radial_poly_c_ * r * r * r; + radial_pdf[i] = radial_emission[i] * std::max(0.0, geometric_factor); + } + + // Check that the refined profile contains positive probability mass before + // constructing the normalized tabular distribution. + double total = 0.0; + for (size_t i = 1; i < radial_grid.size(); ++i) { + total += 0.5 * (radial_pdf[i - 1] + radial_pdf[i]) * + (radial_grid[i] - radial_grid[i - 1]); + } + if (total <= 0.0) { + fatal_error( + "TokamakSource: Integrated emission density is zero or negative. " + "Check emission_density profile."); + } + radial_dist_ = make_unique(radial_grid.data(), radial_pdf.data(), + radial_grid.size(), Interpolation::lin_lin); + + //========================================================================== + // POLOIDAL CDFs (for conditional sampling of alpha given r) + //========================================================================== + // The conditional distribution P(alpha | r) is a mixture: + // P(alpha | r) ~ sum_k w_k(r) * I_hat_k * p_k(alpha) + // where: + // - w_k(r) are the "dynamic" Bernstein weight functions (depend on r) + // - I_hat_k are the "static" normalized integrals (precomputed constants) + // - p_k(alpha) are the normalized basis distributions (precomputed CDFs) + // + // The static weights I_hat_k = I_k / (2*pi*c0) are: + // I_hat_0 = 1 + eps*Dt + // I_hat_1 = 1 + eps*Dt - (3/16)*c1*eps + // I_hat_2 = 1 - (3/8)*c1*eps + // I_hat_3 = 1 + eps*Dt + // I_hat_4 = 1 + (1/2)*eps*Dt - (3/16)*c1*eps + // I_hat_5 = 1 - eps*Dt - (3/8)*c1*eps + + // Compute static weights analytically + poloidal_integrals_[0] = 1.0 + eps * Dt; + poloidal_integrals_[1] = 1.0 + eps * Dt - 0.1875 * c1 * eps; // 3/16 = 0.1875 + poloidal_integrals_[2] = 1.0 - 0.375 * c1 * eps; // 3/8 = 0.375 + poloidal_integrals_[3] = 1.0 + eps * Dt; + poloidal_integrals_[4] = 1.0 + 0.5 * eps * Dt - 0.1875 * c1 * eps; + poloidal_integrals_[5] = 1.0 - eps * Dt - 0.375 * c1 * eps; + + // Build the alpha grid on [0, pi] (half domain due to up-down symmetry) + int n_alpha = n_alpha_; + vector alpha_grid(n_alpha); + double dalpha = PI / (n_alpha - 1); + for (int i = 0; i < n_alpha; ++i) { + alpha_grid[i] = i * dalpha; + } + + // Compute basis function values g_k(alpha) for tabular distributions + // Using Bernstein form: + // R_tilde = b0*(1-r)^2 + 2*b1*r*(1-r) + b2*r^2 + // J_tilde = b3*(1-r) + b4*r + // with: + // b0(alpha) = 1 + eps*Dt + // b1(alpha) = b0 + (eps/2)*cos(psi), psi = alpha + delta*sin(alpha) + // b2(alpha) = 1 + eps*cos(psi) + // b3(alpha) = cos(delta*sin(alpha)) + // + (delta/4)*(cos(alpha - delta*sin(alpha)) + // - cos(3*alpha + delta*sin(alpha))) + // b4(alpha) = b3(alpha) - 2*Dt*cos(alpha) + + array, N_POLOIDAL_BASIS> basis; + for (int k = 0; k < N_POLOIDAL_BASIS; ++k) { + basis[k].resize(n_alpha); + } + + for (int i = 0; i < n_alpha; ++i) { + double alpha = alpha_grid[i]; + double sin_alpha = std::sin(alpha); + double cos_alpha = std::cos(alpha); + double delta_sin_alpha = delta * sin_alpha; + double psi = alpha + delta_sin_alpha; + double cos_psi = std::cos(psi); + + // Bernstein coefficients b0-b4 + double b0 = 1.0 + eps * Dt; + double b1 = b0 + 0.5 * eps * cos_psi; + double b2 = 1.0 + eps * cos_psi; + double b3 = + std::cos(delta_sin_alpha) + 0.25 * delta * + (std::cos(alpha - delta_sin_alpha) - + std::cos(3.0 * alpha + delta_sin_alpha)); + double b4 = b3 - 2.0 * Dt * cos_alpha; + + // 6 basis functions g_k(alpha) = b_i * b_j + basis[0][i] = b0 * b3; // w0 = (1-r)^3 + basis[1][i] = b1 * b3; // w1 = 2*r*(1-r)^2 + basis[2][i] = b2 * b3; // w2 = r^2*(1-r) + basis[3][i] = b0 * b4; // w3 = r*(1-r)^2 + basis[4][i] = b1 * b4; // w4 = 2*r^2*(1-r) + basis[5][i] = b2 * b4; // w5 = r^3 + } + + // Build a linear-linear distribution for each basis function p_k(alpha) + for (int k = 0; k < N_POLOIDAL_BASIS; ++k) { + poloidal_dists_[k] = make_unique( + alpha_grid.data(), basis[k].data(), n_alpha, Interpolation::lin_lin); + } +} + +double TokamakSource::sample_r_over_a(uint64_t* seed) const +{ + return radial_dist_->sample(seed).first; +} + +double TokamakSource::mixture_weight(int k, double r) const +{ + double s = 1.0 - r; + switch (k) { + case 0: + return s * s * s * poloidal_integrals_[0]; + case 1: + return 2.0 * r * s * s * poloidal_integrals_[1]; + case 2: + return r * r * s * poloidal_integrals_[2]; + case 3: + return r * s * s * poloidal_integrals_[3]; + case 4: + return 2.0 * r * r * s * poloidal_integrals_[4]; + case 5: + return r * r * r * poloidal_integrals_[5]; + default: + UNREACHABLE(); + } +} + +double TokamakSource::sample_poloidal_angle(double r_norm, uint64_t* seed) const +{ + // Sample from the conditional distribution P(alpha | r_tilde) using + // mixture sampling with 6 precomputed basis distributions. + // + // The conditional is: P(alpha | r) ~ sum_k w_k(r) * I_hat_k * p_k(alpha) + // where: + // - w_k(r) are the "dynamic" Bernstein weight functions + // - I_hat_k are the "static" normalized integrals (precomputed in + // poloidal_integrals_) + // - p_k(alpha) are the normalized, precomputed basis distributions + // + // The normalization sum_k w_k(r) * I_hat_k equals the radial geometric + // polynomial evaluated at r, which is known analytically. + // + // Algorithm: + // 1. Compute total from analytical normalization + // 2. Lazily evaluate mixture weights with early exit to select component k + // 3. Sample alpha from the selected basis distribution + + // Analytical normalization: sum_k w_k(r) * I_hat_k + double total = + radial_poly_a_ - radial_poly_b_ * r_norm - radial_poly_c_ * r_norm * r_norm; + double xi = prn(seed) * total; + + // Sample component via lazy evaluation with early exit + // Order optimized for peaked emission profiles: 0, 1, 4, 5, 3, 2 + constexpr int order[] = {0, 1, 4, 5, 3, 2}; + double cumsum = 0.0; + int component = order[N_POLOIDAL_BASIS - 1]; + for (int i = 0; i < N_POLOIDAL_BASIS; ++i) { + cumsum += mixture_weight(order[i], r_norm); + if (xi < cumsum) { + component = order[i]; + break; + } + } + + // Sample alpha from [0, pi] + double alpha = poloidal_dists_[component]->sample(seed).first; + + // Exploit up-down symmetry: randomly flip to [pi, 2*pi] with 50% probability + // This is equivalent to flipping the sign of Z in the final position + if (prn(seed) >= 0.5) { + alpha = 2.0 * PI - alpha; + } + return alpha; +} + +std::pair TokamakSource::sample_energy( + double r_norm, uint64_t* seed) const +{ + if (energy_dists_.size() == 1) { + // Single distribution for all r + return energy_dists_[0]->sample(seed); + } + + // Multiple distributions: stochastic selection between bracketing r points + // Find the interval containing r_norm + size_t i = lower_bound_index(r_over_a_.begin(), r_over_a_.end(), r_norm); + + // Handle boundary cases + if (i >= energy_dists_.size() - 1) { + return energy_dists_.back()->sample(seed); + } + + // Stochastic interpolation: randomly select one of the two bracketing + // distributions based on distance to each + double t = (r_norm - r_over_a_[i]) / (r_over_a_[i + 1] - r_over_a_[i]); + size_t idx = (prn(seed) < t) ? i + 1 : i; + return energy_dists_[idx]->sample(seed); +} + +Position TokamakSource::flux_to_cartesian( + double r, double alpha, double phi) const +{ + // Flux surface parameterization: + // R = R0 + r*cos(alpha + delta*sin(alpha)) + Delta*(1 - (r/a)^2) + // Z = kappa * r * sin(alpha) + // x = R * cos(phi) + // y = R * sin(phi) + // z = Z + + double psi = alpha + triangularity_ * std::sin(alpha); + double r_over_a_sq = (r * r) / (minor_radius_ * minor_radius_); + + double R = + major_radius_ + r * std::cos(psi) + shafranov_shift_ * (1.0 - r_over_a_sq); + double Z = elongation_ * r * std::sin(alpha); + + double x = R * std::cos(phi); + double y = R * std::sin(phi); + double z = Z; + + return {x, y, z}; +} + +SourceSite TokamakSource::sample(uint64_t* seed) const +{ + SourceSite site; + site.particle = ParticleType::neutron(); + site.wgt = 1.0; + site.delayed_group = 0; + + // 1. Sample r/a from radial CDF + double r_norm = sample_r_over_a(seed); + double r = r_norm * minor_radius_; + + // 2. Sample poloidal angle from conditional distribution P(alpha|r) + double alpha = sample_poloidal_angle(r_norm, seed); + + // 3. Sample toroidal angle uniformly in [phi_start, phi_start + phi_extent] + double phi = phi_start_ + phi_extent_ * prn(seed); + + // 4. Convert to Cartesian coordinates + site.r = flux_to_cartesian(r, alpha, phi); + + // 4a. Apply vertical shift if non-zero + if (vertical_shift_ != 0.0) { + site.r.z += vertical_shift_; + } + + // 5. Sample isotropic direction + site.u = angle_->sample(seed).first; + + // 6. Sample energy from distribution(s), applying the importance weight so + // that biased distributions are handled correctly + auto [E, E_wgt] = sample_energy(r_norm, seed); + site.E = E; + + // 7. Sample particle creation time + auto [time, time_wgt] = time_->sample(seed); + site.time = time; + + site.wgt *= E_wgt * time_wgt; + + return site; +} + //============================================================================== // Non-member functions //============================================================================== diff --git a/tests/cpp_unit_tests/test_distribution.cpp b/tests/cpp_unit_tests/test_distribution.cpp index e46088a63..479e7c8a7 100644 --- a/tests/cpp_unit_tests/test_distribution.cpp +++ b/tests/cpp_unit_tests/test_distribution.cpp @@ -82,6 +82,31 @@ TEST_CASE("Test alias sampling method for pugixml constructor") } } +TEST_CASE("Test sampling a large linear-linear tabular distribution") +{ + constexpr int n_points = 10001; + constexpr int n_samples = 200000; + openmc::vector x(n_points); + openmc::vector p(n_points); + for (int i = 0; i < n_points; ++i) { + x[i] = static_cast(i) / (n_points - 1); + p[i] = 2.0 * x[i]; + } + + openmc::Tabular dist( + x.data(), p.data(), n_points, openmc::Interpolation::lin_lin); + uint64_t seed = openmc::init_seed(0, 0); + + double mean = 0.0; + for (int i = 0; i < n_samples; ++i) { + mean += dist.sample(&seed).first; + } + mean /= n_samples; + + // The normalized PDF is 2x on [0, 1], which has a mean of 2/3. + REQUIRE_THAT(mean, Catch::Matchers::WithinAbs(2.0 / 3.0, 0.003)); +} + TEST_CASE("Test construction of SpatialBox with parameters") { openmc::Position ll {-1, -2, -3}; diff --git a/tests/cpp_unit_tests/test_math.cpp b/tests/cpp_unit_tests/test_math.cpp index 983322e4b..7467aa1fe 100644 --- a/tests/cpp_unit_tests/test_math.cpp +++ b/tests/cpp_unit_tests/test_math.cpp @@ -46,6 +46,21 @@ TEST_CASE("Test t_percentile") } } +TEST_CASE("Test cylindrical Bessel functions") +{ + constexpr double x = 2.0; + constexpr double expected[] {0.22389077914123567, 0.57672480775687339, + 0.35283402861563773, 0.12894324947440205}; + + for (int n = 0; n < 4; ++n) { + REQUIRE_THAT(openmc::cyl_bessel_j(n, x), + Catch::Matchers::WithinRel(expected[n], 1.0e-14)); + REQUIRE_THAT(openmc::cyl_bessel_j(n, -x), + Catch::Matchers::WithinRel( + (n % 2 == 0 ? 1.0 : -1.0) * expected[n], 1.0e-14)); + } +} + TEST_CASE("Test calc_pn") { // The reference solutions come from scipy.special.eval_legendre diff --git a/tests/unit_tests/test_source_tokamak.py b/tests/unit_tests/test_source_tokamak.py new file mode 100644 index 000000000..f926c2007 --- /dev/null +++ b/tests/unit_tests/test_source_tokamak.py @@ -0,0 +1,263 @@ +import numpy as np +import pytest + +import openmc +import openmc.stats + +from tests.unit_tests import assert_sample_mean + + +def make_source(**kwargs): + """Build a valid TokamakSource, overriding defaults via kwargs.""" + r_over_a = np.linspace(0.0, 1.0, 10) + params = dict( + major_radius=620.0, + minor_radius=200.0, + elongation=1.8, + triangularity=0.45, + shafranov_shift=10.0, + r_over_a=r_over_a, + emission_density=(1.0 - r_over_a**2), + energy=openmc.stats.muir(e0=14.08e6, m_rat=5.0, kt=2.0e4), + ) + params.update(kwargs) + return openmc.TokamakSource(**params) + + +def test_tokamak_source_roundtrip(): + src = make_source( + phi_start=0.1, phi_extent=np.pi, n_alpha=51, vertical_shift=5.0, + strength=2.0, time=openmc.stats.Uniform(0.0, 1e-6)) + + elem = src.to_xml_element() + assert elem.get('type') == 'tokamak' + + new = openmc.SourceBase.from_xml_element(elem) + assert isinstance(new, openmc.TokamakSource) + assert new.major_radius == src.major_radius + assert new.minor_radius == src.minor_radius + assert new.elongation == src.elongation + assert new.triangularity == src.triangularity + assert new.shafranov_shift == src.shafranov_shift + assert new.phi_start == src.phi_start + assert new.phi_extent == src.phi_extent + assert new.n_alpha == src.n_alpha + assert new.vertical_shift == src.vertical_shift + assert new.strength == src.strength + np.testing.assert_allclose(new.r_over_a, src.r_over_a) + np.testing.assert_allclose(new.emission_density, src.emission_density) + assert len(new.energy) == 1 + assert isinstance(new.time, openmc.stats.Uniform) + assert new.time.a == src.time.a + assert new.time.b == src.time.b + + +def test_tokamak_source_default_time(): + src = make_source() + assert src.time is None + + new = openmc.SourceBase.from_xml_element(src.to_xml_element()) + assert new.time is None + + with pytest.raises(TypeError): + make_source(time=1.0) + + +def test_tokamak_source_multiple_energies(): + r_over_a = np.linspace(0.0, 1.0, 5) + energies = [openmc.stats.muir(e0=14.08e6, m_rat=5.0, kt=kt) + for kt in (1.0e4, 1.5e4, 2.0e4, 2.5e4, 3.0e4)] + src = make_source(r_over_a=r_over_a, + emission_density=np.ones_like(r_over_a), + energy=energies) + assert len(src.energy) == len(r_over_a) + + new = openmc.SourceBase.from_xml_element(src.to_xml_element()) + assert len(new.energy) == len(r_over_a) + + +@pytest.mark.parametrize("kwargs, match", [ + (dict(minor_radius=700.0), "smaller than major_radius"), + (dict(shafranov_shift=150.0), "half the minor_radius"), + (dict(emission_density=np.ones(5)), "same length as r_over_a"), + (dict(energy=[openmc.stats.muir(14.08e6, 5.0, 2.0e4)] * 2), + "Number of energy distributions"), + (dict(r_over_a=np.linspace(0.1, 1.0, 10)), "must start at 0"), + (dict(r_over_a=np.linspace(0.0, 0.9, 10)), "must end at 1"), + (dict(emission_density=-np.linspace(0.0, 1.0, 10)), "cannot be negative"), + (dict(emission_density=np.zeros(10)), "must contain a positive value"), +]) +def test_tokamak_source_invalid(kwargs, match): + with pytest.raises(ValueError, match=match): + make_source(**kwargs) + + +@pytest.mark.parametrize("value", [-1.5, 1.5]) +def test_tokamak_source_invalid_triangularity(value): + with pytest.raises(ValueError): + make_source(triangularity=value) + + +def test_tokamak_source_invalid_n_alpha(): + with pytest.raises(ValueError): + make_source(n_alpha=2) + + +@pytest.mark.parametrize(("attribute", "value", "match"), [ + ("minor_radius", 700.0, "smaller than major_radius"), + ("shafranov_shift", 150.0, "half the minor_radius"), + ("emission_density", np.ones(5), "same length as r_over_a"), + ("energy", [openmc.stats.delta_function(1.0)] * 2, + "Number of energy distributions"), +]) +def test_tokamak_source_mutation_validation(attribute, value, match): + src = make_source() + setattr(src, attribute, value) + with pytest.raises(ValueError, match=match): + src.to_xml_element() + + +@pytest.mark.parametrize(("emission_density", "expected_mean"), [ + ([1.0, 1.0], 2.0 / 3.0), + ([1.0, 0.0], 0.5), +]) +def test_tokamak_source_radial_sampling( + run_in_tmpdir, emission_density, expected_mean +): + """Check radial sampling for profiles on the coarsest valid grid.""" + major_radius = 620.0 + minor_radius = 200.0 + src = make_source( + major_radius=major_radius, + minor_radius=minor_radius, + elongation=1.0, + triangularity=0.0, + shafranov_shift=0.0, + r_over_a=[0.0, 1.0], + emission_density=emission_density, + energy=openmc.stats.delta_function(14.07e6), + ) + + sphere = openmc.Sphere(r=2000.0, boundary_type='vacuum') + model = openmc.Model( + geometry=openmc.Geometry([openmc.Cell(region=-sphere)]), + settings=openmc.Settings( + particles=100, batches=1, run_mode='fixed source', source=src), + ) + + sites = model.sample_external_source(20_000) + xyz = np.array([site.r for site in sites]) + major_r = np.hypot(xyz[:, 0], xyz[:, 1]) + r_over_a = np.hypot(major_r - major_radius, xyz[:, 2]) / minor_radius + assert_sample_mean(r_over_a, expected_mean) + + +def test_tokamak_source_poloidal_sampling(run_in_tmpdir): + """Check linear-linear poloidal sampling on a coarse internal grid.""" + major_radius = 620.0 + minor_radius = 200.0 + with pytest.warns(UserWarning, match="below 51"): + src = make_source( + major_radius=major_radius, + minor_radius=minor_radius, + elongation=1.0, + triangularity=0.0, + shafranov_shift=0.0, + emission_density=np.ones(10), + n_alpha=3, + energy=openmc.stats.delta_function(14.07e6), + ) + + sphere = openmc.Sphere(r=2000.0, boundary_type='vacuum') + model = openmc.Model( + geometry=openmc.Geometry([openmc.Cell(region=-sphere)]), + settings=openmc.Settings( + particles=100, batches=1, run_mode='fixed source', source=src), + ) + + sites = model.sample_external_source(100_000) + xyz = np.array([site.r for site in sites]) + major_r = np.hypot(xyz[:, 0], xyz[:, 1]) + + # With three alpha points, linear interpolation of cos(alpha) on [0, pi] + # gives 1 - 2*alpha/pi. Integrating the resulting density gives this mean. + expected_R = ( + major_radius + + 2.0 * minor_radius**2 / (np.pi**2 * major_radius) + ) + assert_sample_mean(major_r, expected_R) + + +@pytest.mark.flaky(reruns=1) +def test_tokamak_source_sampling(run_in_tmpdir): + """Exercise the compiled C++ sampling path and check invariants. + + Sampled moments are compared against direct numerical quadrature of the + exact source density S(r)*R*|J|, where the Jacobian J of the flux-surface + map is computed from analytic partial derivatives. This check is + independent of the Bernstein-mixture factorization used by the + implementation. + """ + R0, a, kappa, delta = 620.0, 200.0, 1.8, -0.5 + shafranov, zshift = 40.0, 25.0 + phi_start, phi_extent = 0.5, np.pi / 2 + + # Fine grids to make discretization error negligible relative to + # statistical uncertainty + r_over_a = np.linspace(0.0, 1.0, 200) + src = make_source( + major_radius=R0, minor_radius=a, elongation=kappa, + triangularity=delta, shafranov_shift=shafranov, vertical_shift=zshift, + r_over_a=r_over_a, emission_density=1.0 - r_over_a**2, + phi_start=phi_start, phi_extent=phi_extent, n_alpha=201, + energy=openmc.stats.delta_function(14.07e6)) + + sphere = openmc.Sphere(r=2000.0, boundary_type='vacuum') + cell = openmc.Cell(region=-sphere) + settings = openmc.Settings( + particles=100, batches=1, run_mode='fixed source', source=src) + model = openmc.Model(geometry=openmc.Geometry([cell]), settings=settings) + + n_samples = 20_000 + sites = model.sample_external_source(n_samples) + + xyz = np.array([s.r for s in sites]) + R = np.hypot(xyz[:, 0], xyz[:, 1]) + z = xyz[:, 2] + + # Energy, weight, and time invariants + assert np.all([s.E == 14.07e6 for s in sites]) + assert np.all([s.wgt == 1.0 for s in sites]) + assert np.all([s.time == 0.0 for s in sites]) + + # Toroidal angle within the requested sector + phi = np.arctan2(xyz[:, 1], xyz[:, 0]) + assert phi.min() >= phi_start + assert phi.max() <= phi_start + phi_extent + + # Positions bounded by the last closed flux surface + assert R.min() >= R0 - a + assert R.max() <= R0 + a + shafranov + assert np.abs(z - zshift).max() <= kappa * a + + # Up-down symmetry about the vertical shift + assert_sample_mean(z, zshift) + + # Reference moments by 2D quadrature of the exact density + r = np.linspace(0.0, 1.0, 1001)[:, np.newaxis] # r/a + alpha = np.linspace(0.0, 2 * np.pi, 2001)[np.newaxis, :] + psi = alpha + delta * np.sin(alpha) + R_map = R0 + a * r * np.cos(psi) + shafranov * (1.0 - r**2) + Z_map = kappa * a * r * np.sin(alpha) + dR_dr = a * np.cos(psi) - 2.0 * shafranov * r + dR_da = -a * r * np.sin(psi) * (1.0 + delta * np.cos(alpha)) + dZ_dr = kappa * a * np.sin(alpha) * np.ones_like(psi) + dZ_da = kappa * a * r * np.cos(alpha) + jac = np.abs(dR_dr * dZ_da - dR_da * dZ_dr) + dens = (1.0 - r**2) * R_map * jac + norm = dens.sum() + expected_R = (R_map * dens).sum() / norm + expected_z2 = (Z_map**2 * dens).sum() / norm + + assert_sample_mean(R, expected_R) + assert_sample_mean((z - zshift)**2, expected_z2)