diff --git a/docs/source/io_formats/settings.rst b/docs/source/io_formats/settings.rst index dcb63d0fc10..5a630cac651 100644 --- a/docs/source/io_formats/settings.rst +++ b/docs/source/io_formats/settings.rst @@ -798,8 +798,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., @@ -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,6 +1110,70 @@ attributes/sub-elements: *Default*: particles are born at :math:`t=0` + 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}`. + + *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 5fad64f3c10..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,72 +299,6 @@ 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`. - -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`. - -Tokamak Plasma Sources ----------------------- - -For fusion applications, the :class:`openmc.TokamakSource` class provides a -native parametric neutron source for tokamak plasmas. Rather than specifying -spatial, angular, and energy distributions separately, the source is defined by -the plasma geometry (using a `Miller-style flux-surface parameterization -`_) and a radial emission profile. Source -sites are sampled directly from the plasma volume without rejection. - -The plasma shape is described by the major radius :math:`R_0`, minor radius -:math:`a`, elongation :math:`\kappa`, triangularity :math:`\delta`, and -Shafranov shift :math:`\Delta`. The neutron birth profile is given as an -emission density :math:`S(r/a)` tabulated on a normalized minor-radius grid that -runs from 0 (magnetic axis) to 1 (last closed flux surface); only the shape of -the profile matters, since it is normalized internally. The emission density is -linearly interpolated between the supplied points and refined internally for -radial sampling. For example:: - - import numpy as np - - r_over_a = np.linspace(0.0, 1.0, 50) - emission = (1.0 - r_over_a**2)**2 # peaked on-axis profile - - source = openmc.TokamakSource( - major_radius=620.0, # cm - minor_radius=200.0, # cm - elongation=1.8, - triangularity=0.45, - shafranov_shift=10.0, # cm - r_over_a=r_over_a, - emission_density=emission, - energy=openmc.stats.muir(e0=14.08e6, m_rat=5.0, kt=20000.0), - ) - - settings.source = source - -The ``energy`` argument accepts either a single -:class:`~openmc.stats.Univariate` distribution applied at all radii, or a -sequence with one distribution per ``r_over_a`` grid point to model a -radially-varying neutron spectrum (energies are then sampled by stochastic -interpolation between the two distributions bracketing the sampled radius). A -time distribution can be given with the ``time`` argument; by default, particles -are born at :math:`t=0`. The toroidal extent can be restricted with -``phi_start`` and ``phi_extent`` to model a sector of the plasma, and -``vertical_shift`` translates the plasma center along the z-axis. - File-based Sources ------------------ @@ -427,6 +388,126 @@ A surface source file is read back like any other source file, with the settings.source = openmc.FileSource('surface_source.h5') +.. _usersguide_tokamak_source: + +Tokamak Plasma Sources +---------------------- + +For fusion applications, the :class:`openmc.TokamakSource` class provides a +native parametric neutron source for tokamak plasmas. Rather than specifying +spatial, angular, and energy distributions separately, the source is defined by +the plasma geometry (using a `Miller-style flux-surface parameterization +`_) and a radial emission profile. Source +sites are sampled directly from the plasma volume without rejection. + +The plasma shape is described by the major radius :math:`R_0`, minor radius +:math:`a`, elongation :math:`\kappa`, triangularity :math:`\delta`, and +Shafranov shift :math:`\Delta`. The neutron birth profile is given as an +emission density :math:`S(r/a)` tabulated on a normalized minor-radius grid that +runs from 0 (magnetic axis) to 1 (last closed flux surface); only the shape of +the profile matters, since it is normalized internally. The emission density is +linearly interpolated between the supplied points and refined internally for +radial sampling. For example:: + + import numpy as np + + r_over_a = np.linspace(0.0, 1.0, 50) + emission = (1.0 - r_over_a**2)**2 # peaked on-axis profile + + source = openmc.TokamakSource( + major_radius=620.0, # cm + minor_radius=200.0, # cm + elongation=1.8, + triangularity=0.45, + shafranov_shift=10.0, # cm + r_over_a=r_over_a, + emission_density=emission, + energy=openmc.stats.muir(e0=14.08e6, m_rat=5.0, kt=20000.0), + ) + + settings.source = source + +The ``energy`` argument accepts either a single +:class:`~openmc.stats.Univariate` distribution applied at all radii, or a +sequence with one distribution per ``r_over_a`` grid point to model a +radially-varying neutron spectrum (energies are then sampled by stochastic +interpolation between the two distributions bracketing the sampled radius). A +time distribution can be given with the ``time`` argument; by default, particles +are born at :math:`t=0`. The toroidal extent can be restricted with +``phi_start`` and ``phi_extent`` to model a sector of the plasma, and +``vertical_shift`` translates the plasma center along the z-axis. + +.. _usersguide_stellarator_source: + +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 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 +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 ``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 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 +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. + .. _compiled_source: Compiled Sources diff --git a/include/openmc/source.h b/include/openmc/source.h index 63e2e244f93..11079a07067 100644 --- a/include/openmc/source.h +++ b/include/openmc/source.h @@ -397,6 +397,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 +//! 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). +//! 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_; + + // Bound-normalized cubic radial PDFs and the distribution of bin masses. + vector> radial_pdf_; + DiscreteIndex radial_bins_; + + // 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..969f79cb744 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -2,14 +2,16 @@ 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 +from typing import Any, Self import lxml.etree as ET import numpy as np import h5py import pandas as pd +from scipy.io import netcdf_file import openmc import openmc.checkvalue as cv @@ -46,7 +48,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 @@ -205,6 +207,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 +1328,626 @@ 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 for the interpolated geometry) by integrating a cubic PDF in each + 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. + + .. 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]): + 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") + 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]): + 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") + 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 + ) -> Self: + """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(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: + 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} + + @classmethod + def from_desc( + cls, + eq, + emission_density, + energy: Univariate | Sequence[Univariate], + n_rho: int = 33, + **kwargs + ) -> Self: + """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. 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 desc.equilibrium.EquilibriaFamily or path-like + A DESC equilibrium object, or a path to a DESC HDF5 output file + (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). + 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 + + """ + 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(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) + 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)) + + 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=eq.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) -> Self: + """Generate stellarator source from an XML element + + Parameters + ---------- + elem : lxml.etree._Element + XML element + + Returns + ------- + openmc.StellaratorSource + Source generated from XML element + + """ + 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( + 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')] + + # 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 = float(get_text(elem, 'strength', 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/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/src/source.cpp b/src/source.cpp index f951722cc18..c944f23f1bf 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 @@ -108,6 +109,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 +1170,507 @@ 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; +} + +// 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]}); + // Add a small relative margin for rounding in the bound and PDF evaluation. + return bound * (1.0 + 1.e-12); +} + +} // 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 (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."); + } + 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 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) { + 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); + } + } + + // 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); + + // 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]; + 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; + } + } + 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; + } + } + }; + + // 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; + 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]); + 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) { + // 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); + 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 = Rh * (Rr * Zth - Rth * Zr); + 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}); + 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]); + // 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; + 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( + "StellaratorSource: the flux-coordinate Jacobian changes sign within " + "the plasma volume; the flux surfaces self-intersect or the Fourier " + "coefficient tables are inconsistent."); + } + + // 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); + radial_pdf_.resize(n_bins); + for (size_t b = 0; b < n_bins; ++b) { + double h = rho_[b + 1] - rho_[b]; + 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]); + 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_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); + + // 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()}); +} + +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]; + 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]); + + // 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) { + 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]; + + // 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]; + 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_) { + // 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]; + 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; + } + } + // 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); +} + +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; + + // 1. Select a radial bin by its integrated mass, then sample within that bin. + int bin = radial_bins_.sample(seed); + double t = prn(seed); + const auto& c = radial_pdf_[bin]; + int64_t n_reject = 0; + while (true) { + double pdf = c[0] + t * (c[1] + t * (c[2] + t * c[3])); + if (prn(seed) < 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; + n_reject = 0; + while (true) { + 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; + 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 //============================================================================== diff --git a/tests/unit_tests/test_source_stellarator.py b/tests/unit_tests/test_source_stellarator.py new file mode 100644 index 00000000000..1d116bdb27a --- /dev/null +++ b/tests/unit_tests/test_source_stellarator.py @@ -0,0 +1,369 @@ +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)) + + +@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, + 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']) + + +@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 + 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]) + + 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', + 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 + + +@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=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, 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 + + +@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( + 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) + 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) 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]'