Store depletion fission yields with common yield matrix

Add two new classes for working with fission yield data
on the Chain. The primary class is FissionYieldDistribution.
This class stores distributions for one nuclide with
potentially many energies, with the assumption that the
products don't change __too__ much across energy.
Looking at the CASL chain and one generated by ENDF data,
this appears to be the case. Most distributions are
"full" in the sense that energies produce the same products.
There are some cases where one or two products may not
exist for a given energy.

The FissionYieldDistribution retains a dictionary-like behavior,
e.g. d[0.0253]["Xe135"] is a valid command and returns the
yield for Xe-135 at a "thermal" spectrum, due to yields provided
at 0.0253 eV. This is done with a helper class, _FissionYield,
implemented first to support this behavior, but also to
allow vector-operations in combining fission yields.

These _FissionYield objects store the same product vector and
a view into the underlying yield_matrix for a single energy.
To support simple weighting of yields, the __mull__ and
__iadd__ methods are implemented. A copy method is provided
to remove the chance of modifying the original yield data.

test_deplete_chain and test_deplete_nuclide have been modified
in order to utilize these classes, without ruining the validity
of the tests. Tests for the view / copy methods and vector
operations on _FissionYield objects are included.

openmc/deplete/nuclide.py

@@ -3,12 +3,19 @@
 Contains the per-nuclide components of a depletion chain.
 """
 
+from numbers import Real
 from collections import namedtuple
+from collections.abc import Iterable, Mapping
+
 try:
     import lxml.etree as ET
 except ImportError:
     import xml.etree.ElementTree as ET
 
+from numpy import asarray, fromiter, empty
+
+from openmc.checkvalue import check_type, check_length
+
 
 DecayTuple = namedtuple('DecayTuple', 'type target branching_ratio')
 DecayTuple.__doc__ = """\
@@ -83,7 +90,7 @@ class Nuclide(object):
         Maps tabulated energy to list of (product, yield) for all
         neutron-induced fission products.
     yield_energies : list of float
