From 2c2d9d7f312e8b7dfcea10de2ca9aee8e8361f3c Mon Sep 17 00:00:00 2001
From: Andrew Johnson <drewej@protonmail.com>
Date: Wed, 18 Mar 2020 20:09:19 -0400
Subject: [PATCH] Provide FissionYieldDistribution.restrict_products

---
 openmc/deplete/nuclide.py                | 32 +++++++++++++++++++++++-
 tests/unit_tests/test_deplete_nuclide.py | 15 +++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/openmc/deplete/nuclide.py b/openmc/deplete/nuclide.py
index 960ba6eff5..7ba72b6d56 100644
--- a/openmc/deplete/nuclide.py
+++ b/openmc/deplete/nuclide.py
@@ -13,7 +13,7 @@ try:
 except ImportError:
     import xml.etree.ElementTree as ET
 
-from numpy import empty
+from numpy import empty, searchsorted
 
 from openmc.checkvalue import check_type
 
@@ -467,6 +467,36 @@ class FissionYieldDistribution(Mapping):
             data_elem = ET.SubElement(yield_element, "data")
             data_elem.text = " ".join(map(str, yield_obj.yields))
 
+    def restrict_products(self, possible_products):
+        """Return a new distribution with select products
+
+        Parameters
+        ----------
+        possible_products : iterable of str
+            Candidate pool of fission products. Existing products
+            not contained here will not exist in the new instance
+
+        Returns
+        -------
+        FissionYieldDistribution or None
+            A value of None indicates no values in
+            ``possible_products`` exist in :attr:`products`
+
+        """
+
+        overlap = set(self.products).intersection(possible_products)
+        if not overlap:
+            return None
+
+        products = sorted(overlap)
+        indices = searchsorted(self.products, products)
+
+        # coerce back to dictionary to pass back to __init__
+        new_yields = {}
+        for ene, yields in zip(self.energies, self.yield_matrix.copy()):
+            new_yields[ene] = dict(zip(products, yields[indices]))
+
+        return type(self)(new_yields)
 
 
 class FissionYield(Mapping):
diff --git a/tests/unit_tests/test_deplete_nuclide.py b/tests/unit_tests/test_deplete_nuclide.py
index b8a258df00..92b004845a 100644
--- a/tests/unit_tests/test_deplete_nuclide.py
+++ b/tests/unit_tests/test_deplete_nuclide.py
@@ -176,6 +176,21 @@ def test_fission_yield_distribution():
     with pytest.raises(TypeError):
         orig_yields *= similar
 
+    # Test restriction of fission products
+    strict_restrict = yield_dist.restrict_products(["Xe135", "Sm149"])
+    with_extras = yield_dist.restrict_products(
+        ["Xe135", "Sm149", "H1", "U235"])
+
+    assert strict_restrict.products == ("Sm149", "Xe135", )
+    assert strict_restrict.energies == yield_dist.energies
+    assert with_extras.products == ("Sm149", "Xe135", )
+    assert with_extras.energies == yield_dist.energies
+    for ene, new_yields in strict_restrict.items():
+        for product in strict_restrict.products:
+            assert new_yields[product] == yield_dist[ene][product]
+            assert with_extras[ene][product] == yield_dist[ene][product]
+
+    assert yield_dist.restrict_products(["U235"]) is None
 
 def test_validate():