From 9aad1ccf0c70820e1af6f23f28212b9fbe3a6bed Mon Sep 17 00:00:00 2001 From: Perry Date: Thu, 11 Jun 2026 11:50:18 -0700 Subject: [PATCH] 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 | 133 ++++--------------- 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, 88 insertions(+), 212 deletions(-) diff --git a/openmc/deplete/microxs.py b/openmc/deplete/microxs.py index b9f1d568b1..29547b806f 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), - ) + list(nuclides), list(reactions), xs_matrix, + np.array(nuc_idx_list, np.int32), np.array(rxn_idx_list, np.int32)) -def _collapse_fluxes( - table: _SparseXSTable, - fluxes: Sequence[np.ndarray], - nuclides: Sequence[str], - reactions: Sequence[str], -) -> list[MicroXS]: +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 - 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)``. - + 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 @@ class MicroXS: 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 b4ed472f45..e1efe1d64f 100644 --- a/openmc/lib/nuclide.py +++ b/openmc/lib/nuclide.py @@ -119,12 +119,9 @@ class Nuclide(_FortranObject): 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 @@ class Nuclide(_FortranObject): 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 fc553f95a4..515c7b998c 100644 --- a/openmc/material.py +++ b/openmc/material.py @@ -2347,18 +2347,14 @@ class Materials(cv.CheckedList): ) 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 @@ class Materials(cv.CheckedList): 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 6a7afb33c1..8f1dbecedb 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 5ab244d27e..9e63a33baf 100644 --- a/tests/unit_tests/test_deplete_ce_collapse.py +++ b/tests/unit_tests/test_deplete_ce_collapse.py @@ -20,16 +20,16 @@ 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" +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) + 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) + for bad in (np.nan, -1.0): + flux = np.ones(N_GROUPS) + flux[0] = bad + with pytest.raises(ValueError): + _collapse_fluxes(table, [flux]) - 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] + 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 0d25a6d02f..238d1a2138 100644 --- a/tests/unit_tests/test_materials.py +++ b/tests/unit_tests/test_materials.py @@ -6,6 +6,7 @@ import pytest 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 test_materials_deplete_groups_by_energy_temperature(monkeypatch): 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 test_materials_deplete_groups_by_energy_temperature(monkeypatch): # 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)