-        Energies at which fission product yiels exist
+        Energies at which fission product yields exist
 
     """
 
@@ -165,12 +172,7 @@ class Nuclide(object):
 
         fpy_elem = element.find('neutron_fission_yields')
         if fpy_elem is not None:
-            for yields_elem in fpy_elem.iter('fission_yields'):
-                E = float(yields_elem.get('energy'))
-                products = yields_elem.find('products').text.split()
-                yields = [float(y) for y in
-                          yields_elem.find('data').text.split()]
-                nuc.yield_data[E] = list(zip(products, yields))
+            nuc.yield_data = FissionYieldDistribution.from_xml_element(fpy_elem)
             nuc.yield_energies = list(sorted(nuc.yield_data.keys()))
 
         return nuc
@@ -211,14 +213,181 @@ class Nuclide(object):
             fpy_elem = ET.SubElement(elem, 'neutron_fission_yields')
             energy_elem = ET.SubElement(fpy_elem, 'energies')
             energy_elem.text = ' '.join(str(E) for E in self.yield_energies)
-
-            for E in self.yield_energies:
-                yields_elem = ET.SubElement(fpy_elem, 'fission_yields')
-                yields_elem.set('energy', str(E))
-
-                products_elem = ET.SubElement(yields_elem, 'products')
-                products_elem.text = ' '.join(x[0] for x in self.yield_data[E])
-                data_elem = ET.SubElement(yields_elem, 'data')
-                data_elem.text = ' '.join(str(x[1]) for x in self.yield_data[E])
+            self.yield_data.to_xml_element(fpy_elem)
 
         return elem
+
+
+class FissionYieldDistribution(Mapping):
+    """Class for storing energy-dependent fission yields for a single nuclide
+
+    Parameters
+    ----------
+    ordered_energies : iterable of real
+        Energies for which fission yield data exist
+    orderded_products : iterable of str
+        Fission products produced by this parent at all energies
+    group_fission_yields : numpy.ndarray or iterable of iterable of float
+        Array of shape ``(n_energy, n_products)`` where
+        ``group_fission_yields[g][j]`` is the yield of
+        ``ordered_products[j]`` due to a fission in energy region ``g``.
+
+    Attributes
+    ----------
+    energies : tuple
+        Energies for which fission yields exist. Converted for
+        indexing
+    products : tuple
+        Fission products produced at all energies. Converted
+        for indexing
+    yield_matrix : numpy.ndarray
+        Array ``(n_energy, n_products)`` where
+        ``yield_matrix[g, j]`` is the fission yield of product
+        ``j`` for energy group ``g``.
+
+    See Also
+    --------
+    :meth:`from_xml`, :meth:`from_dict`
+    """
+
+    def __init__(self, ordered_energies, ordered_products, group_fission_yields):
+        check_type("energies", ordered_energies, Iterable, Real)
+        self.energies = tuple(ordered_energies)
+        check_type("products", ordered_products, Iterable, str)
+        self.products = tuple(ordered_products)
+        yield_matrix = asarray(group_fission_yields, dtype=float)
+        if yield_matrix.shape != (len(self.energies), len(self.products)):
+            raise ValueError(
+                "Shape of yield matrix inconsistent. "
+                "Should be ({}, {}), is {}".format(
+                    len(energy_map), len(ordered_products), yield_matrix.shape))
+        self.yield_matrix = yield_matrix
+
+    def __len__(self):
+        return len(self.energies)
+
+    def __getitem__(self, energy):
+        if energy not in self.energies:
+            raise KeyError(energy)
+        return _FissionYield(
+            self.products, self.yield_matrix[self.energies.index(energy)])
+
+    def __iter__(self):
+        return iter(self.energies)
+
+    @classmethod
+    def from_xml_element(cls, element):
+        """Construct a distribution from a depletion chain xml file
+
+        Parameters
+        ----------
+        element : xml.etree.ElementTree.Element
+            XML element to pull fission yield data from
+
+        Returns
+        -------
+        FissionYieldDistribution
+        """
+        yields = {}
+        for elem_index, yield_elem in enumerate(element.iter("fission_yields")):
+            energy = float(yield_elem.get("energy"))
+            products = yield_elem.find("products").text.split()
+            yield_mapobj = map(float, yield_elem.find("data").text.split())
+            # Get a map of products to their corresponding yield
+            yields[energy] = dict(zip(products, yield_mapobj))
+
+        return cls.from_dict(yields)
+
+    @classmethod
+    def from_dict(cls, fission_yields):
+        """Construct a distribution from a dictionary of yields
+
+        Parameters
+        -----------
+        fission_yields : dict
+            Dictionary ``{energy: {product: yield}}``
+
+        Returns
+        -------
+        FissionYieldDistribution
+        """
+        # mapping {energy: {product: value}}
+        energies = tuple(sorted(fission_yields))
+
+        # Get a consistent set of products to produce a matrix of yields
+        shared_prod = set()
+        for prod_set in map(set, fission_yields.values()):
+            shared_prod |= prod_set
+        ordered_prod = tuple(sorted(shared_prod))
+
+        yield_matrix = empty((len(energies), len(shared_prod)))
+
+        for g_index, energy in enumerate(sorted(energies)):
+            prod_map = fission_yields[energy]
+            for prod_ix, product in enumerate(ordered_prod):
+                yield_val = prod_map.get(product)
+                if yield_val is None:
+                    yield_matrix[g_index, prod_ix] = 0.0
+                else:
+                    yield_matrix[g_index, prod_ix] = yield_val
+
+        return cls(energies, ordered_prod, yield_matrix)
+
+    def to_xml_element(self, root):
+        """Write fission yield data to an xml element
+
+        Parameters
+        ----------
+        root : xml.etree.ElementTree.Element
+            Element to write distribution data to
+        """
+        for energy, yield_obj in self.items():
+            yield_element = ET.SubElement(root, "fission_yields")
+            yield_element.set("energy", str(energy))
+            product_elem = ET.SubElement(yield_element, "products")
+            product_elem.text = " ".join(map(str, yield_obj.products))
+            data_elem = ET.SubElement(yield_element, "data")
+            data_elem.text = " ".join(map(str, yield_obj.yields))
+
+
+class _FissionYield(Mapping):
+    """Mapping for fission yields of a parent at a specific energy
+
+    Abstracted to support nested dictionary-like behavior for
+    :class:`FissionYieldDistribution`, and allowing math operations
+    on a single vector of yields
+
+    Parameters
+    ----------
+    products : tuple of str
+        Products for this specific distribution
+    yields : numpy.ndarray
+        View into associated :attr:`FissionYieldDistribution.yield_matrix`
+    """
+
+    def __init__(self, products, yields):
+        self.products = products
+        self.yields = yields
+
+    def __getitem__(self, product):
+        if product not in self.products:
+            raise KeyError(product)
+        return self.yields[self.products.index(product)]
+
+    def __len__(self):
+        return len(self.products)
+
+    def __iter__(self):
+        return iter(self.products)
+
+    def __iadd__(self, other):
+        """Increment value from other fission yield"""
+        self.yields += other.yields
+        return self
+
+    def __mul__(self, value):
+        return _FissionYield(self.products, self.yields * value)
+
+    def copy(self):
+        """Return an identical yield object, with unique yields"""
+        return _FissionYield(self.products, self.yields.copy())

tests/unit_tests/test_deplete_chain.py

@@ -131,7 +131,8 @@ def test_from_xml(simple_chain):
     # Yield tests
     assert nuc.yield_energies == [0.0253]
     assert list(nuc.yield_data) == [0.0253]
-    assert nuc.yield_data[0.0253] == [("A", 0.0292737), ("B", 0.002566345)]
+    assert nuc.yield_data[0.0253].products == ("A", "B")
+    assert (nuc.yield_data[0.0253].yields == [0.0292737, 0.002566345]).all()
 
 
 def test_export_to_xml(run_in_tmpdir):
@@ -163,7 +164,8 @@ def test_export_to_xml(run_in_tmpdir):
         nuclide.ReactionTuple("(n,gamma)", "B", 0.0, 0.3)
     ]
     C.yield_energies = [0.0253]
-    C.yield_data = {0.0253: [("A", 0.0292737), ("B", 0.002566345)]}
+    C.yield_data = nuclide.FissionYieldDistribution.from_dict({
+        0.0253: {"A": 0.0292737, "B": 0.002566345}})
 
     chain = Chain()
     chain.nuclides = [A, B, C]

tests/unit_tests/test_deplete_nuclide.py

@@ -2,6 +2,8 @@
 
 import xml.etree.ElementTree as ET
 
+import numpy
+
 from openmc.deplete import nuclide
 
 
@@ -69,11 +71,10 @@ def test_from_xml():
         nuclide.ReactionTuple('(n,gamma)', 'U236', 6545200.0, 1.0),
         nuclide.ReactionTuple('fission', None, 193405400.0, 1.0),
     ]
+    expected_yield_data = nuclide.FissionYieldDistribution.from_dict({
+        0.0253: {"Xe138": 0.0481413, "Zr100": 0.0497641, "Te134": 0.062155}})
     assert u235.yield_energies == [0.0253]
-    assert u235.yield_data == {
-        0.0253: [('Te134', 0.062155), ('Zr100', 0.0497641),
-                 ('Xe138', 0.0481413)]
-    }
+    assert u235.yield_data == expected_yield_data
 
 
 def test_to_xml_element():
@@ -91,7 +92,8 @@ def test_to_xml_element():
         nuclide.ReactionTuple('(n,gamma)', 'A', 0.0, 1.0)
     ]
     C.yield_energies = [0.0253]
-    C.yield_data = {0.0253: [("A", 0.0292737), ("B", 0.002566345)]}
+    C.yield_data = nuclide.FissionYieldDistribution.from_dict(
+        {0.0253: {"A": 0.0292737, "B": 0.002566345}})
     element = C.to_xml_element()
 
     assert element.get("half_life") == "0.123"
@@ -114,3 +116,39 @@ def test_to_xml_element():
     assert float(rx_elems[1].get("Q")) == 0.0
 
     assert element.find('neutron_fission_yields') is not None
+
+
+def test_fission_yield_distribution():
+    """Test an energy-dependent yield distribution"""
+    yield_dict = {
+        0.0253: {"Xe135": 7.85e-4, "Gd155": 4.08e-12, "Sm149": 1.71e-12},
+        1.40e7: {"Xe135": 4.54e-3, "Gd155": 5.83e-8, "Sm149": 2.69e-8},
+        5.00e5: {"Xe135": 1.12e-3, "Gd155": 1.32e-12},  # drop Sm149
+    }
+    yield_dist = nuclide.FissionYieldDistribution.from_dict(yield_dict)
+    assert len(yield_dist) == len(yield_dict)
+    assert yield_dist.energies == tuple(sorted(yield_dict.keys()))
+    for exp_ene, exp_dist in yield_dict.items():
+        act_dist = yield_dict[exp_ene]
+        for exp_prod, exp_yield in exp_dist.items():
+            assert act_dist[exp_prod] == exp_yield
+    exp_yield_matrix = numpy.array([
+        [4.08e-12, 1.71e-12, 7.85e-4],
+        [1.32e-12, 0.0, 1.12e-3],
+        [5.83e-8, 2.69e-8, 4.54e-3]])
+    assert numpy.array_equal(yield_dist.yield_matrix, exp_yield_matrix)
+
+    # Test the operations / special methods for fission yield
+    orig_yield_obj = yield_dist[0.0253]
+    # __getitem__ return yields as a view into yield matrix
+    assert orig_yield_obj.yields.base is yield_dist.yield_matrix
+    copied_yield = orig_yield_obj.copy()
+    # copied yields own their own memory -> not a view
+    assert copied_yield.yields.base is None
+
+    # Fission yield feature uses scaled and incremented
+    mod_yields = orig_yield_obj * 2
+    assert numpy.array_equal(orig_yield_obj.yields * 2, mod_yields.yields)
+    mod_yields += orig_yield_obj
+    assert numpy.array_equal(orig_yield_obj.yields * 3, mod_yields.yields)
+