Skip to content

Refactor form_matrix method on depletion chain class #3892

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
paulromano merged 3 commits into from
Apr 3, 2026
+156 −78
Merged

Refactor form_matrix method on depletion chain class #3892

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ccd7fa8
refactor form_matrix method
eepeterson Mar 20, 2026
51ebac3
Add back comment
paulromano Apr 3, 2026
3bdbf5e
Add two unit tests for decay_matrix
paulromano Apr 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
211 changes: 133 additions & 78 deletions openmc/deplete/chain.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,6 +269,7 @@ def __init__(self):
self.reactions = []
self.nuclide_dict = {}
self._fission_yields = None
self._decay_matrix = None

def __contains__(self, nuclide):
return nuclide in self.nuclide_dict
Expand Down Expand Up @@ -604,8 +605,67 @@ def get_default_fission_yields(self):
out[nuc.name] = dict(yield_obj)
return out

def form_matrix(self, rates, fission_yields=None):
"""Forms depletion matrix.
@property
def decay_matrix(self):
"""Sparse CSC decay transmutation matrix.

Contains only terms from radioactive decay: diagonal loss terms
and off-diagonal gain terms (branching ratios, alpha/proton
production). Independent of reaction rates, so computed once and
cached.

See Also
--------
:meth:`form_rxn_matrix`, :meth:`form_matrix`
"""
if self._decay_matrix is None:
n = len(self)
rows, cols, vals = [], [], []

def setval(i, j, val):
rows.append(i)
cols.append(j)
vals.append(val)

for i, nuc in enumerate(self.nuclides):
# Loss from radioactive decay
if nuc.half_life is not None:
decay_constant = math.log(2) / nuc.half_life
if decay_constant != 0.0:
setval(i, i, -decay_constant)

# Gain from radioactive decay
if nuc.n_decay_modes != 0:
for decay_type, target, branching_ratio in nuc.decay_modes:
branch_val = branching_ratio * decay_constant

# Allow for total annihilation for debug purposes
if branch_val != 0.0:
if target is not None:
k = self.nuclide_dict[target]
setval(k, i, branch_val)

# Produce alphas and protons from decay
if 'alpha' in decay_type:
k = self.nuclide_dict.get('He4')
if k is not None:
count = decay_type.count('alpha')
setval(k, i, count * branch_val)
elif 'p' in decay_type:
k = self.nuclide_dict.get('H1')
if k is not None:
count = decay_type.count('p')
setval(k, i, count * branch_val)

self._decay_matrix = csc_array((vals, (rows, cols)), shape=(n, n))
return self._decay_matrix

def form_rxn_matrix(self, rates, fission_yields=None):
"""Form the reaction-rate portion of the transmutation matrix.

Builds only the terms that depend on reaction rates: transmutation
reactions and fission product yields. Does not include radioactive
decay terms (see :attr:`decay_matrix`).

Parameters
----------
Expand All @@ -620,19 +680,19 @@ def form_matrix(self, rates, fission_yields=None):
Returns
-------
scipy.sparse.csc_array
Sparse matrix representing depletion.
Sparse matrix representing reaction-rate terms.

See Also
--------
:meth:`get_default_fission_yields`
:attr:`decay_matrix`, :meth:`form_matrix`
"""
reactions = set()

n = len(self)

# we accumulate indices and value entries for everything and create the matrix
# in one step at the end to avoid expensive index checks scipy otherwise does.
# Accumulate indices/values and then create the matrix at the end to
# avoid expensive index checks scipy otherwise does.
rows, cols, vals = [], [], []

def setval(i, j, val):
rows.append(i)
cols.append(j)
Expand All @@ -642,79 +702,73 @@ def setval(i, j, val):
fission_yields = self.get_default_fission_yields()

for i, nuc in enumerate(self.nuclides):
# Loss from radioactive decay
if nuc.half_life is not None:
decay_constant = math.log(2) / nuc.half_life
if decay_constant != 0.0:
setval(i, i, -decay_constant)

# Gain from radioactive decay
if nuc.n_decay_modes != 0:
for decay_type, target, branching_ratio in nuc.decay_modes:
branch_val = branching_ratio * decay_constant

# Allow for total annihilation for debug purposes
if branch_val != 0.0:
if target is not None:
k = self.nuclide_dict[target]
setval(k, i, branch_val)

# Produce alphas and protons from decay
if 'alpha' in decay_type:
k = self.nuclide_dict.get('He4')
if k is not None:
count = decay_type.count('alpha')
setval(k, i, count * branch_val)
elif 'p' in decay_type:
k = self.nuclide_dict.get('H1')
if k is not None:
count = decay_type.count('p')
setval(k, i, count * branch_val)

if nuc.name in rates.index_nuc:
# Extract all reactions for this nuclide in this cell
nuc_ind = rates.index_nuc[nuc.name]
nuc_rates = rates[nuc_ind, :]

for r_type, target, _, br in nuc.reactions:
# Extract reaction index, and then final reaction rate
r_id = rates.index_rx[r_type]
path_rate = nuc_rates[r_id]

# Loss term -- make sure we only count loss once for
# reactions with branching ratios
if r_type not in reactions:
reactions.add(r_type)
if path_rate != 0.0:
setval(i, i, -path_rate)

# Gain term; allow for total annihilation for debug purposes
if r_type != 'fission':
if target is not None and path_rate != 0.0:
k = self.nuclide_dict[target]
setval(k, i, path_rate * br)
if nuc.name not in rates.index_nuc:
continue

# Determine light nuclide production, e.g., (n,d) should
# produce H2
light_nucs = REACTIONS[r_type].secondaries
for light_nuc in light_nucs:
k = self.nuclide_dict.get(light_nuc)
if k is not None:
setval(k, i, path_rate * br)
nuc_ind = rates.index_nuc[nuc.name]
nuc_rates = rates[nuc_ind, :]

for r_type, target, _, br in nuc.reactions:
r_id = rates.index_rx[r_type]
path_rate = nuc_rates[r_id]

# Loss term -- make sure we only count loss once for
# reactions with branching ratios
if r_type not in reactions:
reactions.add(r_type)
if path_rate != 0.0:
setval(i, i, -path_rate)

# Gain term; allow for total annihilation for debug purposes
if r_type != 'fission':
if target is not None and path_rate != 0.0:
k = self.nuclide_dict[target]
setval(k, i, path_rate * br)

# Determine light nuclide production, e.g., (n,d) should
# produce H2
light_nucs = REACTIONS[r_type].secondaries
for light_nuc in light_nucs:
k = self.nuclide_dict.get(light_nuc)
if k is not None:
setval(k, i, path_rate * br)

else:
for product, y in fission_yields[nuc.name].items():
yield_val = y * path_rate
if yield_val != 0.0:
k = self.nuclide_dict[product]
setval(k, i, yield_val)
else:
for product, y in fission_yields[nuc.name].items():
yield_val = y * path_rate
if yield_val != 0.0:
k = self.nuclide_dict[product]
setval(k, i, yield_val)

# Clear set of reactions
reactions.clear()
reactions.clear()

# Return CSC representation instead of DOK
return csc_array((vals, (rows, cols)), shape=(n, n))

def form_matrix(self, rates, fission_yields=None):
"""Form the full transmutation matrix (decay + reactions).

