From dc6a7f4e12a0967a92ced7e343ede20088173b63 Mon Sep 17 00:00:00 2001 From: Ethan Peterson Date: Sat, 29 Aug 2026 10:53:23 -0400 Subject: [PATCH 01/33] initial stellarator source implementation --- include/openmc/source.h | 119 +++++++ openmc/source.py | 707 ++++++++++++++++++++++++++++++++++++++++ src/source.cpp | 459 ++++++++++++++++++++++++++ 3 files changed, 1285 insertions(+) diff --git a/include/openmc/source.h b/include/openmc/source.h index 51b54a1d106..0501e14c226 100644 --- a/include/openmc/source.h +++ b/include/openmc/source.h @@ -396,6 +396,125 @@ class TokamakSource : public Source { poloidal_integrals_; //!< Integrals of g_k(alpha) over [0, pi] }; +//============================================================================== +//! Parametric stellarator plasma neutron source +//! +//! This source samples neutron positions from a 3-D stellarator plasma using +//! the flux-surface Fourier representation shared by the VMEC and DESC +//! equilibrium codes. The flux coordinates are (rho, theta, zeta), where +//! rho = sqrt(s) is the square root of the normalized toroidal flux +//! (proportional to the average minor radius), theta is the poloidal angle, +//! and zeta is the toroidal angle, which coincides with the cylindrical +//! azimuthal angle phi in both VMEC and DESC. The flux surfaces are +//! +//! R(rho, theta, zeta) = sum_k [ Rc_k(rho) cos(m_k theta - n_k Nfp zeta) +//! + Rs_k(rho) sin(m_k theta - n_k Nfp zeta) ] +//! Z(rho, theta, zeta) = sum_k [ Zs_k(rho) sin(m_k theta - n_k Nfp zeta) +//! + Zc_k(rho) cos(m_k theta - n_k Nfp zeta) ] +//! +//! where Nfp is the number of field periods. For stellarator-symmetric +//! equilibria only Rc and Zs are non-zero. +//! +//! With phi = zeta the volume element is dV = R |tau| drho dtheta dzeta with +//! the poloidal-plane Jacobian +//! +//! tau = dR/drho * dZ/dtheta - dR/dtheta * dZ/drho +//! +//! so for an emission density S(rho) [n/cm^3-s] that is constant on flux +//! surfaces the joint sampling density is +//! +//! p(rho, theta, zeta) ~ S(rho) * R * |tau| +//! +//! The sampling algorithm: +//! 1. Sample rho from the marginal p(rho) ~ S(rho) * V'(rho), where +//! V'(rho) = int R |tau| dtheta dzeta is the differential volume, evaluated +//! exactly at setup via discrete Fourier orthogonality (trapezoidal +//! quadrature of a band-limited trigonometric polynomial) and inverted with +//! a tabulated CDF. +//! 2. Sample (theta, zeta) from the conditional p(theta, zeta | rho) +//! ~ R |tau| by rejection against a precomputed per-radial-bin majorant. +//! 3. Sample energy and time from user-provided distribution(s). +//! 4. Sample an isotropic direction and map to Cartesian coordinates via +//! x = R cos(zeta), y = R sin(zeta), z = Z. +//! +//! Fourier coefficients are interpolated linearly in rho between the surfaces +//! provided on the rho grid; within a radial bin the density R*tau is then +//! exactly quadratic in the interpolation parameter, which allows a rigorous +//! per-bin rejection majorant to be constructed from three evaluations. +//============================================================================== + +class StellaratorSource : public Source { +public: + // Constructors + explicit StellaratorSource(pugi::xml_node node); + + //! Sample from the stellarator source distribution + //! \param[inout] seed Pseudorandom seed pointer + //! \return Sampled site + SourceSite sample(uint64_t* seed) const override; + +private: + //========================================================================== + // Private methods + + //! Precompute the radial CDF and per-bin rejection majorants + void precompute_sampling_distributions(); + + //! Evaluate the (signed) joint density f = R * tau at a point + //! \param bin Radial bin index (between rho grid points bin and bin+1) + //! \param t Interpolation parameter within the bin, in [0, 1] + //! \param theta Poloidal angle [rad] + //! \param zeta Toroidal angle [rad] + //! \param R_out If non-null, receives the major radius R [cm] + //! \param Z_out If non-null, receives the height Z [cm] + //! \return Signed density R * tau (positive after global sign fix) + double eval_density(int bin, double t, double theta, double zeta, + double* R_out = nullptr, double* Z_out = nullptr) const; + + //! Sample energy from the distribution(s) + //! \param rho Normalized radial coordinate (for distribution selection) + //! \param seed Pseudorandom seed pointer + //! \return (Sampled energy [eV], importance weight) + std::pair sample_energy(double rho, uint64_t* seed) const; + + //========================================================================== + // Data members + + // Radial grid and emission profile (input) + vector rho_; //!< Radial grid, rho = sqrt(s) in [0, 1] + vector emission_density_; //!< Emission density S(rho) at grid points + + // Fourier mode description (input) + vector mode_m_; //!< Poloidal mode numbers + vector mode_n_; //!< Toroidal mode numbers (per field period) + int num_field_periods_; //!< Number of field periods Nfp + + // Fourier coefficient tables, flattened [n_rho x n_modes] row-major, [cm] + vector rmnc_; //!< cos coefficients of R + vector zmns_; //!< sin coefficients of Z + vector rmns_; //!< sin coefficients of R (non-symmetric only) + vector zmnc_; //!< cos coefficients of Z (non-symmetric only) + bool asym_ {false}; //!< Whether non-stellarator-symmetric terms present + + // Energy distribution(s): either 1 for all rho, or one per rho point + vector> energy_dists_; + + // Time distribution (defaults to a delta distribution at t=0) + UPtrDist time_; + + // Angular distribution (isotropic) + UPtrAngle angle_; + + // Precomputed distribution for radial sampling + unique_ptr radial_dist_; + + // Per-radial-bin majorant of R*tau for rejection sampling (n_rho - 1) + vector envelope_; + + // Global sign of the flux-coordinate Jacobian (+1 or -1) + double jacobian_sign_ {1.0}; +}; + //============================================================================== // Functions //============================================================================== diff --git a/openmc/source.py b/openmc/source.py index e68bdd5cf25..5c05f2aae3f 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -205,6 +205,8 @@ def from_xml_element(cls, elem: ET.Element, meshes=None) -> SourceBase: return MeshSource.from_xml_element(elem, meshes) elif source_type == 'tokamak': return TokamakSource.from_xml_element(elem) + elif source_type == 'stellarator': + return StellaratorSource.from_xml_element(elem) else: raise ValueError( f'Source type {source_type} is not recognized') @@ -1324,6 +1326,711 @@ def from_xml_element(cls, elem: ET.Element) -> TokamakSource: ) +class StellaratorSource(SourceBase): + r"""A source representing neutron emission from a stellarator plasma. + + This source samples neutron positions from a 3-D stellarator plasma + described by the flux-surface Fourier representation shared by the VMEC + and DESC equilibrium codes. The flux coordinates are + :math:`(\rho, \theta, \zeta)` where :math:`\rho = \sqrt{s}` is the square + root of the normalized toroidal flux (proportional to the average minor + radius), :math:`\theta` is the poloidal angle, and :math:`\zeta` is the + toroidal angle, which coincides with the cylindrical azimuthal angle + :math:`\phi` in both codes. Flux surfaces are given by + + .. math:: + + \begin{aligned} + R(\rho,\theta,\zeta) &= \sum_k \left[ R^c_k(\rho) + \cos(m_k\theta - n_k N_{fp} \zeta) + + R^s_k(\rho) \sin(m_k\theta - n_k N_{fp}\zeta) \right] \\ + Z(\rho,\theta,\zeta) &= \sum_k \left[ Z^s_k(\rho) + \sin(m_k\theta - n_k N_{fp}\zeta) + + Z^c_k(\rho) \cos(m_k\theta - n_k N_{fp}\zeta) \right] + \end{aligned} + + where :math:`N_{fp}` is the number of field periods. For + stellarator-symmetric equilibria only :math:`R^c` and :math:`Z^s` are + non-zero. + + Because :math:`\phi = \zeta`, the volume element is + :math:`dV = R\,|\tau|\, d\rho\, d\theta\, d\zeta` with the poloidal-plane + Jacobian :math:`\tau = \partial_\rho R\, \partial_\theta Z - + \partial_\theta R \,\partial_\rho Z`, so for an emission density + :math:`S(\rho)` constant on flux surfaces the joint density is + :math:`p(\rho,\theta,\zeta) \propto S(\rho) R |\tau|`. The radial + coordinate is sampled from the marginal + :math:`p(\rho) \propto S(\rho) V'(\rho)` (with + :math:`V'(\rho) = \oint\oint R|\tau| \,d\theta\, d\zeta` the differential + volume, evaluated exactly at initialization) via a tabulated CDF, and the + two angles are then sampled from the conditional + :math:`p(\theta,\zeta|\rho) \propto R|\tau|` by rejection against a + precomputed per-radial-bin majorant. + + .. versionadded:: 0.16.1 + + Parameters + ---------- + rho : numpy.ndarray + Radial grid points :math:`\rho = \sqrt{s}`, must start at 0 and end + at 1 and be strictly increasing. + emission_density : numpy.ndarray + Emission density :math:`S(\rho)` at each grid point (arbitrary units, + must be >= 0). Only the shape matters; it is normalized internally. + mode_m : numpy.ndarray of int + Poloidal mode numbers :math:`m_k \ge 0`. + mode_n : numpy.ndarray of int + Toroidal mode numbers :math:`n_k` in units of the number of field + periods (VMEC's ``xn``/``nfp``). + rmnc : numpy.ndarray + Cosine Fourier coefficients of :math:`R` in [cm] with shape + ``(len(rho), len(mode_m))``. + zmns : numpy.ndarray + Sine Fourier coefficients of :math:`Z` in [cm] with shape + ``(len(rho), len(mode_m))``. + energy : openmc.stats.Univariate or Sequence[openmc.stats.Univariate] + Energy distribution(s). Either a single distribution used at all + radii, or one distribution per ``rho`` grid point (selected by + stochastic interpolation between the two bracketing grid points). + num_field_periods : int + Number of field periods :math:`N_{fp}` (default: 1). + rmns : numpy.ndarray, optional + Sine Fourier coefficients of :math:`R` in [cm] for + non-stellarator-symmetric equilibria. + zmnc : numpy.ndarray, optional + Cosine Fourier coefficients of :math:`Z` in [cm] for + non-stellarator-symmetric equilibria. + time : openmc.stats.Univariate, optional + Time distribution of the source. If None, particles are born at + :math:`t = 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 + ---------- + rho : numpy.ndarray + Radial grid points + emission_density : numpy.ndarray + Emission density at each grid point + mode_m : numpy.ndarray + Poloidal mode numbers + mode_n : numpy.ndarray + Toroidal mode numbers (per field period) + rmnc : numpy.ndarray + Cosine coefficients of R in [cm] + zmns : numpy.ndarray + Sine coefficients of Z in [cm] + rmns : numpy.ndarray or None + Sine coefficients of R in [cm] + zmnc : numpy.ndarray or None + Cosine coefficients of Z in [cm] + num_field_periods : int + Number of field periods + energy : list of openmc.stats.Univariate + Energy distribution(s) + time : openmc.stats.Univariate or None + Time distribution of the source + strength : float + Strength of the source + type : str + Indicator of source type: 'stellarator' + constraints : dict + Constraints on sampled source particles + + """ + + def __init__( + self, + rho: Sequence[float], + emission_density: Sequence[float], + mode_m: Sequence[int], + mode_n: Sequence[int], + rmnc: Sequence[Sequence[float]], + zmns: Sequence[Sequence[float]], + energy: Univariate | Sequence[Univariate], + num_field_periods: int = 1, + rmns: Sequence[Sequence[float]] | None = None, + zmnc: Sequence[Sequence[float]] | None = None, + time: Univariate | None = None, + strength: float = 1.0, + constraints: dict[str, Any] | None = None + ): + super().__init__(strength=strength, constraints=constraints) + self.rho = rho + self.emission_density = emission_density + self.mode_m = mode_m + self.mode_n = mode_n + self.rmnc = rmnc + self.zmns = zmns + self.rmns = rmns + self.zmnc = zmnc + self.num_field_periods = num_field_periods + self.energy = energy + self.time = time + + self._validate() + + def _validate(self): + """Validate relationships between stellarator source parameters.""" + n_rho = len(self.rho) + n_modes = len(self.mode_m) + if len(self.emission_density) != n_rho: + raise ValueError( + f"emission_density (length {len(self.emission_density)}) must " + f"have the same length as rho (length {n_rho})") + if not np.any(self.emission_density > 0.0): + raise ValueError("emission_density must contain a positive value") + if len(self.mode_n) != n_modes: + raise ValueError( + f"mode_n (length {len(self.mode_n)}) must have the same " + f"length as mode_m (length {n_modes})") + for name in ('rmnc', 'zmns', 'rmns', 'zmnc'): + coeff = getattr(self, name) + if coeff is not None and coeff.shape != (n_rho, n_modes): + raise ValueError( + f"{name} must have shape (len(rho), len(mode_m)) = " + f"({n_rho}, {n_modes}), got {coeff.shape}") + if (self.rmns is None) != (self.zmnc is None): + raise ValueError( + "rmns and zmnc must both be given for non-stellarator-" + "symmetric equilibria") + if len(self.energy) not in (1, n_rho): + raise ValueError( + f"Number of energy distributions ({len(self.energy)}) must be " + f"either 1 or equal to the number of rho grid points " + f"({n_rho})") + + @property + def type(self) -> str: + return "stellarator" + + @property + def rho(self) -> np.ndarray: + return self._rho + + @rho.setter + def rho(self, value: Sequence[float]): + value = np.asarray(value, dtype=float) + if value.ndim != 1 or len(value) < 2: + raise ValueError("rho must be a 1-D array with at least 2 points") + if value[0] != 0.0: + raise ValueError("rho must start at 0") + if value[-1] != 1.0: + raise ValueError("rho must end at 1") + if not np.all(np.diff(value) > 0): + raise ValueError("rho must be strictly increasing") + self._rho = 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 mode_m(self) -> np.ndarray: + return self._mode_m + + @mode_m.setter + def mode_m(self, value: Sequence[int]): + value = np.asarray(value, dtype=int) + if value.ndim != 1 or len(value) < 1: + raise ValueError("mode_m must be a 1-D array with at least 1 mode") + if np.any(value < 0): + raise ValueError("mode_m values must be >= 0") + self._mode_m = value + + @property + def mode_n(self) -> np.ndarray: + return self._mode_n + + @mode_n.setter + def mode_n(self, value: Sequence[int]): + value = np.asarray(value, dtype=int) + if value.ndim != 1: + raise ValueError("mode_n must be a 1-D array") + self._mode_n = value + + @staticmethod + def _check_coeff(name, value, none_ok=False): + if value is None: + if none_ok: + return None + raise ValueError(f"{name} must be given") + value = np.asarray(value, dtype=float) + if value.ndim != 2: + raise ValueError(f"{name} must be a 2-D array with shape " + "(len(rho), len(mode_m))") + return value + + @property + def rmnc(self) -> np.ndarray: + return self._rmnc + + @rmnc.setter + def rmnc(self, value): + self._rmnc = self._check_coeff('rmnc', value) + + @property + def zmns(self) -> np.ndarray: + return self._zmns + + @zmns.setter + def zmns(self, value): + self._zmns = self._check_coeff('zmns', value) + + @property + def rmns(self) -> np.ndarray | None: + return self._rmns + + @rmns.setter + def rmns(self, value): + self._rmns = self._check_coeff('rmns', value, none_ok=True) + + @property + def zmnc(self) -> np.ndarray | None: + return self._zmnc + + @zmnc.setter + def zmnc(self, value): + self._zmnc = self._check_coeff('zmnc', value, none_ok=True) + + @property + def num_field_periods(self) -> int: + return self._num_field_periods + + @num_field_periods.setter + def num_field_periods(self, value: int): + cv.check_type('num_field_periods', value, Integral) + cv.check_greater_than('num_field_periods', value, 0) + self._num_field_periods = int(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 + + @staticmethod + def _evaluate_emission_density(emission_density, rho): + """Evaluate an emission density callable or validate an array.""" + if callable(emission_density): + return np.asarray(emission_density(rho), dtype=float) + emission_density = np.asarray(emission_density, dtype=float) + if emission_density.shape != rho.shape: + raise ValueError( + f"emission_density array (length {len(emission_density)}) " + f"must have one value per radial surface ({len(rho)}); " + "alternatively provide a callable S(rho)") + return emission_density + + @classmethod + def from_vmec( + cls, + wout: PathLike, + emission_density, + energy: Univariate | Sequence[Univariate], + **kwargs + ) -> StellaratorSource: + """Generate a stellarator source from a VMEC wout file. + + The VMEC radial grid is uniform in the normalized toroidal flux + :math:`s`; it is relabeled here as :math:`\\rho_j = \\sqrt{s_j}`, which + leaves the per-surface Fourier coefficients unchanged. VMEC's ``xn`` + convention (kernel :math:`\\cos(m u - x_n v)` with ``xn`` a multiple of + the number of field periods) is converted to per-field-period mode + numbers. Coefficients are converted from [m] to [cm]. + + Classic-format (NetCDF3) wout files are read with + :func:`scipy.io.netcdf_file` and NetCDF4-format files (which are + HDF5-based) with :mod:`h5py`, so no additional packages are required. + + Parameters + ---------- + wout : path-like + Path to a VMEC ``wout_*.nc`` NetCDF output file. + emission_density : callable or numpy.ndarray + Either a callable ``S(rho)`` evaluated at the radial grid points, + or an array with one value per VMEC radial surface (arbitrary + units, must be >= 0). + energy : openmc.stats.Univariate or Sequence[openmc.stats.Univariate] + Energy distribution(s); see the class docstring. + **kwargs + Additional keyword arguments passed to the constructor + (e.g. ``time``, ``strength``, ``constraints``). + + Returns + ------- + openmc.StellaratorSource + + """ + ds = cls._read_wout(input_path(wout)) + rmnc = ds['rmnc'] + zmns = ds['zmns'] + xm = ds['xm'].astype(int) + xn = ds['xn'].astype(int) + nfp = int(ds['nfp']) + lasym = bool(ds.get('lasym__logical__', 0)) + rmns = ds['rmns'] if lasym else None + zmnc = ds['zmnc'] if lasym else None + + ns = rmnc.shape[0] + s = np.linspace(0.0, 1.0, ns) + rho = np.sqrt(s) + + # m to cm + rmnc = rmnc * 100.0 + zmns = zmns * 100.0 + if lasym: + rmns = rmns * 100.0 + zmnc = zmnc * 100.0 + + return cls( + rho=rho, + emission_density=cls._evaluate_emission_density( + emission_density, rho), + mode_m=xm, + mode_n=xn // nfp, + rmnc=rmnc, + zmns=zmns, + rmns=rmns, + zmnc=zmnc, + num_field_periods=nfp, + energy=energy, + **kwargs + ) + + @staticmethod + def _read_wout(path): + """Read the needed variables from a VMEC wout file. + + Classic (NetCDF3) files are read with scipy; NetCDF4-format files are + HDF5-based and read with h5py. + """ + names = ('rmnc', 'zmns', 'rmns', 'zmnc', 'xm', 'xn', 'nfp', + 'lasym__logical__') + try: + from scipy.io import netcdf_file + with netcdf_file(str(path), mmap=False) as ds: + return {k: np.asarray(v[()], dtype=float) + for k, v in ds.variables.items() if k in names} + except (OSError, TypeError, ValueError): + with h5py.File(path, 'r') as ds: + return {k: np.asarray(ds[k][()], dtype=float) + for k in names if k in ds} + + @staticmethod + def _desc_to_combined(modes_m, modes_n, coeffs, table): + """Fold DESC's product-form double Fourier coefficients into the + combined form sum[c*cos(m*theta - n*Nfp*zeta) + s*sin(...)], m >= 0. + + DESC basis conventions: positive (negative) m selects cos(|m| theta) + (sin(|m| theta)); likewise for n with the toroidal angle. The Ptolemy + identities split each product into the two combined-form harmonics + (m, n) and (m, -n), and modes with m < 0 (or m == 0, n < 0) are folded + using cos(-x) = cos(x), sin(-x) = -sin(x). + """ + def add(m, n, c, s): + if m < 0 or (m == 0 and n < 0): + m, n, s = -m, -n, -s + entry = table.setdefault((m, n), [0.0, 0.0]) + entry[0] += c + entry[1] += s + + for m0, n0, x in zip(modes_m, modes_n, coeffs): + am, an = abs(int(m0)), abs(int(n0)) + h = 0.5 * float(x) + if m0 >= 0 and n0 >= 0: # cos(m theta) * cos(n zeta') + add(am, an, h, 0.0) + add(am, -an, h, 0.0) + elif m0 < 0 and n0 >= 0: # sin(m theta) * cos(n zeta') + add(am, an, 0.0, h) + add(am, -an, 0.0, h) + elif m0 >= 0 and n0 < 0: # cos(m theta) * sin(n zeta') + add(am, an, 0.0, -h) + add(am, -an, 0.0, h) + else: # sin(m theta) * sin(n zeta') + add(am, an, h, 0.0) + add(am, -an, -h, 0.0) + + @staticmethod + def _zernike_radial(rho, l, m): + """Unnormalized Zernike radial polynomial :math:`R_l^{|m|}(\\rho)`.""" + from math import factorial + m = abs(m) + out = np.zeros_like(rho, dtype=float) + for k in range((l - m) // 2 + 1): + c = ((-1)**k * factorial(l - k) + / (factorial(k) * factorial((l + m) // 2 - k) + * factorial((l - m) // 2 - k))) + out += c * rho**(l - 2*k) + return out + + @staticmethod + def _desc_spectral_data(eq): + """Extract Fourier-Zernike modes/coefficients from a DESC equilibrium. + + Accepts either a live ``desc.equilibrium.Equilibrium`` (duck-typed, so + no import of desc is needed) or a path to a DESC HDF5 output file, + which is read directly with h5py. Returns ``(r_modes, r_lmn, z_modes, + z_lmn, nfp)`` where the modes arrays have columns ``(l, m, n)``. + """ + if isinstance(eq, (str, Path)): + with h5py.File(input_path(eq), 'r') as f: + # Output files may contain a family of equilibria; use the last + g = f + if '_equilibria' in f: + idx = sorted((k for k in f['_equilibria'] if k.isdigit()), + key=int) + g = f['_equilibria'][idx[-1]] + return (np.asarray(g['_R_basis/_modes'][()], dtype=int), + np.asarray(g['_R_lmn'][()], dtype=float), + np.asarray(g['_Z_basis/_modes'][()], dtype=int), + np.asarray(g['_Z_lmn'][()], dtype=float), + int(g['_NFP'][()])) + return (np.asarray(eq.R_basis.modes, dtype=int), + np.asarray(eq.R_lmn, dtype=float), + np.asarray(eq.Z_basis.modes, dtype=int), + np.asarray(eq.Z_lmn, dtype=float), + int(eq.NFP)) + + @classmethod + def from_desc( + cls, + eq, + emission_density, + energy: Univariate | Sequence[Univariate], + n_rho: int = 33, + **kwargs + ) -> StellaratorSource: + """Generate a stellarator source from a DESC equilibrium. + + DESC represents :math:`R` and :math:`Z` in a Fourier-Zernike basis, + :math:`X(\\rho,\\theta,\\zeta) = \\sum_{lmn} X_{lmn}\\, + \\mathcal{R}_l^{|m|}(\\rho)\\, \\mathcal{F}_m(\\theta)\\, + \\mathcal{F}_n(N_{fp}\\zeta)`, where :math:`\\mathcal{R}_l^{|m|}` is the + (unnormalized) Zernike radial polynomial and :math:`\\mathcal{F}_m(x)` + is :math:`\\cos(|m|x)` for :math:`m \\ge 0` and :math:`\\sin(|m|x)` + otherwise. The Zernike polynomials are evaluated here on a uniform grid + in :math:`\\rho` (DESC's native radial coordinate) to obtain + per-surface double Fourier coefficients, which are then converted from + DESC's product-form basis to the combined VMEC-style form using the + Ptolemy identities. Coefficients are converted from [m] to [cm]. No + additional packages are required: DESC HDF5 output files are read + directly with :mod:`h5py`, and live equilibrium objects are accessed + through their public attributes only. + + Parameters + ---------- + eq : desc.equilibrium.Equilibrium or path-like + A DESC equilibrium object, or a path to a DESC HDF5 output file + (the last equilibrium is used if the file contains a family). + emission_density : callable or numpy.ndarray + Either a callable ``S(rho)`` evaluated at the radial grid points, + or an array of length ``n_rho`` (arbitrary units, must be >= 0). + energy : openmc.stats.Univariate or Sequence[openmc.stats.Univariate] + Energy distribution(s); see the class docstring. + n_rho : int + Number of radial surfaces to extract (default: 33). + **kwargs + Additional keyword arguments passed to the constructor + (e.g. ``time``, ``strength``, ``constraints``). + + Returns + ------- + openmc.StellaratorSource + + """ + cv.check_greater_than('n_rho', n_rho, 1) + rho = np.linspace(0.0, 1.0, n_rho) + + r_modes, r_lmn, z_modes, z_lmn, nfp = cls._desc_spectral_data(eq) + + # Collapse the Zernike radial dependence onto the rho grid, giving + # product-form double Fourier coefficients for each surface, then + # convert to combined form + surface_tables = [] + radial_r = np.column_stack([ + x * cls._zernike_radial(rho, l, m) + for (l, m, n), x in zip(r_modes, r_lmn)]) + radial_z = np.column_stack([ + x * cls._zernike_radial(rho, l, m) + for (l, m, n), x in zip(z_modes, z_lmn)]) + for i in range(n_rho): + table_r, table_z = {}, {} + cls._desc_to_combined(r_modes[:, 1], r_modes[:, 2], + radial_r[i], table_r) + cls._desc_to_combined(z_modes[:, 1], z_modes[:, 2], + radial_z[i], table_z) + surface_tables.append((table_r, table_z)) + + # Union of modes across R, Z, and all surfaces + modes = sorted({key for table_r, table_z in surface_tables + for key in (*table_r, *table_z)}) + mode_m = np.array([m for m, _ in modes], dtype=int) + mode_n = np.array([n for _, n in modes], dtype=int) + + n_modes = len(modes) + rmnc = np.zeros((n_rho, n_modes)) + rmns = np.zeros((n_rho, n_modes)) + zmnc = np.zeros((n_rho, n_modes)) + zmns = np.zeros((n_rho, n_modes)) + for i, (table_r, table_z) in enumerate(surface_tables): + for k, key in enumerate(modes): + if key in table_r: + rmnc[i, k], rmns[i, k] = table_r[key] + if key in table_z: + zmnc[i, k], zmns[i, k] = table_z[key] + + # m to cm + rmnc *= 100.0 + rmns *= 100.0 + zmnc *= 100.0 + zmns *= 100.0 + + # Drop the asymmetric tables for stellarator-symmetric equilibria + sym = not (np.any(rmns) or np.any(zmnc)) + + return cls( + rho=rho, + emission_density=cls._evaluate_emission_density( + emission_density, rho), + mode_m=mode_m, + mode_n=mode_n, + rmnc=rmnc, + zmns=zmns, + rmns=None if sym else rmns, + zmnc=None if sym else zmnc, + num_field_periods=nfp, + energy=energy, + **kwargs + ) + + def populate_xml_element(self, element): + """Add necessary stellarator source information to an XML element + + Returns + ------- + element : lxml.etree._Element + XML element containing source data + + """ + self._validate() + + ET.SubElement(element, "num_field_periods").text = \ + str(self.num_field_periods) + ET.SubElement(element, "rho").text = \ + ' '.join(str(r) for r in self.rho) + ET.SubElement(element, "emission_density").text = \ + ' '.join(str(s) for s in self.emission_density) + ET.SubElement(element, "mode_m").text = \ + ' '.join(str(m) for m in self.mode_m) + ET.SubElement(element, "mode_n").text = \ + ' '.join(str(n) for n in self.mode_n) + + # Coefficient tables flattened row-major (surface index varies slowest) + for name in ('rmnc', 'zmns', 'rmns', 'zmnc'): + coeff = getattr(self, name) + if coeff is not None: + ET.SubElement(element, name).text = \ + ' '.join(str(c) for c in coeff.ravel()) + + # 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) -> StellaratorSource: + """Generate stellarator source from an XML element + + Parameters + ---------- + elem : lxml.etree._Element + XML element + + Returns + ------- + openmc.StellaratorSource + Source generated from XML element + + """ + nfp_text = get_text(elem, 'num_field_periods') + num_field_periods = int(nfp_text) if nfp_text else 1 + + rho = np.array(get_text(elem, 'rho').split(), dtype=float) + emission_density = np.array( + get_text(elem, 'emission_density').split(), dtype=float) + mode_m = np.array(get_text(elem, 'mode_m').split(), dtype=int) + mode_n = np.array(get_text(elem, 'mode_n').split(), dtype=int) + + shape = (len(rho), len(mode_m)) + coeffs = {} + for name in ('rmnc', 'zmns', 'rmns', 'zmnc'): + text = get_text(elem, name) + coeffs[name] = (np.array(text.split(), dtype=float).reshape(shape) + if text else None) + + # 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( + rho=rho, + emission_density=emission_density, + mode_m=mode_m, + mode_n=mode_n, + rmnc=coeffs['rmnc'], + zmns=coeffs['zmns'], + rmns=coeffs['rmns'], + zmnc=coeffs['zmnc'], + num_field_periods=num_field_periods, + energy=energy, + time=time, + strength=strength, + constraints=constraints + ) + + class SourceParticle: """Source particle diff --git a/src/source.cpp b/src/source.cpp index f951722cc18..531cf63c94c 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -108,6 +108,8 @@ unique_ptr Source::create(pugi::xml_node node) return make_unique(node); } else if (source_type == "tokamak") { return make_unique(node); + } else if (source_type == "stellarator") { + return make_unique(node); } else { fatal_error(fmt::format("Invalid source type '{}' found.", source_type)); } @@ -1167,6 +1169,463 @@ SourceSite TokamakSource::sample(uint64_t* seed) const return site; } +//============================================================================== +// StellaratorSource implementation +//============================================================================== + +namespace { + +//! Maximum over t in [0, 1] of the quadratic through q(0)=f0, q(1/2)=fh, +//! q(1)=f1. Exact, since the sampled density is quadratic in the radial +//! interpolation parameter within a bin. +double quad_max01(double f0, double fh, double f1) +{ + double c = 2.0 * (f0 + f1 - 2.0 * fh); + double b = f1 - f0 - c; + double m = std::max(f0, f1); + if (c < 0.0) { + double t_star = -b / (2.0 * c); + if (t_star > 0.0 && t_star < 1.0) { + m = std::max(m, f0 + t_star * (b + t_star * c)); + } + } + return m; +} + +} // namespace + +StellaratorSource::StellaratorSource(pugi::xml_node node) : Source(node) +{ + // Read number of field periods + if (check_for_node(node, "num_field_periods")) { + num_field_periods_ = std::stoi(get_node_value(node, "num_field_periods")); + } else { + num_field_periods_ = 1; + } + + // Read radial grid and emission profile + rho_ = get_node_array(node, "rho"); + emission_density_ = get_node_array(node, "emission_density"); + + // Read Fourier mode numbers and coefficient tables + mode_m_ = get_node_array(node, "mode_m"); + mode_n_ = get_node_array(node, "mode_n"); + rmnc_ = get_node_array(node, "rmnc"); + zmns_ = get_node_array(node, "zmns"); + if (check_for_node(node, "rmns") || check_for_node(node, "zmnc")) { + if (!check_for_node(node, "rmns") || !check_for_node(node, "zmnc")) { + fatal_error("StellaratorSource: rmns and zmnc must both be given for " + "non-stellarator-symmetric equilibria."); + } + rmns_ = get_node_array(node, "rmns"); + zmnc_ = get_node_array(node, "zmnc"); + asym_ = true; + } + + // 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 + 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 + size_t n_rho = rho_.size(); + size_t n_modes = mode_m_.size(); + if (num_field_periods_ < 1) { + fatal_error("StellaratorSource: num_field_periods must be >= 1."); + } + if (n_rho < 2) { + fatal_error( + "StellaratorSource: At least 2 radial points are required for profiles."); + } + if (rho_.front() != 0.0) { + fatal_error("StellaratorSource: rho must start at 0."); + } + if (rho_.back() != 1.0) { + fatal_error("StellaratorSource: rho must end at 1."); + } + for (size_t i = 1; i < n_rho; ++i) { + if (rho_[i] <= rho_[i - 1]) { + fatal_error("StellaratorSource: rho must be strictly increasing."); + } + } + if (emission_density_.size() != n_rho) { + fatal_error("StellaratorSource: emission_density and rho must have the " + "same length."); + } + for (double s : emission_density_) { + if (s < 0.0) { + fatal_error( + "StellaratorSource: emission_density values cannot be negative."); + } + } + if (n_modes == 0) { + fatal_error("StellaratorSource: At least one Fourier mode is required."); + } + if (mode_n_.size() != n_modes) { + fatal_error( + "StellaratorSource: mode_m and mode_n must have the same length."); + } + for (int m : mode_m_) { + if (m < 0) { + fatal_error("StellaratorSource: poloidal mode numbers must be >= 0 in " + "the combined Fourier form cos/sin(m*theta - n*Nfp*zeta)."); + } + } + if (rmnc_.size() != n_rho * n_modes || zmns_.size() != n_rho * n_modes) { + fatal_error("StellaratorSource: rmnc and zmns must have length " + "len(rho)*len(mode_m)."); + } + if (asym_ && + (rmns_.size() != n_rho * n_modes || zmnc_.size() != n_rho * n_modes)) { + fatal_error("StellaratorSource: rmns and zmnc must have length " + "len(rho)*len(mode_m)."); + } + if (energy_dists_.empty()) { + fatal_error( + "StellaratorSource: At least one energy distribution is required."); + } + if (energy_dists_.size() != 1 && energy_dists_.size() != n_rho) { + fatal_error("StellaratorSource: energy distributions must be either 1 " + "(for all rho) or match the number of rho points."); + } + + // Initialize isotropic angular distribution + angle_ = UPtrAngle {new Isotropic()}; + + precompute_sampling_distributions(); +} + +void StellaratorSource::precompute_sampling_distributions() +{ + size_t n_rho = rho_.size(); + size_t n_modes = mode_m_.size(); + size_t n_bins = n_rho - 1; + + // Angular quadrature grid. The density R*tau is a trigonometric polynomial + // of poloidal degree <= 3*m_max and toroidal degree <= 3*n_max per field + // period, so a uniform trapezoidal product rule with more points than the + // bandwidth integrates it exactly (discrete Fourier orthogonality). The + // denser grid below also controls the discretization error of the rejection + // majorant, which is additionally covered by a safety factor. + int m_max = 0; + int n_max = 0; + for (size_t k = 0; k < n_modes; ++k) { + m_max = std::max(m_max, mode_m_[k]); + n_max = std::max(n_max, std::abs(mode_n_[k])); + } + int n_theta = std::max(64, 8 * m_max); + int n_zeta = std::max(64, 8 * n_max); + size_t n_grid = static_cast(n_theta) * n_zeta; + + // Per-mode trigonometric tables on the angle grid (zeta spans one field + // period; the geometry is periodic with period 2*pi/Nfp) + double dtheta = 2.0 * PI / n_theta; + double dzeta = 2.0 * PI / (num_field_periods_ * n_zeta); + vector cos_mt(n_modes * n_theta), sin_mt(n_modes * n_theta); + vector cos_nz(n_modes * n_zeta), sin_nz(n_modes * n_zeta); + for (size_t k = 0; k < n_modes; ++k) { + for (int i = 0; i < n_theta; ++i) { + double a = mode_m_[k] * i * dtheta; + cos_mt[k * n_theta + i] = std::cos(a); + sin_mt[k * n_theta + i] = std::sin(a); + } + for (int j = 0; j < n_zeta; ++j) { + double a = mode_n_[k] * num_field_periods_ * j * dzeta; + cos_nz[k * n_zeta + j] = std::cos(a); + sin_nz[k * n_zeta + j] = std::sin(a); + } + } + + // Evaluate R, Z, dR/dtheta, dZ/dtheta on the angle grid for every radial + // surface. Radial derivatives within a bin follow from finite differences + // of these grids because the Fourier evaluation is linear in the + // coefficients. + vector R_s(n_rho * n_grid), Z_s(n_rho * n_grid); + vector Rt_s(n_rho * n_grid), Zt_s(n_rho * n_grid); + for (size_t i = 0; i < n_rho; ++i) { + const double* rc = &rmnc_[i * n_modes]; + const double* zs = &zmns_[i * n_modes]; + const double* rs = asym_ ? &rmns_[i * n_modes] : nullptr; + const double* zc = asym_ ? &zmnc_[i * n_modes] : nullptr; + for (int it = 0; it < n_theta; ++it) { + for (int jz = 0; jz < n_zeta; ++jz) { + double R = 0.0, Z = 0.0, Rt = 0.0, Zt = 0.0; + for (size_t k = 0; k < n_modes; ++k) { + // cos/sin(m*theta - n*Nfp*zeta) via angle-addition from tables + double cm = cos_mt[k * n_theta + it]; + double sm = sin_mt[k * n_theta + it]; + double cn = cos_nz[k * n_zeta + jz]; + double sn = sin_nz[k * n_zeta + jz]; + double c = cm * cn + sm * sn; + double s = sm * cn - cm * sn; + double m = mode_m_[k]; + R += rc[k] * c; + Rt -= m * rc[k] * s; + Z += zs[k] * s; + Zt += m * zs[k] * c; + if (asym_) { + R += rs[k] * s; + Rt += m * rs[k] * c; + Z += zc[k] * c; + Zt -= m * zc[k] * s; + } + } + size_t p = i * n_grid + it * n_zeta + jz; + R_s[p] = R; + Z_s[p] = Z; + Rt_s[p] = Rt; + Zt_s[p] = Zt; + } + } + } + + // Sanity check: the map must stay on the +R side of the cylindrical axis + double R_min = *std::min_element(R_s.begin(), R_s.end()); + if (R_min <= 0.0) { + fatal_error("StellaratorSource: flux surfaces reach R <= 0; check the " + "Fourier coefficients (units should be cm)."); + } + + // Pass 1: signed volume integral to fix the global Jacobian sign. Both + // (rho, theta, zeta) handednesses occur in practice (VMEC's Jacobian is + // conventionally negative), so the sign is detected rather than assumed. + double quad_wgt = 4.0 * PI * PI / n_grid; // full-torus angular weight + double v_tot = 0.0; + for (size_t b = 0; b < n_bins; ++b) { + double inv_drho = 1.0 / (rho_[b + 1] - rho_[b]); + const double* R0 = &R_s[b * n_grid]; + const double* R1 = &R_s[(b + 1) * n_grid]; + const double* Z0 = &Z_s[b * n_grid]; + const double* Z1 = &Z_s[(b + 1) * n_grid]; + const double* Rt0 = &Rt_s[b * n_grid]; + const double* Rt1 = &Rt_s[(b + 1) * n_grid]; + const double* Zt0 = &Zt_s[b * n_grid]; + const double* Zt1 = &Zt_s[(b + 1) * n_grid]; + double sum = 0.0; + for (size_t p = 0; p < n_grid; ++p) { + double Rr = (R1[p] - R0[p]) * inv_drho; + double Zr = (Z1[p] - Z0[p]) * inv_drho; + // Midpoint of the bin (all factors are linear in t) + double Rh = 0.5 * (R0[p] + R1[p]); + double Rth = 0.5 * (Rt0[p] + Rt1[p]); + double Zth = 0.5 * (Zt0[p] + Zt1[p]); + sum += Rh * (Rr * Zth - Rth * Zr); + } + v_tot += sum * quad_wgt * (rho_[b + 1] - rho_[b]); + } + jacobian_sign_ = (v_tot >= 0.0) ? 1.0 : -1.0; + + // Pass 2: differential volume V'(rho) at surfaces and bin midpoints, + // rejection majorants, and Jacobian sign-consistency check + vector vp_lo(n_bins), vp_mid(n_bins), vp_hi(n_bins); + envelope_.assign(n_bins, 0.0); + double f_peak = 0.0; + double f_min = 0.0; + for (size_t b = 0; b < n_bins; ++b) { + double inv_drho = 1.0 / (rho_[b + 1] - rho_[b]); + const double* R0 = &R_s[b * n_grid]; + const double* R1 = &R_s[(b + 1) * n_grid]; + const double* Z0 = &Z_s[b * n_grid]; + const double* Z1 = &Z_s[(b + 1) * n_grid]; + const double* Rt0 = &Rt_s[b * n_grid]; + const double* Rt1 = &Rt_s[(b + 1) * n_grid]; + const double* Zt0 = &Zt_s[b * n_grid]; + const double* Zt1 = &Zt_s[(b + 1) * n_grid]; + double s0 = 0.0, sh = 0.0, s1 = 0.0; + double env = 0.0; + for (size_t p = 0; p < n_grid; ++p) { + double Rr = (R1[p] - R0[p]) * inv_drho; + double Zr = (Z1[p] - Z0[p]) * inv_drho; + double f0 = jacobian_sign_ * R0[p] * (Rr * Zt0[p] - Rt0[p] * Zr); + double f1 = jacobian_sign_ * R1[p] * (Rr * Zt1[p] - Rt1[p] * Zr); + double Rh = 0.5 * (R0[p] + R1[p]); + double Rth = 0.5 * (Rt0[p] + Rt1[p]); + double Zth = 0.5 * (Zt0[p] + Zt1[p]); + double fh = jacobian_sign_ * Rh * (Rr * Zth - Rth * Zr); + s0 += f0; + sh += fh; + s1 += f1; + env = std::max(env, quad_max01(f0, fh, f1)); + f_peak = std::max({f_peak, f0, fh, f1}); + f_min = std::min({f_min, f0, fh, f1}); + } + vp_lo[b] = s0 * quad_wgt; + vp_mid[b] = sh * quad_wgt; + vp_hi[b] = s1 * quad_wgt; + // Safety factor covering angular discretization of the majorant + envelope_[b] = 1.05 * env; + } + if (f_min < -1.0e-6 * f_peak) { + fatal_error( + "StellaratorSource: the flux-coordinate Jacobian changes sign within " + "the plasma volume; the flux surfaces self-intersect or the Fourier " + "coefficient tables are inconsistent."); + } + + // Build the marginal radial PDF p(rho) ~ S(rho) * V'(rho) on a refined grid + // containing all surfaces and bin midpoints. At interior surfaces the + // one-sided values from the two adjacent bins are averaged (the piecewise- + // linear coefficient interpolation makes V' one-sided there). + vector grid(2 * n_rho - 1), pdf(2 * n_rho - 1); + for (size_t i = 0; i < n_rho; ++i) { + grid[2 * i] = rho_[i]; + double vp; + if (i == 0) { + vp = vp_lo[0]; + } else if (i == n_rho - 1) { + vp = vp_hi[n_bins - 1]; + } else { + vp = 0.5 * (vp_hi[i - 1] + vp_lo[i]); + } + pdf[2 * i] = emission_density_[i] * std::max(0.0, vp); + } + for (size_t b = 0; b < n_bins; ++b) { + grid[2 * b + 1] = 0.5 * (rho_[b] + rho_[b + 1]); + double s_mid = 0.5 * (emission_density_[b] + emission_density_[b + 1]); + pdf[2 * b + 1] = s_mid * std::max(0.0, vp_mid[b]); + } + + double total = 0.0; + for (size_t i = 1; i < grid.size(); ++i) { + total += 0.5 * (pdf[i - 1] + pdf[i]) * (grid[i] - grid[i - 1]); + } + if (total <= 0.0) { + fatal_error( + "StellaratorSource: Integrated emission density is zero or negative. " + "Check emission_density profile."); + } + radial_dist_ = make_unique( + grid.data(), pdf.data(), grid.size(), Interpolation::lin_lin); +} + +double StellaratorSource::eval_density(int bin, double t, double theta, + double zeta, double* R_out, double* Z_out) const +{ + size_t n_modes = mode_m_.size(); + const double* rc0 = &rmnc_[bin * n_modes]; + const double* rc1 = &rmnc_[(bin + 1) * n_modes]; + const double* zs0 = &zmns_[bin * n_modes]; + const double* zs1 = &zmns_[(bin + 1) * n_modes]; + const double* rs0 = asym_ ? &rmns_[bin * n_modes] : nullptr; + const double* rs1 = asym_ ? &rmns_[(bin + 1) * n_modes] : nullptr; + const double* zc0 = asym_ ? &zmnc_[bin * n_modes] : nullptr; + const double* zc1 = asym_ ? &zmnc_[(bin + 1) * n_modes] : nullptr; + double inv_drho = 1.0 / (rho_[bin + 1] - rho_[bin]); + + double R = 0.0, Rr = 0.0, Rt = 0.0; + double Z = 0.0, Zr = 0.0, Zt = 0.0; + for (size_t k = 0; k < n_modes; ++k) { + double a = mode_m_[k] * theta - mode_n_[k] * num_field_periods_ * zeta; + double c = std::cos(a); + double s = std::sin(a); + double m = mode_m_[k]; + + double drc = rc1[k] - rc0[k]; + double rc = rc0[k] + t * drc; + double dzs = zs1[k] - zs0[k]; + double zs = zs0[k] + t * dzs; + R += rc * c; + Rr += drc * inv_drho * c; + Rt -= m * rc * s; + Z += zs * s; + Zr += dzs * inv_drho * s; + Zt += m * zs * c; + if (asym_) { + double drs = rs1[k] - rs0[k]; + double rs = rs0[k] + t * drs; + double dzc = zc1[k] - zc0[k]; + double zc = zc0[k] + t * dzc; + R += rs * s; + Rr += drs * inv_drho * s; + Rt += m * rs * c; + Z += zc * c; + Zr += dzc * inv_drho * c; + Zt -= m * zc * s; + } + } + if (R_out) + *R_out = R; + if (Z_out) + *Z_out = Z; + return jacobian_sign_ * R * (Rr * Zt - Rt * Zr); +} + +std::pair StellaratorSource::sample_energy( + double rho, uint64_t* seed) const +{ + if (energy_dists_.size() == 1) { + // Single distribution for all rho + return energy_dists_[0]->sample(seed); + } + + // Multiple distributions: stochastic selection between bracketing points + size_t i = lower_bound_index(rho_.begin(), rho_.end(), rho); + if (i >= energy_dists_.size() - 1) { + return energy_dists_.back()->sample(seed); + } + double t = (rho - rho_[i]) / (rho_[i + 1] - rho_[i]); + size_t idx = (prn(seed) < t) ? i + 1 : i; + return energy_dists_[idx]->sample(seed); +} + +SourceSite StellaratorSource::sample(uint64_t* seed) const +{ + SourceSite site; + site.particle = ParticleType::neutron(); + site.wgt = 1.0; + site.delayed_group = 0; + + // 1. Sample rho from the marginal radial CDF + double rho = radial_dist_->sample(seed).first; + int n_bins = static_cast(rho_.size()) - 1; + int bin = static_cast(lower_bound_index(rho_.begin(), rho_.end(), rho)); + bin = std::min(std::max(bin, 0), n_bins - 1); + double t = (rho - rho_[bin]) / (rho_[bin + 1] - rho_[bin]); + + // 2. Rejection-sample (theta, zeta) from p(theta, zeta | rho) ~ R*tau + double env = envelope_[bin]; + double R, Z, zeta; + int64_t n_reject = 0; + while (true) { + double theta = 2.0 * PI * prn(seed); + zeta = 2.0 * PI * prn(seed); + double f = eval_density(bin, t, theta, zeta, &R, &Z); + if (prn(seed) * env < f) + break; + if (++n_reject > MAX_SOURCE_REJECTIONS_PER_SAMPLE) { + fatal_error("StellaratorSource: exceeded the maximum number of " + "rejections while sampling the poloidal/toroidal angles."); + } + } + + // 3. Convert to Cartesian coordinates (zeta is the cylindrical angle) + site.r = {R * std::cos(zeta), R * std::sin(zeta), Z}; + + // 4. Sample isotropic direction + site.u = angle_->sample(seed).first; + + // 5. Sample energy from distribution(s), applying the importance weight + auto [E, E_wgt] = sample_energy(rho, seed); + site.E = E; + + // 6. 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 //============================================================================== From cf616f4322864455ba46e46e4da04a87bd8d94b0 Mon Sep 17 00:00:00 2001 From: Ethan Peterson Date: Sat, 29 Aug 2026 10:54:29 -0400 Subject: [PATCH 02/33] documentation edits --- docs/source/io_formats/settings.rst | 81 ++++++++++++++++++++++++++++- docs/source/pythonapi/base.rst | 1 + docs/source/usersguide/settings.rst | 44 ++++++++++++++++ 3 files changed, 124 insertions(+), 2 deletions(-) diff --git a/docs/source/io_formats/settings.rst b/docs/source/io_formats/settings.rst index 7091757f23b..fdc03f6291f 100644 --- a/docs/source/io_formats/settings.rst +++ b/docs/source/io_formats/settings.rst @@ -778,8 +778,8 @@ attributes/sub-elements: :type: 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. + ``mesh``, ``tokamak``, or ``stellarator``. 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., @@ -1089,6 +1089,83 @@ attributes/sub-elements: *Default*: particles are born at :math:`t=0` + For a source with ``type="stellarator"``, the spatial distribution is + described by the flux-surface Fourier representation shared by the VMEC and + DESC equilibrium codes, using flux coordinates :math:`(\rho, \theta, \zeta)` + where :math:`\rho = \sqrt{s}` is the square root of the normalized toroidal + flux, :math:`\theta` is the poloidal angle, and :math:`\zeta` is the toroidal + (cylindrical) angle: + + .. math:: + + \begin{aligned} + R &= \sum_k \left[ R^c_k(\rho) \cos(m_k\theta - n_k N_{fp} \zeta) + + R^s_k(\rho) \sin(m_k\theta - n_k N_{fp}\zeta) \right] \\ + Z &= \sum_k \left[ Z^s_k(\rho) \sin(m_k\theta - n_k N_{fp}\zeta) + + Z^c_k(\rho) \cos(m_k\theta - n_k N_{fp}\zeta) \right] + \end{aligned} + + The following sub-elements are used instead of the ``space`` element: + + :num_field_periods: + The number of field periods :math:`N_{fp}`. + + *Default*: 1 + + :rho: + A list of radial grid points :math:`\rho = \sqrt{s}`. Must be strictly + increasing, start at 0, and end at 1. + + :emission_density: + A list of neutron emission densities :math:`S(\rho)` evaluated at each + ``rho`` grid point (arbitrary units, must be non-negative). Only the shape + matters, since the profile is normalized internally. Must have the same + length as ``rho`` and contain at least one positive value. + + :mode_m: + A list of poloidal mode numbers :math:`m_k` (must be non-negative). + + :mode_n: + A list of toroidal mode numbers :math:`n_k` in units of the number of + field periods (i.e., VMEC's ``xn`` divided by ``nfp``). Must have the same + length as ``mode_m``. + + :rmnc: + The cosine Fourier coefficients of :math:`R` in [cm], flattened row-major + with the surface index varying slowest (length + ``len(rho)*len(mode_m)``). + + :zmns: + The sine Fourier coefficients of :math:`Z` in [cm], in the same layout as + ``rmnc``. + + :rmns: + The sine Fourier coefficients of :math:`R` in [cm] for + non-stellarator-symmetric equilibria. If given, ``zmnc`` must also be + given. + + *Default*: None + + :zmnc: + The cosine Fourier coefficients of :math:`Z` in [cm] for + non-stellarator-symmetric equilibria. If given, ``rmns`` must also be + given. + + *Default*: None + + :energy: + For a stellarator source, one or more ``energy`` sub-elements specify the + neutron energy distribution(s) with the same semantics as for the tokamak + source: either a single distribution used at all radii, or exactly one + distribution per ``rho`` grid point. + + :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 3ffcac69cda..39283267309 100644 --- a/docs/source/pythonapi/base.rst +++ b/docs/source/pythonapi/base.rst @@ -28,6 +28,7 @@ Simulation Settings openmc.CompiledSource openmc.MeshSource openmc.TokamakSource + openmc.StellaratorSource openmc.SourceParticle openmc.VolumeCalculation openmc.Settings diff --git a/docs/source/usersguide/settings.rst b/docs/source/usersguide/settings.rst index 039a8866265..fddae1d25b5 100644 --- a/docs/source/usersguide/settings.rst +++ b/docs/source/usersguide/settings.rst @@ -338,6 +338,50 @@ 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. +Stellarator Plasma Sources +-------------------------- + +For fully three-dimensional plasmas, the :class:`openmc.StellaratorSource` +class samples neutron positions from the flux-surface Fourier representation +shared by the `VMEC `_ and `DESC +`_ equilibrium codes. The plasma is +described in flux coordinates :math:`(\rho, \theta, \zeta)`, where +:math:`\rho = \sqrt{s}` is the square root of the normalized toroidal flux +(proportional to the average minor radius), :math:`\theta` is the poloidal +angle, and :math:`\zeta` is the toroidal angle (identical to the cylindrical +azimuthal angle in both codes). The user provides the Fourier coefficients of +:math:`R(\rho,\theta,\zeta)` and :math:`Z(\rho,\theta,\zeta)` on a radial grid +together with an emission density :math:`S(\rho)` that is constant on flux +surfaces. The radial coordinate is sampled by inverting a tabulated CDF of the +exact marginal distribution :math:`S(\rho)\, V'(\rho)` and the two angles are +sampled from the conditional distribution +:math:`p(\theta,\zeta|\rho) \propto R\,|\tau|` (with :math:`\tau` the +poloidal-plane Jacobian) by rejection. + +Sources are most conveniently created directly from equilibrium code output:: + + # From a VMEC wout NetCDF file + source = openmc.StellaratorSource.from_vmec( + 'wout_w7x.nc', + emission_density=lambda rho: (1.0 - rho**2)**2, + energy=openmc.stats.muir(e0=14.08e6, m_rat=5.0, kt=20000.0), + ) + + # From a DESC equilibrium object or HDF5 output file + source = openmc.StellaratorSource.from_desc( + eq, # a desc.equilibrium.Equilibrium or path to a DESC HDF5 file + emission_density=lambda rho: (1.0 - rho**2)**2, + energy=openmc.stats.muir(e0=14.08e6, m_rat=5.0, kt=20000.0), + ) + +The ``emission_density`` argument accepts either a callable evaluated on the +radial grid or an array with one value per radial surface; only its shape +matters, since it is normalized internally. The Fourier coefficient tables can +also be passed directly to the constructor (in centimeters) for equilibria from +other sources. As with :class:`openmc.TokamakSource`, the ``energy`` argument +accepts a single distribution or one distribution per radial grid point, and a +``time`` distribution may optionally be given. + File-based Sources ------------------ From 78d58b1f7b20d4467b7416a13bf16958b9b56633 Mon Sep 17 00:00:00 2001 From: Ethan Peterson Date: Thu, 10 Sep 2026 18:21:30 -0400 Subject: [PATCH 03/33] adding tests and code review --- docs/source/io_formats/settings.rst | 14 +- docs/source/usersguide/settings.rst | 7 +- openmc/source.py | 8 +- tests/unit_tests/test_source_stellarator.py | 344 ++++++++++++++++++++ 4 files changed, 348 insertions(+), 25 deletions(-) create mode 100644 tests/unit_tests/test_source_stellarator.py diff --git a/docs/source/io_formats/settings.rst b/docs/source/io_formats/settings.rst index e38f289251e..fdc03f6291f 100644 --- a/docs/source/io_formats/settings.rst +++ b/docs/source/io_formats/settings.rst @@ -626,11 +626,7 @@ found in the :ref:`random ray user guide `. :volume_estimator: Specifies choice of volume estimator for the random ray solver. Options - are 'naive', 'simulation_averaged', 'hybrid', 'adaptive', - 'strict_adaptive', or 'auto'. The default is 'auto', which selects - 'adaptive' for standard solves and 'strict_adaptive' for solves whose - results feed variance reduction (weight window generation and adjoint - workflows). + are 'naive', 'simulation_averaged', or 'hybrid'. The default is 'hybrid'. *Default*: None @@ -1406,13 +1402,6 @@ attributes/sub-elements: *Default*: ``surface_source.h5`` in current working directory -.. deprecated:: - The ```` element is deprecated and will be removed in a - future release. A deprecation warning is emitted when it is present. Use a - file source instead, e.g. ````, - which is equivalent but additionally supports a source strength and source - constraints. - ------------------------------- ```` Element ------------------------------- @@ -1736,7 +1725,6 @@ and 10. The verbosity levels are defined as follows: :5: all of the above + file I/O :6: all of the above + timing statistics and initialization messages :7: all of the above + :math:`k` by generation - :8: all of the above + random ray volume-estimator diagnostics :9: all of the above + indicate when each particle starts :10: all of the above + event information diff --git a/docs/source/usersguide/settings.rst b/docs/source/usersguide/settings.rst index 42232ee7c91..fddae1d25b5 100644 --- a/docs/source/usersguide/settings.rst +++ b/docs/source/usersguide/settings.rst @@ -466,11 +466,6 @@ As an example, to write a maximum of three surface source files::: 'max_source_files': 3 } -A surface source file is read back like any other source file, with the -:class:`openmc.FileSource` class:: - - settings.source = openmc.FileSource('surface_source.h5') - .. _compiled_source: Compiled Sources @@ -536,7 +531,7 @@ OpenMC shared library. This can be done by writing a CMakeLists.txt file: cmake_minimum_required(VERSION 3.3 FATAL_ERROR) project(openmc_sources CXX) - add_library(source MODULE source_ring.cpp) + add_library(source SHARED source_ring.cpp) find_package(OpenMC REQUIRED HINTS ) target_link_libraries(source OpenMC::libopenmc) diff --git a/openmc/source.py b/openmc/source.py index 5c05f2aae3f..d5e6122d880 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -1984,8 +1984,7 @@ def from_xml_element(cls, elem: ET.Element) -> StellaratorSource: Source generated from XML element """ - nfp_text = get_text(elem, 'num_field_periods') - num_field_periods = int(nfp_text) if nfp_text else 1 + num_field_periods = int(get_text(elem, 'num_field_periods', 1)) rho = np.array(get_text(elem, 'rho').split(), dtype=float) emission_density = np.array( @@ -2002,8 +2001,6 @@ def from_xml_element(cls, elem: ET.Element) -> StellaratorSource: # 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') @@ -2011,8 +2008,7 @@ def from_xml_element(cls, elem: ET.Element) -> StellaratorSource: # 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 + strength = float(get_text(elem, 'strength', 1.0)) return cls( rho=rho, diff --git a/tests/unit_tests/test_source_stellarator.py b/tests/unit_tests/test_source_stellarator.py new file mode 100644 index 00000000000..023dc5c6a24 --- /dev/null +++ b/tests/unit_tests/test_source_stellarator.py @@ -0,0 +1,344 @@ +import numpy as np +import pytest + +import openmc +import openmc.stats + +from tests.unit_tests import assert_sample_mean + + +def circular_tokamak_modes(r0=600.0, a=150.0, n_rho=17): + """Axisymmetric circular torus expressed in stellarator Fourier form.""" + rho = np.linspace(0.0, 1.0, n_rho) + mode_m = [0, 1] + mode_n = [0, 0] + rmnc = np.column_stack([np.full(n_rho, r0), a * rho]) + zmns = np.column_stack([np.zeros(n_rho), a * rho]) + return rho, mode_m, mode_n, rmnc, zmns + + +def rotating_ellipse_modes(r0=600.0, a=150.0, e=50.0, n_rho=17, nfp=5): + """Classic rotating-ellipse stellarator: + + R = R0 + a*rho*cos(theta) + e*rho*cos(theta - nfp*zeta) + Z = a*rho*sin(theta) - e*rho*sin(theta - nfp*zeta) + """ + rho = np.linspace(0.0, 1.0, n_rho) + mode_m = [0, 1, 1] + mode_n = [0, 0, 1] + rmnc = np.column_stack([np.full(n_rho, r0), a * rho, e * rho]) + zmns = np.column_stack([np.zeros(n_rho), a * rho, -e * rho]) + return rho, mode_m, mode_n, rmnc, zmns + + +def make_source(**kwargs): + """Build a valid StellaratorSource, overriding defaults via kwargs.""" + rho, mode_m, mode_n, rmnc, zmns = rotating_ellipse_modes() + params = dict( + rho=rho, + emission_density=(1.0 - rho**2), + mode_m=mode_m, + mode_n=mode_n, + rmnc=rmnc, + zmns=zmns, + num_field_periods=5, + energy=openmc.stats.muir(e0=14.08e6, m_rat=5.0, kt=2.0e4), + ) + params.update(kwargs) + return openmc.StellaratorSource(**params) + + +def make_model(src): + sphere = openmc.Sphere(r=2000.0, boundary_type='vacuum') + return openmc.Model( + geometry=openmc.Geometry([openmc.Cell(region=-sphere)]), + settings=openmc.Settings( + particles=100, batches=1, run_mode='fixed source', source=src), + ) + + +def reference_moments(rho, mode_m, mode_n, rmnc, zmns, nfp, emission_density, + n_theta=128, n_zeta=128): + """Compute moments of the exact density S(rho)*R*|tau| by quadrature. + + The Fourier coefficients here are linear in rho, so the piecewise-linear + coefficient interpolation used by the implementation is exact and the + continuous density can be used as a reference. + """ + rq = np.linspace(0.0, 1.0, 201)[1:] # exclude axis (zero measure) + theta = np.linspace(0.0, 2 * np.pi, n_theta, endpoint=False) + zeta = np.linspace(0.0, 2 * np.pi / nfp, n_zeta, endpoint=False) + rg, tg, zg = np.meshgrid(rq, theta, zeta, indexing='ij') + + # Coefficients are c_k(rho) = slope_k * rho + const_k + slopes_r = (rmnc[-1] - rmnc[0]) # over rho in [0, 1] + consts_r = rmnc[0] + slopes_z = (zmns[-1] - zmns[0]) + consts_z = zmns[0] + + R = np.zeros_like(rg) + Rr = np.zeros_like(rg) + Rt = np.zeros_like(rg) + Z = np.zeros_like(rg) + Zr = np.zeros_like(rg) + Zt = np.zeros_like(rg) + for k, (m, n) in enumerate(zip(mode_m, mode_n)): + ang = m * tg - n * nfp * zg + c, s = np.cos(ang), np.sin(ang) + rc = consts_r[k] + slopes_r[k] * rg + zs = consts_z[k] + slopes_z[k] * rg + R += rc * c + Rr += slopes_r[k] * c + Rt += -m * rc * s + Z += zs * s + Zr += slopes_z[k] * s + Zt += m * zs * c + + S = np.interp(rg, rho, emission_density) + f = S * np.abs(R * (Rr * Zt - Rt * Zr)) + w = f / f.sum() + return { + 'rho': (w * rg).sum(), + 'R': (w * R).sum(), + 'Z': (w * Z).sum(), + 'Z2': (w * Z**2).sum(), + } + + +def test_stellarator_source_roundtrip(): + src = make_source(strength=2.0, time=openmc.stats.Uniform(0.0, 1e-6)) + + elem = src.to_xml_element() + assert elem.get('type') == 'stellarator' + + new = openmc.SourceBase.from_xml_element(elem) + assert isinstance(new, openmc.StellaratorSource) + assert new.num_field_periods == src.num_field_periods + assert new.strength == src.strength + np.testing.assert_allclose(new.rho, src.rho) + np.testing.assert_allclose(new.emission_density, src.emission_density) + np.testing.assert_array_equal(new.mode_m, src.mode_m) + np.testing.assert_array_equal(new.mode_n, src.mode_n) + np.testing.assert_allclose(new.rmnc, src.rmnc) + np.testing.assert_allclose(new.zmns, src.zmns) + assert new.rmns is None + assert new.zmnc is None + 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_stellarator_source_asymmetric_roundtrip(): + rho, mode_m, mode_n, rmnc, zmns = rotating_ellipse_modes() + src = make_source(rmns=0.01 * rmnc, zmnc=0.01 * zmns) + new = openmc.SourceBase.from_xml_element(src.to_xml_element()) + np.testing.assert_allclose(new.rmns, src.rmns) + np.testing.assert_allclose(new.zmnc, src.zmnc) + + +def test_stellarator_source_multiple_energies(): + rho, mode_m, mode_n, rmnc, zmns = rotating_ellipse_modes(n_rho=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(rho=rho, emission_density=np.ones_like(rho), + rmnc=rmnc, zmns=zmns, energy=energies) + assert len(src.energy) == len(rho) + assert src.time is None + + new = openmc.SourceBase.from_xml_element(src.to_xml_element()) + assert len(new.energy) == len(rho) + assert new.time is None + + +@pytest.mark.parametrize("kwargs, match", [ + (dict(rho=np.linspace(0.1, 1.0, 17)), "must start at 0"), + (dict(rho=np.linspace(0.0, 0.9, 17)), "must end at 1"), + (dict(emission_density=np.ones(5)), "same length as rho"), + (dict(emission_density=np.zeros(17)), "must contain a positive value"), + (dict(mode_m=[0, 1]), "same length as mode_m"), + (dict(mode_m=[0, -1, 1]), "must be >= 0"), + (dict(rmnc=np.ones((17, 2))), "must have shape"), + (dict(rmns=np.ones((17, 3))), "must both be given"), + (dict(energy=[openmc.stats.muir(14.08e6, 5.0, 2.0e4)] * 2), + "Number of energy distributions"), + (dict(num_field_periods=0), "num_field_periods"), +]) +def test_stellarator_source_invalid(kwargs, match): + with pytest.raises(ValueError, match=match): + make_source(**kwargs) + + +def test_stellarator_source_axisymmetric_sampling(run_in_tmpdir): + """Circular axisymmetric plasma has analytic moments. + + With S = 1 the joint density ~ (R0 + a*rho*cos(theta)) * a^2 * rho gives + E[rho] = 2/3 and E[R] = R0 + a^2/(4*R0). + """ + r0, a = 600.0, 150.0 + rho, mode_m, mode_n, rmnc, zmns = circular_tokamak_modes(r0, a) + src = make_source( + rho=rho, emission_density=np.ones_like(rho), mode_m=mode_m, + mode_n=mode_n, rmnc=rmnc, zmns=zmns, num_field_periods=1, + energy=openmc.stats.delta_function(14.07e6), + ) + model = make_model(src) + + sites = model.sample_external_source(30_000) + xyz = np.array([site.r for site in sites]) + major_r = np.hypot(xyz[:, 0], xyz[:, 1]) + minor_rho = np.hypot(major_r - r0, xyz[:, 2]) / a + + assert minor_rho.max() < 1.0 + 1e-9 + assert_sample_mean(minor_rho, 2.0 / 3.0) + assert_sample_mean(major_r, r0 + a**2 / (4.0 * r0)) + + +def test_stellarator_source_rotating_ellipse_sampling(run_in_tmpdir): + """Compare sampled moments of a 3-D configuration against quadrature of + the exact density S(rho) * R * |tau|.""" + r0, a, e, nfp = 600.0, 150.0, 50.0, 5 + rho, mode_m, mode_n, rmnc, zmns = rotating_ellipse_modes( + r0, a, e, nfp=nfp) + emission_density = 1.0 - rho**2 + src = make_source( + rho=rho, emission_density=emission_density, mode_m=mode_m, + mode_n=mode_n, rmnc=rmnc, zmns=zmns, num_field_periods=nfp, + energy=openmc.stats.delta_function(14.07e6), + ) + model = make_model(src) + + sites = model.sample_external_source(30_000) + xyz = np.array([site.r for site in sites]) + major_r = np.hypot(xyz[:, 0], xyz[:, 1]) + z = xyz[:, 2] + + ref = reference_moments( + rho, mode_m, mode_n, rmnc, zmns, nfp, emission_density) + assert_sample_mean(major_r, ref['R']) + assert_sample_mean(z, ref['Z']) + assert_sample_mean(z**2, ref['Z2']) + + +def test_stellarator_source_from_vmec(run_in_tmpdir): + """Round-trip a synthetic rotating-ellipse wout file through from_vmec.""" + from scipy.io import netcdf_file + + ns, nfp = 16, 5 + r0_m, a_m, e_m = 6.0, 1.5, 0.5 # VMEC files are in meters + s = np.linspace(0.0, 1.0, ns) + rho = np.sqrt(s) + xm = np.array([0, 1, 1]) + xn = np.array([0, 0, nfp]) # xn includes the nfp factor + rmnc = np.column_stack([np.full(ns, r0_m), a_m * rho, e_m * rho]) + zmns = np.column_stack([np.zeros(ns), a_m * rho, -e_m * rho]) + + with netcdf_file('wout_test.nc', 'w') as ds: + ds.createDimension('radius', ns) + ds.createDimension('mn_mode', len(xm)) + ds.createVariable('rmnc', 'd', ('radius', 'mn_mode'))[:] = rmnc + ds.createVariable('zmns', 'd', ('radius', 'mn_mode'))[:] = zmns + ds.createVariable('xm', 'd', ('mn_mode',))[:] = xm + ds.createVariable('xn', 'd', ('mn_mode',))[:] = xn + ds.createVariable('nfp', 'i', ()).data.fill(nfp) + ds.createVariable('lasym__logical__', 'i', ()).data.fill(0) + + src = openmc.StellaratorSource.from_vmec( + 'wout_test.nc', + emission_density=lambda r: 1.0 - r**2, + energy=openmc.stats.delta_function(14.07e6), + ) + assert src.num_field_periods == nfp + np.testing.assert_allclose(src.rho, rho) + np.testing.assert_array_equal(src.mode_m, xm) + np.testing.assert_array_equal(src.mode_n, [0, 0, 1]) + np.testing.assert_allclose(src.rmnc, 100.0 * rmnc) # m to cm + np.testing.assert_allclose(src.zmns, 100.0 * zmns) + np.testing.assert_allclose(src.emission_density, 1.0 - rho**2) + assert src.rmns is None + + +def test_stellarator_source_desc_ptolemy(): + """Check the DESC product-form to combined-form conversion identities.""" + rng = np.random.default_rng(42) + theta = rng.uniform(0.0, 2 * np.pi, 50) + zp = rng.uniform(0.0, 2 * np.pi, 50) # zeta' = nfp * zeta + + # One mode of each sign class with random coefficients + modes_m = np.array([2, -2, 2, -2, 0, 3]) + modes_n = np.array([1, 1, -1, -1, -2, 0]) + coeffs = rng.normal(size=modes_m.size) + + # Direct evaluation in DESC's product-form convention + direct = np.zeros_like(theta) + for m0, n0, x in zip(modes_m, modes_n, coeffs): + pol = np.cos(abs(m0) * theta) if m0 >= 0 else np.sin(abs(m0) * theta) + tor = np.cos(abs(n0) * zp) if n0 >= 0 else np.sin(abs(n0) * zp) + direct += x * pol * tor + + # Combined-form evaluation + table = {} + openmc.StellaratorSource._desc_to_combined(modes_m, modes_n, coeffs, table) + combined = np.zeros_like(theta) + for (m, n), (c, s) in table.items(): + ang = m * theta - n * zp + combined += c * np.cos(ang) + s * np.sin(ang) + + np.testing.assert_allclose(combined, direct, atol=1e-12) + + +def test_stellarator_source_zernike_radial(): + """Check the Zernike radial polynomial against known closed forms.""" + zr = openmc.StellaratorSource._zernike_radial + rho = np.linspace(0.0, 1.0, 11) + np.testing.assert_allclose(zr(rho, 0, 0), np.ones_like(rho)) + np.testing.assert_allclose(zr(rho, 1, 1), rho) + np.testing.assert_allclose(zr(rho, 2, 0), 2 * rho**2 - 1) + np.testing.assert_allclose(zr(rho, 2, 2), rho**2) + np.testing.assert_allclose(zr(rho, 3, 1), 3 * rho**3 - 2 * rho) + np.testing.assert_allclose(zr(rho, 4, 0), 6 * rho**4 - 6 * rho**2 + 1) + np.testing.assert_allclose(zr(rho, 3, -1), zr(rho, 3, 1)) # |m| is used + + +def test_stellarator_source_from_desc(run_in_tmpdir): + """Read a synthetic rotating-ellipse DESC HDF5 file with from_desc. + + The equilibrium R = R0 + a*rho*cos(theta - Nfp*zeta') expands in DESC's + product-form Fourier-Zernike basis (Zernike R_1^1(rho) = rho) as modes + (l=1, m=1, n=1) and (l=1, m=-1, n=-1), and likewise for Z. + """ + import h5py + + nfp = 5 + r0_m, a_m = 6.0, 1.5 + r_modes = np.array([[0, 0, 0], [1, 1, 1], [1, -1, -1]]) + r_lmn = np.array([r0_m, a_m, a_m]) + z_modes = np.array([[1, 1, -1], [1, -1, 1]]) + z_lmn = np.array([-a_m, a_m]) # Z = a*rho*sin(theta - Nfp*zeta') + + with h5py.File('desc_test.h5', 'w') as f: + g = f.create_group('_equilibria').create_group('0') + g['_R_basis/_modes'] = r_modes + g['_R_lmn'] = r_lmn + g['_Z_basis/_modes'] = z_modes + g['_Z_lmn'] = z_lmn + g['_NFP'] = nfp + + src = openmc.StellaratorSource.from_desc( + 'desc_test.h5', + emission_density=lambda r: 1.0 - r**2, + energy=openmc.stats.delta_function(14.07e6), + n_rho=9, + ) + assert src.num_field_periods == nfp + rho = np.linspace(0.0, 1.0, 9) + np.testing.assert_allclose(src.rho, rho) + assert src.rmns is None # symmetric equilibrium + + # Combined form: R = R0 + a*rho*cos(theta - Nfp*zeta), Z = a*rho*sin(...) + k00 = np.flatnonzero((src.mode_m == 0) & (src.mode_n == 0))[0] + k11 = np.flatnonzero((src.mode_m == 1) & (src.mode_n == 1))[0] + np.testing.assert_allclose(src.rmnc[:, k00], 100.0 * r0_m) + np.testing.assert_allclose(src.rmnc[:, k11], 100.0 * a_m * rho, atol=1e-12) + np.testing.assert_allclose(src.zmns[:, k11], 100.0 * a_m * rho, atol=1e-12) + np.testing.assert_allclose(src.zmns[:, k00], 0.0, atol=1e-12) From 0e4bd03ce79a549e54508f9c1ef1990db0b55520 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Mon, 28 Sep 2026 10:06:53 -0500 Subject: [PATCH 04/33] Stabilize high-order DESC Zernike evaluation Addresses original StellaratorSource review finding 4. --- openmc/source.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/openmc/source.py b/openmc/source.py index d5e6122d880..12aa71ae6a4 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -10,6 +10,7 @@ import numpy as np import h5py import pandas as pd +from scipy.special import eval_jacobi import openmc import openmc.checkvalue as cv @@ -1780,17 +1781,13 @@ def add(m, n, c, s): add(am, -an, -h, 0.0) @staticmethod - def _zernike_radial(rho, l, m): + def _zernike_radial(rho, l, m): # noqa: E741 """Unnormalized Zernike radial polynomial :math:`R_l^{|m|}(\\rho)`.""" - from math import factorial m = abs(m) - out = np.zeros_like(rho, dtype=float) - for k in range((l - m) // 2 + 1): - c = ((-1)**k * factorial(l - k) - / (factorial(k) * factorial((l + m) // 2 - k) - * factorial((l - m) // 2 - k))) - out += c * rho**(l - 2*k) - return out + if l < m or (l - m) % 2: + raise ValueError("Zernike l - abs(m) must be nonnegative and even") + rho = np.asarray(rho, dtype=float) + return rho**m * eval_jacobi((l - m)//2, 0, m, 2*rho**2 - 1) @staticmethod def _desc_spectral_data(eq): @@ -1878,10 +1875,10 @@ def from_desc( surface_tables = [] radial_r = np.column_stack([ x * cls._zernike_radial(rho, l, m) - for (l, m, n), x in zip(r_modes, r_lmn)]) + for (l, m, n), x in zip(r_modes, r_lmn)]) # noqa: E741 radial_z = np.column_stack([ x * cls._zernike_radial(rho, l, m) - for (l, m, n), x in zip(z_modes, z_lmn)]) + for (l, m, n), x in zip(z_modes, z_lmn)]) # noqa: E741 for i in range(n_rho): table_r, table_z = {}, {} cls._desc_to_combined(r_modes[:, 1], r_modes[:, 2], From b197eeded90ee13eec9f4b4f194f274d407d1881 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Mon, 28 Sep 2026 10:06:54 -0500 Subject: [PATCH 05/33] Restore unrelated settings documentation reverted by the branch Addresses original StellaratorSource review finding 5. --- docs/source/io_formats/settings.rst | 14 +++++++++++++- docs/source/usersguide/settings.rst | 7 ++++++- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/docs/source/io_formats/settings.rst b/docs/source/io_formats/settings.rst index fdc03f6291f..e38f289251e 100644 --- a/docs/source/io_formats/settings.rst +++ b/docs/source/io_formats/settings.rst @@ -626,7 +626,11 @@ found in the :ref:`random ray user guide `. :volume_estimator: Specifies choice of volume estimator for the random ray solver. Options - are 'naive', 'simulation_averaged', or 'hybrid'. The default is 'hybrid'. + are 'naive', 'simulation_averaged', 'hybrid', 'adaptive', + 'strict_adaptive', or 'auto'. The default is 'auto', which selects + 'adaptive' for standard solves and 'strict_adaptive' for solves whose + results feed variance reduction (weight window generation and adjoint + workflows). *Default*: None @@ -1402,6 +1406,13 @@ attributes/sub-elements: *Default*: ``surface_source.h5`` in current working directory +.. deprecated:: + The ```` element is deprecated and will be removed in a + future release. A deprecation warning is emitted when it is present. Use a + file source instead, e.g. ````, + which is equivalent but additionally supports a source strength and source + constraints. + ------------------------------- ```` Element ------------------------------- @@ -1725,6 +1736,7 @@ and 10. The verbosity levels are defined as follows: :5: all of the above + file I/O :6: all of the above + timing statistics and initialization messages :7: all of the above + :math:`k` by generation + :8: all of the above + random ray volume-estimator diagnostics :9: all of the above + indicate when each particle starts :10: all of the above + event information diff --git a/docs/source/usersguide/settings.rst b/docs/source/usersguide/settings.rst index fddae1d25b5..42232ee7c91 100644 --- a/docs/source/usersguide/settings.rst +++ b/docs/source/usersguide/settings.rst @@ -466,6 +466,11 @@ As an example, to write a maximum of three surface source files::: 'max_source_files': 3 } +A surface source file is read back like any other source file, with the +:class:`openmc.FileSource` class:: + + settings.source = openmc.FileSource('surface_source.h5') + .. _compiled_source: Compiled Sources @@ -531,7 +536,7 @@ OpenMC shared library. This can be done by writing a CMakeLists.txt file: cmake_minimum_required(VERSION 3.3 FATAL_ERROR) project(openmc_sources CXX) - add_library(source SHARED source_ring.cpp) + add_library(source MODULE source_ring.cpp) find_package(OpenMC REQUIRED HINTS ) target_link_libraries(source OpenMC::libopenmc) From a937da8203c0aaf230d08e9a7cb324fb2d491f2d Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Mon, 28 Sep 2026 10:57:02 -0500 Subject: [PATCH 06/33] Prune exactly zero Fourier modes from DESC sources --- openmc/source.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/openmc/source.py b/openmc/source.py index 12aa71ae6a4..58deb8669e2 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -1911,6 +1911,11 @@ def from_desc( zmnc *= 100.0 zmns *= 100.0 + # Drop modes whose R and Z coefficients are exactly zero on every surface. + keep = np.any((rmnc != 0) | (rmns != 0) | (zmnc != 0) | (zmns != 0), axis=0) + mode_m, mode_n = mode_m[keep], mode_n[keep] + rmnc, rmns, zmnc, zmns = (a[:, keep] for a in (rmnc, rmns, zmnc, zmns)) + # Drop the asymmetric tables for stellarator-symmetric equilibria sym = not (np.any(rmns) or np.any(zmnc)) From dd69bc2095f4342f8c2854780f85a2ce146f4f9c Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Mon, 28 Sep 2026 11:00:37 -0500 Subject: [PATCH 07/33] Integrate and sample the stellarator radial marginal consistently Retain the original angular setup and rejection sampler. Integrate each radial interval independently, preserving one-sided geometry derivatives, and invert the resulting quartic CDF. --- docs/source/usersguide/settings.rst | 6 +- include/openmc/source.h | 10 +-- include/openmc/stellarator.h | 47 +++++++++++ openmc/source.py | 3 +- src/source.cpp | 90 +++++++++++++-------- tests/cpp_unit_tests/CMakeLists.txt | 1 + tests/cpp_unit_tests/test_stellarator.cpp | 24 ++++++ tests/unit_tests/test_source_stellarator.py | 33 ++++++++ 8 files changed, 171 insertions(+), 43 deletions(-) create mode 100644 include/openmc/stellarator.h create mode 100644 tests/cpp_unit_tests/test_stellarator.cpp diff --git a/docs/source/usersguide/settings.rst b/docs/source/usersguide/settings.rst index 42232ee7c91..88460bfb8d0 100644 --- a/docs/source/usersguide/settings.rst +++ b/docs/source/usersguide/settings.rst @@ -352,8 +352,10 @@ angle, and :math:`\zeta` is the toroidal angle (identical to the cylindrical azimuthal angle in both codes). The user provides the Fourier coefficients of :math:`R(\rho,\theta,\zeta)` and :math:`Z(\rho,\theta,\zeta)` on a radial grid together with an emission density :math:`S(\rho)` that is constant on flux -surfaces. The radial coordinate is sampled by inverting a tabulated CDF of the -exact marginal distribution :math:`S(\rho)\, V'(\rho)` and the two angles are +surfaces. Fourier coefficients and emission are interpolated linearly in +:math:`\rho`. The cubic marginal density :math:`S(\rho)\, V'(\rho)` is integrated +analytically within each interval and its quartic CDF is inverted numerically. +The two angles are sampled from the conditional distribution :math:`p(\theta,\zeta|\rho) \propto R\,|\tau|` (with :math:`\tau` the poloidal-plane Jacobian) by rejection. diff --git a/include/openmc/source.h b/include/openmc/source.h index 0501e14c226..36ba35bf2b5 100644 --- a/include/openmc/source.h +++ b/include/openmc/source.h @@ -428,9 +428,8 @@ class TokamakSource : public Source { //! The sampling algorithm: //! 1. Sample rho from the marginal p(rho) ~ S(rho) * V'(rho), where //! V'(rho) = int R |tau| dtheta dzeta is the differential volume, evaluated -//! exactly at setup via discrete Fourier orthogonality (trapezoidal -//! quadrature of a band-limited trigonometric polynomial) and inverted with -//! a tabulated CDF. +//! by Fourier quadrature. Each radial bin has a cubic PDF whose integral +//! is evaluated analytically and inverted with safeguarded Newton iteration. //! 2. Sample (theta, zeta) from the conditional p(theta, zeta | rho) //! ~ R |tau| by rejection against a precomputed per-radial-bin majorant. //! 3. Sample energy and time from user-provided distribution(s). @@ -505,8 +504,9 @@ class StellaratorSource : public Source { // Angular distribution (isotropic) UPtrAngle angle_; - // Precomputed distribution for radial sampling - unique_ptr radial_dist_; + // Cubic radial PDFs in local bin coordinates and normalized bin masses. + vector> radial_pdf_; + vector radial_cdf_; // Per-radial-bin majorant of R*tau for rejection sampling (n_rho - 1) vector envelope_; diff --git a/include/openmc/stellarator.h b/include/openmc/stellarator.h new file mode 100644 index 00000000000..927321dd7ef --- /dev/null +++ b/include/openmc/stellarator.h @@ -0,0 +1,47 @@ +#ifndef OPENMC_STELLARATOR_H +#define OPENMC_STELLARATOR_H + +#include +#include + +#include "openmc/array.h" + +namespace openmc { + +// Integral on [0,t] of a cubic density, evaluated without the radial bin width. +inline double stellarator_integral(const array& c, double t) +{ + return t * (c[0] + t * (c[1] / 2 + t * (c[2] / 3 + t * c[3] / 4))); +} + +// Invert the normalized integral of a nonnegative cubic. Bisection guarantees +// progress even at a zero of the PDF. Stop at a bracket width of 8 epsilon; +// this does not bound rounding error in evaluating an ill-conditioned CDF. +inline double stellarator_invert(const array& c, double u) +{ + if (u <= 0) + return 0; + if (u >= 1) + return 1; + double target = u * stellarator_integral(c, 1); + double lo = 0, hi = 1, t = u; + for (int iteration = 0; iteration < 512; ++iteration) { + double residual = stellarator_integral(c, t) - target; + if (residual > 0) + hi = t; + else + lo = t; + if (hi - lo <= 8 * std::numeric_limits::epsilon()) + break; + double pdf = c[0] + t * (c[1] + t * (c[2] + t * c[3])); + double next = pdf > 0 ? t - residual / pdf : (lo + hi) / 2; + // Keep Newton steps away from the bracket edges to bound iteration count. + t = next > lo + 0.1 * (hi - lo) && next < hi - 0.1 * (hi - lo) + ? next + : (lo + hi) / 2; + } + return (lo + hi) / 2; +} + +} // namespace openmc +#endif diff --git a/openmc/source.py b/openmc/source.py index 58deb8669e2..2eaa4fca0b2 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -1363,7 +1363,8 @@ class StellaratorSource(SourceBase): coordinate is sampled from the marginal :math:`p(\rho) \propto S(\rho) V'(\rho)` (with :math:`V'(\rho) = \oint\oint R|\tau| \,d\theta\, d\zeta` the differential - volume, evaluated exactly at initialization) via a tabulated CDF, and the + volume for the interpolated geometry) by integrating a cubic PDF in each + radial interval and numerically inverting its quartic CDF, and the two angles are then sampled from the conditional :math:`p(\theta,\zeta|\rho) \propto R|\tau|` by rejection against a precomputed per-radial-bin majorant. diff --git a/src/source.cpp b/src/source.cpp index 531cf63c94c..69768d300b2 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -37,6 +37,7 @@ #include "openmc/settings.h" #include "openmc/simulation.h" #include "openmc/state_point.h" +#include "openmc/stellarator.h" #include "openmc/string_utils.h" #include "openmc/surface.h" #include "openmc/xml_interface.h" @@ -1471,40 +1472,57 @@ void StellaratorSource::precompute_sampling_distributions() "coefficient tables are inconsistent."); } - // Build the marginal radial PDF p(rho) ~ S(rho) * V'(rho) on a refined grid - // containing all surfaces and bin midpoints. At interior surfaces the - // one-sided values from the two adjacent bins are averaged (the piecewise- - // linear coefficient interpolation makes V' one-sided there). - vector grid(2 * n_rho - 1), pdf(2 * n_rho - 1); - for (size_t i = 0; i < n_rho; ++i) { - grid[2 * i] = rho_[i]; - double vp; - if (i == 0) { - vp = vp_lo[0]; - } else if (i == n_rho - 1) { - vp = vp_hi[n_bins - 1]; - } else { - vp = 0.5 * (vp_hi[i - 1] + vp_lo[i]); - } - pdf[2 * i] = emission_density_[i] * std::max(0.0, vp); - } + // Integrate S(rho)*V'(rho) independently on each radial interval. + double s_scale = + *std::max_element(emission_density_.begin(), emission_density_.end()); + if (!(s_scale > 0)) + fatal_error( + "StellaratorSource: integrated emission density must be positive."); + vector masses(n_bins); + radial_pdf_.resize(n_bins); + radial_cdf_.assign(n_bins + 1, 0); for (size_t b = 0; b < n_bins; ++b) { - grid[2 * b + 1] = 0.5 * (rho_[b] + rho_[b + 1]); - double s_mid = 0.5 * (emission_density_[b] + emission_density_[b + 1]); - pdf[2 * b + 1] = s_mid * std::max(0.0, vp_mid[b]); - } - - double total = 0.0; - for (size_t i = 1; i < grid.size(); ++i) { - total += 0.5 * (pdf[i - 1] + pdf[i]) * (grid[i] - grid[i - 1]); + double h = rho_[b + 1] - rho_[b]; + double v0 = vp_lo[b]; + double vh = vp_mid[b]; + double v1 = vp_hi[b]; + array a {v0, 4 * vh - 3 * v0 - v1, 2 * (v0 + v1 - 2 * vh)}; + double minimum = std::min(v0, v1); + if (a[2] > 0) { + double t = -a[1] / (2 * a[2]); + if (t > 0 && t < 1) + minimum = std::min(minimum, a[0] + t * (a[1] + t * a[2])); + } + if (!std::isfinite(v0 + vh + v1) || minimum < 0) + fatal_error("StellaratorSource: invalid differential volume polynomial."); + double s0 = emission_density_[b] / s_scale; + double ds = emission_density_[b + 1] / s_scale - s0; + auto& c = radial_pdf_[b]; + c = {s0 * a[0], s0 * a[1] + ds * a[0], s0 * a[2] + ds * a[1], ds * a[2]}; + double scale = 0; + for (double v : c) + scale = std::max(scale, std::abs(v)); + if (scale > 0) { + for (double& v : c) + v /= scale; + masses[b] = + static_cast(h) * scale * stellarator_integral(c, 1); + if (!std::isfinite(masses[b]) || masses[b] <= 0) + fatal_error("StellaratorSource: invalid radial interval mass."); + } } - if (total <= 0.0) { + long double total = 0; + for (auto mass : masses) + total += mass; + if (!(total > 0) || !std::isfinite(total)) fatal_error( - "StellaratorSource: Integrated emission density is zero or negative. " - "Check emission_density profile."); + "StellaratorSource: integrated emission must be finite and positive."); + long double cumulative = 0; + for (size_t b = 0; b < n_bins; ++b) { + cumulative += masses[b]; + radial_cdf_[b + 1] = cumulative / total; } - radial_dist_ = make_unique( - grid.data(), pdf.data(), grid.size(), Interpolation::lin_lin); + radial_cdf_.back() = 1; } double StellaratorSource::eval_density(int bin, double t, double theta, @@ -1585,11 +1603,13 @@ SourceSite StellaratorSource::sample(uint64_t* seed) const site.delayed_group = 0; // 1. Sample rho from the marginal radial CDF - double rho = radial_dist_->sample(seed).first; - int n_bins = static_cast(rho_.size()) - 1; - int bin = static_cast(lower_bound_index(rho_.begin(), rho_.end(), rho)); - bin = std::min(std::max(bin, 0), n_bins - 1); - double t = (rho - rho_[bin]) / (rho_[bin + 1] - rho_[bin]); + double u = prn(seed); + int bin = std::upper_bound(radial_cdf_.begin(), radial_cdf_.end(), u) - + radial_cdf_.begin() - 1; + double local = + (u - radial_cdf_[bin]) / (radial_cdf_[bin + 1] - radial_cdf_[bin]); + double t = stellarator_invert(radial_pdf_[bin], local); + double rho = rho_[bin] + t * (rho_[bin + 1] - rho_[bin]); // 2. Rejection-sample (theta, zeta) from p(theta, zeta | rho) ~ R*tau double env = envelope_[bin]; diff --git a/tests/cpp_unit_tests/CMakeLists.txt b/tests/cpp_unit_tests/CMakeLists.txt index b6892549bde..c9f9b5a55b8 100644 --- a/tests/cpp_unit_tests/CMakeLists.txt +++ b/tests/cpp_unit_tests/CMakeLists.txt @@ -10,6 +10,7 @@ set(TEST_NAMES test_ray test_region test_surface + test_stellarator test_tensor test_geometry # Add additional unit test files here diff --git a/tests/cpp_unit_tests/test_stellarator.cpp b/tests/cpp_unit_tests/test_stellarator.cpp new file mode 100644 index 00000000000..ba99ae2332c --- /dev/null +++ b/tests/cpp_unit_tests/test_stellarator.cpp @@ -0,0 +1,24 @@ +#include "openmc/constants.h" +#include "openmc/stellarator.h" + +#include +#include + +TEST_CASE("Stellarator cubic CDF inversion", "[stellarator]") +{ + for (int degree = 0; degree < 4; ++degree) { + openmc::array c {}; + c[degree] = 1; + for (double u : + {0.0, 1.e-30, 1.e-12, 0.01, 0.5, 0.99, std::nextafter(1.0, 0.0), 1.0}) { + double t = openmc::stellarator_invert(c, u); + REQUIRE(std::abs(t - std::pow(u, 1.0 / (degree + 1))) < 4.e-15); + } + } + // PDF vanishes at both endpoints, and has a negative monomial coefficient. + openmc::array c {0, 1, -1, 0}; + for (double t : {0.0, 0.001, 0.2, 0.5, 0.99, 1.0}) { + double u = 3 * t * t - 2 * t * t * t; + REQUIRE(std::abs(openmc::stellarator_invert(c, u) - t) < 1.e-13); + } +} diff --git a/tests/unit_tests/test_source_stellarator.py b/tests/unit_tests/test_source_stellarator.py index 023dc5c6a24..7e5a4f8a79a 100644 --- a/tests/unit_tests/test_source_stellarator.py +++ b/tests/unit_tests/test_source_stellarator.py @@ -342,3 +342,36 @@ def test_stellarator_source_from_desc(run_in_tmpdir): np.testing.assert_allclose(src.rmnc[:, k11], 100.0 * a_m * rho, atol=1e-12) np.testing.assert_allclose(src.zmns[:, k11], 100.0 * a_m * rho, atol=1e-12) np.testing.assert_allclose(src.zmns[:, k00], 0.0, atol=1e-12) + + +@pytest.mark.parametrize('case', ['emission', 'geometry', 'zero_intervals']) +def test_stellarator_radial_cdf(case, run_in_tmpdir): + from scipy.stats import kstest + + if case == 'emission': + rho, q, emission = np.array([0., 1.]), np.array([0., 1.]), [0., 1.] + def cdf(x): + return np.clip(x, 0, 1)**3 + elif case == 'geometry': + rho, q, emission = np.array([0., .5, 1.]), np.array([0., .1, 1.]), [1.] * 3 + def cdf(x): + return np.clip(x, 0, 1)**2 + else: + rho = q = np.array([0., .1, .2, .4, .7, 1.]) + emission = np.array([0., 0., 1., 0., 0., 0.]) + # Independent integration of rho*S(rho), retaining exact breakpoints. + from scipy.integrate import quad + norm = quad(lambda x: x*np.interp(x, rho, emission), 0, 1, points=rho)[0] + def cdf(x): + return np.array([quad(lambda t: t*np.interp(t, rho, emission), 0, y, + points=rho[rho < y])[0]/norm for y in x]) + src = openmc.StellaratorSource( + rho, emission, [0, 1], [0, 0], + np.column_stack([np.full(len(rho), 600.), 150*q]), + np.column_stack([np.zeros(len(rho)), 150*q]), + openmc.stats.Discrete([14e6], [1.])) + n = 3000 if case == 'zero_intervals' else 30000 + sites = make_model(src).sample_external_source(n, prn_seed=42, as_array=True) + xyz = sites['r'] + radial = np.hypot(np.hypot(xyz[:, 0], xyz[:, 1])-600, xyz[:, 2])/150 + assert kstest(radial, cdf).statistic < 2.2/np.sqrt(n) From 9c4804e89b56104014669d05692f0967338d2cdb Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Sat, 3 Oct 2026 13:12:16 -0500 Subject: [PATCH 08/33] Change long double to double --- src/source.cpp | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/source.cpp b/src/source.cpp index 69768d300b2..172d1a7aa74 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -1478,7 +1478,7 @@ void StellaratorSource::precompute_sampling_distributions() if (!(s_scale > 0)) fatal_error( "StellaratorSource: integrated emission density must be positive."); - vector masses(n_bins); + vector masses(n_bins); radial_pdf_.resize(n_bins); radial_cdf_.assign(n_bins + 1, 0); for (size_t b = 0; b < n_bins; ++b) { @@ -1505,19 +1505,18 @@ void StellaratorSource::precompute_sampling_distributions() if (scale > 0) { for (double& v : c) v /= scale; - masses[b] = - static_cast(h) * scale * stellarator_integral(c, 1); + masses[b] = h * scale * stellarator_integral(c, 1); if (!std::isfinite(masses[b]) || masses[b] <= 0) fatal_error("StellaratorSource: invalid radial interval mass."); } } - long double total = 0; + double total = 0; for (auto mass : masses) total += mass; if (!(total > 0) || !std::isfinite(total)) fatal_error( "StellaratorSource: integrated emission must be finite and positive."); - long double cumulative = 0; + double cumulative = 0; for (size_t b = 0; b < n_bins; ++b) { cumulative += masses[b]; radial_cdf_[b + 1] = cumulative / total; From c6b4774cea052b86e5a2887d974a32b79dfbca62 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Sat, 3 Oct 2026 13:17:14 -0500 Subject: [PATCH 09/33] Sample stellarator radial intervals with local rejection --- docs/source/usersguide/settings.rst | 5 +-- include/openmc/source.h | 4 +-- include/openmc/stellarator.h | 37 ++++++--------------- openmc/source.py | 4 +-- src/source.cpp | 22 +++++++++--- tests/cpp_unit_tests/test_stellarator.cpp | 23 +++++++------ tests/unit_tests/test_source_stellarator.py | 9 ++++- 7 files changed, 55 insertions(+), 49 deletions(-) diff --git a/docs/source/usersguide/settings.rst b/docs/source/usersguide/settings.rst index 88460bfb8d0..a82d089fa04 100644 --- a/docs/source/usersguide/settings.rst +++ b/docs/source/usersguide/settings.rst @@ -354,8 +354,9 @@ azimuthal angle in both codes). The user provides the Fourier coefficients of together with an emission density :math:`S(\rho)` that is constant on flux surfaces. Fourier coefficients and emission are interpolated linearly in :math:`\rho`. The cubic marginal density :math:`S(\rho)\, V'(\rho)` is integrated -analytically within each interval and its quartic CDF is inverted numerically. -The two angles are +analytically to select an interval, then sampled within that interval by +rejection against a bound derived from its Bernstein coefficients. Rejected +radial proposals are resampled within the same interval. The two angles are sampled from the conditional distribution :math:`p(\theta,\zeta|\rho) \propto R\,|\tau|` (with :math:`\tau` the poloidal-plane Jacobian) by rejection. diff --git a/include/openmc/source.h b/include/openmc/source.h index 36ba35bf2b5..b3c271d994a 100644 --- a/include/openmc/source.h +++ b/include/openmc/source.h @@ -428,8 +428,8 @@ class TokamakSource : public Source { //! The sampling algorithm: //! 1. Sample rho from the marginal p(rho) ~ S(rho) * V'(rho), where //! V'(rho) = int R |tau| dtheta dzeta is the differential volume, evaluated -//! by Fourier quadrature. Each radial bin has a cubic PDF whose integral -//! is evaluated analytically and inverted with safeguarded Newton iteration. +//! by Fourier quadrature. Each radial bin is selected using its analytically +//! integrated cubic PDF, then rho is sampled within the bin by rejection. //! 2. Sample (theta, zeta) from the conditional p(theta, zeta | rho) //! ~ R |tau| by rejection against a precomputed per-radial-bin majorant. //! 3. Sample energy and time from user-provided distribution(s). diff --git a/include/openmc/stellarator.h b/include/openmc/stellarator.h index 927321dd7ef..cdf3f8ddd88 100644 --- a/include/openmc/stellarator.h +++ b/include/openmc/stellarator.h @@ -1,6 +1,7 @@ #ifndef OPENMC_STELLARATOR_H #define OPENMC_STELLARATOR_H +#include #include #include @@ -14,33 +15,17 @@ inline double stellarator_integral(const array& c, double t) return t * (c[0] + t * (c[1] / 2 + t * (c[2] / 3 + t * c[3] / 4))); } -// Invert the normalized integral of a nonnegative cubic. Bisection guarantees -// progress even at a zero of the PDF. Stop at a bracket width of 8 epsilon; -// this does not bound rounding error in evaluating an ill-conditioned CDF. -inline double stellarator_invert(const array& c, double u) +// Bernstein basis functions are nonnegative and sum to one on [0,1], so the +// largest Bernstein coefficient bounds the cubic for rejection sampling. +inline double stellarator_rejection_bound(const array& c) { - if (u <= 0) - return 0; - if (u >= 1) - return 1; - double target = u * stellarator_integral(c, 1); - double lo = 0, hi = 1, t = u; - for (int iteration = 0; iteration < 512; ++iteration) { - double residual = stellarator_integral(c, t) - target; - if (residual > 0) - hi = t; - else - lo = t; - if (hi - lo <= 8 * std::numeric_limits::epsilon()) - break; - double pdf = c[0] + t * (c[1] + t * (c[2] + t * c[3])); - double next = pdf > 0 ? t - residual / pdf : (lo + hi) / 2; - // Keep Newton steps away from the bracket edges to bound iteration count. - t = next > lo + 0.1 * (hi - lo) && next < hi - 0.1 * (hi - lo) - ? next - : (lo + hi) / 2; - } - return (lo + hi) / 2; + double bound = std::max({c[0], c[0] + c[1] / 3, + c[0] + 2 * c[1] / 3 + c[2] / 3, c[0] + c[1] + c[2] + c[3]}); + // Allow for rounding in both the bound and the polynomial evaluation. + double magnitude = 0; + for (double value : c) + magnitude += std::abs(value); + return bound + 64 * std::numeric_limits::epsilon() * magnitude; } } // namespace openmc diff --git a/openmc/source.py b/openmc/source.py index 2eaa4fca0b2..5a2ae917535 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -1364,8 +1364,8 @@ class StellaratorSource(SourceBase): :math:`p(\rho) \propto S(\rho) V'(\rho)` (with :math:`V'(\rho) = \oint\oint R|\tau| \,d\theta\, d\zeta` the differential volume for the interpolated geometry) by integrating a cubic PDF in each - radial interval and numerically inverting its quartic CDF, and the - two angles are then sampled from the conditional + radial interval and using rejection sampling within the selected interval. + The two angles are then sampled from the conditional :math:`p(\theta,\zeta|\rho) \propto R|\tau|` by rejection against a precomputed per-radial-bin majorant. diff --git a/src/source.cpp b/src/source.cpp index 172d1a7aa74..363ea08f189 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -1601,19 +1601,31 @@ SourceSite StellaratorSource::sample(uint64_t* seed) const site.wgt = 1.0; site.delayed_group = 0; - // 1. Sample rho from the marginal radial CDF + // 1. Select a radial bin by its integrated mass, then sample within that bin. double u = prn(seed); int bin = std::upper_bound(radial_cdf_.begin(), radial_cdf_.end(), u) - radial_cdf_.begin() - 1; - double local = - (u - radial_cdf_[bin]) / (radial_cdf_[bin + 1] - radial_cdf_[bin]); - double t = stellarator_invert(radial_pdf_[bin], local); + double t = (u - radial_cdf_[bin]) / (radial_cdf_[bin + 1] - radial_cdf_[bin]); + const auto& c = radial_pdf_[bin]; + double bound = stellarator_rejection_bound(c); + int64_t n_reject = 0; + while (true) { + double pdf = c[0] + t * (c[1] + t * (c[2] + t * c[3])); + if (prn(seed) * bound < pdf) + break; + if (++n_reject > MAX_SOURCE_REJECTIONS_PER_SAMPLE) { + fatal_error("StellaratorSource: exceeded the maximum number of " + "rejections while sampling within a radial interval."); + } + // Retain the selected bin so rejection does not change its probability. + t = prn(seed); + } double rho = rho_[bin] + t * (rho_[bin + 1] - rho_[bin]); // 2. Rejection-sample (theta, zeta) from p(theta, zeta | rho) ~ R*tau double env = envelope_[bin]; double R, Z, zeta; - int64_t n_reject = 0; + n_reject = 0; while (true) { double theta = 2.0 * PI * prn(seed); zeta = 2.0 * PI * prn(seed); diff --git a/tests/cpp_unit_tests/test_stellarator.cpp b/tests/cpp_unit_tests/test_stellarator.cpp index ba99ae2332c..7630ed354f2 100644 --- a/tests/cpp_unit_tests/test_stellarator.cpp +++ b/tests/cpp_unit_tests/test_stellarator.cpp @@ -1,24 +1,25 @@ -#include "openmc/constants.h" #include "openmc/stellarator.h" #include #include -TEST_CASE("Stellarator cubic CDF inversion", "[stellarator]") +TEST_CASE("Stellarator cubic integration and rejection bounds", "[stellarator]") { for (int degree = 0; degree < 4; ++degree) { openmc::array c {}; c[degree] = 1; - for (double u : - {0.0, 1.e-30, 1.e-12, 0.01, 0.5, 0.99, std::nextafter(1.0, 0.0), 1.0}) { - double t = openmc::stellarator_invert(c, u); - REQUIRE(std::abs(t - std::pow(u, 1.0 / (degree + 1))) < 4.e-15); - } + REQUIRE(openmc::stellarator_integral(c, 1) == 1.0 / (degree + 1)); + REQUIRE(openmc::stellarator_rejection_bound(c) >= 1); } // PDF vanishes at both endpoints, and has a negative monomial coefficient. openmc::array c {0, 1, -1, 0}; - for (double t : {0.0, 0.001, 0.2, 0.5, 0.99, 1.0}) { - double u = 3 * t * t - 2 * t * t * t; - REQUIRE(std::abs(openmc::stellarator_invert(c, u) - t) < 1.e-13); - } + REQUIRE(std::abs(openmc::stellarator_integral(c, 1) - 1.0 / 6) < 1.e-15); + REQUIRE(openmc::stellarator_rejection_bound(c) >= 0.25); + // t*(1-t)^2 has its maximum at t=1/3, despite zero endpoint values. + c = {0, 1, -2, 1}; + REQUIRE(std::abs(openmc::stellarator_integral(c, 1) - 1.0 / 12) < 1.e-15); + REQUIRE(openmc::stellarator_rejection_bound(c) >= 4.0 / 27); + // A cubic may have negative Bernstein coefficients and still be nonnegative. + c = {1, -4, 4, 0}; + REQUIRE(openmc::stellarator_rejection_bound(c) >= 1); } diff --git a/tests/unit_tests/test_source_stellarator.py b/tests/unit_tests/test_source_stellarator.py index 7e5a4f8a79a..ef23f6e24ac 100644 --- a/tests/unit_tests/test_source_stellarator.py +++ b/tests/unit_tests/test_source_stellarator.py @@ -344,7 +344,8 @@ def test_stellarator_source_from_desc(run_in_tmpdir): np.testing.assert_allclose(src.zmns[:, k00], 0.0, atol=1e-12) -@pytest.mark.parametrize('case', ['emission', 'geometry', 'zero_intervals']) +@pytest.mark.parametrize('case', [ + 'emission', 'decreasing_emission', 'geometry', 'zero_intervals']) def test_stellarator_radial_cdf(case, run_in_tmpdir): from scipy.stats import kstest @@ -352,6 +353,12 @@ def test_stellarator_radial_cdf(case, run_in_tmpdir): rho, q, emission = np.array([0., 1.]), np.array([0., 1.]), [0., 1.] def cdf(x): return np.clip(x, 0, 1)**3 + elif case == 'decreasing_emission': + # The PDF vanishes at both endpoints and peaks inside the interval. + rho, q, emission = np.array([0., 1.]), np.array([0., 1.]), [1., 0.] + def cdf(x): + x = np.clip(x, 0, 1) + return 3*x**2 - 2*x**3 elif case == 'geometry': rho, q, emission = np.array([0., .5, 1.]), np.array([0., .1, 1.]), [1.] * 3 def cdf(x): From 4ee49347d5621026a2b6194826c6a45bc6dc8040 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Sat, 3 Oct 2026 13:25:09 -0500 Subject: [PATCH 10/33] Keep stellarator radial polynomial helpers private --- include/openmc/stellarator.h | 32 ----------------------- src/source.cpp | 25 +++++++++++++++--- tests/cpp_unit_tests/CMakeLists.txt | 1 - tests/cpp_unit_tests/test_stellarator.cpp | 25 ------------------ 4 files changed, 22 insertions(+), 61 deletions(-) delete mode 100644 include/openmc/stellarator.h delete mode 100644 tests/cpp_unit_tests/test_stellarator.cpp diff --git a/include/openmc/stellarator.h b/include/openmc/stellarator.h deleted file mode 100644 index cdf3f8ddd88..00000000000 --- a/include/openmc/stellarator.h +++ /dev/null @@ -1,32 +0,0 @@ -#ifndef OPENMC_STELLARATOR_H -#define OPENMC_STELLARATOR_H - -#include -#include -#include - -#include "openmc/array.h" - -namespace openmc { - -// Integral on [0,t] of a cubic density, evaluated without the radial bin width. -inline double stellarator_integral(const array& c, double t) -{ - return t * (c[0] + t * (c[1] / 2 + t * (c[2] / 3 + t * c[3] / 4))); -} - -// Bernstein basis functions are nonnegative and sum to one on [0,1], so the -// largest Bernstein coefficient bounds the cubic for rejection sampling. -inline double stellarator_rejection_bound(const array& c) -{ - double bound = std::max({c[0], c[0] + c[1] / 3, - c[0] + 2 * c[1] / 3 + c[2] / 3, c[0] + c[1] + c[2] + c[3]}); - // Allow for rounding in both the bound and the polynomial evaluation. - double magnitude = 0; - for (double value : c) - magnitude += std::abs(value); - return bound + 64 * std::numeric_limits::epsilon() * magnitude; -} - -} // namespace openmc -#endif diff --git a/src/source.cpp b/src/source.cpp index 363ea08f189..ad63b950eba 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -6,6 +6,7 @@ #include // for max #include // for sin, cos, abs +#include // for numeric_limits #include // for move #ifdef HAS_DYNAMIC_LINKING @@ -37,7 +38,6 @@ #include "openmc/settings.h" #include "openmc/simulation.h" #include "openmc/state_point.h" -#include "openmc/stellarator.h" #include "openmc/string_utils.h" #include "openmc/surface.h" #include "openmc/xml_interface.h" @@ -1193,6 +1193,25 @@ double quad_max01(double f0, double fh, double f1) return m; } +// Integral of a cubic density on [0, 1], excluding the radial bin width. +double cubic_integral01(const array& c) +{ + return c[0] + (c[1] / 2 + (c[2] / 3 + c[3] / 4)); +} + +// Bernstein basis functions are nonnegative and sum to one on [0,1], so the +// largest Bernstein coefficient bounds the cubic for rejection sampling. +double cubic_rejection_bound(const array& c) +{ + double bound = std::max({c[0], c[0] + c[1] / 3, + c[0] + 2 * c[1] / 3 + c[2] / 3, c[0] + c[1] + c[2] + c[3]}); + // Allow for rounding in both the bound and the polynomial evaluation. + double magnitude = 0; + for (double value : c) + magnitude += std::abs(value); + return bound + 64 * std::numeric_limits::epsilon() * magnitude; +} + } // namespace StellaratorSource::StellaratorSource(pugi::xml_node node) : Source(node) @@ -1505,7 +1524,7 @@ void StellaratorSource::precompute_sampling_distributions() if (scale > 0) { for (double& v : c) v /= scale; - masses[b] = h * scale * stellarator_integral(c, 1); + masses[b] = h * scale * cubic_integral01(c); if (!std::isfinite(masses[b]) || masses[b] <= 0) fatal_error("StellaratorSource: invalid radial interval mass."); } @@ -1607,7 +1626,7 @@ SourceSite StellaratorSource::sample(uint64_t* seed) const radial_cdf_.begin() - 1; double t = (u - radial_cdf_[bin]) / (radial_cdf_[bin + 1] - radial_cdf_[bin]); const auto& c = radial_pdf_[bin]; - double bound = stellarator_rejection_bound(c); + double bound = cubic_rejection_bound(c); int64_t n_reject = 0; while (true) { double pdf = c[0] + t * (c[1] + t * (c[2] + t * c[3])); diff --git a/tests/cpp_unit_tests/CMakeLists.txt b/tests/cpp_unit_tests/CMakeLists.txt index c9f9b5a55b8..b6892549bde 100644 --- a/tests/cpp_unit_tests/CMakeLists.txt +++ b/tests/cpp_unit_tests/CMakeLists.txt @@ -10,7 +10,6 @@ set(TEST_NAMES test_ray test_region test_surface - test_stellarator test_tensor test_geometry # Add additional unit test files here diff --git a/tests/cpp_unit_tests/test_stellarator.cpp b/tests/cpp_unit_tests/test_stellarator.cpp deleted file mode 100644 index 7630ed354f2..00000000000 --- a/tests/cpp_unit_tests/test_stellarator.cpp +++ /dev/null @@ -1,25 +0,0 @@ -#include "openmc/stellarator.h" - -#include -#include - -TEST_CASE("Stellarator cubic integration and rejection bounds", "[stellarator]") -{ - for (int degree = 0; degree < 4; ++degree) { - openmc::array c {}; - c[degree] = 1; - REQUIRE(openmc::stellarator_integral(c, 1) == 1.0 / (degree + 1)); - REQUIRE(openmc::stellarator_rejection_bound(c) >= 1); - } - // PDF vanishes at both endpoints, and has a negative monomial coefficient. - openmc::array c {0, 1, -1, 0}; - REQUIRE(std::abs(openmc::stellarator_integral(c, 1) - 1.0 / 6) < 1.e-15); - REQUIRE(openmc::stellarator_rejection_bound(c) >= 0.25); - // t*(1-t)^2 has its maximum at t=1/3, despite zero endpoint values. - c = {0, 1, -2, 1}; - REQUIRE(std::abs(openmc::stellarator_integral(c, 1) - 1.0 / 12) < 1.e-15); - REQUIRE(openmc::stellarator_rejection_bound(c) >= 4.0 / 27); - // A cubic may have negative Bernstein coefficients and still be nonnegative. - c = {1, -4, 4, 0}; - REQUIRE(openmc::stellarator_rejection_bound(c) >= 1); -} From 9247ad5e0373dc9a31a7c50571fe0c9b1c91e936 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Sat, 3 Oct 2026 13:34:43 -0500 Subject: [PATCH 11/33] Simplify the radial rejection bound rounding margin --- src/source.cpp | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/source.cpp b/src/source.cpp index ad63b950eba..cb252c903e0 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -6,7 +6,6 @@ #include // for max #include // for sin, cos, abs -#include // for numeric_limits #include // for move #ifdef HAS_DYNAMIC_LINKING @@ -1205,11 +1204,8 @@ double cubic_rejection_bound(const array& c) { double bound = std::max({c[0], c[0] + c[1] / 3, c[0] + 2 * c[1] / 3 + c[2] / 3, c[0] + c[1] + c[2] + c[3]}); - // Allow for rounding in both the bound and the polynomial evaluation. - double magnitude = 0; - for (double value : c) - magnitude += std::abs(value); - return bound + 64 * std::numeric_limits::epsilon() * magnitude; + // Add a small relative margin for rounding in the bound and PDF evaluation. + return bound * (1.0 + 1.e-12); } } // namespace From 4a653ad891e0a3f3f6176a1f0f7f0140786d4af1 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Sat, 3 Oct 2026 14:04:55 -0500 Subject: [PATCH 12/33] Use std::accumulate --- src/source.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/source.cpp b/src/source.cpp index cb252c903e0..12c88cce772 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -6,6 +6,7 @@ #include // for max #include // for sin, cos, abs +#include // for accumulate #include // for move #ifdef HAS_DYNAMIC_LINKING @@ -1525,9 +1526,7 @@ void StellaratorSource::precompute_sampling_distributions() fatal_error("StellaratorSource: invalid radial interval mass."); } } - double total = 0; - for (auto mass : masses) - total += mass; + double total = std::accumulate(masses.begin(), masses.end(), 0.0); if (!(total > 0) || !std::isfinite(total)) fatal_error( "StellaratorSource: integrated emission must be finite and positive."); From d1335e4ff219fb97e593fc789df832989d1450af Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Sat, 3 Oct 2026 14:16:49 -0500 Subject: [PATCH 13/33] Simplify stellarator radial sampling --- include/openmc/source.h | 4 ++-- src/source.cpp | 41 +++++++++++++---------------------------- 2 files changed, 15 insertions(+), 30 deletions(-) diff --git a/include/openmc/source.h b/include/openmc/source.h index b3c271d994a..9dddc7479e8 100644 --- a/include/openmc/source.h +++ b/include/openmc/source.h @@ -504,9 +504,9 @@ class StellaratorSource : public Source { // Angular distribution (isotropic) UPtrAngle angle_; - // Cubic radial PDFs in local bin coordinates and normalized bin masses. + // Bound-normalized cubic radial PDFs and the distribution of bin masses. vector> radial_pdf_; - vector radial_cdf_; + DiscreteIndex radial_bins_; // Per-radial-bin majorant of R*tau for rejection sampling (n_rho - 1) vector envelope_; diff --git a/src/source.cpp b/src/source.cpp index 12c88cce772..fa3e7684f7e 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -1193,12 +1193,6 @@ double quad_max01(double f0, double fh, double f1) return m; } -// Integral of a cubic density on [0, 1], excluding the radial bin width. -double cubic_integral01(const array& c) -{ - return c[0] + (c[1] / 2 + (c[2] / 3 + c[3] / 4)); -} - // Bernstein basis functions are nonnegative and sum to one on [0,1], so the // largest Bernstein coefficient bounds the cubic for rejection sampling. double cubic_rejection_bound(const array& c) @@ -1496,7 +1490,6 @@ void StellaratorSource::precompute_sampling_distributions() "StellaratorSource: integrated emission density must be positive."); vector masses(n_bins); radial_pdf_.resize(n_bins); - radial_cdf_.assign(n_bins + 1, 0); for (size_t b = 0; b < n_bins; ++b) { double h = rho_[b + 1] - rho_[b]; double v0 = vp_lo[b]; @@ -1512,30 +1505,25 @@ void StellaratorSource::precompute_sampling_distributions() if (!std::isfinite(v0 + vh + v1) || minimum < 0) fatal_error("StellaratorSource: invalid differential volume polynomial."); double s0 = emission_density_[b] / s_scale; - double ds = emission_density_[b + 1] / s_scale - s0; + double s1 = emission_density_[b + 1] / s_scale; + double ds = s1 - s0; + // Simpson's rule exactly integrates the cubic S(rho)*V'(rho). + masses[b] = h * ((s0 / 6) * v0 + ((s0 + s1) / 3) * vh + (s1 / 6) * v1); auto& c = radial_pdf_[b]; c = {s0 * a[0], s0 * a[1] + ds * a[0], s0 * a[2] + ds * a[1], ds * a[2]}; - double scale = 0; - for (double v : c) - scale = std::max(scale, std::abs(v)); - if (scale > 0) { - for (double& v : c) - v /= scale; - masses[b] = h * scale * cubic_integral01(c); - if (!std::isfinite(masses[b]) || masses[b] <= 0) + double bound = cubic_rejection_bound(c); + if (bound > 0) { + if (!std::isfinite(bound) || !std::isfinite(masses[b]) || masses[b] <= 0) fatal_error("StellaratorSource: invalid radial interval mass."); + for (double& v : c) + v /= bound; } } double total = std::accumulate(masses.begin(), masses.end(), 0.0); if (!(total > 0) || !std::isfinite(total)) fatal_error( "StellaratorSource: integrated emission must be finite and positive."); - double cumulative = 0; - for (size_t b = 0; b < n_bins; ++b) { - cumulative += masses[b]; - radial_cdf_[b + 1] = cumulative / total; - } - radial_cdf_.back() = 1; + radial_bins_.assign({masses.data(), masses.size()}); } double StellaratorSource::eval_density(int bin, double t, double theta, @@ -1616,16 +1604,13 @@ SourceSite StellaratorSource::sample(uint64_t* seed) const site.delayed_group = 0; // 1. Select a radial bin by its integrated mass, then sample within that bin. - double u = prn(seed); - int bin = std::upper_bound(radial_cdf_.begin(), radial_cdf_.end(), u) - - radial_cdf_.begin() - 1; - double t = (u - radial_cdf_[bin]) / (radial_cdf_[bin + 1] - radial_cdf_[bin]); + int bin = radial_bins_.sample(seed); + double t = prn(seed); const auto& c = radial_pdf_[bin]; - double bound = cubic_rejection_bound(c); int64_t n_reject = 0; while (true) { double pdf = c[0] + t * (c[1] + t * (c[2] + t * c[3])); - if (prn(seed) * bound < pdf) + if (prn(seed) < pdf) break; if (++n_reject > MAX_SOURCE_REJECTIONS_PER_SAMPLE) { fatal_error("StellaratorSource: exceeded the maximum number of " From d07c4f5bfaca780c656616866ad18752eedd92a8 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Sat, 3 Oct 2026 14:24:23 -0500 Subject: [PATCH 14/33] Fix style in settings.rst --- docs/source/usersguide/settings.rst | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/docs/source/usersguide/settings.rst b/docs/source/usersguide/settings.rst index a82d089fa04..87b491c5eb3 100644 --- a/docs/source/usersguide/settings.rst +++ b/docs/source/usersguide/settings.rst @@ -341,25 +341,24 @@ are born at :math:`t=0`. The toroidal extent can be restricted with Stellarator Plasma Sources -------------------------- -For fully three-dimensional plasmas, the :class:`openmc.StellaratorSource` -class samples neutron positions from the flux-surface Fourier representation -shared by the `VMEC `_ and `DESC +For fully three-dimensional plasmas, the :class:`openmc.StellaratorSource` class +samples neutron positions from the flux-surface Fourier representation shared by +the `VMEC `_ and `DESC `_ equilibrium codes. The plasma is -described in flux coordinates :math:`(\rho, \theta, \zeta)`, where -:math:`\rho = \sqrt{s}` is the square root of the normalized toroidal flux -(proportional to the average minor radius), :math:`\theta` is the poloidal -angle, and :math:`\zeta` is the toroidal angle (identical to the cylindrical -azimuthal angle in both codes). The user provides the Fourier coefficients of +described in flux coordinates :math:`(\rho, \theta, \zeta)`, where :math:`\rho = +\sqrt{s}` is the square root of the normalized toroidal flux (proportional to +the average minor radius), :math:`\theta` is the poloidal angle, and +:math:`\zeta` is the toroidal angle (identical to the cylindrical azimuthal +angle in both codes). The user provides the Fourier coefficients of :math:`R(\rho,\theta,\zeta)` and :math:`Z(\rho,\theta,\zeta)` on a radial grid together with an emission density :math:`S(\rho)` that is constant on flux surfaces. Fourier coefficients and emission are interpolated linearly in -:math:`\rho`. The cubic marginal density :math:`S(\rho)\, V'(\rho)` is integrated -analytically to select an interval, then sampled within that interval by -rejection against a bound derived from its Bernstein coefficients. Rejected +:math:`\rho`. The cubic marginal density :math:`S(\rho)\, V'(\rho)` is +integrated analytically to select an interval, then sampled within that interval +by rejection against a bound derived from its Bernstein coefficients. Rejected radial proposals are resampled within the same interval. The two angles are -sampled from the conditional distribution -:math:`p(\theta,\zeta|\rho) \propto R\,|\tau|` (with :math:`\tau` the -poloidal-plane Jacobian) by rejection. +sampled from the conditional distribution :math:`p(\theta,\zeta|\rho) \propto +R\,|\tau|` (with :math:`\tau` the poloidal-plane Jacobian) by rejection. Sources are most conveniently created directly from equilibrium code output:: From 70c9db5e4ac7d6b85cc75345de60739b42e0b3a2 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Sat, 3 Oct 2026 14:26:21 -0500 Subject: [PATCH 15/33] Remove CDF test --- tests/unit_tests/test_source_stellarator.py | 40 --------------------- 1 file changed, 40 deletions(-) diff --git a/tests/unit_tests/test_source_stellarator.py b/tests/unit_tests/test_source_stellarator.py index ef23f6e24ac..023dc5c6a24 100644 --- a/tests/unit_tests/test_source_stellarator.py +++ b/tests/unit_tests/test_source_stellarator.py @@ -342,43 +342,3 @@ def test_stellarator_source_from_desc(run_in_tmpdir): np.testing.assert_allclose(src.rmnc[:, k11], 100.0 * a_m * rho, atol=1e-12) np.testing.assert_allclose(src.zmns[:, k11], 100.0 * a_m * rho, atol=1e-12) np.testing.assert_allclose(src.zmns[:, k00], 0.0, atol=1e-12) - - -@pytest.mark.parametrize('case', [ - 'emission', 'decreasing_emission', 'geometry', 'zero_intervals']) -def test_stellarator_radial_cdf(case, run_in_tmpdir): - from scipy.stats import kstest - - if case == 'emission': - rho, q, emission = np.array([0., 1.]), np.array([0., 1.]), [0., 1.] - def cdf(x): - return np.clip(x, 0, 1)**3 - elif case == 'decreasing_emission': - # The PDF vanishes at both endpoints and peaks inside the interval. - rho, q, emission = np.array([0., 1.]), np.array([0., 1.]), [1., 0.] - def cdf(x): - x = np.clip(x, 0, 1) - return 3*x**2 - 2*x**3 - elif case == 'geometry': - rho, q, emission = np.array([0., .5, 1.]), np.array([0., .1, 1.]), [1.] * 3 - def cdf(x): - return np.clip(x, 0, 1)**2 - else: - rho = q = np.array([0., .1, .2, .4, .7, 1.]) - emission = np.array([0., 0., 1., 0., 0., 0.]) - # Independent integration of rho*S(rho), retaining exact breakpoints. - from scipy.integrate import quad - norm = quad(lambda x: x*np.interp(x, rho, emission), 0, 1, points=rho)[0] - def cdf(x): - return np.array([quad(lambda t: t*np.interp(t, rho, emission), 0, y, - points=rho[rho < y])[0]/norm for y in x]) - src = openmc.StellaratorSource( - rho, emission, [0, 1], [0, 0], - np.column_stack([np.full(len(rho), 600.), 150*q]), - np.column_stack([np.zeros(len(rho)), 150*q]), - openmc.stats.Discrete([14e6], [1.])) - n = 3000 if case == 'zero_intervals' else 30000 - sites = make_model(src).sample_external_source(n, prn_seed=42, as_array=True) - xyz = sites['r'] - radial = np.hypot(np.hypot(xyz[:, 0], xyz[:, 1])-600, xyz[:, 2])/150 - assert kstest(radial, cdf).statistic < 2.2/np.sqrt(n) From bed8ee65b34df47616cae8737e8c938efac14172 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Sun, 4 Oct 2026 17:05:24 -0500 Subject: [PATCH 16/33] Address @eepeterson comments --- src/source.cpp | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/src/source.cpp b/src/source.cpp index fa3e7684f7e..a98f36ebd8d 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -1278,6 +1278,11 @@ StellaratorSource::StellaratorSource(pugi::xml_node node) : Source(node) "StellaratorSource: emission_density values cannot be negative."); } } + if (std::none_of(emission_density_.begin(), emission_density_.end(), + [](double s) { return s > 0.0; })) { + fatal_error("StellaratorSource: emission_density must contain at least " + "one positive value."); + } if (n_modes == 0) { fatal_error("StellaratorSource: At least one Fourier mode is required."); } @@ -1483,11 +1488,8 @@ void StellaratorSource::precompute_sampling_distributions() } // Integrate S(rho)*V'(rho) independently on each radial interval. - double s_scale = + double s_max = *std::max_element(emission_density_.begin(), emission_density_.end()); - if (!(s_scale > 0)) - fatal_error( - "StellaratorSource: integrated emission density must be positive."); vector masses(n_bins); radial_pdf_.resize(n_bins); for (size_t b = 0; b < n_bins; ++b) { @@ -1504,8 +1506,8 @@ void StellaratorSource::precompute_sampling_distributions() } if (!std::isfinite(v0 + vh + v1) || minimum < 0) fatal_error("StellaratorSource: invalid differential volume polynomial."); - double s0 = emission_density_[b] / s_scale; - double s1 = emission_density_[b + 1] / s_scale; + double s0 = emission_density_[b] / s_max; + double s1 = emission_density_[b + 1] / s_max; double ds = s1 - s0; // Simpson's rule exactly integrates the cubic S(rho)*V'(rho). masses[b] = h * ((s0 / 6) * v0 + ((s0 + s1) / 3) * vh + (s1 / 6) * v1); @@ -1520,9 +1522,6 @@ void StellaratorSource::precompute_sampling_distributions() } } double total = std::accumulate(masses.begin(), masses.end(), 0.0); - if (!(total > 0) || !std::isfinite(total)) - fatal_error( - "StellaratorSource: integrated emission must be finite and positive."); radial_bins_.assign({masses.data(), masses.size()}); } From 5a85f829333e8ec89af284434eed4aec1b34cded Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Tue, 13 Oct 2026 14:53:59 -0500 Subject: [PATCH 17/33] Remove unused total variable --- src/source.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/source.cpp b/src/source.cpp index a98f36ebd8d..d297ea902e3 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -1521,7 +1521,6 @@ void StellaratorSource::precompute_sampling_distributions() v /= bound; } } - double total = std::accumulate(masses.begin(), masses.end(), 0.0); radial_bins_.assign({masses.data(), masses.size()}); } From 0cf3df1417ce466eed9aaaf7d95dbadc2dc5a017 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Tue, 13 Oct 2026 14:54:48 -0500 Subject: [PATCH 18/33] List stellarator in the SourceBase type docstring The new source type was registered everywhere it needed to be except the SourceBase.type attribute docstring, which still enumerated only independent, file, compiled, mesh and tokamak. The parallel list in the settings I/O format reference was updated by this branch, so the Python API documentation was the one place left inconsistent with the StellaratorSource.type property and with SourceBase.from_xml_element, both of which already handle 'stellarator'. Co-Authored-By: Claude Opus 5 --- openmc/source.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openmc/source.py b/openmc/source.py index 5a2ae917535..3e5b781dd92 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -47,7 +47,7 @@ class SourceBase(ABC): Attributes ---------- - type : {'independent', 'file', 'compiled', 'mesh', 'tokamak'} + type : {'independent', 'file', 'compiled', 'mesh', 'tokamak', 'stellarator'} Indicator of source type. strength : float Strength of the source From ed3bb2bced680db1fd7061c6f9497b7df0577c51 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Wed, 14 Oct 2026 17:38:04 -0500 Subject: [PATCH 19/33] Check type against PathLike, not Path --- openmc/source.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/openmc/source.py b/openmc/source.py index 3e5b781dd92..807d2cae680 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -2,6 +2,7 @@ from abc import ABC, abstractmethod from collections.abc import Iterable, Sequence from numbers import Integral, Real +import os from pathlib import Path import warnings from typing import Any @@ -1799,7 +1800,7 @@ def _desc_spectral_data(eq): which is read directly with h5py. Returns ``(r_modes, r_lmn, z_modes, z_lmn, nfp)`` where the modes arrays have columns ``(l, m, n)``. """ - if isinstance(eq, (str, Path)): + if isinstance(eq, (str, os.PathLike)): with h5py.File(input_path(eq), 'r') as f: # Output files may contain a family of equilibria; use the last g = f From 13add724a983f33ad572fddb78bc2e9617f6947d Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Wed, 14 Oct 2026 23:33:56 -0500 Subject: [PATCH 20/33] Use optional DESC APIs for stellarator source conversion --- docs/source/usersguide/settings.rst | 6 + openmc/source.py | 183 ++++++------------- pyproject.toml | 1 + tests/unit_tests/test_source_stellarator.py | 185 +++++++++++--------- 4 files changed, 165 insertions(+), 210 deletions(-) diff --git a/docs/source/usersguide/settings.rst b/docs/source/usersguide/settings.rst index 87b491c5eb3..28d5413915b 100644 --- a/docs/source/usersguide/settings.rst +++ b/docs/source/usersguide/settings.rst @@ -376,6 +376,12 @@ Sources are most conveniently created directly from equilibrium code output:: energy=openmc.stats.muir(e0=14.08e6, m_rat=5.0, kt=20000.0), ) +The ``from_desc`` constructor requires the optional ``desc-opt`` package, +available through ``pip install openmc[desc]`` or +``conda install -c conda-forge desc-opt``. DESC loads its own files and performs +the spectral conversion. An equilibrium family uses its last member. DESC is +not required to import OpenMC, read the resulting source XML, or run transport. + The ``emission_density`` argument accepts either a callable evaluated on the radial grid or an array with one value per radial surface; only its shape matters, since it is normalized internally. The Fourier coefficient tables can diff --git a/openmc/source.py b/openmc/source.py index 807d2cae680..866fd951993 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -11,7 +11,6 @@ import numpy as np import h5py import pandas as pd -from scipy.special import eval_jacobi import openmc import openmc.checkvalue as cv @@ -1748,77 +1747,6 @@ def _read_wout(path): return {k: np.asarray(ds[k][()], dtype=float) for k in names if k in ds} - @staticmethod - def _desc_to_combined(modes_m, modes_n, coeffs, table): - """Fold DESC's product-form double Fourier coefficients into the - combined form sum[c*cos(m*theta - n*Nfp*zeta) + s*sin(...)], m >= 0. - - DESC basis conventions: positive (negative) m selects cos(|m| theta) - (sin(|m| theta)); likewise for n with the toroidal angle. The Ptolemy - identities split each product into the two combined-form harmonics - (m, n) and (m, -n), and modes with m < 0 (or m == 0, n < 0) are folded - using cos(-x) = cos(x), sin(-x) = -sin(x). - """ - def add(m, n, c, s): - if m < 0 or (m == 0 and n < 0): - m, n, s = -m, -n, -s - entry = table.setdefault((m, n), [0.0, 0.0]) - entry[0] += c - entry[1] += s - - for m0, n0, x in zip(modes_m, modes_n, coeffs): - am, an = abs(int(m0)), abs(int(n0)) - h = 0.5 * float(x) - if m0 >= 0 and n0 >= 0: # cos(m theta) * cos(n zeta') - add(am, an, h, 0.0) - add(am, -an, h, 0.0) - elif m0 < 0 and n0 >= 0: # sin(m theta) * cos(n zeta') - add(am, an, 0.0, h) - add(am, -an, 0.0, h) - elif m0 >= 0 and n0 < 0: # cos(m theta) * sin(n zeta') - add(am, an, 0.0, -h) - add(am, -an, 0.0, h) - else: # sin(m theta) * sin(n zeta') - add(am, an, h, 0.0) - add(am, -an, -h, 0.0) - - @staticmethod - def _zernike_radial(rho, l, m): # noqa: E741 - """Unnormalized Zernike radial polynomial :math:`R_l^{|m|}(\\rho)`.""" - m = abs(m) - if l < m or (l - m) % 2: - raise ValueError("Zernike l - abs(m) must be nonnegative and even") - rho = np.asarray(rho, dtype=float) - return rho**m * eval_jacobi((l - m)//2, 0, m, 2*rho**2 - 1) - - @staticmethod - def _desc_spectral_data(eq): - """Extract Fourier-Zernike modes/coefficients from a DESC equilibrium. - - Accepts either a live ``desc.equilibrium.Equilibrium`` (duck-typed, so - no import of desc is needed) or a path to a DESC HDF5 output file, - which is read directly with h5py. Returns ``(r_modes, r_lmn, z_modes, - z_lmn, nfp)`` where the modes arrays have columns ``(l, m, n)``. - """ - if isinstance(eq, (str, os.PathLike)): - with h5py.File(input_path(eq), 'r') as f: - # Output files may contain a family of equilibria; use the last - g = f - if '_equilibria' in f: - idx = sorted((k for k in f['_equilibria'] if k.isdigit()), - key=int) - g = f['_equilibria'][idx[-1]] - return (np.asarray(g['_R_basis/_modes'][()], dtype=int), - np.asarray(g['_R_lmn'][()], dtype=float), - np.asarray(g['_Z_basis/_modes'][()], dtype=int), - np.asarray(g['_Z_lmn'][()], dtype=float), - int(g['_NFP'][()])) - return (np.asarray(eq.R_basis.modes, dtype=int), - np.asarray(eq.R_lmn, dtype=float), - np.asarray(eq.Z_basis.modes, dtype=int), - np.asarray(eq.Z_lmn, dtype=float), - int(eq.NFP)) - @classmethod def from_desc( cls, @@ -1836,20 +1764,17 @@ def from_desc( \\mathcal{F}_n(N_{fp}\\zeta)`, where :math:`\\mathcal{R}_l^{|m|}` is the (unnormalized) Zernike radial polynomial and :math:`\\mathcal{F}_m(x)` is :math:`\\cos(|m|x)` for :math:`m \\ge 0` and :math:`\\sin(|m|x)` - otherwise. The Zernike polynomials are evaluated here on a uniform grid - in :math:`\\rho` (DESC's native radial coordinate) to obtain - per-surface double Fourier coefficients, which are then converted from - DESC's product-form basis to the combined VMEC-style form using the - Ptolemy identities. Coefficients are converted from [m] to [cm]. No - additional packages are required: DESC HDF5 output files are read - directly with :mod:`h5py`, and live equilibrium objects are accessed - through their public attributes only. + otherwise. DESC evaluates the radial basis on a uniform grid in + :math:`\\rho` and converts the product-form coefficients to the + combined VMEC-style form. Coefficients are converted from [m] to [cm]. + This method requires the optional ``desc-opt`` package, which is + imported only when the method is called. Parameters ---------- - eq : desc.equilibrium.Equilibrium or path-like + eq : desc.equilibrium.Equilibrium or desc.equilibrium.EquilibriaFamily or path-like A DESC equilibrium object, or a path to a DESC HDF5 output file - (the last equilibrium is used if the file contains a family). + (the last equilibrium is used for an equilibrium family). emission_density : callable or numpy.ndarray Either a callable ``S(rho)`` evaluated at the radial grid points, or an array of length ``n_rho`` (arbitrary units, must be >= 0). @@ -1866,52 +1791,56 @@ def from_desc( openmc.StellaratorSource """ + try: + from desc.basis import zernike_radial + from desc.equilibrium import Equilibrium, EquilibriaFamily + from desc.io import load + from desc.vmec_utils import ptolemy_identity_rev + except ModuleNotFoundError as exc: + if exc.name != 'desc': + raise + raise ImportError( + "StellaratorSource.from_desc requires DESC. " + "Install it with 'pip install openmc[desc]' or " + "'pip install desc-opt'.") from exc + + cv.check_type('n_rho', n_rho, Integral) cv.check_greater_than('n_rho', n_rho, 1) rho = np.linspace(0.0, 1.0, n_rho) - - r_modes, r_lmn, z_modes, z_lmn, nfp = cls._desc_spectral_data(eq) - - # Collapse the Zernike radial dependence onto the rho grid, giving - # product-form double Fourier coefficients for each surface, then - # convert to combined form - surface_tables = [] - radial_r = np.column_stack([ - x * cls._zernike_radial(rho, l, m) - for (l, m, n), x in zip(r_modes, r_lmn)]) # noqa: E741 - radial_z = np.column_stack([ - x * cls._zernike_radial(rho, l, m) - for (l, m, n), x in zip(z_modes, z_lmn)]) # noqa: E741 - for i in range(n_rho): - table_r, table_z = {}, {} - cls._desc_to_combined(r_modes[:, 1], r_modes[:, 2], - radial_r[i], table_r) - cls._desc_to_combined(z_modes[:, 1], z_modes[:, 2], - radial_z[i], table_z) - surface_tables.append((table_r, table_z)) - - # Union of modes across R, Z, and all surfaces - modes = sorted({key for table_r, table_z in surface_tables - for key in (*table_r, *table_z)}) - mode_m = np.array([m for m, _ in modes], dtype=int) - mode_n = np.array([n for _, n in modes], dtype=int) - - n_modes = len(modes) - rmnc = np.zeros((n_rho, n_modes)) - rmns = np.zeros((n_rho, n_modes)) - zmnc = np.zeros((n_rho, n_modes)) - zmns = np.zeros((n_rho, n_modes)) - for i, (table_r, table_z) in enumerate(surface_tables): - for k, key in enumerate(modes): - if key in table_r: - rmnc[i, k], rmns[i, k] = table_r[key] - if key in table_z: - zmnc[i, k], zmns[i, k] = table_z[key] - - # m to cm - rmnc *= 100.0 - rmns *= 100.0 - zmnc *= 100.0 - zmns *= 100.0 + if isinstance(eq, (str, os.PathLike)): + eq = load(input_path(eq), file_format='hdf5') + if isinstance(eq, EquilibriaFamily): + if not len(eq): + raise ValueError("The DESC equilibrium family is empty") + eq = eq[-1] + cv.check_type('eq', eq, Equilibrium) + + def convert(name): + basis = getattr(eq, f'{name}_basis') + l, m, _ = basis.modes.T # noqa: E741 + modes, inverse = np.unique(basis.modes[:, 1:], axis=0, + return_inverse=True) + radial = np.asarray(zernike_radial(rho[:, None], l, m)) + values = np.zeros((n_rho, len(modes))) + # Sum radial modes sharing the same angular mode on each surface. + np.add.at(values.T, inverse, + (radial * np.asarray(getattr(eq, f'{name}_lmn'))).T) + m, n, sine, cosine = ptolemy_identity_rev( + modes[:, 0], modes[:, 1], values) + return list(zip(m, n)), sine, cosine + + r_modes, r_sine, r_cosine = convert('R') + z_modes, z_sine, z_cosine = convert('Z') + modes = sorted(set(r_modes) | set(z_modes)) + mode_m, mode_n = np.asarray(modes, dtype=int).T + indices = {mode: k for k, mode in enumerate(modes)} + rmnc, rmns, zmnc, zmns = ( + np.zeros((n_rho, len(modes))) for _ in range(4)) + # Align the R/Z mode lists and convert meters to centimeters. + r_indices = [indices[mode] for mode in r_modes] + z_indices = [indices[mode] for mode in z_modes] + rmnc[:, r_indices], rmns[:, r_indices] = 100*r_cosine, 100*r_sine + zmnc[:, z_indices], zmns[:, z_indices] = 100*z_cosine, 100*z_sine # Drop modes whose R and Z coefficients are exactly zero on every surface. keep = np.any((rmnc != 0) | (rmns != 0) | (zmnc != 0) | (zmns != 0), axis=0) @@ -1931,7 +1860,7 @@ def from_desc( zmns=zmns, rmns=None if sym else rmns, zmnc=None if sym else zmnc, - num_field_periods=nfp, + num_field_periods=eq.NFP, energy=energy, **kwargs ) diff --git a/pyproject.toml b/pyproject.toml index 7133b7bb866..1ebd1f940a2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ dependencies = [ ] [project.optional-dependencies] +desc = ["desc-opt>=0.17.3"] depletion-mpi = ["mpi4py"] docs = [ "breathe", diff --git a/tests/unit_tests/test_source_stellarator.py b/tests/unit_tests/test_source_stellarator.py index 023dc5c6a24..fc8e0d2b20d 100644 --- a/tests/unit_tests/test_source_stellarator.py +++ b/tests/unit_tests/test_source_stellarator.py @@ -258,87 +258,106 @@ def test_stellarator_source_from_vmec(run_in_tmpdir): assert src.rmns is None -def test_stellarator_source_desc_ptolemy(): - """Check the DESC product-form to combined-form conversion identities.""" - rng = np.random.default_rng(42) - theta = rng.uniform(0.0, 2 * np.pi, 50) - zp = rng.uniform(0.0, 2 * np.pi, 50) # zeta' = nfp * zeta - - # One mode of each sign class with random coefficients - modes_m = np.array([2, -2, 2, -2, 0, 3]) - modes_n = np.array([1, 1, -1, -1, -2, 0]) - coeffs = rng.normal(size=modes_m.size) - - # Direct evaluation in DESC's product-form convention - direct = np.zeros_like(theta) - for m0, n0, x in zip(modes_m, modes_n, coeffs): - pol = np.cos(abs(m0) * theta) if m0 >= 0 else np.sin(abs(m0) * theta) - tor = np.cos(abs(n0) * zp) if n0 >= 0 else np.sin(abs(n0) * zp) - direct += x * pol * tor - - # Combined-form evaluation - table = {} - openmc.StellaratorSource._desc_to_combined(modes_m, modes_n, coeffs, table) - combined = np.zeros_like(theta) - for (m, n), (c, s) in table.items(): - ang = m * theta - n * zp - combined += c * np.cos(ang) + s * np.sin(ang) - - np.testing.assert_allclose(combined, direct, atol=1e-12) - - -def test_stellarator_source_zernike_radial(): - """Check the Zernike radial polynomial against known closed forms.""" - zr = openmc.StellaratorSource._zernike_radial - rho = np.linspace(0.0, 1.0, 11) - np.testing.assert_allclose(zr(rho, 0, 0), np.ones_like(rho)) - np.testing.assert_allclose(zr(rho, 1, 1), rho) - np.testing.assert_allclose(zr(rho, 2, 0), 2 * rho**2 - 1) - np.testing.assert_allclose(zr(rho, 2, 2), rho**2) - np.testing.assert_allclose(zr(rho, 3, 1), 3 * rho**3 - 2 * rho) - np.testing.assert_allclose(zr(rho, 4, 0), 6 * rho**4 - 6 * rho**2 + 1) - np.testing.assert_allclose(zr(rho, 3, -1), zr(rho, 3, 1)) # |m| is used - - -def test_stellarator_source_from_desc(run_in_tmpdir): - """Read a synthetic rotating-ellipse DESC HDF5 file with from_desc. - - The equilibrium R = R0 + a*rho*cos(theta - Nfp*zeta') expands in DESC's - product-form Fourier-Zernike basis (Zernike R_1^1(rho) = rho) as modes - (l=1, m=1, n=1) and (l=1, m=-1, n=-1), and likewise for Z. - """ - import h5py - - nfp = 5 - r0_m, a_m = 6.0, 1.5 - r_modes = np.array([[0, 0, 0], [1, 1, 1], [1, -1, -1]]) - r_lmn = np.array([r0_m, a_m, a_m]) - z_modes = np.array([[1, 1, -1], [1, -1, 1]]) - z_lmn = np.array([-a_m, a_m]) # Z = a*rho*sin(theta - Nfp*zeta') - - with h5py.File('desc_test.h5', 'w') as f: - g = f.create_group('_equilibria').create_group('0') - g['_R_basis/_modes'] = r_modes - g['_R_lmn'] = r_lmn - g['_Z_basis/_modes'] = z_modes - g['_Z_lmn'] = z_lmn - g['_NFP'] = nfp - +@pytest.fixture +def desc_equilibrium(): + pytest.importorskip('desc') + from desc.equilibrium import Equilibrium + + # Circular torus plus a nontrivial l=3 radial contribution. + eq = Equilibrium(L=3, M=2, N=1, NFP=3, sym=False) + r = np.zeros(eq.R_basis.num_modes) + z = np.zeros(eq.Z_basis.num_modes) + r[eq.R_basis.get_idx(0, 0, 0)] = 6.0 + r[eq.R_basis.get_idx(1, 1, 0)] = 1.5 + z[eq.Z_basis.get_idx(1, -1, 0)] = 1.5 + r[eq.R_basis.get_idx(3, 1, -1)] = 0.02 + z[eq.Z_basis.get_idx(3, -1, -1)] = 0.03 + eq.R_lmn, eq.Z_lmn = r, z + return eq + + +@pytest.mark.parametrize('sym', [False, True]) +def test_stellarator_source_from_desc(desc_equilibrium, sym): + from desc.grid import Grid + + eq = desc_equilibrium + if sym: + eq.change_resolution(sym=True) src = openmc.StellaratorSource.from_desc( - 'desc_test.h5', - emission_density=lambda r: 1.0 - r**2, - energy=openmc.stats.delta_function(14.07e6), - n_rho=9, - ) - assert src.num_field_periods == nfp - rho = np.linspace(0.0, 1.0, 9) - np.testing.assert_allclose(src.rho, rho) - assert src.rmns is None # symmetric equilibrium - - # Combined form: R = R0 + a*rho*cos(theta - Nfp*zeta), Z = a*rho*sin(...) - k00 = np.flatnonzero((src.mode_m == 0) & (src.mode_n == 0))[0] - k11 = np.flatnonzero((src.mode_m == 1) & (src.mode_n == 1))[0] - np.testing.assert_allclose(src.rmnc[:, k00], 100.0 * r0_m) - np.testing.assert_allclose(src.rmnc[:, k11], 100.0 * a_m * rho, atol=1e-12) - np.testing.assert_allclose(src.zmns[:, k11], 100.0 * a_m * rho, atol=1e-12) - np.testing.assert_allclose(src.zmns[:, k00], 0.0, atol=1e-12) + eq, lambda r: 1-r**2, openmc.stats.delta_function(14.07e6), n_rho=7) + assert src.num_field_periods == eq.NFP + np.testing.assert_allclose(src.emission_density, 1-src.rho**2) + assert (src.rmns is None) == sym + + # Compare physical surfaces against DESC's full basis evaluation, which + # independently exercises radial evaluation, mode signs, and units. + rng = np.random.default_rng(640) + theta, zeta = rng.uniform(0, 2*np.pi, (2, 30)) + phase = src.mode_m[:, None]*theta - src.mode_n[:, None]*eq.NFP*zeta + for i, rho in enumerate(src.rho): + nodes = np.column_stack([np.full(theta.size, rho), theta, zeta]) + # Supplying a Grid preserves rho=0 without the default axis shift. + grid = Grid(nodes) + expected_r = np.asarray(eq.R_basis.evaluate(grid) @ eq.R_lmn)*100 + expected_z = np.asarray(eq.Z_basis.evaluate(grid) @ eq.Z_lmn)*100 + r = src.rmnc[i] @ np.cos(phase) + z = src.zmns[i] @ np.sin(phase) + if not sym: + r += src.rmns[i] @ np.sin(phase) + z += src.zmnc[i] @ np.cos(phase) + np.testing.assert_allclose(r, expected_r, atol=1e-10, rtol=1e-12) + np.testing.assert_allclose(z, expected_z, atol=1e-10, rtol=1e-12) + + +@pytest.mark.parametrize('family', [False, True]) +def test_stellarator_source_desc_file(desc_equilibrium, tmp_path, family): + from desc.equilibrium import EquilibriaFamily + + eq = desc_equilibrium + obj = EquilibriaFamily(eq.copy(), eq) if family else eq + if family: + obj[0].R_lmn = obj[0].R_lmn * 2 + path = tmp_path/'equilibrium.h5' + obj.save(path) + args = dict(emission_density=lambda r: 1-r**2, + energy=openmc.stats.delta_function(14.07e6), n_rho=5) + reference = openmc.StellaratorSource.from_desc(eq, **args) + for value in (obj, path, str(path)): + source = openmc.StellaratorSource.from_desc(value, **args) + assert source.num_field_periods == reference.num_field_periods + for name in ('rho', 'mode_m', 'mode_n', 'rmnc', 'rmns', 'zmnc', 'zmns'): + np.testing.assert_allclose(getattr(source, name), + getattr(reference, name)) + + +def test_stellarator_source_desc_invalid(desc_equilibrium, tmp_path): + from desc.equilibrium import EquilibriaFamily + + args = dict(emission_density=lambda r: 1-r**2, + energy=openmc.stats.delta_function(14.07e6)) + with pytest.raises(ValueError, match='empty'): + openmc.StellaratorSource.from_desc(EquilibriaFamily(), **args) + with pytest.raises(TypeError, match='eq'): + openmc.StellaratorSource.from_desc(object(), **args) + with pytest.raises(TypeError, match='n_rho'): + openmc.StellaratorSource.from_desc(desc_equilibrium, n_rho=3.5, **args) + with pytest.raises(ValueError, match='n_rho'): + openmc.StellaratorSource.from_desc(desc_equilibrium, n_rho=1, **args) + with pytest.raises(FileNotFoundError): + openmc.StellaratorSource.from_desc(tmp_path/'missing.h5', **args) + + +def test_stellarator_source_desc_missing_dependency(monkeypatch): + import builtins + + original_import = builtins.__import__ + + def without_desc(name, *args, **kwargs): + if name == 'desc' or name.startswith('desc.'): + raise ModuleNotFoundError("No module named 'desc'", name='desc') + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, '__import__', without_desc) + with pytest.raises(ImportError, match='pip install openmc'): + openmc.StellaratorSource.from_desc( + 'unused.h5', [1, 1], openmc.stats.delta_function(14.07e6)) From 16420f793a5aec4d1435685f46b6307dc95360d6 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Wed, 14 Oct 2026 23:39:34 -0500 Subject: [PATCH 21/33] Expand DESC adapter coverage and document startup overhead --- docs/source/usersguide/settings.rst | 2 ++ tests/unit_tests/test_source_stellarator.py | 21 ++++++++++++++++----- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/docs/source/usersguide/settings.rst b/docs/source/usersguide/settings.rst index 28d5413915b..9439543750a 100644 --- a/docs/source/usersguide/settings.rst +++ b/docs/source/usersguide/settings.rst @@ -381,6 +381,8 @@ available through ``pip install openmc[desc]`` or ``conda install -c conda-forge desc-opt``. DESC loads its own files and performs the spectral conversion. An equilibrium family uses its last member. DESC is not required to import OpenMC, read the resulting source XML, or run transport. +The first call can incur DESC/JAX initialization and compilation overhead; +reusing an already-loaded equilibrium avoids repeated file loading. The ``emission_density`` argument accepts either a callable evaluated on the radial grid or an array with one value per radial surface; only its shape diff --git a/tests/unit_tests/test_source_stellarator.py b/tests/unit_tests/test_source_stellarator.py index fc8e0d2b20d..a1343b4a0b5 100644 --- a/tests/unit_tests/test_source_stellarator.py +++ b/tests/unit_tests/test_source_stellarator.py @@ -258,20 +258,27 @@ def test_stellarator_source_from_vmec(run_in_tmpdir): assert src.rmns is None -@pytest.fixture -def desc_equilibrium(): +@pytest.fixture(params=[0, 1], ids=['axisymmetric', 'stellarator']) +def desc_equilibrium(request): pytest.importorskip('desc') from desc.equilibrium import Equilibrium # Circular torus plus a nontrivial l=3 radial contribution. - eq = Equilibrium(L=3, M=2, N=1, NFP=3, sym=False) + eq = Equilibrium(L=3, M=2, N=request.param, NFP=3, sym=False) r = np.zeros(eq.R_basis.num_modes) z = np.zeros(eq.Z_basis.num_modes) r[eq.R_basis.get_idx(0, 0, 0)] = 6.0 r[eq.R_basis.get_idx(1, 1, 0)] = 1.5 z[eq.Z_basis.get_idx(1, -1, 0)] = 1.5 - r[eq.R_basis.get_idx(3, 1, -1)] = 0.02 - z[eq.Z_basis.get_idx(3, -1, -1)] = 0.03 + r[eq.R_basis.get_idx(3, -1, 0)] = 0.02 + z[eq.Z_basis.get_idx(2, 0, 0)] = 0.03 + if eq.N: + # Exercise all product-basis sign combinations and m=0 harmonics. + for m, n in [(1, 1), (1, -1), (-1, 1), (-1, -1)]: + r[eq.R_basis.get_idx(3, m, n)] = 0.02 + z[eq.Z_basis.get_idx(3, m, n)] = 0.03 + r[eq.R_basis.get_idx(0, 0, -1)] = 0.01 + z[eq.Z_basis.get_idx(0, 0, 1)] = 0.01 eq.R_lmn, eq.Z_lmn = r, z return eq @@ -345,6 +352,10 @@ def test_stellarator_source_desc_invalid(desc_equilibrium, tmp_path): openmc.StellaratorSource.from_desc(desc_equilibrium, n_rho=1, **args) with pytest.raises(FileNotFoundError): openmc.StellaratorSource.from_desc(tmp_path/'missing.h5', **args) + path = tmp_path/'surface.h5' + desc_equilibrium.get_surface_at(rho=1.0).save(path) + with pytest.raises(TypeError, match='eq'): + openmc.StellaratorSource.from_desc(path, **args) def test_stellarator_source_desc_missing_dependency(monkeypatch): From 910f058cbe78016b8e71fb82461f735b66d1456a Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Oct 2026 00:07:29 -0500 Subject: [PATCH 22/33] Extend coverage to NetCDF4 wout files --- tests/unit_tests/test_source_stellarator.py | 29 ++++++++++++++------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/tests/unit_tests/test_source_stellarator.py b/tests/unit_tests/test_source_stellarator.py index a1343b4a0b5..2055f53da02 100644 --- a/tests/unit_tests/test_source_stellarator.py +++ b/tests/unit_tests/test_source_stellarator.py @@ -220,8 +220,10 @@ def test_stellarator_source_rotating_ellipse_sampling(run_in_tmpdir): assert_sample_mean(z**2, ref['Z2']) -def test_stellarator_source_from_vmec(run_in_tmpdir): +@pytest.mark.parametrize('file_format', ['netcdf3', 'hdf5']) +def test_stellarator_source_from_vmec(run_in_tmpdir, file_format): """Round-trip a synthetic rotating-ellipse wout file through from_vmec.""" + import h5py from scipy.io import netcdf_file ns, nfp = 16, 5 @@ -233,15 +235,22 @@ def test_stellarator_source_from_vmec(run_in_tmpdir): rmnc = np.column_stack([np.full(ns, r0_m), a_m * rho, e_m * rho]) zmns = np.column_stack([np.zeros(ns), a_m * rho, -e_m * rho]) - with netcdf_file('wout_test.nc', 'w') as ds: - ds.createDimension('radius', ns) - ds.createDimension('mn_mode', len(xm)) - ds.createVariable('rmnc', 'd', ('radius', 'mn_mode'))[:] = rmnc - ds.createVariable('zmns', 'd', ('radius', 'mn_mode'))[:] = zmns - ds.createVariable('xm', 'd', ('mn_mode',))[:] = xm - ds.createVariable('xn', 'd', ('mn_mode',))[:] = xn - ds.createVariable('nfp', 'i', ()).data.fill(nfp) - ds.createVariable('lasym__logical__', 'i', ()).data.fill(0) + if file_format == 'netcdf3': + with netcdf_file('wout_test.nc', 'w') as ds: + ds.createDimension('radius', ns) + ds.createDimension('mn_mode', len(xm)) + ds.createVariable('rmnc', 'd', ('radius', 'mn_mode'))[:] = rmnc + ds.createVariable('zmns', 'd', ('radius', 'mn_mode'))[:] = zmns + ds.createVariable('xm', 'd', ('mn_mode',))[:] = xm + ds.createVariable('xn', 'd', ('mn_mode',))[:] = xn + ds.createVariable('nfp', 'i', ()).data.fill(nfp) + ds.createVariable('lasym__logical__', 'i', ()).data.fill(0) + else: + # Exercise the HDF5 fallback without constructing a full NetCDF4 file. + with h5py.File('wout_test.nc', 'w') as ds: + for name, value in dict(rmnc=rmnc, zmns=zmns, xm=xm, xn=xn, + nfp=nfp, lasym__logical__=0).items(): + ds[name] = value src = openmc.StellaratorSource.from_vmec( 'wout_test.nc', From 755d0f42f1b380c865f9df8bc1395b531ab1e1fc Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Oct 2026 00:27:14 -0500 Subject: [PATCH 23/33] Install DESC in CI --- tools/ci/gha-install.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/ci/gha-install.sh b/tools/ci/gha-install.sh index 8a125d908b2..37807e381e8 100755 --- a/tools/ci/gha-install.sh +++ b/tools/ci/gha-install.sh @@ -53,4 +53,4 @@ fi python tools/ci/gha-install.py # Install Python API in editable mode -pip install -e .[test,vtk,ci] +pip install -e '.[test,vtk,ci,desc]' From 581cf595bb76f804d6a3b98950365bb340507940 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Oct 2026 16:15:30 -0500 Subject: [PATCH 24/33] Documentation improvements --- docs/source/io_formats/settings.rst | 26 +-- docs/source/usersguide/settings.rst | 255 +++++++++++++++------------- 2 files changed, 148 insertions(+), 133 deletions(-) diff --git a/docs/source/io_formats/settings.rst b/docs/source/io_formats/settings.rst index 81cd2d9d213..5a630cac651 100644 --- a/docs/source/io_formats/settings.rst +++ b/docs/source/io_formats/settings.rst @@ -1036,8 +1036,9 @@ attributes/sub-elements: ```` 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: + a Miller-style flux-surface parameterization (see + :ref:`usersguide_tokamak_source`). The following sub-elements are used instead + of the ``space`` element: :major_radius: The major radius :math:`R_0` of the plasma in [cm]. @@ -1109,23 +1110,10 @@ attributes/sub-elements: *Default*: particles are born at :math:`t=0` - For a source with ``type="stellarator"``, the spatial distribution is - described by the flux-surface Fourier representation shared by the VMEC and - DESC equilibrium codes, using flux coordinates :math:`(\rho, \theta, \zeta)` - where :math:`\rho = \sqrt{s}` is the square root of the normalized toroidal - flux, :math:`\theta` is the poloidal angle, and :math:`\zeta` is the toroidal - (cylindrical) angle: - - .. math:: - - \begin{aligned} - R &= \sum_k \left[ R^c_k(\rho) \cos(m_k\theta - n_k N_{fp} \zeta) - + R^s_k(\rho) \sin(m_k\theta - n_k N_{fp}\zeta) \right] \\ - Z &= \sum_k \left[ Z^s_k(\rho) \sin(m_k\theta - n_k N_{fp}\zeta) - + Z^c_k(\rho) \cos(m_k\theta - n_k N_{fp}\zeta) \right] - \end{aligned} - - The following sub-elements are used instead of the ``space`` element: + For a source with ``type="stellarator"``, the spatial distribution uses the + flux-surface Fourier representation described in + :ref:`usersguide_stellarator_source`. The following sub-elements are used + instead of the ``space`` element: :num_field_periods: The number of field periods :math:`N_{fp}`. diff --git a/docs/source/usersguide/settings.rst b/docs/source/usersguide/settings.rst index 9439543750a..c793f815955 100644 --- a/docs/source/usersguide/settings.rst +++ b/docs/source/usersguide/settings.rst @@ -170,14 +170,19 @@ External Source Distributions External source distributions can be specified through the :attr:`Settings.source` attribute. If you have a single external source, you can create an instance of any of the subclasses of :class:`openmc.SourceBase` -(:class:`openmc.IndependentSource`, :class:`openmc.FileSource`, -:class:`openmc.CompiledSource`) and use it to set the :attr:`Settings.source` -attribute. If you have multiple external sources with varying source strengths, -:attr:`Settings.source` should be set to a list of :class:`openmc.SourceBase` -objects. - -The :class:`openmc.IndependentSource` class is the primary class for defining -source distributions and has four main attributes that one can set: +(:class:`openmc.IndependentSource`, :class:`openmc.MeshSource`, +:class:`openmc.FileSource`, :class:`openmc.TokamakSource`, +:class:`openmc.StellaratorSource`, or :class:`openmc.CompiledSource`) and use it +to set the :attr:`Settings.source` attribute. If you have multiple external +sources with varying source strengths, :attr:`Settings.source` should be set to +a list of :class:`openmc.SourceBase` objects. + +Sources with Independent Distributions +-------------------------------------- + +The :class:`openmc.IndependentSource` class samples position, direction, energy, +and emission time independently of one another. Each distribution is specified +separately through one of four attributes: :attr:`IndependentSource.space`, which defines the spatial distribution, :attr:`IndependentSource.angle`, which defines the angular distribution, :attr:`IndependentSource.energy`, which defines the energy distribution, and @@ -239,6 +244,28 @@ would run:: source.time = openmc.stats.Uniform(0, 1e-6) settings.source = source +Sampling from an :class:`openmc.IndependentSource` may be biased +for local or global variance reduction by modifying the +:attr:`~openmc.IndependentSource.bias` attribute of each of its four main +distributions. Further discussion of source biasing can be found in +:ref:`source_biasing`. + +Finally, the :attr:`IndependentSource.particle` attribute can be used to +indicate the source should be composed of particles other than neutrons. For +example, the following would generate a photon source:: + + source = openmc.IndependentSource() + source.particle = 'photon' + ... + + settings.source = source + +For a full list of all classes related to statistical distributions, see +:ref:`pythonapi_stats`. + +Combining Sources +----------------- + All subclasses of :class:`openmc.SourceBase` have a :attr:`SourceBase.strength` attribute that indicates the relative strength of a source distribution if multiple are used. For example, to create two sources, one that should be @@ -272,24 +299,96 @@ option:: settings.source = [src1, src2] settings.uniform_source_sampling = True -Additionally, sampling from an :class:`openmc.IndependentSource` may be biased -for local or global variance reduction by modifying the -:attr:`~openmc.IndependentSource.bias` attribute of each of its four main -distributions. Further discussion of source biasing can be found in -:ref:`source_biasing`. +File-based Sources +------------------ -Finally, the :attr:`IndependentSource.particle` attribute can be used to -indicate the source should be composed of particles other than neutrons. For -example, the following would generate a photon source:: +OpenMC can use a pregenerated HDF5 source file through the +:class:`openmc.FileSource` class:: - source = openmc.IndependentSource() - source.particle = 'photon' - ... + settings.source = openmc.FileSource('source.h5') - settings.source = source +Statepoint and source files are generated automatically when a simulation is run +and can be used as the starting source in a new simulation. Alternatively, a +source file can be manually generated with the :func:`openmc.write_source_file` +function. This is particularly useful for coupling OpenMC with another program +that generates a source to be used in OpenMC. -For a full list of all classes related to statistical distributions, see -:ref:`pythonapi_stats`. +Surface Sources ++++++++++++++++ + +A source file based on particles that cross one or more surfaces can be +generated during a simulation using the :attr:`Settings.surf_source_write` +attribute:: + + settings.surf_source_write = { + 'surfaces_ids': [1, 2, 3], + 'max_particles': 10000 + } + +In this example, at most 10,000 source particles are stored when particles cross +surfaces with IDs of 1, 2, or 3. If no surface IDs are declared, particles +crossing any surface of the model will be banked:: + + settings.surf_source_write = {'max_particles': 10000} + +A cell ID can also be used to bank particles that are crossing any surface of +a cell that particles are either coming from or going to:: + + settings.surf_source_write = {'cell': 1, 'max_particles': 10000} + +In this example, particles that are crossing any surface that bounds cell 1 will +be banked excluding any surface that does not use a 'transmission' or 'vacuum' +boundary condition. + +.. note:: Surfaces with boundary conditions that are not "transmission" or "vacuum" + are not eligible to store any particles when using ``cell``, ``cellfrom`` + or ``cellto`` attributes. It is recommended to use surface IDs instead. + +Surface IDs can be used in combination with a cell ID:: + + settings.surf_source_write = { + 'cell': 1, + 'surfaces_ids': [1, 2, 3], + 'max_particles': 10000 + } + +In that case, only particles that are crossing the declared surfaces coming from +cell 1 or going to cell 1 will be banked. To account specifically for particles +leaving or entering a given cell, ``cellfrom`` and ``cellto`` are also available +to respectively account for particles coming from a cell:: + + settings.surf_source_write = { + 'cellfrom': 1, + 'max_particles': 10000 + } + +or particles going to a cell:: + + settings.surf_source_write = { + 'cellto': 1, + 'max_particles': 10000 + } + +.. note:: The ``cell``, ``cellfrom`` and ``cellto`` attributes cannot be + used simultaneously. + +To generate more than one surface source files when the maximum number of stored +particles is reached, ``max_source_files`` is available. The surface source bank +will be cleared in simulation memory each time a surface source file is written. +As an example, to write a maximum of three surface source files::: + + settings.surf_source_write = { + 'surfaces_ids': [1, 2, 3], + 'max_particles': 10000, + 'max_source_files': 3 + } + +A surface source file is read back like any other source file, with the +:class:`openmc.FileSource` class:: + + settings.source = openmc.FileSource('surface_source.h5') + +.. _usersguide_tokamak_source: Tokamak Plasma Sources ---------------------- @@ -338,6 +437,8 @@ 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. +.. _usersguide_stellarator_source: + Stellarator Plasma Sources -------------------------- @@ -352,9 +453,24 @@ the average minor radius), :math:`\theta` is the poloidal angle, and angle in both codes). The user provides the Fourier coefficients of :math:`R(\rho,\theta,\zeta)` and :math:`Z(\rho,\theta,\zeta)` on a radial grid together with an emission density :math:`S(\rho)` that is constant on flux -surfaces. Fourier coefficients and emission are interpolated linearly in -:math:`\rho`. The cubic marginal density :math:`S(\rho)\, V'(\rho)` is -integrated analytically to select an interval, then sampled within that interval +surfaces. The cylindrical coordinates are represented as + +.. math:: + + \begin{aligned} + R &= \sum_k \left[ R^c_k(\rho) \cos(m_k\theta - n_k N_{fp} \zeta) + + R^s_k(\rho) \sin(m_k\theta - n_k N_{fp}\zeta) \right] \\ + Z &= \sum_k \left[ Z^s_k(\rho) \sin(m_k\theta - n_k N_{fp}\zeta) + + Z^c_k(\rho) \cos(m_k\theta - n_k N_{fp}\zeta) \right] + \end{aligned} + +where :math:`k` indexes Fourier modes, :math:`m_k` and :math:`n_k` are the +poloidal and toroidal mode numbers, and :math:`N_{fp}` is the number of field +periods. Superscripts :math:`c` and :math:`s` denote cosine and sine coefficients. + +Fourier coefficients and emission are interpolated linearly in :math:`\rho`. +The cubic marginal density :math:`S(\rho)\, V'(\rho)` is integrated analytically +to select an interval, then sampled within that interval by rejection against a bound derived from its Bernstein coefficients. Rejected radial proposals are resampled within the same interval. The two angles are sampled from the conditional distribution :math:`p(\theta,\zeta|\rho) \propto @@ -392,95 +508,6 @@ other sources. As with :class:`openmc.TokamakSource`, the ``energy`` argument accepts a single distribution or one distribution per radial grid point, and a ``time`` distribution may optionally be given. -File-based Sources ------------------- - -OpenMC can use a pregenerated HDF5 source file through the -:class:`openmc.FileSource` class:: - - settings.source = openmc.FileSource('source.h5') - -Statepoint and source files are generated automatically when a simulation is run -and can be used as the starting source in a new simulation. Alternatively, a -source file can be manually generated with the :func:`openmc.write_source_file` -function. This is particularly useful for coupling OpenMC with another program -that generates a source to be used in OpenMC. - -Surface Sources -+++++++++++++++ - -A source file based on particles that cross one or more surfaces can be -generated during a simulation using the :attr:`Settings.surf_source_write` -attribute:: - - settings.surf_source_write = { - 'surfaces_ids': [1, 2, 3], - 'max_particles': 10000 - } - -In this example, at most 10,000 source particles are stored when particles cross -surfaces with IDs of 1, 2, or 3. If no surface IDs are declared, particles -crossing any surface of the model will be banked:: - - settings.surf_source_write = {'max_particles': 10000} - -A cell ID can also be used to bank particles that are crossing any surface of -a cell that particles are either coming from or going to:: - - settings.surf_source_write = {'cell': 1, 'max_particles': 10000} - -In this example, particles that are crossing any surface that bounds cell 1 will -be banked excluding any surface that does not use a 'transmission' or 'vacuum' -boundary condition. - -.. note:: Surfaces with boundary conditions that are not "transmission" or "vacuum" - are not eligible to store any particles when using ``cell``, ``cellfrom`` - or ``cellto`` attributes. It is recommended to use surface IDs instead. - -Surface IDs can be used in combination with a cell ID:: - - settings.surf_source_write = { - 'cell': 1, - 'surfaces_ids': [1, 2, 3], - 'max_particles': 10000 - } - -In that case, only particles that are crossing the declared surfaces coming from -cell 1 or going to cell 1 will be banked. To account specifically for particles -leaving or entering a given cell, ``cellfrom`` and ``cellto`` are also available -to respectively account for particles coming from a cell:: - - settings.surf_source_write = { - 'cellfrom': 1, - 'max_particles': 10000 - } - -or particles going to a cell:: - - settings.surf_source_write = { - 'cellto': 1, - 'max_particles': 10000 - } - -.. note:: The ``cell``, ``cellfrom`` and ``cellto`` attributes cannot be - used simultaneously. - -To generate more than one surface source files when the maximum number of stored -particles is reached, ``max_source_files`` is available. The surface source bank -will be cleared in simulation memory each time a surface source file is written. -As an example, to write a maximum of three surface source files::: - - settings.surf_source_write = { - 'surfaces_ids': [1, 2, 3], - 'max_particles': 10000, - 'max_source_files': 3 - } - -A surface source file is read back like any other source file, with the -:class:`openmc.FileSource` class:: - - settings.source = openmc.FileSource('surface_source.h5') - .. _compiled_source: Compiled Sources From b05c5cfc8c4ca3b5fb6d5a8e0abacbc651bc1a63 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Oct 2026 16:24:52 -0500 Subject: [PATCH 25/33] Small updates in StellaratorSource --- openmc/source.py | 48 +++++++++++++++++------------------------------- 1 file changed, 17 insertions(+), 31 deletions(-) diff --git a/openmc/source.py b/openmc/source.py index 866fd951993..6a66ae2fdbd 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -5,7 +5,7 @@ import os from pathlib import Path import warnings -from typing import Any +from typing import Any, Self import lxml.etree as ET import numpy as np @@ -1659,7 +1659,7 @@ def from_vmec( emission_density, energy: Univariate | Sequence[Univariate], **kwargs - ) -> StellaratorSource: + ) -> Self: """Generate a stellarator source from a VMEC wout file. The VMEC radial grid is uniform in the normalized toroidal flux @@ -1755,7 +1755,7 @@ def from_desc( energy: Univariate | Sequence[Univariate], n_rho: int = 33, **kwargs - ) -> StellaratorSource: + ) -> Self: """Generate a stellarator source from a DESC equilibrium. DESC represents :math:`R` and :math:`Z` in a Fourier-Zernike basis, @@ -1791,25 +1791,17 @@ def from_desc( openmc.StellaratorSource """ - try: - from desc.basis import zernike_radial - from desc.equilibrium import Equilibrium, EquilibriaFamily - from desc.io import load - from desc.vmec_utils import ptolemy_identity_rev - except ModuleNotFoundError as exc: - if exc.name != 'desc': - raise - raise ImportError( - "StellaratorSource.from_desc requires DESC. " - "Install it with 'pip install openmc[desc]' or " - "'pip install desc-opt'.") from exc + from desc.basis import zernike_radial + from desc.equilibrium import Equilibrium, EquilibriaFamily + from desc.io import load + from desc.vmec_utils import ptolemy_identity_rev cv.check_type('n_rho', n_rho, Integral) cv.check_greater_than('n_rho', n_rho, 1) rho = np.linspace(0.0, 1.0, n_rho) if isinstance(eq, (str, os.PathLike)): eq = load(input_path(eq), file_format='hdf5') - if isinstance(eq, EquilibriaFamily): + elif isinstance(eq, EquilibriaFamily): if not len(eq): raise ValueError("The DESC equilibrium family is empty") eq = eq[-1] @@ -1818,15 +1810,13 @@ def from_desc( def convert(name): basis = getattr(eq, f'{name}_basis') l, m, _ = basis.modes.T # noqa: E741 - modes, inverse = np.unique(basis.modes[:, 1:], axis=0, - return_inverse=True) + modes, inverse = np.unique(basis.modes[:, 1:], axis=0, return_inverse=True) radial = np.asarray(zernike_radial(rho[:, None], l, m)) values = np.zeros((n_rho, len(modes))) # Sum radial modes sharing the same angular mode on each surface. np.add.at(values.T, inverse, (radial * np.asarray(getattr(eq, f'{name}_lmn'))).T) - m, n, sine, cosine = ptolemy_identity_rev( - modes[:, 0], modes[:, 1], values) + m, n, sine, cosine = ptolemy_identity_rev(modes[:, 0], modes[:, 1], values) return list(zip(m, n)), sine, cosine r_modes, r_sine, r_cosine = convert('R') @@ -1876,16 +1866,12 @@ def populate_xml_element(self, element): """ self._validate() - ET.SubElement(element, "num_field_periods").text = \ - str(self.num_field_periods) - ET.SubElement(element, "rho").text = \ - ' '.join(str(r) for r in self.rho) - ET.SubElement(element, "emission_density").text = \ - ' '.join(str(s) for s in self.emission_density) - ET.SubElement(element, "mode_m").text = \ - ' '.join(str(m) for m in self.mode_m) - ET.SubElement(element, "mode_n").text = \ - ' '.join(str(n) for n in self.mode_n) + ET.SubElement(element, "num_field_periods").text = str(self.num_field_periods) + ET.SubElement(element, "rho").text = ' '.join(str(r) for r in self.rho) + ET.SubElement(element, "emission_density").text = ' '.join( + str(s) for s in self.emission_density) + ET.SubElement(element, "mode_m").text = ' '.join(str(m) for m in self.mode_m) + ET.SubElement(element, "mode_n").text = ' '.join(str(n) for n in self.mode_n) # Coefficient tables flattened row-major (surface index varies slowest) for name in ('rmnc', 'zmns', 'rmns', 'zmnc'): @@ -1903,7 +1889,7 @@ def populate_xml_element(self, element): element.append(self.time.to_xml_element('time')) @classmethod - def from_xml_element(cls, elem: ET.Element) -> StellaratorSource: + def from_xml_element(cls, elem: ET.Element) -> Self: """Generate stellarator source from an XML element Parameters From f785ad1f792becf1883f61a024d11aec2e5ff6d7 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Oct 2026 16:47:57 -0500 Subject: [PATCH 26/33] Remove missing dependency test --- tests/unit_tests/test_source_stellarator.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/tests/unit_tests/test_source_stellarator.py b/tests/unit_tests/test_source_stellarator.py index 2055f53da02..5d24c68e189 100644 --- a/tests/unit_tests/test_source_stellarator.py +++ b/tests/unit_tests/test_source_stellarator.py @@ -365,19 +365,3 @@ def test_stellarator_source_desc_invalid(desc_equilibrium, tmp_path): desc_equilibrium.get_surface_at(rho=1.0).save(path) with pytest.raises(TypeError, match='eq'): openmc.StellaratorSource.from_desc(path, **args) - - -def test_stellarator_source_desc_missing_dependency(monkeypatch): - import builtins - - original_import = builtins.__import__ - - def without_desc(name, *args, **kwargs): - if name == 'desc' or name.startswith('desc.'): - raise ModuleNotFoundError("No module named 'desc'", name='desc') - return original_import(name, *args, **kwargs) - - monkeypatch.setattr(builtins, '__import__', without_desc) - with pytest.raises(ImportError, match='pip install openmc'): - openmc.StellaratorSource.from_desc( - 'unused.h5', [1, 1], openmc.stats.delta_function(14.07e6)) From 64ead4c311c8e6537881d9118a0c70430c0625f9 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Oct 2026 16:49:18 -0500 Subject: [PATCH 27/33] Fix 'if' statement I missed up --- openmc/source.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openmc/source.py b/openmc/source.py index 6a66ae2fdbd..f850795256d 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -1801,7 +1801,7 @@ def from_desc( rho = np.linspace(0.0, 1.0, n_rho) if isinstance(eq, (str, os.PathLike)): eq = load(input_path(eq), file_format='hdf5') - elif isinstance(eq, EquilibriaFamily): + if isinstance(eq, EquilibriaFamily): if not len(eq): raise ValueError("The DESC equilibrium family is empty") eq = eq[-1] From 636ae8d55e9f4fe840a631edc0144909a85028a2 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Oct 2026 17:02:57 -0500 Subject: [PATCH 28/33] Avoid potentially high temporary memory use during initialization --- src/source.cpp | 102 +++++++++----------- tests/unit_tests/test_source_stellarator.py | 4 +- 2 files changed, 49 insertions(+), 57 deletions(-) diff --git a/src/source.cpp b/src/source.cpp index d297ea902e3..bff4166a717 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -1361,13 +1361,11 @@ void StellaratorSource::precompute_sampling_distributions() } } - // Evaluate R, Z, dR/dtheta, dZ/dtheta on the angle grid for every radial - // surface. Radial derivatives within a bin follow from finite differences - // of these grids because the Fourier evaluation is linear in the - // coefficients. - vector R_s(n_rho * n_grid), Z_s(n_rho * n_grid); - vector Rt_s(n_rho * n_grid), Zt_s(n_rho * n_grid); - for (size_t i = 0; i < n_rho; ++i) { + // Keep only adjacent surfaces, reusing the older surface's storage after + // processing each interval. Each surface is evaluated once. + vector R_s(2 * n_grid), Z_s(2 * n_grid); + vector Rt_s(2 * n_grid), Zt_s(2 * n_grid); + auto evaluate_surface = [&](size_t i) { const double* rc = &rmnc_[i * n_modes]; const double* zs = &zmns_[i * n_modes]; const double* rs = asym_ ? &rmns_[i * n_modes] : nullptr; @@ -1395,90 +1393,82 @@ void StellaratorSource::precompute_sampling_distributions() Zt -= m * zc[k] * s; } } - size_t p = i * n_grid + it * n_zeta + jz; + if (R <= 0.0) { + fatal_error( + "StellaratorSource: flux surfaces reach R <= 0; check the " + "Fourier coefficients (units should be cm)."); + } + size_t p = (i % 2) * n_grid + it * n_zeta + jz; R_s[p] = R; Z_s[p] = Z; Rt_s[p] = Rt; Zt_s[p] = Zt; } } - } + }; - // Sanity check: the map must stay on the +R side of the cylindrical axis - double R_min = *std::min_element(R_s.begin(), R_s.end()); - if (R_min <= 0.0) { - fatal_error("StellaratorSource: flux surfaces reach R <= 0; check the " - "Fourier coefficients (units should be cm)."); - } - - // Pass 1: signed volume integral to fix the global Jacobian sign. Both - // (rho, theta, zeta) handednesses occur in practice (VMEC's Jacobian is - // conventionally negative), so the sign is detected rather than assumed. + // Accumulate signed integrals and bounds for both orientations. The total + // volume determines the Jacobian sign after all intervals are processed. double quad_wgt = 4.0 * PI * PI / n_grid; // full-torus angular weight double v_tot = 0.0; - for (size_t b = 0; b < n_bins; ++b) { - double inv_drho = 1.0 / (rho_[b + 1] - rho_[b]); - const double* R0 = &R_s[b * n_grid]; - const double* R1 = &R_s[(b + 1) * n_grid]; - const double* Z0 = &Z_s[b * n_grid]; - const double* Z1 = &Z_s[(b + 1) * n_grid]; - const double* Rt0 = &Rt_s[b * n_grid]; - const double* Rt1 = &Rt_s[(b + 1) * n_grid]; - const double* Zt0 = &Zt_s[b * n_grid]; - const double* Zt1 = &Zt_s[(b + 1) * n_grid]; - double sum = 0.0; - for (size_t p = 0; p < n_grid; ++p) { - double Rr = (R1[p] - R0[p]) * inv_drho; - double Zr = (Z1[p] - Z0[p]) * inv_drho; - // Midpoint of the bin (all factors are linear in t) - double Rh = 0.5 * (R0[p] + R1[p]); - double Rth = 0.5 * (Rt0[p] + Rt1[p]); - double Zth = 0.5 * (Zt0[p] + Zt1[p]); - sum += Rh * (Rr * Zth - Rth * Zr); - } - v_tot += sum * quad_wgt * (rho_[b + 1] - rho_[b]); - } - jacobian_sign_ = (v_tot >= 0.0) ? 1.0 : -1.0; - - // Pass 2: differential volume V'(rho) at surfaces and bin midpoints, - // rejection majorants, and Jacobian sign-consistency check vector vp_lo(n_bins), vp_mid(n_bins), vp_hi(n_bins); + vector negative_envelope(n_bins); envelope_.assign(n_bins, 0.0); double f_peak = 0.0; double f_min = 0.0; + evaluate_surface(0); for (size_t b = 0; b < n_bins; ++b) { + evaluate_surface(b + 1); double inv_drho = 1.0 / (rho_[b + 1] - rho_[b]); - const double* R0 = &R_s[b * n_grid]; - const double* R1 = &R_s[(b + 1) * n_grid]; - const double* Z0 = &Z_s[b * n_grid]; - const double* Z1 = &Z_s[(b + 1) * n_grid]; - const double* Rt0 = &Rt_s[b * n_grid]; - const double* Rt1 = &Rt_s[(b + 1) * n_grid]; - const double* Zt0 = &Zt_s[b * n_grid]; - const double* Zt1 = &Zt_s[(b + 1) * n_grid]; + size_t lo = (b % 2) * n_grid; + size_t hi = ((b + 1) % 2) * n_grid; + const double* R0 = &R_s[lo]; + const double* R1 = &R_s[hi]; + const double* Z0 = &Z_s[lo]; + const double* Z1 = &Z_s[hi]; + const double* Rt0 = &Rt_s[lo]; + const double* Rt1 = &Rt_s[hi]; + const double* Zt0 = &Zt_s[lo]; + const double* Zt1 = &Zt_s[hi]; double s0 = 0.0, sh = 0.0, s1 = 0.0; double env = 0.0; + double negative_env = 0.0; for (size_t p = 0; p < n_grid; ++p) { double Rr = (R1[p] - R0[p]) * inv_drho; double Zr = (Z1[p] - Z0[p]) * inv_drho; - double f0 = jacobian_sign_ * R0[p] * (Rr * Zt0[p] - Rt0[p] * Zr); - double f1 = jacobian_sign_ * R1[p] * (Rr * Zt1[p] - Rt1[p] * Zr); + double f0 = R0[p] * (Rr * Zt0[p] - Rt0[p] * Zr); + double f1 = R1[p] * (Rr * Zt1[p] - Rt1[p] * Zr); double Rh = 0.5 * (R0[p] + R1[p]); double Rth = 0.5 * (Rt0[p] + Rt1[p]); double Zth = 0.5 * (Zt0[p] + Zt1[p]); - double fh = jacobian_sign_ * Rh * (Rr * Zth - Rth * Zr); + double fh = Rh * (Rr * Zth - Rth * Zr); s0 += f0; sh += fh; s1 += f1; env = std::max(env, quad_max01(f0, fh, f1)); + negative_env = std::max(negative_env, quad_max01(-f0, -fh, -f1)); f_peak = std::max({f_peak, f0, fh, f1}); f_min = std::min({f_min, f0, fh, f1}); } vp_lo[b] = s0 * quad_wgt; vp_mid[b] = sh * quad_wgt; vp_hi[b] = s1 * quad_wgt; + v_tot += sh * quad_wgt * (rho_[b + 1] - rho_[b]); // Safety factor covering angular discretization of the majorant envelope_[b] = 1.05 * env; + negative_envelope[b] = 1.05 * negative_env; + } + jacobian_sign_ = (v_tot >= 0.0) ? 1.0 : -1.0; + if (jacobian_sign_ < 0.0) { + envelope_.swap(negative_envelope); + double old_peak = f_peak; + f_peak = -f_min; + f_min = -old_peak; + for (size_t b = 0; b < n_bins; ++b) { + vp_lo[b] = -vp_lo[b]; + vp_mid[b] = -vp_mid[b]; + vp_hi[b] = -vp_hi[b]; + } } if (f_min < -1.0e-6 * f_peak) { fatal_error( diff --git a/tests/unit_tests/test_source_stellarator.py b/tests/unit_tests/test_source_stellarator.py index 5d24c68e189..1d116bdb27a 100644 --- a/tests/unit_tests/test_source_stellarator.py +++ b/tests/unit_tests/test_source_stellarator.py @@ -194,12 +194,14 @@ def test_stellarator_source_axisymmetric_sampling(run_in_tmpdir): assert_sample_mean(major_r, r0 + a**2 / (4.0 * r0)) -def test_stellarator_source_rotating_ellipse_sampling(run_in_tmpdir): +@pytest.mark.parametrize('orientation', [1, -1]) +def test_stellarator_source_rotating_ellipse_sampling(run_in_tmpdir, orientation): """Compare sampled moments of a 3-D configuration against quadrature of the exact density S(rho) * R * |tau|.""" r0, a, e, nfp = 600.0, 150.0, 50.0, 5 rho, mode_m, mode_n, rmnc, zmns = rotating_ellipse_modes( r0, a, e, nfp=nfp) + zmns *= orientation emission_density = 1.0 - rho**2 src = make_source( rho=rho, emission_density=emission_density, mode_m=mode_m, From 2e20ce8375bc1698f396adfeee3246139b24a5c6 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Oct 2026 18:11:04 -0500 Subject: [PATCH 29/33] Remove use of input_path --- openmc/source.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/openmc/source.py b/openmc/source.py index f850795256d..ca2a330da0b 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -1692,7 +1692,7 @@ def from_vmec( openmc.StellaratorSource """ - ds = cls._read_wout(input_path(wout)) + ds = cls._read_wout(wout) rmnc = ds['rmnc'] zmns = ds['zmns'] xm = ds['xm'].astype(int) @@ -1800,7 +1800,7 @@ def from_desc( cv.check_greater_than('n_rho', n_rho, 1) rho = np.linspace(0.0, 1.0, n_rho) if isinstance(eq, (str, os.PathLike)): - eq = load(input_path(eq), file_format='hdf5') + eq = load(eq, file_format='hdf5') if isinstance(eq, EquilibriaFamily): if not len(eq): raise ValueError("The DESC equilibrium family is empty") From c2b003b6bd4cfee56c1894c616fa232200267496 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Oct 2026 18:16:48 -0500 Subject: [PATCH 30/33] Move scipy import to top-level --- openmc/source.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openmc/source.py b/openmc/source.py index ca2a330da0b..a964e6d89f1 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -11,6 +11,7 @@ import numpy as np import h5py import pandas as pd +from scipy.io import netcdf_file import openmc import openmc.checkvalue as cv @@ -1738,7 +1739,6 @@ def _read_wout(path): names = ('rmnc', 'zmns', 'rmns', 'zmnc', 'xm', 'xn', 'nfp', 'lasym__logical__') try: - from scipy.io import netcdf_file with netcdf_file(str(path), mmap=False) as ds: return {k: np.asarray(v[()], dtype=float) for k, v in ds.variables.items() if k in names} From 8c3214ee7a9ef2a012d47f3cdeb6d7d808f7398f Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Oct 2026 18:26:21 -0500 Subject: [PATCH 31/33] Add additional comments in source.cpp --- src/source.cpp | 38 ++++++++++++++++++++++++++++++++++---- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/src/source.cpp b/src/source.cpp index bff4166a717..02010ec5f0a 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -1330,8 +1330,8 @@ void StellaratorSource::precompute_sampling_distributions() // of poloidal degree <= 3*m_max and toroidal degree <= 3*n_max per field // period, so a uniform trapezoidal product rule with more points than the // bandwidth integrates it exactly (discrete Fourier orthogonality). The - // denser grid below also controls the discretization error of the rejection - // majorant, which is additionally covered by a safety factor. + // denser grid below also estimates the angular maximum for rejection + // sampling; a heuristic safety factor allows for peaks between grid points. int m_max = 0; int n_max = 0; for (size_t k = 0; k < n_modes; ++k) { @@ -1365,6 +1365,8 @@ void StellaratorSource::precompute_sampling_distributions() // processing each interval. Each surface is evaluated once. vector R_s(2 * n_grid), Z_s(2 * n_grid); vector Rt_s(2 * n_grid), Zt_s(2 * n_grid); + + // Evaluate positions and poloidal derivatives at each angular grid point. auto evaluate_surface = [&](size_t i) { const double* rc = &rmnc_[i * n_modes]; const double* zs = &zmns_[i * n_modes]; @@ -1434,6 +1436,8 @@ void StellaratorSource::precompute_sampling_distributions() double env = 0.0; double negative_env = 0.0; for (size_t p = 0; p < n_grid; ++p) { + // Linear interpolation makes radial derivatives constant within a bin. + // R and the poloidal derivatives are linear, so R*tau is quadratic. double Rr = (R1[p] - R0[p]) * inv_drho; double Zr = (Z1[p] - Z0[p]) * inv_drho; double f0 = R0[p] * (Rr * Zt0[p] - Rt0[p] * Zr); @@ -1445,6 +1449,9 @@ void StellaratorSource::precompute_sampling_distributions() s0 += f0; sh += fh; s1 += f1; + + // Endpoints and midpoint determine the radial quadratic at this angle. + // Bound both signs until the global Jacobian orientation is known. env = std::max(env, quad_max01(f0, fh, f1)); negative_env = std::max(negative_env, quad_max01(-f0, -fh, -f1)); f_peak = std::max({f_peak, f0, fh, f1}); @@ -1454,12 +1461,13 @@ void StellaratorSource::precompute_sampling_distributions() vp_mid[b] = sh * quad_wgt; vp_hi[b] = s1 * quad_wgt; v_tot += sh * quad_wgt * (rho_[b + 1] - rho_[b]); - // Safety factor covering angular discretization of the majorant + // Heuristic allowance for angular peaks between quadrature points envelope_[b] = 1.05 * env; negative_envelope[b] = 1.05 * negative_env; } jacobian_sign_ = (v_tot >= 0.0) ? 1.0 : -1.0; if (jacobian_sign_ < 0.0) { + // Reverse the integrals and extrema to use a positive volume element. envelope_.swap(negative_envelope); double old_peak = f_peak; f_peak = -f_min; @@ -1477,7 +1485,8 @@ void StellaratorSource::precompute_sampling_distributions() "coefficient tables are inconsistent."); } - // Integrate S(rho)*V'(rho) independently on each radial interval. + // Integrate S(rho)*V'(rho) independently on each radial interval. Scaling + // the emission profile by its maximum preserves the normalized distribution. double s_max = *std::max_element(emission_density_.begin(), emission_density_.end()); vector masses(n_bins); @@ -1487,7 +1496,12 @@ void StellaratorSource::precompute_sampling_distributions() double v0 = vp_lo[b]; double vh = vp_mid[b]; double v1 = vp_hi[b]; + + // Recover V'(t) = a[0] + a[1]*t + a[2]*t^2 for t in [0, 1]. Each + // interval retains its own one-sided radial derivatives at the endpoints. array a {v0, 4 * vh - 3 * v0 - v1, 2 * (v0 + v1 - 2 * vh)}; + + // A convex quadratic can have its minimum inside the interval. double minimum = std::min(v0, v1); if (a[2] > 0) { double t = -a[1] / (2 * a[2]); @@ -1501,16 +1515,20 @@ void StellaratorSource::precompute_sampling_distributions() double ds = s1 - s0; // Simpson's rule exactly integrates the cubic S(rho)*V'(rho). masses[b] = h * ((s0 / 6) * v0 + ((s0 + s1) / 3) * vh + (s1 / 6) * v1); + + // Multiply V'(t) by the linear emission profile to form the radial PDF. auto& c = radial_pdf_[b]; c = {s0 * a[0], s0 * a[1] + ds * a[0], s0 * a[2] + ds * a[1], ds * a[2]}; double bound = cubic_rejection_bound(c); if (bound > 0) { if (!std::isfinite(bound) || !std::isfinite(masses[b]) || masses[b] <= 0) fatal_error("StellaratorSource: invalid radial interval mass."); + // Store the acceptance probability for uniform proposals within the bin. for (double& v : c) v /= bound; } } + // Select intervals by integrated probability, including their radial widths. radial_bins_.assign({masses.data(), masses.size()}); } @@ -1518,6 +1536,7 @@ double StellaratorSource::eval_density(int bin, double t, double theta, double zeta, double* R_out, double* Z_out) const { size_t n_modes = mode_m_.size(); + // Coefficients at the two surfaces bracketing the sampled radial position const double* rc0 = &rmnc_[bin * n_modes]; const double* rc1 = &rmnc_[(bin + 1) * n_modes]; const double* zs0 = &zmns_[bin * n_modes]; @@ -1528,6 +1547,8 @@ double StellaratorSource::eval_density(int bin, double t, double theta, const double* zc1 = asym_ ? &zmnc_[(bin + 1) * n_modes] : nullptr; double inv_drho = 1.0 / (rho_[bin + 1] - rho_[bin]); + // Accumulate positions and derivatives; suffixes r and t denote rho and + // theta. double R = 0.0, Rr = 0.0, Rt = 0.0; double Z = 0.0, Zr = 0.0, Zt = 0.0; for (size_t k = 0; k < n_modes; ++k) { @@ -1536,6 +1557,9 @@ double StellaratorSource::eval_density(int bin, double t, double theta, double s = std::sin(a); double m = mode_m_[k]; + // Interpolate coefficients in t; their rho derivatives are the surface + // differences divided by the bin width. Angular derivatives act on the + // Fourier basis and introduce a factor of m. double drc = rc1[k] - rc0[k]; double rc = rc0[k] + t * drc; double dzs = zs1[k] - zs0[k]; @@ -1547,6 +1571,7 @@ double StellaratorSource::eval_density(int bin, double t, double theta, Zr += dzs * inv_drho * s; Zt += m * zs * c; if (asym_) { + // Add the sine-R and cosine-Z terms for non-stellarator-symmetric shapes. double drs = rs1[k] - rs0[k]; double rs = rs0[k] + t * drs; double dzc = zc1[k] - zc0[k]; @@ -1559,10 +1584,15 @@ double StellaratorSource::eval_density(int bin, double t, double theta, Zt -= m * zc * s; } } + // Reuse the evaluated position when an angular proposal is accepted. if (R_out) *R_out = R; if (Z_out) *Z_out = Z; + + // Since zeta is the cylindrical azimuth, the volume Jacobian is R*tau, + // with tau = R_rho*Z_theta - R_theta*Z_rho. S(rho) cancels from the + // conditional angular density; the global sign accounts for handedness. return jacobian_sign_ * R * (Rr * Zt - Rt * Zr); } From b12871d5d5e018550bdf9471b9a98777ec074846 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 15 Oct 2026 23:42:32 -0500 Subject: [PATCH 32/33] Updates in StellaratorSource::sample --- src/source.cpp | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/source.cpp b/src/source.cpp index 02010ec5f0a..c944f23f1bf 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -1617,9 +1617,6 @@ std::pair StellaratorSource::sample_energy( SourceSite StellaratorSource::sample(uint64_t* seed) const { SourceSite site; - site.particle = ParticleType::neutron(); - site.wgt = 1.0; - site.delayed_group = 0; // 1. Select a radial bin by its integrated mass, then sample within that bin. int bin = radial_bins_.sample(seed); @@ -1644,8 +1641,8 @@ SourceSite StellaratorSource::sample(uint64_t* seed) const double R, Z, zeta; n_reject = 0; while (true) { - double theta = 2.0 * PI * prn(seed); - zeta = 2.0 * PI * prn(seed); + double theta = uniform_distribution(0.0, 2.0 * PI, seed); + zeta = uniform_distribution(0.0, 2.0 * PI, seed); double f = eval_density(bin, t, theta, zeta, &R, &Z); if (prn(seed) * env < f) break; From cd5a960b9eb1aacefe60489c1acf7f7b462b0b24 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Sat, 17 Oct 2026 01:41:05 -0500 Subject: [PATCH 33/33] Add type checks to mode_m and mode_n --- openmc/source.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/openmc/source.py b/openmc/source.py index a964e6d89f1..969f79cb744 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -1546,6 +1546,7 @@ def mode_m(self) -> np.ndarray: @mode_m.setter def mode_m(self, value: Sequence[int]): + cv.check_type('mode_m', value, Iterable, Integral) value = np.asarray(value, dtype=int) if value.ndim != 1 or len(value) < 1: raise ValueError("mode_m must be a 1-D array with at least 1 mode") @@ -1559,6 +1560,7 @@ def mode_n(self) -> np.ndarray: @mode_n.setter def mode_n(self, value: Sequence[int]): + cv.check_type('mode_n', value, Iterable, Integral) value = np.asarray(value, dtype=int) if value.ndim != 1: raise ValueError("mode_n must be a 1-D array")