Add base class for TallyHelpers

Interface is designed such that one could make a tally
helper work with unique materials, mainly for pulling
fission Q values on a per material basis

openmc/deplete/operator.py

@@ -154,7 +154,7 @@ class Operator(TransportOperator):
             self.local_mats, self._burnable_nucs, self.chain.reactions)
 
         # Get class to assist working with tallies
-        self._tally_helper = ChainFissTallyHelper(len(self.local_mats))
+        self._tally_helper = ChainFissTallyHelper()
 
 
     def __call__(self, vec, power, print_out=True):
@@ -185,7 +185,7 @@ class Operator(TransportOperator):
 
         # Update material compositions and tally nuclides
         self._update_materials()
-        self._tally_helper.reaction_tally.nuclides = self._get_tally_nuclides()
+        self._tally_helper.nuclides = self._get_tally_nuclides()
 
         # Run OpenMC
         openmc.capi.reset()
@@ -515,7 +515,7 @@ class Operator(TransportOperator):
 
         # Extract tally bins
         materials = self.burnable_mats
-        nuclides = self._tally_helper.reaction_tally.nuclides
+        nuclides = self._tally_helper.nuclides
 
         # Form fast map
         nuc_ind = [rates.index_nuc[nuc] for nuc in nuclides]
@@ -560,7 +560,8 @@ class Operator(TransportOperator):
                     j += 1
 
             # Accumulate energy from fission
-            energy += self._tally_helper.get_fiss_energy(rates_expanded[:, fission_ind], i)
+            energy += self._tally_helper.get_fission_energy(
+                rates_expanded[:, fission_ind], i)
 
             # Divide by total number and store
             for i_nuc_results in nuc_ind:

openmc/deplete/tally_helpers.py

@@ -1,35 +1,26 @@
 """
 Class for normalizing fission energy deposition
 """
+from abc import ABC, abstractmethod
 
 from numpy import dot, zeros
 
+from openmc.checkvalue import check_type
 from openmc.capi import Tally, MaterialFilter
 
 
-class ChainFissTallyHelper(object):
-    """Fission Q-values are pulled from chain"""
+class TallyHelperBase(ABC):
+    """Base class for working with tallies for depletion"""
 
-    def __init__(self, n_materials):
+    def __init__(self):
         self._fiss_q = None
-        self.n_materials = n_materials
         self._rx_tally = None
+        self._nuclides = []
 
+    @abstractmethod
     def set_fission_q(self, chain_nucs, rate_index):
-        if (self._fiss_q is not None
-                and self._fiss_q.shape == (len(rate_index), )):
-            return
-
-        fq = zeros(len(rate_index))
-
-        for nuclide in chain_nucs:
-            if nuclide.name in rate_index:
-                for rx in nuclide.reactions:
-                    if rx.type == "fission":
-                        fq[rate_index[nuclide.name]] = rx.Q
-                        break
-
-        self._fiss_q = fq
+        """Populate the energy released per fission Q value array"""
+        pass
 
     @property
     def reaction_tally(self):
@@ -46,5 +37,74 @@ class ChainFissTallyHelper(object):
         self._rx_tally.scores = scores
         self._rx_tally.filters = [MaterialFilter(materials)]
 
-    def get_fiss_energy(self, fiss_rates, _mat_index):
+    @property
+    def nuclides(self):
+        """List of nuclides with requested reaction rates"""
+        return self._nuclides
+
+    @nuclides.setter
+    def nuclides(self, nuclides):
+        check_type("nuclides", nuclides, list, str)
+        self._nuclides = nuclides
+        self._rx_tally.nuclides = nuclides
+
+    @abstractmethod
+    def get_fission_energy(self, fission_rates, mat_index):
+        """return a vector of the isotopic fission energy for this material
+
+        parameters
+        ----------
+        fission_rates: numpy.ndarray
+            fission reaction rate for each isotope in the specified
+            material. should be ordered corresponding to initial
+            ``rate_index`` used in :meth:`set_fission_q`
+        mat_index: int
+            index for the material requested.
+        """
+
+
+class ChainFissTallyHelper(TallyHelperBase):
+    """Fission Q-values are pulled from chain"""
+
+    def set_fission_q(self, chain_nucs, rate_index):
+        """Populate the fission Q value vector from a chain.
+
+        Paramters
+        ---------
+        chain_nucs: iterable of :class:`openmc.deplete.Nuclide`
+            Nuclides used in this depletion chain. Do not need
+            to be ordered
+        rate_index: dict of str to int
+            Dictionary mapping names of nuclides, e.g. ``"U235"``,
+            to a corresponding index in the desired fission Q
+            vector.
+        """
+        if (self._fiss_q is not None
+                and self._fiss_q.shape == (len(rate_index), )):
+            return
+
+        fq = zeros(len(rate_index))
+
+        for nuclide in chain_nucs:
+            if nuclide.name in rate_index:
+                for rx in nuclide.reactions:
+                    if rx.type == "fission":
+                        fq[rate_index[nuclide.name]] = rx.Q
+                        break
+
+        self._fiss_q = fq
+
+    def get_fission_energy(self, fiss_rates, _mat_index):
+        """Return a vector of the isotopic fission energy for this material
+
+        parameters
+        ----------
+        fission_rates: numpy.ndarray
+            fission reaction rate for each isotope in the specified
+            material. should be ordered corresponding to initial
+            ``rate_index`` used in :meth:`set_fission_q`
+        mat_index: int
+            index for the material requested. Unused, as all
+            isotopes in all materials have the same Q value.
+        """
         return dot(fiss_rates, self._fiss_q)