From f3e7c7c623d2aba1a524b98a89467d1002b17fc1 Mon Sep 17 00:00:00 2001
From: Andrew Johnson <drewej@protonmail.com>
Date: Wed, 24 Jul 2019 11:58:54 -0500
Subject: [PATCH] Support for passing fission yields to Chain.form_matrix

In support of using energy dependent fission yields, this
commit allows a dictionary of fission yields to be passed
into Chain.form_matrix. The dictionary is expected to
be of the form
        {source_name: {target_name: yield_fraction}}
where source_name and target_name are string GND names
for fissionable nuclide and fission yield products, respectively.

Currently, the fission yield dictionary does not have to be passed,
defaulting to using the thermal yield values. This is done to
make testing easier, and for back compatability. The goal of this
feature, however, is to pass material [and thus spectrum] specific
yields into this method.

The test test_deplete_chain::test_form_matrix has been modified
to pass the exact fission yield dictionary into Chain.form_matrix.
The resulting matrix is compared against the matrix when
no yields are passed.
---
 openmc/deplete/chain.py                | 25 ++++++++++++++++++++++---
 tests/unit_tests/test_deplete_chain.py |  7 +++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/openmc/deplete/chain.py b/openmc/deplete/chain.py
index 1f16d9cafa..865f533ceb 100644
--- a/openmc/deplete/chain.py
+++ b/openmc/deplete/chain.py
@@ -129,6 +129,7 @@ class Chain(object):
         self.nuclides = []
         self.reactions = []
         self.nuclide_dict = OrderedDict()
+        self._default_fsn_yields = None
 
     def __contains__(self, nuclide):
         return nuclide in self.nuclide_dict
@@ -374,13 +375,27 @@ class Chain(object):
             clean_indentation(root_elem)
             tree.write(str(filename), encoding='utf-8')
 
-    def form_matrix(self, rates):
+    def _build_default_yields(self):
+        """Return dictionary {str: {str: float}}"""
+        # Take lowest energy for back compatability
+        # Should be removed by end of this feature
+        out = {}
+        for nuc in self.nuclides:
+            if len(nuc.yield_data) == 0:
+                continue
+            _energy, yield_data = sorted(nuc.yield_data.items())[0]
+            out[nuc.name] = {prod: frac for prod, frac in yield_data}
+        return out
+
+    def form_matrix(self, rates, fission_yields=None):
         """Forms depletion matrix.
 
         Parameters
         ----------
         rates : numpy.ndarray
             2D array indexed by (nuclide, reaction)
+        fission_yields : dictionary of tuple to float, optional
+            Option to use a custom set of fission yields
 
         Returns
         -------
@@ -391,6 +406,11 @@ class Chain(object):
         matrix = defaultdict(float)
         reactions = set()
 
+        if fission_yields is None:
+            if self._default_fsn_yields is None:
+                self._default_fsn_yields = self._build_default_yields()
+            fission_yields = self._default_fsn_yields
+
         for i, nuc in enumerate(self.nuclides):
 
             if nuc.n_decay_modes != 0:
@@ -438,8 +458,7 @@ class Chain(object):
                             # Assume that we should always use thermal fission
                             # yields. At some point it would be nice to account
                             # for the energy-dependence..
-                            energy, data = sorted(nuc.yield_data.items())[0]
-                            for product, y in data:
+                            for product, y in fission_yields[nuc.name].items():
                                 yield_val = y * path_rate
                                 if yield_val != 0.0:
                                     k = self.nuclide_dict[product]
diff --git a/tests/unit_tests/test_deplete_chain.py b/tests/unit_tests/test_deplete_chain.py
index 33b13b596d..df82eab1b0 100644
--- a/tests/unit_tests/test_deplete_chain.py
+++ b/tests/unit_tests/test_deplete_chain.py
@@ -3,6 +3,7 @@
 from collections.abc import Mapping
 import os
 from pathlib import Path
+from itertools import product
 
 import numpy as np
 from openmc.data import zam, ATOMIC_SYMBOL
@@ -220,6 +221,12 @@ def test_form_matrix(simple_chain):
     assert mat[1, 2] == mat12
     assert mat[2, 2] == mat22
 
+    # Pass equivalent fission yields directly
+    # Ensure identical matrix is formed
+    f_yields = {"C": {"A": 0.0292737, "B": 0.002566345}}
+    new_mat = chain.form_matrix(react[0], f_yields)
+    for r, c in product(range(3), range(3)):
+        assert new_mat[r, c] == mat[r, c]
 
 def test_getitem():
     """ Test nuc_by_ind converter function. """