Add set/get capture branching ratios method for Chain

Related to #1237.

openmc.deplete.Chain.set_capture_branches takes in
a dictionary of parent nuclide names to {daughter: ratio}, e.g.
{"Am241": {"Am242": 0.9, "Am242_m1": 0.1}}
that will be used to overwrite existing capture
reactions with new branching ratios.

openmc.deplete.Chain.get_capture_branches returns a
similar dictionary of capture reactions with multiple targets
and their branching ratios.

openmc/deplete/chain.py

@@ -11,7 +11,8 @@ import re
 from collections import OrderedDict, defaultdict
 from collections.abc import Mapping
 
-from openmc.checkvalue import check_type
+from openmc.checkvalue import check_type, check_less_than
+from openmc.data import gnd_name, zam
 
 # Try to use lxml if it is available. It preserves the order of attributes and
 # provides a pretty-printer by default. If not available,
@@ -451,3 +452,137 @@ class Chain(object):
         matrix_dok = sp.dok_matrix((n, n))
         dict.update(matrix_dok, matrix)
         return matrix_dok.tocsr()
+
+    def get_capture_branches(self):
+        """Return a dictionary with capture branching ratios
+
+        Returns
+        -------
+        capt :
+            nested dict of parent nuclide keys with capture targets and
+            branching ratios::
+
+                {"Am241": {"Am242": 0.91, "Am242_m1": 0.09}}
+
+        See Also
+        --------
+        :meth:`set_capture_branches`
+
+        """
+
+        capt = {}
+        for nuclide in self.nuclides:
+            nuc_capt = {}
+            for rx in nuclide.reactions:
+                if rx.type == "(n,gamma)" and rx.branching_ratio != 1.0:
+                    nuc_capt[rx.target] = rx.branching_ratio
+            if len(nuc_capt) > 0:
+                capt[nuclide.name] = nuc_capt
+        return capt
+
+    def set_capture_branches(self, branch_ratios):
+        """Set the capture branching ratios
+
+        ``branch_ratios`` may be modified in place, only to
+        insert missing ground state reactions. These will be
+        inserted only if:
+
+            1) There is no branch directly to a ground state
+               target, and
+            2) The sum of all ratios on this branch does not
+               equal 1.
+
+        Parameters
+        ----------
+        branch_ratios : dict of {str: {str: float}}
+            Capture branching ratios to be inserted.
+            First layer keys are names of parent nuclides, e.g.
+            ``"Am241"``. The capture branching ratios for these
+            parents will be modified. Corresponding values are
+            dictionaries of ``{target: branching_ratio}``
+
+        See Also
+        --------
+        :meth:`get_capture_branches`
+        """
+
+        # Store some useful information through the validation stage
+
+        sums = {}
+        capt_ix_map = {}
+        grounds = {}
+
+        missing = set()
+        no_capture = set()
+
+        # Check for validity before manipulation
+
+        check_type("branch_ratios", branch_ratios, dict, str)
+
+        for parent, sub in branch_ratios.items():
+            if parent not in self:
+                # TODO How to handle missing branching ratios
+                missing.add(parent)
+                continue
+
+            # Make sure this nuclide has capture reactions
+
+            indexes = []
+            for ix, rx in enumerate(self[parent].reactions):
+                if rx.type == "(n,gamma)":
+                    indexes.append(ix)
+                    if "_m" not in rx.target:
+                        grounds[parent] = rx.target
+
+            if len(indexes) == 0:
+                no_capture.add(parent)
+                continue
+
+            capt_ix_map[parent] = indexes
+
+            check_type(parent, sub, dict, str)
+            check_type(parent + " ratios", list(sub.values()), list, float)
+            this_sum = sum(sub.values())
+            check_less_than(parent + " ratios", this_sum, 1.0, True)
+            sums[parent] = this_sum
+
+        if len(missing) > 0:
+            print("The following nuclides were not found in {}: {}".format(
+                self.__class__.__name__, ", ".join(sorted(missing))))
+
+        if len(no_capture) > 0:
+            print("The following nuclides did not have capture reactions: "
+                  "{}".format(", ".join(sorted(no_capture))))
+
+        # Insert new ReactionTuples with updated branch ratios
+
+        for parent_name, capt_index in capt_ix_map.items():
+
+            parent = self[parent_name]
+            new_ratios = branch_ratios[parent_name]
+            capt_index = capt_ix_map[parent_name]
+
+            # Assume Q value is independent of target state
+            capt_Q = parent.reactions[capt_index[0]].Q
+
+            # Remove existing capture reactions
+
+            for ix in reversed(capt_index):
+                parent.reactions.pop(ix)
+
+            all_meta = False
+
+            for tgt, br in new_ratios.items():
+                all_meta |= ("_m" in tgt)
+                parent.reactions.append(ReactionTuple(
+                    "(n,gamma)", tgt, capt_Q, br))
+
+            if all_meta and sums[parent_name] != 1.0:
+                ground_br = 1.0 - sums[parent_name]
+                ground_tgt = grounds.get(parent_name, None)
+                if ground_tgt is None:
+                    pz, pa, pm = zam(parent_name)
+                    ground_tgt = gnd_name(pz, pa + 1, 0)
+                new_ratios[ground_tgt] = ground_br
+                parent.reactions.append(ReactionTuple(
+                    "(n,gamma)", ground_tgt, capt_Q, ground_br))