Parameters
----------
rates : numpy.ndarray
2D array indexed by (nuclide, reaction)
fission_yields : dict, optional
Option to use a custom set of fission yields. Expected
to be of the form ``{parent : {product : f_yield}}``
with string nuclide names for ``parent`` and ``product``,
and ``f_yield`` as the respective fission yield

Returns
-------
scipy.sparse.csc_array
Sparse matrix representing depletion.

See Also
--------
:attr:`decay_matrix`, :meth:`form_rxn_matrix`,
:meth:`get_default_fission_yields`
"""
return self.decay_matrix + self.form_rxn_matrix(rates, fission_yields)

def add_redox_term(self, matrix, buffer, oxidation_states):
r"""Adds a redox term to the depletion matrix from data contained in
the matrix itself and a few user-inputs.
Expand Down Expand Up @@ -807,7 +861,7 @@ def form_rr_term(self, tr_rates, current_timestep, mats):
# Use DOK as intermediate representation
n = len(self)
matrix = dok_array((n, n))

check_type("mats", mats, (tuple, str))
if not isinstance(mats, str):
check_type("mats", mats, tuple, str)
Expand All @@ -816,8 +870,8 @@ def form_rr_term(self, tr_rates, current_timestep, mats):
else:
mat = mats
dest_mat = None
# Build transfer term

# Build transfer term
components = tr_rates.get_components(mat, current_timestep, dest_mat)

for i, nuc in enumerate(self.nuclides):
Expand All @@ -829,7 +883,7 @@ def form_rr_term(self, tr_rates, current_timestep, mats):
else:
continue
matrix[i, i] = sum(tr_rates.get_external_rate(mat, key, current_timestep, dest_mat))

# Return CSC instead of DOK
return matrix.tocsc()

Expand Down Expand Up @@ -1363,6 +1417,7 @@ def _get_chain(

def _invalidate_chain_cache(chain):
"""Invalidate the cache for a specific Chain (when it is modifed)."""
chain._decay_matrix = None
if hasattr(chain, '_xml_path'):
# Remove all entries with the same path as self._xml_path
for key in list(_CHAIN_CACHE.keys()):
Expand Down
23 changes: 23 additions & 0 deletions tests/unit_tests/test_deplete_chain.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -246,6 +246,29 @@ def test_form_matrix(simple_chain):
assert new_mat[r, c] == mat[r, c]


def test_decay_matrix(simple_chain):
"""Test that decay_matrix contains only radioactive decay terms."""
# Nuclide order: H1(0), A(1), B(2), C(3)
decay_A = log(2) / 2.36520E+04
decay_B = log(2) / 3.29040E+04

expected = np.zeros((4, 4))
expected[1, 1] = -decay_A # Loss: A decays
expected[2, 1] = decay_A * 0.6 # A -> B (branching ratio 0.6)
expected[3, 1] = decay_A * 0.4 # A -> C (branching ratio 0.4)
expected[1, 2] = decay_B # B -> A (branching ratio 1.0)
expected[2, 2] = -decay_B # Loss: B decays

assert np.allclose(expected, simple_chain.decay_matrix.toarray())


def test_decay_matrix_cached(simple_chain):
"""Test that decay_matrix is lazily computed and returns the same object."""
m1 = simple_chain.decay_matrix
m2 = simple_chain.decay_matrix
assert m1 is m2


def test_getitem():
"""Test nuc_by_ind converter function."""
chain = Chain()
Expand Down
Loading