From a694b3c34c04b63a9e9760793cf015ec2431ef70 Mon Sep 17 00:00:00 2001 From: Perry Date: Tue, 9 Jun 2026 06:16:26 -0700 Subject: [PATCH 1/9] Add group_xs C API for group-averaged cross sections Add Reaction::group_xs and Nuclide::group_xs, exposed via a new openmc_nuclide_group_xs C API, to compute the group-averaged cross section in each energy group -- a flat-in-bin (constant intra-group flux) average. A group cross section table can then be built once and reused to collapse the flux across many domains. Reaction::collapse_rate makes the same flat-in-bin assumption, so the two share one union-grid panel-walk kernel (for_each_panel, after NJOY/GROUPR's panel) rather than duplicating the integration loop; collapse_rate stays bit-identical. --- include/openmc/nuclide.h | 9 +++++++ include/openmc/reaction.h | 10 ++++++++ src/nuclide.cpp | 52 +++++++++++++++++++++++++++++++++++++++ src/reaction.cpp | 49 +++++++++++++++++++++++++++--------- 4 files changed, 108 insertions(+), 12 deletions(-) diff --git a/include/openmc/nuclide.h b/include/openmc/nuclide.h index ae39a53ddf5..7a2e631cf37 100644 --- a/include/openmc/nuclide.h +++ b/include/openmc/nuclide.h @@ -84,6 +84,15 @@ class Nuclide { double collapse_rate(int MT, double temperature, span energy, span flux) const; + //! Calculate group-averaged cross sections for a reaction + // + //! \param[in] MT ENDF MT value for desired reaction + //! \param[in] temperature Temperature in [K] + //! \param[in] energy Energy group boundaries in [eV] (size n_groups+1) + //! \param[out] xs Group-averaged cross section in each energy group + void group_xs(int MT, double temperature, span energy, + span xs) const; + //! Return a ParticleType object representing this nuclide ParticleType particle_type() const { return {Z_, A_, metastable_}; } diff --git a/include/openmc/reaction.h b/include/openmc/reaction.h index e95db1665d7..4116aff433e 100644 --- a/include/openmc/reaction.h +++ b/include/openmc/reaction.h @@ -52,6 +52,16 @@ class Reaction { double collapse_rate(int64_t i_temp, span energy, span flux, const vector& grid) const; + //! Calculate group-averaged cross sections over a group structure + // + //! \param[in] i_temp Temperature index + //! \param[in] energy Energy group boundaries in [eV] (size n_groups+1) + //! \param[in] grid Nuclide energy grid at i_temp + //! \param[out] xs Group-averaged cross section in each energy group; + //! must be zeroed on entry + void group_xs(int64_t i_temp, span energy, + const vector& grid, span xs) const; + //! Cross section at a single temperature struct TemperatureXS { int threshold; diff --git a/src/nuclide.cpp b/src/nuclide.cpp index 17d6e952c39..6a7afb33c18 100644 --- a/src/nuclide.cpp +++ b/src/nuclide.cpp @@ -1072,6 +1072,40 @@ double Nuclide::collapse_rate(int MT, double temperature, } } +void Nuclide::group_xs( + int MT, double temperature, span energy, span xs) const +{ + assert(MT > 0); + assert(energy.size() > 0); + assert(energy.size() == xs.size() + 1); + + // Zero the output; group_xs only writes groups that have data + std::fill(xs.data(), xs.data() + xs.size(), 0.0); + + int i_rx = reaction_index_[MT]; + if (i_rx < 0) + return; + const auto& rx = reactions_[i_rx]; + + // Determine temperature index + int64_t i_temp; + double f; + std::tie(i_temp, f) = this->find_temperature(temperature); + + // Get group-averaged cross sections at lower temperature + const auto& grid_low = grid_[i_temp].energy; + rx->group_xs(i_temp, energy, grid_low, xs); + + if (f > 0.0) { + // Interpolate element-wise between lower and higher temperature + const auto& grid_high = grid_[i_temp + 1].energy; + vector xs_high(xs.size(), 0.0); + rx->group_xs(i_temp + 1, energy, grid_high, xs_high); + for (std::size_t g = 0; g < xs.size(); ++g) + xs[g] += f * (xs_high[g] - xs[g]); + } +} + //============================================================================== // Non-member functions //============================================================================== @@ -1215,6 +1249,24 @@ extern "C" int openmc_nuclide_collapse_rate(int index, int MT, return 0; } +extern "C" int openmc_nuclide_group_xs(int index, int MT, double temperature, + const double* energy, int n_groups, double* group_xs) +{ + if (index < 0 || index >= data::nuclides.size()) { + set_errmsg("Index in nuclides vector is out of bounds."); + return OPENMC_E_OUT_OF_BOUNDS; + } + + try { + data::nuclides[index]->group_xs(MT, temperature, + {energy, energy + n_groups + 1}, {group_xs, group_xs + n_groups}); + } catch (const std::out_of_range& e) { + set_errmsg(e.what()); + return OPENMC_E_OUT_OF_BOUNDS; + } + return 0; +} + void nuclides_clear() { data::nuclides.clear(); diff --git a/src/reaction.cpp b/src/reaction.cpp index c02e0cc407e..5e0241c61cf 100644 --- a/src/reaction.cpp +++ b/src/reaction.cpp @@ -3,7 +3,7 @@ #include // for remove_if #include #include -#include // for move +#include // for move, pair #include @@ -112,31 +112,33 @@ double Reaction::xs(const NuclideMicroXS& micro) const return this->xs(micro.index_temp, micro.index_grid, micro.interp_factor); } -double Reaction::collapse_rate(int64_t i_temp, span energy, - span flux, const vector& grid) const +namespace { + +//! Call f(group, xs_avg, dE) for each panel (union-grid interval); returns +//! the inclusive range of groups visited ({0, -1} if none). +template +std::pair for_each_panel(span energy, + const vector& grid, const vector& xs, int i_threshold, F&& f) { // Find index corresponding to first energy - const auto& xs = xs_[i_temp].value; int i_low = lower_bound_index(grid.cbegin(), grid.cend(), energy.front()); // Check for threshold and adjust starting point if necessary int j_start = 0; - int i_threshold = xs_[i_temp].threshold; if (i_low < i_threshold) { i_low = i_threshold; while (energy[j_start + 1] < grid[i_low]) { ++j_start; if (j_start + 1 == energy.size()) - return 0.0; + return {0, -1}; } } - double xs_flux_sum = 0.0; - - for (int j = j_start; j < flux.size(); ++j) { + int n_groups = energy.size() - 1; + int j = j_start; + for (; j < n_groups; ++j) { double E_group_low = energy[j]; double E_group_high = energy[j + 1]; - double flux_per_eV = flux[j] / (E_group_high - E_group_low); // Determine energy grid index corresponding to group high int i_high = i_low; @@ -165,8 +167,7 @@ double Reaction::collapse_rate(int64_t i_temp, span energy, double xs_avg = 0.5 * (xs_low + xs_high); // Add contribution from segment - double dE = (E_high - E_low); - xs_flux_sum += flux_per_eV * xs_avg * dE; + f(j, xs_avg, E_high - E_low); } i_low = i_high; @@ -175,10 +176,34 @@ double Reaction::collapse_rate(int64_t i_temp, span energy, if (i_low + 1 == grid.size()) break; } + return {j_start, std::min(j, n_groups - 1)}; +} + +} // namespace +double Reaction::collapse_rate(int64_t i_temp, span energy, + span flux, const vector& grid) const +{ + double xs_flux_sum = 0.0; + for_each_panel(energy, grid, xs_[i_temp].value, xs_[i_temp].threshold, + [&](int j, double xs_avg, double dE) { + double flux_per_eV = flux[j] / (energy[j + 1] - energy[j]); + xs_flux_sum += flux_per_eV * xs_avg * dE; + }); return xs_flux_sum; } +void Reaction::group_xs(int64_t i_temp, span energy, + const vector& grid, span xs) const +{ + // Sum XS integrals into the zeroed output, then divide by group width + auto [j_first, j_last] = + for_each_panel(energy, grid, xs_[i_temp].value, xs_[i_temp].threshold, + [&](int j, double xs_avg, double dE) { xs[j] += xs_avg * dE; }); + for (int j = j_first; j <= j_last; ++j) + xs[j] /= (energy[j + 1] - energy[j]); +} + //============================================================================== // Non-member functions //============================================================================== From 81e033d2e2ceeb442ec0eb5a09f5c3955704cc7f Mon Sep 17 00:00:00 2001 From: Perry Date: Tue, 9 Jun 2026 07:07:03 -0700 Subject: [PATCH 2/9] Add Nuclide.group_xs Python binding Wrap the openmc_nuclide_group_xs C API in openmc.lib.Nuclide.group_xs, returning the group-averaged cross section in each energy group as a NumPy array of length len(energy) - 1. --- openmc/lib/nuclide.py | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/openmc/lib/nuclide.py b/openmc/lib/nuclide.py index ef1287cf34a..b4ed472f458 100644 --- a/openmc/lib/nuclide.py +++ b/openmc/lib/nuclide.py @@ -29,6 +29,10 @@ _array_1d_dble, _array_1d_dble, c_int, POINTER(c_double)] _dll.openmc_nuclide_collapse_rate.restype = c_int _dll.openmc_nuclide_collapse_rate.errcheck = _error_handler +_dll.openmc_nuclide_group_xs.argtypes = [c_int, c_int, c_double, + _array_1d_dble, c_int, _array_1d_dble] +_dll.openmc_nuclide_group_xs.restype = c_int +_dll.openmc_nuclide_group_xs.errcheck = _error_handler _dll.nuclides_size.restype = c_size_t @@ -112,6 +116,38 @@ def collapse_rate(self, MT, temperature, energy, flux): flux, len(flux), xs) return xs.value + def group_xs(self, MT, temperature, energy): + """Calculate group-averaged microscopic cross sections. + + Computes the cross section averaged over each energy group, + ``integral(sigma dE) / dE_group``, for a given reaction. This is a + flat-in-bin average -- it assumes the flux is constant within each + group -- so the group cross sections do not depend on the per-group + flux magnitudes and the collapsed reaction rate ``sum(flux * group_xs)`` + matches :meth:`collapse_rate`, which makes the same assumption. + + Parameters + ---------- + MT : int + ENDF MT value of the desired reaction + temperature : float + Temperature in [K] at which to evaluate cross sections + energy : iterable of float + Energy group boundaries in [eV] + + Returns + ------- + numpy.ndarray + Group-averaged cross section in [b] for each of the + ``len(energy) - 1`` energy groups + + """ + energy = np.asarray(energy, dtype=float) + xs = np.zeros(energy.size - 1) + _dll.openmc_nuclide_group_xs(self._index, MT, temperature, energy, + len(xs), xs) + return xs + class _NuclideMapping(Mapping): """Provide mapping from nuclide name to index in nuclides array.""" From d92af26e397e52a353b9e5b921471af9e70ec2b2 Mon Sep 17 00:00:00 2001 From: Perry Date: Tue, 9 Jun 2026 07:26:28 -0700 Subject: [PATCH 3/9] Build the group cross section table once for flux collapse In get_microxs_and_flux(reaction_rate_mode='flux'), build the group-averaged cross sections once into a _SparseXSTable via _build_xs_table_ce and collapse every domain with a single matrix-vector product (_collapse_fluxes), instead of re-deriving the group cross sections from continuous-energy data for each domain. Because the group cross sections are domain-independent -- a flat-in-bin average of the continuous-energy data -- this is ~N times faster for N domains and reproduces the previous result to within floating-point summation order. The collapse loads cross sections from the same library the model resolved, matching the transport solve. MicroXS.from_multigroup_flux is now the single-flux case of the same engine (_flux_collapse_micros); it additionally validates that the flux is finite and non-negative, raising ValueError. The 'direct' mode and reaction_rate_opts paths are untouched. Only Independent Operator microxs, for now. --- openmc/deplete/microxs.py | 268 ++++++++++++++++++++++++++++++++------ 1 file changed, 231 insertions(+), 37 deletions(-) diff --git a/openmc/deplete/microxs.py b/openmc/deplete/microxs.py index 42bb958caf7..cd3f31fe29c 100644 --- a/openmc/deplete/microxs.py +++ b/openmc/deplete/microxs.py @@ -6,6 +6,8 @@ from __future__ import annotations from collections.abc import Sequence +from contextlib import nullcontext +from dataclasses import dataclass import shutil from tempfile import TemporaryDirectory from typing import Union, TypeAlias, Self @@ -140,6 +142,8 @@ def get_microxs_and_flux( chain = _get_chain(chain_file) if reactions is None: reactions = chain.reactions + nuclides_with_data = None + cross_sections = None if not nuclides: cross_sections = _find_cross_sections(model) nuclides_with_data = _get_nuclides_with_data(cross_sections) @@ -309,14 +313,15 @@ def get_microxs_and_flux( MicroXS(xs_i, rr_nuclides, rr_reactions) for xs_i in xs) if reaction_rate_mode == 'flux': - # Compute flux-collapsed microscopic XS - flux_micros = [MicroXS.from_multigroup_flux( - energies=collapse_energies, - multigroup_flux=flux_i, - chain_file=chain_file, - nuclides=nuclides, - reactions=reactions - ) for flux_i in fluxes] + # Resolve data availability if the caller passed nuclides explicitly + if nuclides_with_data is None: + cross_sections = _find_cross_sections(model) + nuclides_with_data = _get_nuclides_with_data(cross_sections) + # Build the XS table once and collapse every domain against it; looping + # from_multigroup_flux per domain would rebuild the table each call. + flux_micros = _flux_collapse_micros( + nuclides, reactions, collapse_energies, fluxes, + nuclides_with_data, cross_sections=cross_sections) # We need to return one-group fluxes to match the microscopic cross # sections, which are always one-group by virtue of the collapse @@ -336,6 +341,215 @@ def get_microxs_and_flux( return fluxes, micros +@dataclass +class _SparseXSTable: + """Sparse storage of group cross sections for vectorized flux collapse. + + Stores only the non-zero ``(nuclide, reaction)`` pairs. ``xs_matrix`` has + shape ``(nnz, n_groups)``; ``nuc_indices`` and ``rxn_indices`` map each row + to its position in the dense ``(n_nuclides, n_reactions)`` result. + + Parameters + ---------- + nuclides : list of str + Nuclide names defining the result's nuclide axis. + reactions : list of str + Reaction names defining the result's reaction axis. + n_groups : int + Number of energy groups. + xs_matrix : numpy.ndarray + Group cross sections for the stored rows, shape ``(nnz, n_groups)``. + nuc_indices : numpy.ndarray + Row-to-nuclide index map, shape ``(nnz,)``. + rxn_indices : numpy.ndarray + Row-to-reaction index map, shape ``(nnz,)``. + + Notes + ----- + Provider-agnostic: rows may be filled from continuous-energy data + (:func:`_build_xs_table_ce`) or any other group cross section source. + + """ + nuclides: list[str] + reactions: list[str] + n_groups: int + xs_matrix: np.ndarray + nuc_indices: np.ndarray + rxn_indices: np.ndarray + + def collapse(self, phi_norm: np.ndarray) -> np.ndarray: + """Collapse the table to one-group cross sections. + + Parameters + ---------- + phi_norm : numpy.ndarray + Normalized group flux, shape ``(n_groups,)``. Should sum to 1. + + Returns + ------- + numpy.ndarray + Dense ``(n_nuclides, n_reactions)`` one-group cross sections. + + """ + if len(phi_norm) != self.n_groups: + raise ValueError( + f'phi_norm has length {len(phi_norm)} but table expects ' + f'{self.n_groups} groups') + collapsed_sparse = self.xs_matrix @ phi_norm + result = np.zeros((len(self.nuclides), len(self.reactions))) + result[self.nuc_indices, self.rxn_indices] = collapsed_sparse + return result + + +def _build_xs_table_ce( + nuclides: Sequence[str], + reactions: Sequence[str], + energies: Sequence[float], + temperature: float, + nuclides_with_data: set, + cross_sections=None, + **init_kwargs, +) -> _SparseXSTable: + """Build a sparse group cross section table from continuous-energy data. + + For each requested ``(nuclide, reaction)`` the group-averaged cross sections + are computed once via the :meth:`openmc.lib.Nuclide.group_xs` C API and + stored as one sparse row, so every domain can later be collapsed with a + single matrix-vector product. Rows that are entirely zero (reaction absent, + or threshold above the whole group structure) are skipped. Cross sections + are evaluated inside a single :class:`openmc.lib.TemporarySession`, loading + each nuclide once. + + Parameters + ---------- + nuclides : sequence of str + Nuclide names defining the result's nuclide axis. + reactions : sequence of str + Reaction names defining the result's reaction axis. + energies : sequence of float + Energy group boundaries in [eV], ascending, length ``n_groups + 1``. + temperature : float + Temperature in [K] for cross section evaluation. + nuclides_with_data : set + Nuclides available in the cross section library; others are skipped. + cross_sections : PathLike, optional + Cross section library for the session, so it matches the one + ``nuclides_with_data`` was resolved from. Defaults to ``openmc.config``. + **init_kwargs : dict + Keyword arguments passed to :func:`openmc.lib.init` via the temporary + session. + + Returns + ------- + _SparseXSTable + + """ + mts = [REACTION_MT[name] for name in reactions] + energies = np.asarray(energies, dtype=float) + n_groups = len(energies) - 1 + + rows = [] + nuc_idx_list = [] + rxn_idx_list = [] + # Load against the same library nuclides_with_data was resolved from + library = (openmc.config.patch('cross_sections', cross_sections) + if cross_sections is not None else nullcontext()) + with library, openmc.lib.TemporarySession(**init_kwargs): + for nuc_idx, nuc in enumerate(nuclides): + if nuc not in nuclides_with_data: + continue + lib_nuc = openmc.lib.load_nuclide(nuc) + # Index by reaction, not MT, so fission/(n,fission) stay separate + for rxn_idx, mt in enumerate(mts): + xs_g = lib_nuc.group_xs(mt, temperature, energies) + if xs_g.any(): + rows.append(xs_g) + nuc_idx_list.append(nuc_idx) + rxn_idx_list.append(rxn_idx) + + if rows: + xs_matrix = np.vstack(rows) + else: + xs_matrix = np.empty((0, n_groups)) + + return _SparseXSTable( + nuclides=list(nuclides), + reactions=list(reactions), + n_groups=n_groups, + xs_matrix=xs_matrix, + nuc_indices=np.array(nuc_idx_list, dtype=np.int32), + rxn_indices=np.array(rxn_idx_list, dtype=np.int32), + ) + + +def _collapse_fluxes( + table: _SparseXSTable, + fluxes: Sequence[np.ndarray], + nuclides: Sequence[str], + reactions: Sequence[str], +) -> list[MicroXS]: + """Collapse each domain's multigroup flux against a built XS table. + + Each flux is validated (finite, non-negative) and normalized to sum 1 before + being collapsed; a zero-sum flux yields an all-zero :class:`MicroXS`. + + Parameters + ---------- + table : _SparseXSTable + Pre-built group cross section table. + fluxes : sequence of numpy.ndarray + One ``(n_groups,)`` group-flux vector per domain. + nuclides : sequence of str + Nuclide axis of the output MicroXS objects. + reactions : sequence of str + Reaction axis of the output MicroXS objects. + + Returns + ------- + list of MicroXS + One per domain, each of shape ``(n_nuclides, n_reactions, 1)``. + + """ + micros = [] + for flux in fluxes: + flux = np.asarray(flux, dtype=float) + if not np.isfinite(flux).all(): + raise ValueError('Multigroup flux contains non-finite values') + if (flux < 0).any(): + raise ValueError('Multigroup flux contains negative values') + flux_sum = flux.sum() + if flux_sum == 0.0: + micros.append(MicroXS( + np.zeros((len(nuclides), len(reactions), 1)), + nuclides, reactions)) + continue + collapsed = table.collapse(flux / flux_sum) + micros.append(MicroXS(collapsed[:, :, np.newaxis], + nuclides, reactions)) + return micros + + +def _flux_collapse_micros( + nuclides: Sequence[str], + reactions: Sequence[str], + energies: Sequence[float], + fluxes: Sequence[np.ndarray], + nuclides_with_data: set, + temperature: float = 293.6, + cross_sections=None, + **init_kwargs, +) -> list[MicroXS]: + """Build the CE group XS table once, then collapse each domain's flux. + + Returns one :class:`MicroXS` per entry in ``fluxes``. + + """ + table = _build_xs_table_ce( + nuclides, reactions, energies, temperature, nuclides_with_data, + cross_sections=cross_sections, **init_kwargs) + return _collapse_fluxes(table, fluxes, nuclides, reactions) + + class MicroXS: """Microscopic cross section data for use in transport-independent depletion. @@ -407,7 +621,8 @@ def from_multigroup_flux( energies : iterable of float or str Energy group boundaries in [eV] or the name of the group structure multigroup_flux : iterable of float - Energy-dependent multigroup flux values + Energy-dependent multigroup flux values. Must be finite and + non-negative. chain_file : PathLike or Chain, optional Path to the depletion chain XML file or an instance of openmc.deplete.Chain. Defaults to ``openmc.config['chain_file']``. @@ -450,36 +665,15 @@ def from_multigroup_flux( nuclides = chain.nuclides nuclides = [nuc.name for nuc in nuclides] - # Get reaction MT values. If no reactions specified, default to the - # reactions available in the chain file + # If no reactions specified, default to those in the chain file if reactions is None: reactions = chain.reactions - mts = [REACTION_MT[name] for name in reactions] - - # Create 3D array for microscopic cross sections - microxs_arr = np.zeros((len(nuclides), len(mts), 1)) - - # If flux is zero, safely return zero cross sections - multigroup_flux = np.array(multigroup_flux) - if (flux_sum := multigroup_flux.sum()) == 0.0: - return cls(microxs_arr, nuclides, reactions) - - # Normalize multigroup flux - multigroup_flux /= flux_sum - - # Compute microscopic cross sections within a temporary session - with openmc.lib.TemporarySession(**init_kwargs): - # For each nuclide and reaction, compute the flux-averaged xs - for nuc_index, nuc in enumerate(nuclides): - if nuc not in nuclides_with_data: - continue - lib_nuc = openmc.lib.load_nuclide(nuc) - for mt_index, mt in enumerate(mts): - microxs_arr[nuc_index, mt_index, 0] = lib_nuc.collapse_rate( - mt, temperature, energies, multigroup_flux - ) - - return cls(microxs_arr, nuclides, reactions) + + # Delegate to the shared build-and-collapse engine (single-flux case) + return _flux_collapse_micros( + nuclides, reactions, energies, [multigroup_flux], + nuclides_with_data, temperature, cross_sections=cross_sections, + **init_kwargs)[0] @classmethod def from_csv(cls, csv_file, **kwargs): From 5e0f0044733b3183372a61233288715f35d3bd66 Mon Sep 17 00:00:00 2001 From: Perry Date: Tue, 9 Jun 2026 09:04:04 -0700 Subject: [PATCH 4/9] Add unit tests for the group_xs flux-collapse path Test the new continuous-energy group cross section path: - group_xs reproduces collapse_rate: sum_g(flux_g * group_xs_g) matches collapse_rate over several independent fluxes, with a scale guard that the values are per-group averages (barns), not the raw integral. - _build_xs_table_ce + _collapse_fluxes (and the delegated from_multigroup_flux) reproduce the old per-domain collapse_rate result, including a threshold reaction and a zero-flux domain. - _collapse_fluxes rejects non-finite and negative flux and returns a zero MicroXS for a zero-sum flux. --- tests/unit_tests/test_deplete_ce_collapse.py | 155 +++++++++++++++++++ 1 file changed, 155 insertions(+) create mode 100644 tests/unit_tests/test_deplete_ce_collapse.py diff --git a/tests/unit_tests/test_deplete_ce_collapse.py b/tests/unit_tests/test_deplete_ce_collapse.py new file mode 100644 index 00000000000..ececaedd672 --- /dev/null +++ b/tests/unit_tests/test_deplete_ce_collapse.py @@ -0,0 +1,155 @@ +"""Unit tests for the continuous-energy group cross section / flux-collapse path: +``openmc.lib.Nuclide.group_xs``, ``_SparseXSTable``, ``_build_xs_table_ce`` and +``_collapse_fluxes`` in :mod:`openmc.deplete.microxs`. + +Requires nuclear data (``OPENMC_CROSS_SECTIONS``) and the compiled library. +""" + +from pathlib import Path + +import numpy as np +import pytest + +import openmc +import openmc.lib +from openmc.data import REACTION_MT +from openmc.mgxs import GROUP_STRUCTURES +from openmc.deplete import MicroXS +from openmc.deplete.microxs import _build_xs_table_ce, _collapse_fluxes +from openmc.deplete.coupled_operator import ( + _find_cross_sections, _get_nuclides_with_data) + +CHAIN_FILE = Path(__file__).parents[1] / "chain_simple.xml" +GROUP_STRUCTURE = "CASMO-40" +TEMPERATURE = 293.6 + + +def _energies(): + return np.asarray(GROUP_STRUCTURES[GROUP_STRUCTURE], dtype=float) + + +def _reference_micros(nuclides, reactions, mts, energies, fluxes, temperature): + """Old per-domain reference: collapse_rate(mt, T, E, flux/sum) per (nuc, mt).""" + out = [] + with openmc.lib.TemporarySession(): + lib_nucs = {n: openmc.lib.load_nuclide(n) for n in nuclides} + for flux in fluxes: + flux = np.asarray(flux, dtype=float) + arr = np.zeros((len(nuclides), len(reactions), 1)) + flux_sum = flux.sum() + if flux_sum != 0.0: + phi = flux / flux_sum + for i, nuc in enumerate(nuclides): + for j, mt in enumerate(mts): + arr[i, j, 0] = lib_nucs[nuc].collapse_rate( + mt, temperature, energies, phi) + out.append(arr) + return out + + +def test_group_xs_equivalence(): + """group_xs reproduces collapse_rate: sum_g(phi_g * group_xs_g) == collapse_rate. + + A single dot product is under-determined, so equivalence is checked over + several independent fluxes (random + the flat-per-eV bin-width flux). A scale + guard catches a forgotten division by the group width (the undivided integral + would be ~1e5-1e7 eV-b, far outside the barns range). + """ + energies = _energies() + n_groups = energies.size - 1 + dE = np.diff(energies) + rng = np.random.default_rng(0) + + # (nuclide, MT): include a threshold reaction (U238 (n,2n)) plus non-threshold + cases = [ + ('U238', REACTION_MT['(n,2n)']), + ('U235', REACTION_MT['fission']), + ('O16', REACTION_MT['(n,gamma)']), + ] + with openmc.lib.TemporarySession(): + for name, mt in cases: + nuc = openmc.lib.load_nuclide(name) + group_xs = nuc.group_xs(mt, TEMPERATURE, energies) + assert group_xs.shape == (n_groups,) + + for flux in (rng.random(n_groups), rng.random(n_groups), dE): + assert np.dot(flux, group_xs) == pytest.approx( + nuc.collapse_rate(mt, TEMPERATURE, energies, flux), rel=1e-10) + + # Per-group AVERAGE cross section, in barns -- not the raw integral. + assert group_xs.max() < 1e6 + + if mt == REACTION_MT['(n,2n)']: + n2n = group_xs + + # The threshold reaction must leave below-threshold groups at zero. + assert (n2n == 0.0).any() and (n2n != 0.0).any() + + +def test_group_xs_structure_below_threshold(): + """A structure entirely below a reaction threshold collapses to all zeros.""" + energies = np.array([1.0e-5, 1.0, 1.0e2]) # U238 (n,2n) threshold ~6.2 MeV + flux = np.array([1.0, 2.0]) + with openmc.lib.TemporarySession(): + nuc = openmc.lib.load_nuclide('U238') + mt = REACTION_MT['(n,2n)'] + assert np.all(nuc.group_xs(mt, TEMPERATURE, energies) == 0.0) + assert nuc.collapse_rate(mt, TEMPERATURE, energies, flux) == 0.0 + + +def test_ce_collapse_matches_collapse_rate(): + """_build_xs_table_ce + _collapse_fluxes (and the delegated from_multigroup_flux) + reproduce the old per-domain collapse_rate result, with a threshold reaction + and a zero-flux domain. + """ + nuclides = ['U235', 'U238', 'O16'] + reactions = ['fission', '(n,gamma)', '(n,2n)'] # (n,2n) is a threshold reaction + mts = [REACTION_MT[r] for r in reactions] + energies = _energies() + n_groups = energies.size - 1 + nuclides_with_data = _get_nuclides_with_data(_find_cross_sections(model=None)) + + rng = np.random.default_rng(1) + fluxes = [rng.random(n_groups), rng.random(n_groups), np.zeros(n_groups)] + + table = _build_xs_table_ce( + nuclides, reactions, energies, TEMPERATURE, nuclides_with_data) + new = _collapse_fluxes(table, fluxes, nuclides, reactions) + ref = _reference_micros(nuclides, reactions, mts, energies, fluxes, TEMPERATURE) + + for micro, ref_data in zip(new, ref): + assert micro.data == pytest.approx(ref_data, rel=1e-9, abs=1e-12) + # Zero-flux domain -> all-zero MicroXS of the right shape. + assert new[-1].data.shape == (len(nuclides), len(reactions), 1) + assert np.all(new[-1].data == 0.0) + + # from_multigroup_flux delegates to the same engine -> identical result. + micro = MicroXS.from_multigroup_flux( + energies=energies, multigroup_flux=fluxes[0], chain_file=CHAIN_FILE, + nuclides=nuclides, reactions=reactions) + assert micro.data == pytest.approx(ref[0], rel=1e-9, abs=1e-12) + + +def test_collapse_fluxes_guards(): + """_collapse_fluxes rejects non-finite / negative flux and zero-fills zero flux.""" + nuclides = ['U235'] + reactions = ['fission'] + energies = _energies() + n_groups = energies.size - 1 + nuclides_with_data = _get_nuclides_with_data(_find_cross_sections(model=None)) + table = _build_xs_table_ce( + nuclides, reactions, energies, TEMPERATURE, nuclides_with_data) + + nan_flux = np.ones(n_groups) + nan_flux[0] = np.nan + with pytest.raises(ValueError): + _collapse_fluxes(table, [nan_flux], nuclides, reactions) + + neg_flux = np.ones(n_groups) + neg_flux[0] = -1.0 + with pytest.raises(ValueError): + _collapse_fluxes(table, [neg_flux], nuclides, reactions) + + zero = _collapse_fluxes(table, [np.zeros(n_groups)], nuclides, reactions)[0] + assert zero.data.shape == (1, 1, 1) + assert np.all(zero.data == 0.0) From 1911ba4e1c110108c9523054219d047b6a159bd0 Mon Sep 17 00:00:00 2001 From: Perry Date: Wed, 10 Jun 2026 08:06:04 -0700 Subject: [PATCH 5/9] Make from_multigroup_flux the multi-flux collapse engine Overload from_multigroup_flux to accept a 2-D batch (or list of 1-D fluxes) and build the group XS table once for all of them, returning a list of MicroXS. Add a keyword-only cross_sections argument and resolve the chain lazily. Delete _flux_collapse_micros and rewire get_microxs_and_flux to call from_multigroup_flux directly. --- openmc/deplete/microxs.py | 116 ++++++++++--------- tests/unit_tests/test_deplete_ce_collapse.py | 90 ++++++++++++++ 2 files changed, 149 insertions(+), 57 deletions(-) diff --git a/openmc/deplete/microxs.py b/openmc/deplete/microxs.py index cd3f31fe29c..b9f1d568b1e 100644 --- a/openmc/deplete/microxs.py +++ b/openmc/deplete/microxs.py @@ -142,8 +142,6 @@ def get_microxs_and_flux( chain = _get_chain(chain_file) if reactions is None: reactions = chain.reactions - nuclides_with_data = None - cross_sections = None if not nuclides: cross_sections = _find_cross_sections(model) nuclides_with_data = _get_nuclides_with_data(cross_sections) @@ -313,15 +311,12 @@ def get_microxs_and_flux( MicroXS(xs_i, rr_nuclides, rr_reactions) for xs_i in xs) if reaction_rate_mode == 'flux': - # Resolve data availability if the caller passed nuclides explicitly - if nuclides_with_data is None: - cross_sections = _find_cross_sections(model) - nuclides_with_data = _get_nuclides_with_data(cross_sections) - # Build the XS table once and collapse every domain against it; looping - # from_multigroup_flux per domain would rebuild the table each call. - flux_micros = _flux_collapse_micros( - nuclides, reactions, collapse_energies, fluxes, - nuclides_with_data, cross_sections=cross_sections) + # Resolve the library from the model (from_multigroup_flux defaults to config) + cross_sections = _find_cross_sections(model) + # Collapse all domains against one table, built once + flux_micros = MicroXS.from_multigroup_flux( + collapse_energies, fluxes, chain_file=chain, nuclides=nuclides, + reactions=reactions, cross_sections=cross_sections) # We need to return one-group fluxes to match the microscopic cross # sections, which are always one-group by virtue of the collapse @@ -529,27 +524,6 @@ def _collapse_fluxes( return micros -def _flux_collapse_micros( - nuclides: Sequence[str], - reactions: Sequence[str], - energies: Sequence[float], - fluxes: Sequence[np.ndarray], - nuclides_with_data: set, - temperature: float = 293.6, - cross_sections=None, - **init_kwargs, -) -> list[MicroXS]: - """Build the CE group XS table once, then collapse each domain's flux. - - Returns one :class:`MicroXS` per entry in ``fluxes``. - - """ - table = _build_xs_table_ce( - nuclides, reactions, energies, temperature, nuclides_with_data, - cross_sections=cross_sections, **init_kwargs) - return _collapse_fluxes(table, fluxes, nuclides, reactions) - - class MicroXS: """Microscopic cross section data for use in transport-independent depletion. @@ -597,32 +571,44 @@ def __init__(self, data: np.ndarray, nuclides: list[str], reactions: list[str]): def from_multigroup_flux( cls, energies: Sequence[float] | str, - multigroup_flux: Sequence[float], + multigroup_flux: Sequence[float] | Sequence[Sequence[float]], chain_file: PathLike | None = None, temperature: float = 293.6, nuclides: Sequence[str] | None = None, reactions: Sequence[str] | None = None, + *, + cross_sections: PathLike | None = None, **init_kwargs: dict, - ) -> MicroXS: + ) -> MicroXS | list[MicroXS]: """Generated microscopic cross sections from a known flux. The size of the MicroXS matrix depends on the chain file and cross sections available. MicroXS entry will be 0 if the nuclide cross section is not found. + Multiple fluxes can be collapsed at once by passing a 2-D array (or a + list of 1-D arrays); the group cross section table is then built only + once and reused for every flux. + It is recommended to make repeated calls to this method within a context manager using the :class:`openmc.lib.TemporarySession` class to avoid re-initializing OpenMC and loading cross sections each time. .. versionadded:: 0.15.0 + .. versionchanged:: 0.15.4 + ``multigroup_flux`` may be 2-D (or a list of 1-D arrays) to collapse + several fluxes against a single shared cross section table, returning + a list of :class:`MicroXS`. Added the ``cross_sections`` argument. + Parameters ---------- energies : iterable of float or str Energy group boundaries in [eV] or the name of the group structure - multigroup_flux : iterable of float + multigroup_flux : iterable of float or iterable of iterable of float Energy-dependent multigroup flux values. Must be finite and - non-negative. + non-negative. A 1-D input is a single flux; a 2-D input (or a list of + 1-D arrays) is a batch of fluxes that share the same group structure. chain_file : PathLike or Chain, optional Path to the depletion chain XML file or an instance of openmc.deplete.Chain. Defaults to ``openmc.config['chain_file']``. @@ -634,12 +620,18 @@ def from_multigroup_flux( reactions : list of str, optional Reactions to get cross sections for. If not specified, all neutron reactions listed in the depletion chain file are used. + cross_sections : PathLike, optional + Cross section library used to resolve nuclide data availability and + evaluate cross sections. Defaults to ``openmc.config['cross_sections']``. **init_kwargs : dict Keyword arguments passed to :func:`openmc.lib.init` Returns ------- - MicroXS + MicroXS or list of MicroXS + A single :class:`MicroXS` for a 1-D ``multigroup_flux``; a list with + one entry per flux for a 2-D input (a 1-row batch returns a + 1-element list, not an unwrapped :class:`MicroXS`). """ check_type("temperature", temperature, (int, float)) @@ -652,28 +644,38 @@ def from_multigroup_flux( if not np.all(np.diff(energies) > 0): raise ValueError('Energy group boundaries must be in ascending order') - # check dimension consistency - if len(multigroup_flux) != len(energies) - 1: - raise ValueError('Length of flux array should be len(energies)-1') - - chain = _get_chain(chain_file) - cross_sections = _find_cross_sections(model=None) + # A 1-D flux is a single domain; 2-D (or a list of 1-D) is a batch + flux_array = np.asarray(multigroup_flux, dtype=float) + if flux_array.ndim not in (1, 2): + raise ValueError('multigroup_flux must be 1-D or 2-D') + single = flux_array.ndim == 1 + fluxes = [flux_array] if single else list(flux_array) + + # check dimension consistency per flux + n_groups = len(energies) - 1 + for flux in fluxes: + if len(flux) != n_groups: + raise ValueError('Length of flux array should be len(energies)-1') + + # Resolve the library once; data availability is derived from it + if cross_sections is None: + cross_sections = _find_cross_sections(model=None) nuclides_with_data = _get_nuclides_with_data(cross_sections) - # If no nuclides were specified, default to all nuclides from the chain - if not nuclides: - nuclides = chain.nuclides - nuclides = [nuc.name for nuc in nuclides] - - # If no reactions specified, default to those in the chain file - if reactions is None: - reactions = chain.reactions - - # Delegate to the shared build-and-collapse engine (single-flux case) - return _flux_collapse_micros( - nuclides, reactions, energies, [multigroup_flux], - nuclides_with_data, temperature, cross_sections=cross_sections, - **init_kwargs)[0] + # Default nuclides/reactions from the chain only when needed + if not nuclides or reactions is None: + chain = _get_chain(chain_file) + if not nuclides: + nuclides = [nuc.name for nuc in chain.nuclides] + if reactions is None: + reactions = chain.reactions + + # Build the XS table once and collapse every flux against it + table = _build_xs_table_ce( + nuclides, reactions, energies, temperature, nuclides_with_data, + cross_sections=cross_sections, **init_kwargs) + micros = _collapse_fluxes(table, fluxes, nuclides, reactions) + return micros[0] if single else micros @classmethod def from_csv(cls, csv_file, **kwargs): diff --git a/tests/unit_tests/test_deplete_ce_collapse.py b/tests/unit_tests/test_deplete_ce_collapse.py index ececaedd672..5ab244d27e8 100644 --- a/tests/unit_tests/test_deplete_ce_collapse.py +++ b/tests/unit_tests/test_deplete_ce_collapse.py @@ -153,3 +153,93 @@ def test_collapse_fluxes_guards(): zero = _collapse_fluxes(table, [np.zeros(n_groups)], nuclides, reactions)[0] assert zero.data.shape == (1, 1, 1) assert np.all(zero.data == 0.0) + + +def test_from_multigroup_flux_batch(): + """A 2-D batch builds one shared table and matches the per-flux singular + result and the old per-domain collapse_rate reference, including a threshold + reaction and a zero-flux row, without mutating the input. + """ + nuclides = ['U235', 'U238', 'O16'] + reactions = ['fission', '(n,gamma)', '(n,2n)'] # (n,2n) is a threshold reaction + mts = [REACTION_MT[r] for r in reactions] + energies = _energies() + n_groups = energies.size - 1 + + rng = np.random.default_rng(2) + fluxes = np.vstack([rng.random(n_groups), rng.random(n_groups), + np.zeros(n_groups)]) + flux_copy = fluxes.copy() + + batch = MicroXS.from_multigroup_flux( + energies=energies, multigroup_flux=fluxes, chain_file=CHAIN_FILE, + nuclides=nuclides, reactions=reactions) + assert isinstance(batch, list) + assert len(batch) == len(fluxes) + assert all(isinstance(m, MicroXS) for m in batch) + + # The batch call must not mutate its input + assert np.array_equal(fluxes, flux_copy) + + # Each batch element equals the per-flux singular call ... + for row, micro in zip(fluxes, batch): + single = MicroXS.from_multigroup_flux( + energies=energies, multigroup_flux=row, chain_file=CHAIN_FILE, + nuclides=nuclides, reactions=reactions) + assert isinstance(single, MicroXS) + assert micro.data == pytest.approx(single.data, rel=1e-9, abs=1e-12) + + # ... and the old per-domain collapse_rate reference + ref = _reference_micros(nuclides, reactions, mts, energies, fluxes, TEMPERATURE) + for micro, ref_data in zip(batch, ref): + assert micro.data == pytest.approx(ref_data, rel=1e-9, abs=1e-12) + + # Zero-flux row -> all-zero MicroXS + assert np.all(batch[-1].data == 0.0) + + +def test_from_multigroup_flux_one_row_batch(): + """A 1-row 2-D batch returns a 1-element list, not an unwrapped MicroXS.""" + energies = _energies() + n_groups = energies.size - 1 + flux = np.random.default_rng(3).random((1, n_groups)) + + result = MicroXS.from_multigroup_flux( + energies=energies, multigroup_flux=flux, chain_file=CHAIN_FILE, + nuclides=['U235'], reactions=['fission']) + assert isinstance(result, list) + assert len(result) == 1 + assert isinstance(result[0], MicroXS) + + +def test_from_multigroup_flux_invalid_shape(): + """ndim 0/3 and ragged fluxes raise ValueError.""" + energies = _energies() + n_groups = energies.size - 1 + kwargs = dict(energies=energies, chain_file=CHAIN_FILE, + nuclides=['U235'], reactions=['fission']) + + with pytest.raises(ValueError): + MicroXS.from_multigroup_flux(multigroup_flux=1.0, **kwargs) + + with pytest.raises(ValueError): + MicroXS.from_multigroup_flux( + multigroup_flux=np.ones((2, 1, n_groups)), **kwargs) + + with pytest.raises(ValueError): + MicroXS.from_multigroup_flux( + multigroup_flux=[np.ones(n_groups), np.ones(n_groups - 1)], **kwargs) + + +def test_from_multigroup_flux_explicit_chain_no_config(monkeypatch): + """An explicit chain_file works even when openmc.config has no chain_file.""" + monkeypatch.delitem(openmc.config, 'chain_file', raising=False) + energies = _energies() + n_groups = energies.size - 1 + flux = np.random.default_rng(4).random(n_groups) + + # With no nuclides/reactions the chain must be resolved from chain_file + micro = MicroXS.from_multigroup_flux( + energies=energies, multigroup_flux=flux, chain_file=CHAIN_FILE) + assert isinstance(micro, MicroXS) + assert len(micro.nuclides) > 0 From 0b5b8072b9f861e658520e63079f043187d329a7 Mon Sep 17 00:00:00 2001 From: Perry Date: Wed, 10 Jun 2026 08:09:40 -0700 Subject: [PATCH 6/9] Group Materials.deplete by energy and temperature Group materials sharing an (energy structure, temperature) and collapse each group's fluxes in a single batched from_multigroup_flux call, so the cross section table is built once per distinct group instead of once per material. Output order is preserved. --- openmc/material.py | 29 ++++++++++--- tests/unit_tests/test_materials.py | 70 ++++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+), 7 deletions(-) diff --git a/openmc/material.py b/openmc/material.py index eded9c3faa6..fc553f95a40 100644 --- a/openmc/material.py +++ b/openmc/material.py @@ -2332,27 +2332,42 @@ def deplete( chain = _get_chain(chain_file) # Create MicroXS objects for all materials - micros = [] + micros = [None] * len(self) fluxes = [] with openmc.lib.TemporarySession(): - for material, flux, energy in zip( + # Group by (energy, temperature) so the table is built once per group + groups = {} + for i, (material, flux, energy) in enumerate(zip( self, multigroup_fluxes, energy_group_structures - ): + )): if material.volume is None: raise ValueError( f"Material {material.id} has no volume; cannot deplete" ) + fluxes.append(material.volume) temperature = material.temperature or 293.6 - micro_xs = openmc.deplete.MicroXS.from_multigroup_flux( + if isinstance(energy, str): + energy_key = energy + else: + energy = np.asarray(energy, dtype=float) + energy_key = energy.tobytes() + groups.setdefault((energy_key, temperature), []).append( + (i, energy, flux)) + + # Collapse each group's fluxes against a single shared table + for (_, temperature), members in groups.items(): + energy = members[0][1] + group_fluxes = [flux for _, _, flux in members] + micros_grp = openmc.deplete.MicroXS.from_multigroup_flux( energies=energy, - multigroup_flux=flux, + multigroup_flux=group_fluxes, chain_file=chain, temperature=temperature, reactions=reactions, ) - micros.append(micro_xs) - fluxes.append(material.volume) + for (i, _, _), micro in zip(members, micros_grp): + micros[i] = micro # Create a single operator for all materials operator = openmc.deplete.IndependentOperator( diff --git a/tests/unit_tests/test_materials.py b/tests/unit_tests/test_materials.py index 6620402fdcd..0d25a6d02f0 100644 --- a/tests/unit_tests/test_materials.py +++ b/tests/unit_tests/test_materials.py @@ -1,8 +1,12 @@ from pathlib import Path +from unittest import mock +import numpy as np import pytest import openmc +import openmc.deplete +import openmc.lib from openmc.deplete import Chain @@ -127,3 +131,69 @@ def __exit__(self, exc_type, exc, tb): source_rates=1.0, chain_file=chain, ) + + +def test_materials_deplete_groups_by_energy_temperature(monkeypatch): + """Materials.deplete groups materials by (energy, temperature): the group + cross section table is built once per distinct group, the micros keep + material order, and each equals the ungrouped from_multigroup_flux result. + """ + import openmc.deplete.microxs as microxs_mod + + chain = Path(__file__).parents[1] / "chain_ni.xml" + energy = "VITAMIN-J-42" + n_groups = 42 + + rng = np.random.default_rng(0) + fluxes = [rng.random(n_groups) for _ in range(3)] + temps = [293.6, 293.6, 600.0] + + def make_mat(nuclide, temperature): + m = openmc.Material() + m.add_nuclide(nuclide, 1.0) + m.set_density("g/cm3", 7.87) + m.depletable = True + m.temperature = temperature + m.volume = 1.0 + return m + + # mat0 and mat1 share (energy, T=293.6); mat2 differs only in temperature + mats = openmc.Materials([ + make_mat("Ni58", temps[0]), + make_mat("Ni60", temps[1]), + make_mat("Ni62", temps[2]), + ]) + + # Spy on the table build and capture the micros the operator receives + build_spy = mock.MagicMock(wraps=microxs_mod._build_xs_table_ce) + monkeypatch.setattr(microxs_mod, "_build_xs_table_ce", build_spy) + + captured = {} + def fake_operator(*args, **kwargs): + captured["micros"] = kwargs["micros"] + raise StopIteration + monkeypatch.setattr(openmc.deplete, "IndependentOperator", fake_operator) + + with pytest.raises(StopIteration): + mats.deplete( + multigroup_fluxes=fluxes, + energy_group_structures=[energy, energy, energy], + timesteps=[1.0], + source_rates=1.0, + chain_file=chain, + ) + + # Two distinct (energy, temperature) groups -> two builds, not three + assert build_spy.call_count == 2 + + micros = captured["micros"] + assert len(micros) == 3 + assert all(m is not None for m in micros) + + # Each per-material micro equals the ungrouped single-flux result + with openmc.lib.TemporarySession(): + for flux, temperature, micro in zip(fluxes, temps, micros): + ref = openmc.deplete.MicroXS.from_multigroup_flux( + energies=energy, multigroup_flux=flux, chain_file=chain, + temperature=temperature) + assert micro.data == pytest.approx(ref.data) From 9aad1ccf0c70820e1af6f23f28212b9fbe3a6bed Mon Sep 17 00:00:00 2001 From: Perry Date: Thu, 11 Jun 2026 11:50:18 -0700 Subject: [PATCH 7/9] Streamline the flux-collapse additions Docstrings trimmed to the package's sibling density and mechanical compactions (tuple init, conditional empty-matrix, dict-membership cache check, single batched collapse assignment). No functional change; the redundant collapse() length guard is dropped in favor of numpy's own shape error. --- openmc/deplete/microxs.py | 137 ++++--------------- openmc/lib/nuclide.py | 12 +- openmc/material.py | 18 +-- src/nuclide.cpp | 10 +- tests/unit_tests/test_deplete_ce_collapse.py | 115 ++++++---------- tests/unit_tests/test_materials.py | 12 +- 6 files changed, 90 insertions(+), 214 deletions(-) diff --git a/openmc/deplete/microxs.py b/openmc/deplete/microxs.py index b9f1d568b1e..29547b806fb 100644 --- a/openmc/deplete/microxs.py +++ b/openmc/deplete/microxs.py @@ -338,61 +338,28 @@ def get_microxs_and_flux( @dataclass class _SparseXSTable: - """Sparse storage of group cross sections for vectorized flux collapse. - - Stores only the non-zero ``(nuclide, reaction)`` pairs. ``xs_matrix`` has - shape ``(nnz, n_groups)``; ``nuc_indices`` and ``rxn_indices`` map each row - to its position in the dense ``(n_nuclides, n_reactions)`` result. - - Parameters - ---------- - nuclides : list of str - Nuclide names defining the result's nuclide axis. - reactions : list of str - Reaction names defining the result's reaction axis. - n_groups : int - Number of energy groups. - xs_matrix : numpy.ndarray - Group cross sections for the stored rows, shape ``(nnz, n_groups)``. - nuc_indices : numpy.ndarray - Row-to-nuclide index map, shape ``(nnz,)``. - rxn_indices : numpy.ndarray - Row-to-reaction index map, shape ``(nnz,)``. - - Notes - ----- - Provider-agnostic: rows may be filled from continuous-energy data - (:func:`_build_xs_table_ce`) or any other group cross section source. + """Sparse group cross sections for vectorized flux collapse. + Only non-zero ``(nuclide, reaction)`` pairs are stored: ``xs_matrix`` holds + one ``(n_groups,)`` row per pair and ``nuc_indices``/``rxn_indices`` map each + row into the dense ``(n_nuclides, n_reactions)`` result. Rows may come from + any group cross section source (e.g. :func:`_build_xs_table_ce`). """ nuclides: list[str] reactions: list[str] - n_groups: int xs_matrix: np.ndarray nuc_indices: np.ndarray rxn_indices: np.ndarray def collapse(self, phi_norm: np.ndarray) -> np.ndarray: - """Collapse the table to one-group cross sections. - - Parameters - ---------- - phi_norm : numpy.ndarray - Normalized group flux, shape ``(n_groups,)``. Should sum to 1. - - Returns - ------- - numpy.ndarray - Dense ``(n_nuclides, n_reactions)`` one-group cross sections. + """Collapse the table against a group flux with one matrix-vector product. + A normalized flux (summing to 1) yields one-group cross sections, a raw + flux yields reaction rates. Returns a dense + ``(n_nuclides, n_reactions)`` array. """ - if len(phi_norm) != self.n_groups: - raise ValueError( - f'phi_norm has length {len(phi_norm)} but table expects ' - f'{self.n_groups} groups') - collapsed_sparse = self.xs_matrix @ phi_norm result = np.zeros((len(self.nuclides), len(self.reactions))) - result[self.nuc_indices, self.rxn_indices] = collapsed_sparse + result[self.nuc_indices, self.rxn_indices] = self.xs_matrix @ phi_norm return result @@ -407,13 +374,10 @@ def _build_xs_table_ce( ) -> _SparseXSTable: """Build a sparse group cross section table from continuous-energy data. - For each requested ``(nuclide, reaction)`` the group-averaged cross sections - are computed once via the :meth:`openmc.lib.Nuclide.group_xs` C API and - stored as one sparse row, so every domain can later be collapsed with a - single matrix-vector product. Rows that are entirely zero (reaction absent, - or threshold above the whole group structure) are skipped. Cross sections - are evaluated inside a single :class:`openmc.lib.TemporarySession`, loading - each nuclide once. + Group-averaged cross sections for each requested ``(nuclide, reaction)`` are + computed once via :meth:`openmc.lib.Nuclide.group_xs` inside a single + :class:`openmc.lib.TemporarySession`; all-zero rows (reaction absent, or + threshold above the group structure) are skipped. Parameters ---------- @@ -422,30 +386,22 @@ def _build_xs_table_ce( reactions : sequence of str Reaction names defining the result's reaction axis. energies : sequence of float - Energy group boundaries in [eV], ascending, length ``n_groups + 1``. + Ascending energy group boundaries in [eV], length ``n_groups + 1``. temperature : float Temperature in [K] for cross section evaluation. nuclides_with_data : set Nuclides available in the cross section library; others are skipped. cross_sections : PathLike, optional - Cross section library for the session, so it matches the one + Cross section library for the session, matching the one ``nuclides_with_data`` was resolved from. Defaults to ``openmc.config``. **init_kwargs : dict - Keyword arguments passed to :func:`openmc.lib.init` via the temporary - session. - - Returns - ------- - _SparseXSTable - + Keyword arguments passed to :func:`openmc.lib.init`. """ mts = [REACTION_MT[name] for name in reactions] energies = np.asarray(energies, dtype=float) n_groups = len(energies) - 1 - rows = [] - nuc_idx_list = [] - rxn_idx_list = [] + rows, nuc_idx_list, rxn_idx_list = [], [], [] # Load against the same library nuclides_with_data was resolved from library = (openmc.config.patch('cross_sections', cross_sections) if cross_sections is not None else nullcontext()) @@ -462,48 +418,19 @@ def _build_xs_table_ce( nuc_idx_list.append(nuc_idx) rxn_idx_list.append(rxn_idx) - if rows: - xs_matrix = np.vstack(rows) - else: - xs_matrix = np.empty((0, n_groups)) + xs_matrix = np.vstack(rows) if rows else np.empty((0, n_groups)) return _SparseXSTable( - nuclides=list(nuclides), - reactions=list(reactions), - n_groups=n_groups, - xs_matrix=xs_matrix, - nuc_indices=np.array(nuc_idx_list, dtype=np.int32), - rxn_indices=np.array(rxn_idx_list, dtype=np.int32), - ) - - -def _collapse_fluxes( - table: _SparseXSTable, - fluxes: Sequence[np.ndarray], - nuclides: Sequence[str], - reactions: Sequence[str], -) -> list[MicroXS]: - """Collapse each domain's multigroup flux against a built XS table. + list(nuclides), list(reactions), xs_matrix, + np.array(nuc_idx_list, np.int32), np.array(rxn_idx_list, np.int32)) - Each flux is validated (finite, non-negative) and normalized to sum 1 before - being collapsed; a zero-sum flux yields an all-zero :class:`MicroXS`. - Parameters - ---------- - table : _SparseXSTable - Pre-built group cross section table. - fluxes : sequence of numpy.ndarray - One ``(n_groups,)`` group-flux vector per domain. - nuclides : sequence of str - Nuclide axis of the output MicroXS objects. - reactions : sequence of str - Reaction axis of the output MicroXS objects. - - Returns - ------- - list of MicroXS - One per domain, each of shape ``(n_nuclides, n_reactions, 1)``. +def _collapse_fluxes(table: _SparseXSTable, fluxes: Sequence[np.ndarray]) -> list[MicroXS]: + """Collapse each domain's multigroup flux against a built XS table. + Each flux is validated (finite, non-negative) and normalized to sum 1 before + collapse; a zero-sum flux yields an all-zero MicroXS. Returns one + ``(n_nuclides, n_reactions, 1)`` :class:`MicroXS` per domain. """ micros = [] for flux in fluxes: @@ -513,14 +440,10 @@ def _collapse_fluxes( if (flux < 0).any(): raise ValueError('Multigroup flux contains negative values') flux_sum = flux.sum() - if flux_sum == 0.0: - micros.append(MicroXS( - np.zeros((len(nuclides), len(reactions), 1)), - nuclides, reactions)) - continue - collapsed = table.collapse(flux / flux_sum) + # Zero-sum flux (all zeros, given the checks above) collapses to zeros + collapsed = table.collapse(flux / flux_sum if flux_sum else flux) micros.append(MicroXS(collapsed[:, :, np.newaxis], - nuclides, reactions)) + table.nuclides, table.reactions)) return micros @@ -674,7 +597,7 @@ def from_multigroup_flux( table = _build_xs_table_ce( nuclides, reactions, energies, temperature, nuclides_with_data, cross_sections=cross_sections, **init_kwargs) - micros = _collapse_fluxes(table, fluxes, nuclides, reactions) + micros = _collapse_fluxes(table, fluxes) return micros[0] if single else micros @classmethod diff --git a/openmc/lib/nuclide.py b/openmc/lib/nuclide.py index b4ed472f458..e1efe1d64fd 100644 --- a/openmc/lib/nuclide.py +++ b/openmc/lib/nuclide.py @@ -119,12 +119,9 @@ def collapse_rate(self, MT, temperature, energy, flux): def group_xs(self, MT, temperature, energy): """Calculate group-averaged microscopic cross sections. - Computes the cross section averaged over each energy group, - ``integral(sigma dE) / dE_group``, for a given reaction. This is a - flat-in-bin average -- it assumes the flux is constant within each - group -- so the group cross sections do not depend on the per-group - flux magnitudes and the collapsed reaction rate ``sum(flux * group_xs)`` - matches :meth:`collapse_rate`, which makes the same assumption. + The average over each group, ``integral(sigma dE) / dE_group``, + assumes a flat flux within the group, so the collapsed rate + ``sum(flux * group_xs)`` matches :meth:`collapse_rate`. Parameters ---------- @@ -138,8 +135,7 @@ def group_xs(self, MT, temperature, energy): Returns ------- numpy.ndarray - Group-averaged cross section in [b] for each of the - ``len(energy) - 1`` energy groups + Group-averaged cross section in [b] for each energy group """ energy = np.asarray(energy, dtype=float) diff --git a/openmc/material.py b/openmc/material.py index fc553f95a40..515c7b998ce 100644 --- a/openmc/material.py +++ b/openmc/material.py @@ -2347,18 +2347,14 @@ def deplete( ) fluxes.append(material.volume) temperature = material.temperature or 293.6 - if isinstance(energy, str): - energy_key = energy - else: - energy = np.asarray(energy, dtype=float) - energy_key = energy.tobytes() - groups.setdefault((energy_key, temperature), []).append( - (i, energy, flux)) + key = (energy if isinstance(energy, str) + else np.asarray(energy, dtype=float).tobytes()) + groups.setdefault((key, temperature), (energy, []))[1].append( + (i, flux)) # Collapse each group's fluxes against a single shared table - for (_, temperature), members in groups.items(): - energy = members[0][1] - group_fluxes = [flux for _, _, flux in members] + for (_, temperature), (energy, members) in groups.items(): + indices, group_fluxes = zip(*members) micros_grp = openmc.deplete.MicroXS.from_multigroup_flux( energies=energy, multigroup_flux=group_fluxes, @@ -2366,7 +2362,7 @@ def deplete( temperature=temperature, reactions=reactions, ) - for (i, _, _), micro in zip(members, micros_grp): + for i, micro in zip(indices, micros_grp): micros[i] = micro # Create a single operator for all materials diff --git a/src/nuclide.cpp b/src/nuclide.cpp index 6a7afb33c18..8f1dbecedbe 100644 --- a/src/nuclide.cpp +++ b/src/nuclide.cpp @@ -1088,19 +1088,15 @@ void Nuclide::group_xs( const auto& rx = reactions_[i_rx]; // Determine temperature index - int64_t i_temp; - double f; - std::tie(i_temp, f) = this->find_temperature(temperature); + auto [i_temp, f] = this->find_temperature(temperature); // Get group-averaged cross sections at lower temperature - const auto& grid_low = grid_[i_temp].energy; - rx->group_xs(i_temp, energy, grid_low, xs); + rx->group_xs(i_temp, energy, grid_[i_temp].energy, xs); if (f > 0.0) { // Interpolate element-wise between lower and higher temperature - const auto& grid_high = grid_[i_temp + 1].energy; vector xs_high(xs.size(), 0.0); - rx->group_xs(i_temp + 1, energy, grid_high, xs_high); + rx->group_xs(i_temp + 1, energy, grid_[i_temp + 1].energy, xs_high); for (std::size_t g = 0; g < xs.size(); ++g) xs[g] += f * (xs_high[g] - xs[g]); } diff --git a/tests/unit_tests/test_deplete_ce_collapse.py b/tests/unit_tests/test_deplete_ce_collapse.py index 5ab244d27e8..9e63a33bafe 100644 --- a/tests/unit_tests/test_deplete_ce_collapse.py +++ b/tests/unit_tests/test_deplete_ce_collapse.py @@ -20,16 +20,16 @@ _find_cross_sections, _get_nuclides_with_data) CHAIN_FILE = Path(__file__).parents[1] / "chain_simple.xml" -GROUP_STRUCTURE = "CASMO-40" +ENERGIES = np.asarray(GROUP_STRUCTURES["CASMO-40"], dtype=float) +N_GROUPS = ENERGIES.size - 1 TEMPERATURE = 293.6 +NUCLIDES = ['U235', 'U238', 'O16'] +REACTIONS = ['fission', '(n,gamma)', '(n,2n)'] # (n,2n) is a threshold reaction -def _energies(): - return np.asarray(GROUP_STRUCTURES[GROUP_STRUCTURE], dtype=float) - - -def _reference_micros(nuclides, reactions, mts, energies, fluxes, temperature): +def _reference_micros(nuclides, reactions, fluxes): """Old per-domain reference: collapse_rate(mt, T, E, flux/sum) per (nuc, mt).""" + mts = [REACTION_MT[r] for r in reactions] out = [] with openmc.lib.TemporarySession(): lib_nucs = {n: openmc.lib.load_nuclide(n) for n in nuclides} @@ -42,7 +42,7 @@ def _reference_micros(nuclides, reactions, mts, energies, fluxes, temperature): for i, nuc in enumerate(nuclides): for j, mt in enumerate(mts): arr[i, j, 0] = lib_nucs[nuc].collapse_rate( - mt, temperature, energies, phi) + mt, TEMPERATURE, ENERGIES, phi) out.append(arr) return out @@ -55,9 +55,7 @@ def test_group_xs_equivalence(): guard catches a forgotten division by the group width (the undivided integral would be ~1e5-1e7 eV-b, far outside the barns range). """ - energies = _energies() - n_groups = energies.size - 1 - dE = np.diff(energies) + dE = np.diff(ENERGIES) rng = np.random.default_rng(0) # (nuclide, MT): include a threshold reaction (U238 (n,2n)) plus non-threshold @@ -69,12 +67,12 @@ def test_group_xs_equivalence(): with openmc.lib.TemporarySession(): for name, mt in cases: nuc = openmc.lib.load_nuclide(name) - group_xs = nuc.group_xs(mt, TEMPERATURE, energies) - assert group_xs.shape == (n_groups,) + group_xs = nuc.group_xs(mt, TEMPERATURE, ENERGIES) + assert group_xs.shape == (N_GROUPS,) - for flux in (rng.random(n_groups), rng.random(n_groups), dE): + for flux in (rng.random(N_GROUPS), rng.random(N_GROUPS), dE): assert np.dot(flux, group_xs) == pytest.approx( - nuc.collapse_rate(mt, TEMPERATURE, energies, flux), rel=1e-10) + nuc.collapse_rate(mt, TEMPERATURE, ENERGIES, flux), rel=1e-10) # Per-group AVERAGE cross section, in barns -- not the raw integral. assert group_xs.max() < 1e6 @@ -102,31 +100,26 @@ def test_ce_collapse_matches_collapse_rate(): reproduce the old per-domain collapse_rate result, with a threshold reaction and a zero-flux domain. """ - nuclides = ['U235', 'U238', 'O16'] - reactions = ['fission', '(n,gamma)', '(n,2n)'] # (n,2n) is a threshold reaction - mts = [REACTION_MT[r] for r in reactions] - energies = _energies() - n_groups = energies.size - 1 nuclides_with_data = _get_nuclides_with_data(_find_cross_sections(model=None)) rng = np.random.default_rng(1) - fluxes = [rng.random(n_groups), rng.random(n_groups), np.zeros(n_groups)] + fluxes = [rng.random(N_GROUPS), rng.random(N_GROUPS), np.zeros(N_GROUPS)] table = _build_xs_table_ce( - nuclides, reactions, energies, TEMPERATURE, nuclides_with_data) - new = _collapse_fluxes(table, fluxes, nuclides, reactions) - ref = _reference_micros(nuclides, reactions, mts, energies, fluxes, TEMPERATURE) + NUCLIDES, REACTIONS, ENERGIES, TEMPERATURE, nuclides_with_data) + new = _collapse_fluxes(table, fluxes) + ref = _reference_micros(NUCLIDES, REACTIONS, fluxes) for micro, ref_data in zip(new, ref): assert micro.data == pytest.approx(ref_data, rel=1e-9, abs=1e-12) # Zero-flux domain -> all-zero MicroXS of the right shape. - assert new[-1].data.shape == (len(nuclides), len(reactions), 1) + assert new[-1].data.shape == (len(NUCLIDES), len(REACTIONS), 1) assert np.all(new[-1].data == 0.0) # from_multigroup_flux delegates to the same engine -> identical result. micro = MicroXS.from_multigroup_flux( - energies=energies, multigroup_flux=fluxes[0], chain_file=CHAIN_FILE, - nuclides=nuclides, reactions=reactions) + energies=ENERGIES, multigroup_flux=fluxes[0], chain_file=CHAIN_FILE, + nuclides=NUCLIDES, reactions=REACTIONS) assert micro.data == pytest.approx(ref[0], rel=1e-9, abs=1e-12) @@ -134,23 +127,17 @@ def test_collapse_fluxes_guards(): """_collapse_fluxes rejects non-finite / negative flux and zero-fills zero flux.""" nuclides = ['U235'] reactions = ['fission'] - energies = _energies() - n_groups = energies.size - 1 nuclides_with_data = _get_nuclides_with_data(_find_cross_sections(model=None)) table = _build_xs_table_ce( - nuclides, reactions, energies, TEMPERATURE, nuclides_with_data) - - nan_flux = np.ones(n_groups) - nan_flux[0] = np.nan - with pytest.raises(ValueError): - _collapse_fluxes(table, [nan_flux], nuclides, reactions) + nuclides, reactions, ENERGIES, TEMPERATURE, nuclides_with_data) - neg_flux = np.ones(n_groups) - neg_flux[0] = -1.0 - with pytest.raises(ValueError): - _collapse_fluxes(table, [neg_flux], nuclides, reactions) + for bad in (np.nan, -1.0): + flux = np.ones(N_GROUPS) + flux[0] = bad + with pytest.raises(ValueError): + _collapse_fluxes(table, [flux]) - zero = _collapse_fluxes(table, [np.zeros(n_groups)], nuclides, reactions)[0] + zero = _collapse_fluxes(table, [np.zeros(N_GROUPS)])[0] assert zero.data.shape == (1, 1, 1) assert np.all(zero.data == 0.0) @@ -160,20 +147,14 @@ def test_from_multigroup_flux_batch(): result and the old per-domain collapse_rate reference, including a threshold reaction and a zero-flux row, without mutating the input. """ - nuclides = ['U235', 'U238', 'O16'] - reactions = ['fission', '(n,gamma)', '(n,2n)'] # (n,2n) is a threshold reaction - mts = [REACTION_MT[r] for r in reactions] - energies = _energies() - n_groups = energies.size - 1 - rng = np.random.default_rng(2) - fluxes = np.vstack([rng.random(n_groups), rng.random(n_groups), - np.zeros(n_groups)]) + fluxes = np.vstack([rng.random(N_GROUPS), rng.random(N_GROUPS), + np.zeros(N_GROUPS)]) flux_copy = fluxes.copy() batch = MicroXS.from_multigroup_flux( - energies=energies, multigroup_flux=fluxes, chain_file=CHAIN_FILE, - nuclides=nuclides, reactions=reactions) + energies=ENERGIES, multigroup_flux=fluxes, chain_file=CHAIN_FILE, + nuclides=NUCLIDES, reactions=REACTIONS) assert isinstance(batch, list) assert len(batch) == len(fluxes) assert all(isinstance(m, MicroXS) for m in batch) @@ -184,13 +165,13 @@ def test_from_multigroup_flux_batch(): # Each batch element equals the per-flux singular call ... for row, micro in zip(fluxes, batch): single = MicroXS.from_multigroup_flux( - energies=energies, multigroup_flux=row, chain_file=CHAIN_FILE, - nuclides=nuclides, reactions=reactions) + energies=ENERGIES, multigroup_flux=row, chain_file=CHAIN_FILE, + nuclides=NUCLIDES, reactions=REACTIONS) assert isinstance(single, MicroXS) assert micro.data == pytest.approx(single.data, rel=1e-9, abs=1e-12) # ... and the old per-domain collapse_rate reference - ref = _reference_micros(nuclides, reactions, mts, energies, fluxes, TEMPERATURE) + ref = _reference_micros(NUCLIDES, REACTIONS, fluxes) for micro, ref_data in zip(batch, ref): assert micro.data == pytest.approx(ref_data, rel=1e-9, abs=1e-12) @@ -200,12 +181,10 @@ def test_from_multigroup_flux_batch(): def test_from_multigroup_flux_one_row_batch(): """A 1-row 2-D batch returns a 1-element list, not an unwrapped MicroXS.""" - energies = _energies() - n_groups = energies.size - 1 - flux = np.random.default_rng(3).random((1, n_groups)) + flux = np.random.default_rng(3).random((1, N_GROUPS)) result = MicroXS.from_multigroup_flux( - energies=energies, multigroup_flux=flux, chain_file=CHAIN_FILE, + energies=ENERGIES, multigroup_flux=flux, chain_file=CHAIN_FILE, nuclides=['U235'], reactions=['fission']) assert isinstance(result, list) assert len(result) == 1 @@ -214,32 +193,22 @@ def test_from_multigroup_flux_one_row_batch(): def test_from_multigroup_flux_invalid_shape(): """ndim 0/3 and ragged fluxes raise ValueError.""" - energies = _energies() - n_groups = energies.size - 1 - kwargs = dict(energies=energies, chain_file=CHAIN_FILE, + kwargs = dict(energies=ENERGIES, chain_file=CHAIN_FILE, nuclides=['U235'], reactions=['fission']) - with pytest.raises(ValueError): - MicroXS.from_multigroup_flux(multigroup_flux=1.0, **kwargs) - - with pytest.raises(ValueError): - MicroXS.from_multigroup_flux( - multigroup_flux=np.ones((2, 1, n_groups)), **kwargs) - - with pytest.raises(ValueError): - MicroXS.from_multigroup_flux( - multigroup_flux=[np.ones(n_groups), np.ones(n_groups - 1)], **kwargs) + for bad in (1.0, np.ones((2, 1, N_GROUPS)), + [np.ones(N_GROUPS), np.ones(N_GROUPS - 1)]): + with pytest.raises(ValueError): + MicroXS.from_multigroup_flux(multigroup_flux=bad, **kwargs) def test_from_multigroup_flux_explicit_chain_no_config(monkeypatch): """An explicit chain_file works even when openmc.config has no chain_file.""" monkeypatch.delitem(openmc.config, 'chain_file', raising=False) - energies = _energies() - n_groups = energies.size - 1 - flux = np.random.default_rng(4).random(n_groups) + flux = np.random.default_rng(4).random(N_GROUPS) # With no nuclides/reactions the chain must be resolved from chain_file micro = MicroXS.from_multigroup_flux( - energies=energies, multigroup_flux=flux, chain_file=CHAIN_FILE) + energies=ENERGIES, multigroup_flux=flux, chain_file=CHAIN_FILE) assert isinstance(micro, MicroXS) assert len(micro.nuclides) > 0 diff --git a/tests/unit_tests/test_materials.py b/tests/unit_tests/test_materials.py index 0d25a6d02f0..238d1a2138a 100644 --- a/tests/unit_tests/test_materials.py +++ b/tests/unit_tests/test_materials.py @@ -6,6 +6,7 @@ import openmc import openmc.deplete +import openmc.deplete.microxs as microxs_mod import openmc.lib from openmc.deplete import Chain @@ -138,8 +139,6 @@ def test_materials_deplete_groups_by_energy_temperature(monkeypatch): cross section table is built once per distinct group, the micros keep material order, and each equals the ungrouped from_multigroup_flux result. """ - import openmc.deplete.microxs as microxs_mod - chain = Path(__file__).parents[1] / "chain_ni.xml" energy = "VITAMIN-J-42" n_groups = 42 @@ -168,11 +167,8 @@ def make_mat(nuclide, temperature): build_spy = mock.MagicMock(wraps=microxs_mod._build_xs_table_ce) monkeypatch.setattr(microxs_mod, "_build_xs_table_ce", build_spy) - captured = {} - def fake_operator(*args, **kwargs): - captured["micros"] = kwargs["micros"] - raise StopIteration - monkeypatch.setattr(openmc.deplete, "IndependentOperator", fake_operator) + op_spy = mock.MagicMock(side_effect=StopIteration) + monkeypatch.setattr(openmc.deplete, "IndependentOperator", op_spy) with pytest.raises(StopIteration): mats.deplete( @@ -186,7 +182,7 @@ def fake_operator(*args, **kwargs): # Two distinct (energy, temperature) groups -> two builds, not three assert build_spy.call_count == 2 - micros = captured["micros"] + micros = op_spy.call_args.kwargs["micros"] assert len(micros) == 3 assert all(m is not None for m in micros) From 302e301061d72c7d13921acee5d11e7984e7b459 Mon Sep 17 00:00:00 2001 From: Perry Date: Wed, 10 Jun 2026 08:21:17 -0700 Subject: [PATCH 8/9] Optional/Bonus: Cache per-temperature XS table in FluxCollapseHelper Collapse each material's raw flux against a group cross section table built once per distinct temperature and reused across materials and timesteps, replacing the per-(nuclide, reaction) collapse_rate calls. The nuclides setter clears the cache when the nuclide set changes; the direct-tally override is unchanged. Optional because the cross-timestep reuse assumes the tallied nuclide set is stable. Today _get_reaction_nuclides returns a fixed set, but it documents an intent to filter on nonzero density; with that, the set would grow step to step and each change drops the cache back to a once-per-step rebuild. --- openmc/deplete/helpers.py | 45 +++++--- tests/unit_tests/test_deplete_activation.py | 107 ++++++++++++++++++++ 2 files changed, 139 insertions(+), 13 deletions(-) diff --git a/openmc/deplete/helpers.py b/openmc/deplete/helpers.py index 14780d61a20..612e4c70edc 100644 --- a/openmc/deplete/helpers.py +++ b/openmc/deplete/helpers.py @@ -9,7 +9,7 @@ from numbers import Real import sys -from numpy import dot, zeros, newaxis, asarray +from numpy import dot, zeros, newaxis, asarray, ix_ from openmc.mpi import comm from openmc.checkvalue import check_type, check_greater_than @@ -235,6 +235,10 @@ class FluxCollapseHelper(ReactionRateHelper): .. versionadded:: 0.12.1 + .. versionchanged:: 0.15.4 + Flux-collapsed rates reuse a group cross section table built once per + material temperature instead of collapsing each pair every step. + Parameters ---------- n_nucs : int @@ -255,15 +259,24 @@ class FluxCollapseHelper(ReactionRateHelper): nuclides : list of str All nuclides with desired reaction rates. + Notes + ----- + One group cross section table is cached per distinct material temperature, + which suits a handful of temperatures but not hundreds of distinct ones. + """ def __init__(self, n_nucs, n_reacts, energies, reactions=None, nuclides=None): super().__init__(n_nucs, n_reacts) self._energies = asarray(energies) + self._xs_tables = {} # group XS tables cached by temperature self._reactions_direct = list(reactions) if reactions is not None else [] self._nuclides_direct = list(nuclides) if nuclides is not None else None @ReactionRateHelper.nuclides.setter def nuclides(self, nuclides): + # Invalidate cached tables when the nuclide set changes (rare) + if nuclides != self._nuclides: + self._xs_tables = {} ReactionRateHelper.nuclides.fset(self, nuclides) if self._reactions_direct and self._nuclides_direct is None: self._rate_tally.nuclides = nuclides @@ -347,6 +360,16 @@ def reset_tally_means(self): if self._reactions_direct: self._rate_tally_means_cache = None + def _table(self, temperature): + """Return the cached group cross section table for a temperature.""" + if temperature not in self._xs_tables: + # Lazy import avoids the microxs -> coupled_operator -> helpers cycle + from .microxs import _build_xs_table_ce + self._xs_tables[temperature] = _build_xs_table_ce( + self.nuclides, self._scores, self._energies, temperature, + set(self.nuclides)) + return self._xs_tables[temperature] + def get_material_rates(self, mat_index, nuc_index, react_index): """Return an array of reaction rates for a material @@ -384,20 +407,16 @@ def get_material_rates(self, mat_index, nuc_index, react_index): mat = self._materials[mat_index] + # One matvec gives all flux-mode rates, in (nuclides, scores) order + self._results_cache[ix_(nuc_index, react_index)] = \ + self._table(mat.temperature).collapse(flux) + + # Overwrite the (nuclide, reaction) pairs tallied directly for name, i_nuc in zip(self.nuclides, nuc_index): - for mt, score, i_rx in zip(self._mts, self._scores, react_index): + for score, i_rx in zip(self._scores, react_index): if score in self._reactions_direct and name in nuclides_direct: - # Get reaction rate from tally - i_rx_direct = direct_rx_index[score] - i_nuc_direct = direct_nuc_index[name] - self._results_cache[i_nuc, i_rx] = rx_rates[i_nuc_direct, i_rx_direct] - else: - # Use flux to collapse reaction rate (per N) - nuc = openmc.lib.nuclides[name] - rate_per_nuc = nuc.collapse_rate( - mt, mat.temperature, self._energies, flux) - - self._results_cache[i_nuc, i_rx] = rate_per_nuc + self._results_cache[i_nuc, i_rx] = rx_rates[ + direct_nuc_index[name], direct_rx_index[score]] return self._results_cache diff --git a/tests/unit_tests/test_deplete_activation.py b/tests/unit_tests/test_deplete_activation.py index 2f171388050..6c5ca5a20c7 100644 --- a/tests/unit_tests/test_deplete_activation.py +++ b/tests/unit_tests/test_deplete_activation.py @@ -1,11 +1,18 @@ from math import pi, log, log10 from random import uniform, normalvariate +from types import SimpleNamespace +from unittest import mock import numpy as np import openmc.deplete +import openmc.deplete.microxs as microxs_mod import openmc +import openmc.lib import pytest +from openmc.data import REACTION_MT +from openmc.deplete.helpers import FluxCollapseHelper +from openmc.mgxs import GROUP_STRUCTURES @pytest.fixture @@ -193,3 +200,103 @@ def test_flux_rr_missing_nuclide(run_in_tmpdir, model): op, [100.0], source_rates=[10.0] ) integrator.integrate() + + +CASMO40 = np.asarray(GROUP_STRUCTURES['CASMO-40'], dtype=float) +N_GROUPS = CASMO40.size - 1 +TEMPERATURE = 293.6 + + +def _wire_helper(nuclides, scores, flux, n_nucs=None, + reactions_direct=None, nuclides_direct=None): + """Wire a FluxCollapseHelper for one material without building tallies.""" + helper = FluxCollapseHelper( + n_nucs or len(nuclides), len(scores), CASMO40, + reactions=reactions_direct, nuclides=nuclides_direct) + helper._materials = [SimpleNamespace(temperature=TEMPERATURE)] + helper._xs_tables = {} + helper._mts = [REACTION_MT[s] for s in scores] + helper._scores = list(scores) + helper._flux_tally_means_cache = np.asarray(flux, dtype=float) + helper.nuclides = list(nuclides) + return helper + + +def test_flux_collapse_helper_cache_lifecycle(monkeypatch): + """The FluxCollapseHelper per-temperature cache lifecycle: flux-path rates + equal collapse_rate; the table is built once per temperature; re-setting the + same nuclides (as happens each timestep) does not rebuild it; and growing the + nuclide set clears the cache, so the table is rebuilt and the newly-added + nuclide gets correct (non-stale) rates. + """ + nuclides = ['U235', 'U238', 'O16'] + grown = nuclides + ['Pu239'] # larger set, added later + scores = ['fission', '(n,gamma)', '(n,2n)'] # (n,2n) is a threshold reaction + flux = np.random.default_rng(0).random(N_GROUPS) + react_index = list(range(len(scores))) + + build_spy = mock.MagicMock(wraps=microxs_mod._build_xs_table_ce) + monkeypatch.setattr(microxs_mod, '_build_xs_table_ce', build_spy) + + with openmc.lib.TemporarySession(): + # Size the result cache for the larger (grown) set from the start + helper = _wire_helper(nuclides, scores, flux, n_nucs=len(grown)) + rates = helper.get_material_rates( + 0, list(range(len(nuclides))), react_index).copy() + + # Flux-path rate equals collapse_rate for every (nuclide, score) ... + for i, name in enumerate(nuclides): + nuc = openmc.lib.nuclides[name] + for j, s in enumerate(scores): + expected = nuc.collapse_rate( + REACTION_MT[s], TEMPERATURE, CASMO40, flux) + assert rates[i, j] == pytest.approx(expected, rel=1e-10) + + # ... built once for the single temperature ... + assert build_spy.call_count == 1 + + # ... and re-setting the same nuclide list (each step) does not rebuild + helper.nuclides = list(nuclides) + helper.get_material_rates(0, list(range(len(nuclides))), react_index) + assert build_spy.call_count == 1 + + # Growing the set clears the cache -> exactly one rebuild at the same T + helper.nuclides = list(grown) + rates = helper.get_material_rates( + 0, list(range(len(grown))), react_index).copy() + assert build_spy.call_count == 2 + + # The newly-added nuclide gets correct rates, not stale zeros + pu = openmc.lib.nuclides['Pu239'] + for j, s in enumerate(scores): + expected = pu.collapse_rate( + REACTION_MT[s], TEMPERATURE, CASMO40, flux) + assert rates[len(nuclides), j] == pytest.approx(expected, rel=1e-10) + assert rates[len(nuclides), 0] > 0.0 # Pu239 fission nonzero -> not stale + + +def test_flux_collapse_helper_direct_override(): + """A direct-tally (nuclide, reaction) pair overrides the flux-collapsed value.""" + nuclides = ['U235', 'U238'] + scores = ['fission', '(n,gamma)'] + flux = np.random.default_rng(1).random(N_GROUPS) + + with openmc.lib.TemporarySession(): + # Direct-tally U235 fission; everything else via flux collapse + helper = _wire_helper(nuclides, scores, flux, + reactions_direct=['fission'], nuclides_direct=['U235']) + direct_value = 1.2345e-3 + helper._rate_tally = SimpleNamespace(nuclides=['U235']) + helper._rate_tally_means_cache = np.array([[direct_value]]) + + rates = helper.get_material_rates( + 0, list(range(len(nuclides))), list(range(len(scores)))).copy() + + # U235 fission comes from the direct tally, not the flux collapse + assert rates[0, 0] == pytest.approx(direct_value) + + # U235 (n,gamma) still comes from the flux collapse + nuc = openmc.lib.nuclides['U235'] + expected = nuc.collapse_rate( + REACTION_MT['(n,gamma)'], TEMPERATURE, CASMO40, flux) + assert rates[0, 1] == pytest.approx(expected, rel=1e-10) From 30f5d1a6b855657ce67a787bb67b0b5249c8c5af Mon Sep 17 00:00:00 2001 From: Perry Date: Wed, 17 Jun 2026 21:02:54 -0700 Subject: [PATCH 9/9] Avoid copying the whole flux batch in from_multigroup_flux Detect single vs batch from the first flux's ndim and pass the batch through, instead of copying the whole (N, G) array just to read ndim. --- openmc/deplete/microxs.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/openmc/deplete/microxs.py b/openmc/deplete/microxs.py index 29547b806fb..37105ae7e0d 100644 --- a/openmc/deplete/microxs.py +++ b/openmc/deplete/microxs.py @@ -567,12 +567,12 @@ def from_multigroup_flux( if not np.all(np.diff(energies) > 0): raise ValueError('Energy group boundaries must be in ascending order') - # A 1-D flux is a single domain; 2-D (or a list of 1-D) is a batch - flux_array = np.asarray(multigroup_flux, dtype=float) - if flux_array.ndim not in (1, 2): - raise ValueError('multigroup_flux must be 1-D or 2-D') - single = flux_array.ndim == 1 - fluxes = [flux_array] if single else list(flux_array) + # A 1-D flux is a single domain; 2-D (or a list of 1-D arrays) is a batch + try: + single = {1: True, 2: False}[1 + np.ndim(multigroup_flux[0])] + except (TypeError, IndexError, KeyError): + raise ValueError('multigroup_flux must be 1-D or 2-D') from None + fluxes = [np.asarray(multigroup_flux, dtype=float)] if single else multigroup_flux # check dimension consistency per flux n_groups = len(energies) - 1