Provide Chain.reduce for following paths

openmc/deplete/chain.py

@@ -10,7 +10,7 @@ import math
 import re
 from collections import OrderedDict, defaultdict
 from collections.abc import Mapping, Iterable
-from numbers import Real
+from numbers import Real, Integral
 from warnings import warn
 
 from openmc.checkvalue import check_type, check_greater_than
@@ -820,3 +820,143 @@ class Chain(object):
                 return stat
             valid = valid and stat
         return valid
+
+    def reduce(self, initial_isotopes, level=None):
+        """Reduce the size of the chain by follwing transmutation paths
+
+        Parameters
+        ----------
+        initial_isotopes : iterable of str
+            Start the search based on the contents of these isotopes
+        level : int, optional
+            Depth of transmuation path to follow. Must be greater than
+            or equal to zero. A value of zero returns a chain with
+            ``initial_isotopes``. The default value of None implies
+            that all isotopes that appear in the transmutation paths
+            of the initial isotopes and their progeny should be
+            explored
+
+        Returns
+        -------
+        Chain
+
+        """
+        check_type("initial_isotopes", initial_isotopes, Iterable, str)
+        if level is None:
+            level = math.inf
+        else:
+            check_type("level", level, Integral)
+            check_greater_than("level", level, 0, equality=True)
+
+        all_isotopes = self._follow(set(initial_isotopes), level)
+
+        # Avoid re-sorting for fission yields
+        name_sort = sorted(all_isotopes)
+
+        nuclides = []
+        nuclide_dict = {}
+        reactions = set()
+
+        for idx, iso in enumerate(sorted(all_isotopes, key=openmc.data.zam)):
+            previous = self[iso]
+            new_nuclide = Nuclide(previous.name)
+            new_nuclide.half_life = previous.half_life
+            new_nuclide.decay_energy = new_nuclide.decay_energy
+
+            new_decay = []
+            for mode in previous.decay_modes:
+                if mode.target in all_isotopes:
+                    new_decay.append(mode)
+                else:
+                    new_decay.append(DecayTuple(
+                        mode.type, None, mode.branching_ratio))
+            new_nuclide.decay_modes = new_decay
+
+            new_reactions = []
+            for rxn in previous.reactions:
+                if rxn.target in all_isotopes:
+                    new_reactions.append(rxn)
+                    reactions.add(rxn.type)
+                elif rxn.type == "fission":
+                    new_yields = new_nuclide.yield_data = (
+                        previous.yield_data.restrict_products(name_sort))
+                    if new_yields is not None:
+                        new_reactions.append(rxn)
+                        reactions.add("fission")
+                # Maintain total destruction rates but set no target
+                else:
+                    new_reactions.append(ReactionTuple(
+                        rxn.type, None, rxn.Q, rxn.branching_ratio))
+                    reactions.add(rxn.type)
+
+            new_nuclide.reactions = new_reactions
+
+            nuclides.append(new_nuclide)
+            nuclide_dict[iso] = idx
+
+        new_chain = type(self)()
+        new_chain.nuclides = nuclides
+        new_chain.nuclide_dict = nuclide_dict
+
+        # Doesn't appear that the ordering matters for the reactions,
+        # just the contents
+        new_chain.reactions = sorted(reactions)
+
+        return new_chain
+
+    def _follow(self, isotopes, level):
+        """Return all isotopes present up to depth level"""
+        found = set(isotopes)
+        remaining = set(self.nuclide_dict)
+        if not found.issubset(remaining):
+            raise IndexError(
+                "The following isotopes were not found in the chain: "
+                "{}".format(", ".join(found - remaining)))
+
+        if level == 0:
+            return found
+
+        remaining.difference_update(found)
+
+        depth = 0
+        next_iso = set()
+
+        while depth < level and remaining:
+            # Exhaust all isotopes at this level
+            while isotopes:
+                iso = isotopes.pop()
+                found.add(iso)
+                nuclide = self[iso]
+
+                # Follow all transmutation paths for this nuclide
+                for rxn in nuclide.reactions + nuclide.decay_modes:
+                    if rxn.type == "fission" or rxn.target is None:
+                        continue
+                    # Skip if we've already come across this isotope
+                    elif (rxn.target in next_iso
+                          or rxn.target in found or rxn.target in isotopes):
+                        continue
+                    next_iso.add(rxn.target)
+
+                if nuclide.yield_data is not None:
+                    for product in nuclide.yield_data.products:
+                        if (product in next_iso
+                                or product in found or product in isotopes):
+                            continue
+                        next_iso.add(product)
+
+            if not next_iso:
+                # No additional isotopes to process, nor to update the
+                # current set of discovered isotopes
+                return found
+
+            # Prepare for next dig
+            depth += 1
+            isotopes.update(next_iso)
+            remaining.difference_update(next_iso)
+            next_iso.clear()
+
+        # Process isotope that would have started next depth
+        found.update(isotopes)
+
+        return found