From 49b896b0eb339d4cd6cbba549ee75feb7ee130f4 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Mon, 2 Mar 2026 21:06:52 -0600 Subject: [PATCH] Enable "hybrid" tallies in `get_microxs_and_flux` (#3831) --- openmc/deplete/microxs.py | 147 ++++++++++++++++-- tests/regression_tests/microxs/test.py | 12 +- .../test_reference_materials_hybrid.csv | 7 + tests/unit_tests/test_deplete_microxs.py | 129 ++++++++++++++- 4 files changed, 279 insertions(+), 16 deletions(-) create mode 100644 tests/regression_tests/microxs/test_reference_materials_hybrid.csv diff --git a/openmc/deplete/microxs.py b/openmc/deplete/microxs.py index 879a2d4ee..e6e2dbce2 100644 --- a/openmc/deplete/microxs.py +++ b/openmc/deplete/microxs.py @@ -50,7 +50,8 @@ def get_microxs_and_flux( chain_file: PathLike | Chain | None = None, path_statepoint: PathLike | None = None, path_input: PathLike | None = None, - run_kwargs=None + run_kwargs=None, + reaction_rate_opts: dict | None = None, ) -> tuple[list[np.ndarray], list[MicroXS]]: """Generate microscopic cross sections and fluxes for multiple domains. @@ -80,10 +81,12 @@ def get_microxs_and_flux( Energy group boundaries in [eV] or the name of the group structure. If left as None energies will default to [0.0, 100e6] reaction_rate_mode : {"direct", "flux"}, optional - Indicate how reaction rates should be calculated. The "direct" method - tallies reaction rates directly. The "flux" method tallies a multigroup - flux spectrum and then collapses multigroup reaction rates after a - transport solve (with an option to tally some reaction rates directly). + The "direct" method tallies reaction rates directly (per energy + group). The "flux" method tallies a multigroup flux spectrum and then + collapses reaction rates after a transport solve. When + `reaction_rate_opts` is provided with `reaction_rate_mode='flux'`, the + specified nuclide/reaction pairs are tallied directly and those values + override the flux-collapsed values. chain_file : PathLike or Chain, optional Path to the depletion chain XML file or an instance of openmc.deplete.Chain. Used to determine cross sections for materials not @@ -99,6 +102,10 @@ def get_microxs_and_flux( not kept. run_kwargs : dict, optional Keyword arguments passed to :meth:`openmc.Model.run` + reaction_rate_opts : dict, optional + When `reaction_rate_mode="flux"`, allows selecting a subset of + nuclide/reaction pairs to be computed via direct reaction-rate tallies + (per energy group). Supported keys: "nuclides", "reactions". Returns ------- @@ -153,12 +160,36 @@ def get_microxs_and_flux( flux_tally.scores = ['flux'] model.tallies = [flux_tally] + # Prepare reaction-rate tally for 'direct' or subset for 'flux' with opts + rr_tally = None + rr_nuclides: list[str] = [] + rr_reactions: list[str] = [] if reaction_rate_mode == 'direct': + rr_nuclides = list(nuclides) + rr_reactions = list(reactions) + elif reaction_rate_mode == 'flux' and reaction_rate_opts: + opts = reaction_rate_opts or {} + rr_nuclides = list(opts.get('nuclides', [])) + rr_reactions = list(opts.get('reactions', [])) + # Keep only requested pairs within overall sets + if rr_nuclides: + rr_nuclides = [n for n in rr_nuclides if n in set(nuclides)] + if rr_reactions: + rr_reactions = [r for r in rr_reactions if r in set(reactions)] + + # Only construct tally if both lists are non-empty + if rr_nuclides and rr_reactions: rr_tally = openmc.Tally(name='MicroXS RR') - rr_tally.filters = [domain_filter, energy_filter] - rr_tally.nuclides = nuclides + # Use 1-group energy filter for RR in flux mode + if reaction_rate_mode == 'flux': + rr_energy_filter = openmc.EnergyFilter( + [energy_filter.values[0], energy_filter.values[-1]]) + else: + rr_energy_filter = energy_filter + rr_tally.filters = [domain_filter, rr_energy_filter] + rr_tally.nuclides = rr_nuclides rr_tally.multiply_density = False - rr_tally.scores = reactions + rr_tally.scores = rr_reactions model.tallies.append(rr_tally) if openmc.lib.is_initialized: @@ -196,7 +227,7 @@ def get_microxs_and_flux( # Read in tally results (on all ranks) with StatePoint(statepoint_path) as sp: - if reaction_rate_mode == 'direct': + if rr_tally is not None: rr_tally = sp.tallies[rr_tally.id] rr_tally._read_results() flux_tally = sp.tallies[flux_tally.id] @@ -209,24 +240,31 @@ def get_microxs_and_flux( # Create list where each item corresponds to one domain fluxes = list(flux.squeeze((1, 2))) - if reaction_rate_mode == 'direct': + # If we built a reaction-rate tally, compute microscopic cross sections + if rr_tally is not None: # Get reaction rates reaction_rates = rr_tally.get_reshaped_data() # (domains, groups, nuclides, reactions) # Make energy groups last dimension reaction_rates = np.moveaxis(reaction_rates, 1, -1) # (domains, nuclides, reactions, groups) + # If RR is 1-group, sum flux over groups + if reaction_rate_mode == "flux": + flux = flux.sum(axis=-1, keepdims=True) # (domains, 1, 1, 1) + # Divide RR by flux to get microscopic cross sections. The indexing # ensures that only non-zero flux values are used, and broadcasting is # applied to align the shapes of reaction_rates and flux for division. - xs = np.empty_like(reaction_rates) # (domains, nuclides, reactions, groups) + xs = np.zeros_like(reaction_rates) # (domains, nuclides, reactions, groups) d, _, _, g = np.nonzero(flux) xs[d, ..., g] = reaction_rates[d, ..., g] / flux[d, :, :, g] # Create lists where each item corresponds to one domain - micros = [MicroXS(xs_i, nuclides, reactions) for xs_i in xs] - else: - micros = [MicroXS.from_multigroup_flux( + direct_micros = [MicroXS(xs_i, rr_nuclides, rr_reactions) for xs_i in xs] + + # If using flux mode, compute flux-collapsed microscopic XS + if reaction_rate_mode == 'flux': + flux_micros = [MicroXS.from_multigroup_flux( energies=energies, multigroup_flux=flux_i, chain_file=chain_file, @@ -234,6 +272,14 @@ def get_microxs_and_flux( reactions=reactions ) for flux_i in fluxes] + # Decide which micros to use and merge if needed + if reaction_rate_mode == 'flux' and rr_tally is not None: + micros = [m1.merge(m2) for m1, m2 in zip(flux_micros, direct_micros)] + elif rr_tally is not None: + micros = direct_micros + else: + micros = flux_micros + # Reset tallies model.tallies = original_tallies @@ -484,6 +530,79 @@ class MicroXS: return cls(data, nuclides, reactions) + def merge(self, other: Self, prefer: str = 'other') -> Self: + """Merge two MicroXS objects by taking the union of nuclides/reactions. + + If the two objects contain overlapping nuclide/reaction entries, values + from `other` will overwrite values from `self` when `prefer='other'`. + When `prefer='self'`, values from `self` are retained for overlapping + entries, and values from `other` are used only for non-overlapping + entries. + + Parameters + ---------- + other : MicroXS + Other MicroXS instance to merge with this one. + prefer : {"other", "self"} + Which instance's data should take precedence on overlap. + + Returns + ------- + MicroXS + New instance containing the merged data. + """ + check_value('prefer', prefer, {'other', 'self'}) + + # Require same number of energy groups + if self.data.shape[2] != other.data.shape[2]: + raise ValueError( + 'Cannot merge MicroXS with different number of energy groups: ' + f"{self.data.shape[2]} vs {other.data.shape[2]}. Ensure that " + 'both were generated with consistent group structures and ' + 'treatments (e.g., both multigroup or both collapsed).' + ) + + # Build unified axes preserving order (self first, then other's new) + new_nuclides = list(self.nuclides) + for nuc in other.nuclides: + if nuc not in self._index_nuc: + new_nuclides.append(nuc) + new_reactions = list(self.reactions) + for rx in other.reactions: + if rx not in self._index_rx: + new_reactions.append(rx) + + # Allocate and fill from self (self's nuclides/reactions map to the + # first indices of new_nuclides/new_reactions by construction) + groups = self.data.shape[2] + data = np.zeros((len(new_nuclides), len(new_reactions), groups)) + idx_n = {nuc: i for i, nuc in enumerate(new_nuclides)} + idx_r = {rx: i for i, rx in enumerate(new_reactions)} + + n_self = len(self.nuclides) + r_self = len(self.reactions) + data[:n_self, :r_self] = self.data + + # Build destination index arrays for other's nuclides/reactions + dst_n = np.array([idx_n[nuc] for nuc in other.nuclides]) + dst_r = np.array([idx_r[rx] for rx in other.reactions]) + + # Copy from other, respecting precedence + if prefer == 'other': + data[np.ix_(dst_n, dst_r)] = other.data + else: + # Copy only entries where nuc or rx is absent from self + nuc_is_new = np.array( + [nuc not in self._index_nuc for nuc in other.nuclides]) + rx_is_new = np.array( + [rx not in self._index_rx for rx in other.reactions]) + mask = nuc_is_new[:, np.newaxis] | rx_is_new[np.newaxis, :] + src_i, src_j = np.where(mask) + if src_i.size: + data[dst_n[src_i], dst_r[src_j]] = other.data[src_i, src_j] + + return MicroXS(data, new_nuclides, new_reactions) + def write_microxs_hdf5( micros: Sequence[MicroXS], diff --git a/tests/regression_tests/microxs/test.py b/tests/regression_tests/microxs/test.py index 70833bb39..781be3ef7 100644 --- a/tests/regression_tests/microxs/test.py +++ b/tests/regression_tests/microxs/test.py @@ -34,6 +34,7 @@ def model(): [ ("materials", "direct"), ("materials", "flux"), + ("materials", "hybrid"), ("mesh", "direct"), ("mesh", "flux"), ] @@ -49,12 +50,21 @@ def test_from_model(model, domain_type, rr_mode): domains = mesh nuclides = ['U235', 'O16', 'Xe135'] kwargs = { - 'reaction_rate_mode': rr_mode, 'chain_file': CHAIN_FILE, 'path_statepoint': 'neutron_transport.h5', } if rr_mode == 'flux': + kwargs['reaction_rate_mode'] = 'flux' kwargs['energies'] = 'CASMO-40' + elif rr_mode == 'hybrid': + kwargs['reaction_rate_mode'] = 'flux' + kwargs['energies'] = 'CASMO-40' + kwargs['reaction_rate_opts'] = { + 'nuclides': ['U235'], + 'reactions': ['fission'], + } + else: + kwargs['reaction_rate_mode'] = rr_mode _, test_xs = get_microxs_and_flux(model, domains, nuclides, **kwargs) if config['update']: test_xs[0].to_csv(f'test_reference_{domain_type}_{rr_mode}.csv') diff --git a/tests/regression_tests/microxs/test_reference_materials_hybrid.csv b/tests/regression_tests/microxs/test_reference_materials_hybrid.csv new file mode 100644 index 000000000..6d17a5e2b --- /dev/null +++ b/tests/regression_tests/microxs/test_reference_materials_hybrid.csv @@ -0,0 +1,7 @@ +nuclides,reactions,groups,xs +U235,"(n,gamma)",1,0.1500301670375847 +U235,fission,1,1.2578145727734291 +O16,"(n,gamma)",1,0.00012069778439640312 +O16,fission,1,0.0 +Xe135,"(n,gamma)",1,0.014820264774863558 +Xe135,fission,1,0.0 diff --git a/tests/unit_tests/test_deplete_microxs.py b/tests/unit_tests/test_deplete_microxs.py index 5762a8511..340ba6163 100644 --- a/tests/unit_tests/test_deplete_microxs.py +++ b/tests/unit_tests/test_deplete_microxs.py @@ -6,12 +6,15 @@ to a custom file with new depletion_chain node from os import remove from pathlib import Path +from unittest.mock import patch import pytest -from openmc.deplete import MicroXS +import openmc +from openmc.deplete import MicroXS, get_microxs_and_flux import numpy as np ONE_GROUP_XS = Path(__file__).parents[1] / "micro_xs_simple.csv" +CHAIN_FILE = Path(__file__).parents[1] / "chain_simple.xml" def test_from_array(): @@ -124,3 +127,127 @@ def test_microxs_zero_flux(): # All microscopic cross sections should be zero assert np.all(microxs.data == 0.0) + + +def test_hybrid_tally_setup(): + """In hybrid mode a 1-group RR tally is added alongside the flux tally.""" + # Create a simple model with one material and a few nuclides for testing + model = openmc.Model() + mat = openmc.Material(components={'U235': 1.0, 'O16': 2.0}) + sphere = openmc.Sphere(r=10.0, boundary_type='vacuum') + cell = openmc.Cell(region=-sphere, fill=mat) + model.geometry = openmc.Geometry([cell]) + model.settings.batches = 2 + model.settings.particles = 10 + + # Define 2-group energy structure for the test + energies = [0., 0.625, 2.0e7] + + # Function to replace Model.run and capture the tallies that were created + captured = {} + def capture_run(**kwargs): + captured['tallies'] = list(model.tallies) + raise StopIteration + + # Call get_microxs_and_flux but replace Model.run with a function that + # captures the tallies and raises StopIteration to exit early + with patch.object(model, 'run', side_effect=capture_run): + with pytest.raises(StopIteration): + get_microxs_and_flux( + model, [mat], + nuclides=['U235', 'O16'], + reactions=['fission', '(n,gamma)'], + energies=energies, + reaction_rate_mode='flux', + reaction_rate_opts={'nuclides': ['U235'], 'reactions': ['fission']}, + chain_file=CHAIN_FILE, + ) + + # Check that both tallies were created with the expected properties + tally_names = [t.name for t in captured['tallies']] + assert 'MicroXS flux' in tally_names + assert 'MicroXS RR' in tally_names + + # Check that the RR tally has the expected nuclides and reactions + rr = next(t for t in captured['tallies'] if t.name == 'MicroXS RR') + assert rr.nuclides == ['U235'] + assert rr.scores == ['fission'] + + # RR tally must use a 1-group energy filter spanning the full energy range + ef = next(f for f in rr.filters if isinstance(f, openmc.EnergyFilter)) + assert len(ef.values) == 2 + assert ef.values[0] == pytest.approx(energies[0]) + assert ef.values[-1] == pytest.approx(energies[-1]) + +# --------------------------------------------------------------------------- +# Tests for MicroXS.merge() +# --------------------------------------------------------------------------- + +def _make_microxs(nuclides, reactions, values, groups=1): + """Helper: build a MicroXS from a flat list of values (nuclide-major order).""" + data = np.array(values, dtype=float).reshape( + len(nuclides), len(reactions), groups) + return MicroXS(data, nuclides, reactions) + + +def test_merge_disjoint(): + """Merging two MicroXS with no overlapping nuclides or reactions.""" + m1 = _make_microxs(['U235', 'U238'], ['fission', '(n,gamma)'], + [1., 2., 3., 4.]) + m2 = _make_microxs(['Pu239'], ['(n,2n)'], [5.]) + + merged = m1.merge(m2) + + assert merged.nuclides == ['U235', 'U238', 'Pu239'] + assert merged.reactions == ['fission', '(n,gamma)', '(n,2n)'] + assert merged.data.shape == (3, 3, 1) + + # Self data preserved + assert merged['U235', 'fission'] == pytest.approx([1.]) + assert merged['U238', '(n,gamma)'] == pytest.approx([4.]) + # New data from other + assert merged['Pu239', '(n,2n)'] == pytest.approx([5.]) + # Cross-terms that had no data should be zero + assert merged['U235', '(n,2n)'] == pytest.approx([0.]) + assert merged['Pu239', 'fission'] == pytest.approx([0.]) + + +def test_merge_prefer_other(): + """prefer='other': other overwrites shared entries, adds new ones.""" + # m1: U235 and U238, reactions fission and (n,gamma) + m1 = _make_microxs(['U235', 'U238'], ['fission', '(n,gamma)'], + [1., 2., 3., 4.]) + # m2: only U235, reactions fission (overlap) and (n,2n) (new) + m2 = _make_microxs(['U235'], ['fission', '(n,2n)'], [9., 5.]) + + merged = m1.merge(m2) + + # Nuclide/reaction sets + assert set(merged.nuclides) == {'U235', 'U238'} + assert set(merged.reactions) == {'fission', '(n,gamma)', '(n,2n)'} + + # U235/fission overwritten by m2 + assert merged['U235', 'fission'] == pytest.approx([9.]) + # U235/(n,gamma) untouched + assert merged['U235', '(n,gamma)'] == pytest.approx([2.]) + # New reaction added from m2 + assert merged['U235', '(n,2n)'] == pytest.approx([5.]) + # U238 data preserved + assert merged['U238', 'fission'] == pytest.approx([3.]) + assert merged['U238', '(n,gamma)'] == pytest.approx([4.]) + # U238/(n,2n) not in either → zero + assert merged['U238', '(n,2n)'] == pytest.approx([0.]) + + +def test_merge_prefer_self(): + """prefer='self': shared pairs keep self's value; new entries use other's value.""" + m1 = _make_microxs(['U235', 'U238'], ['fission', '(n,gamma)'], + [1., 2., 3., 4.]) + m2 = _make_microxs(['U235'], ['fission', '(n,2n)'], [9., 5.]) + + merged = m1.merge(m2, prefer='self') + + # U235/fission: self wins + assert merged['U235', 'fission'] == pytest.approx([1.]) + # U235/(n,2n): other used + assert merged['U235', '(n,2n)'] == pytest.approx([5.])