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

@@ -9,7 +9,7 @@ from itertools import product
 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 @@ class FluxCollapseHelper(ReactionRateHelper):
         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 @@ class FluxCollapseHelper(ReactionRateHelper):
 
         mat = self._materials[mat_index]
 
-        for name, i_nuc in zip(self.nuclides, nuc_index):
-            for mt, score, i_rx in zip(self._mts, 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)
+        # 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)
 
-                    self._results_cache[i_nuc, i_rx] = rate_per_nuc
+        # Overwrite the (nuclide, reaction) pairs tallied directly
+        for name, i_nuc in zip(self.nuclides, nuc_index):
+            for score, i_rx in zip(self._scores, react_index):
+                if score in self._reactions_direct and name in nuclides_direct:
+                    self._results_cache[i_nuc, i_rx] = rx_rates[
+                        direct_nuc_index[name], direct_rx_index[score]]
 
         return self._results_cache
 

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)