From f82dc83473cd27c80820077beb5b013444187aaf Mon Sep 17 00:00:00 2001
From: Andrew Johnson <drewej@protonmail.com>
Date: Tue, 13 Aug 2019 16:38:24 -0500
Subject: [PATCH] Add FissionYieldCutoffHelper for weighting yields by a cutoff
 energy

This helper performs the following tasks to weight fission yields:
1) Determine a set of "fast" and "thermal" yields
based on a user specified cutoff energy,
2) Tally fission rates above and below the cutoff, and
3) Compute the effective yield by weighting the fast and
thermal yields by the fraction of fast and thermal fissions
for all nuclides in all burnable materials.

The user is allowed to specify the cutoff energy and energies
preferred for the fast and thermal yields. By default, the
cutoff is 112 eV, chosen as it is the logarithmic mean of
0.0253 eV and 500 keV. These two values are the default energies
for thermal and fast yields.

Tests are added to check for failure in constructing this
helper (thermal energy > cutoff, non-real values passed, etc.),
selection of yields given a range of user data, and
a proxy class that emulates a variety of thermal and fast splits
to check the eventual computed yield libraries.
---
 docs/source/pythonapi/deplete.rst             |   1 +
 openmc/deplete/helpers.py                     | 189 +++++++++++++++---
 .../unit_tests/test_deplete_fission_yields.py | 143 +++++++++++++
 3 files changed, 306 insertions(+), 27 deletions(-)

diff --git a/docs/source/pythonapi/deplete.rst b/docs/source/pythonapi/deplete.rst
index 7efc8c639..44f8b70e0 100644
--- a/docs/source/pythonapi/deplete.rst
+++ b/docs/source/pythonapi/deplete.rst
@@ -80,6 +80,7 @@ data, such as number densities and reaction rates for each material.
    ChainFissionHelper
    ConstantFissionYieldHelper
    DirectReactionRateHelper
+   FissionYieldCutoffHelper
    OperatorResult
    ReactionRates
    Results
diff --git a/openmc/deplete/helpers.py b/openmc/deplete/helpers.py
index 44ed9293d..37d7359e8 100644
--- a/openmc/deplete/helpers.py
+++ b/openmc/deplete/helpers.py
@@ -3,18 +3,20 @@ Class for normalizing fission energy deposition
 """
 from itertools import product
 from numbers import Real
+import operator
 
-from numpy import dot, zeros, newaxis, asarray, empty_like, where
+from numpy import dot, zeros, newaxis
 
 from openmc.checkvalue import check_type, check_greater_than
 from openmc.capi import (
     Tally, MaterialFilter, EnergyFilter)
 from .abc import (
-    ReactionRateHelper, EnergyHelper, FissionYieldHelper)
+    ReactionRateHelper, EnergyHelper, FissionYieldHelper,
+    TalliedFissionYieldHelper)
 
 __all__ = (
     "DirectReactionRateHelper", "ChainFissionHelper",
-    "ConstantFissionYieldHelper")
+    "ConstantFissionYieldHelper", "FissionYieldCutoffHelper")
 
 # -------------------------------------
 # Helpers for generating reaction rates
@@ -220,41 +222,174 @@ class ConstantFissionYieldHelper(FissionYieldHelper):
         """
         return self.constant_yields
 
-    def compute_yields(self, local_mat_index):
-        """Compute single fission yields using :attr:`results`
 
-        Produces a new library in :attr:`libraries`
+class FissionYieldCutoffHelper(TalliedFissionYieldHelper):
+    """Helper that computes fission yields based on a cutoff energy
+
+    Tally fission rates above and below the cutoff energy.
+    Assume that all fissions below cutoff energy have use thermal fission
+    product yield distributions, while all fissions above use a faster
+    set of yield distributions.
+
+    Uses a limit of 20 MeV for tallying fission.
+
+    Parameters
+    ----------
+    chain_nuclides : iterable of openmc.deplete.Nuclide
+        Nuclides tracked in the depletion chain. Not necessary
+        that all have yield data.
+    n_bmats : int
+        Number of burnable materials tracked in the problem
+    cutoff : float, optional
+        Cutoff energy in [eV] below which all fissions will be
+        use thermal yields. All other fissions will use a
+        faster set of yields. Default: 112 [eV]
+    thermal_energy : float, optional
+        Energy of yield data corresponding to thermal yields.
+        Default: 0.0253 [eV]
+    fast_energy : float, optional
+        Energy of yield data corresponding to fast yields.
+        Default: 500 [kev]
+
+    Attributes
+    ----------
+    n_bmats : int
+        Number of burnable materials tracked in the problem
+    thermal_yields : dict
+        Dictionary of the form ``{parent: {product: yield}}``
+        with thermal yields
+    fast_yields : dict
+        Dictionary of the form ``{parent: {product: yield}}``
+        with fast yields
+    results : numpy.ndarray
+        Array of fission rate fractions with shape
+        ``(n_mats, 2, n_nucs)``. ``results[:, 0]``
+        corresponds to the fraction of all fissions
+        that occured below ``cutoff``. The number
+        of materials in the first axis corresponds
+        to the number of materials burned by the
+        :class:`openmc.deplete.Operator`
+    """
+
+    def __init__(self, chain_nuclides, n_bmats, cutoff=112.0,
+                 thermal_energy=0.0253, fast_energy=500.0e3):
+        check_type("cutoff", cutoff, Real)
+        check_type("thermal_energy", thermal_energy, Real)
+        check_type("fast_energy", fast_energy, Real)
+        check_greater_than("thermal_energy", thermal_energy, 0.0, equality=True)
+        check_greater_than("cutoff", cutoff, thermal_energy, equality=False)
+        check_greater_than("fast_energy", fast_energy, cutoff, equality=False)
+        super().__init__(chain_nuclides, n_bmats)
+        self._cutoff = cutoff
+        self._thermal_yields = {}
+        self._fast_yields = {}
+        for name, nuc in self._chain_nuclides.items():
+            yields = nuc.yield_data
+            energies = nuc.yield_energies
+            thermal = yields.get(thermal_energy)
+            if thermal is None:
+                # find first index >= cutoff
+                ix = self._find_fallback_energy(
+                    name, energies, cutoff, True)
+                thermal = yields[energies[ix - 1]]
+            fast = yields.get(fast_energy)
+            if fast is None:
+                # find first index <= cutoff
+                rev_ix = self._find_fallback_energy(
+                    name, list(reversed(energies)), cutoff, False)
+                fast = yields[energies[-rev_ix]]
+            self._thermal_yields[name] = thermal
+            self._fast_yields[name] = fast
+
+    @staticmethod
+    def _find_fallback_energy(name, energies, cutoff, check_under):
+        cutoff_func = operator.ge if check_under else operator.le
+        found = False
+        for ix, ene in enumerate(energies):
+            if cutoff_func(ene, cutoff):
+                found = True
+                break
+        if found and ix != 0:
+            return ix
+        domain = "thermal" if check_under else "fast"
+        raise ValueError("Could not find {} yields for {} "
+                         "with cutoff {} eV".format(domain, name, cutoff))
+
+    def generate_tallies(self, materials, mat_indexes):
+        """Use C API to produce a fission rate tally in burnable materials
+
+        Include a :class:`openmc.capi.EnergyFilter` to tally fission rates
+        above and below cutoff energy.
+
+        Parameters
+        ----------
+        materials : iterable of :class:`openmc.capi.Material`
+            Materials to be used in :class:`openmc.capi.MaterialFilter`
+        mat_indexes : iterable of int
+            Indexes for materials in ``materials`` tracked on this
+            process
+        """
+        super().generate_tallies(materials, mat_indexes)
+        energy_filter = EnergyFilter()
+        energy_filter.bins = (0.0, self._cutoff, self._upper_energy)
+        self._fission_rate_tally.filters.append(energy_filter)
+
+    def unpack(self):
+        """Obtain fast and thermal fission fractions from tally"""
+        fission_rates = self._fission_rate_tally.results[..., 1].reshape(
+            self.n_bmats, 2, len(self._tally_index))
+        self.results = fission_rates[self._local_indexes]
+        total_fission = self.results.sum(axis=1)
+        nz_mat, nz_nuc = total_fission.nonzero()
+        self.results[nz_mat, :, nz_nuc] /= total_fission[nz_mat, newaxis, nz_nuc]
+
+    def weighted_yields(self, local_mat_index):
+        """Return fission yields for a specific material
+
+        For nuclides with both yield data above and below
+        the cutoff energy, the effective yield for nuclide ``A``
+        will be a weighted sum of fast and thermal yields. The
+        weights will be the fraction of ``A``s fission events
+        in the above and below the cutoff energy.
+
+        If ``A`` has fission product distribution ``F``
+        for fast fissions and ``T`` for thermal fissions, and
+        70% of ``A``'s fissions are considered thermal, then
+        the effective fission product yield distributions
+        for ``A`` is ``0.7 * T + 0.3 * F``
 
         Parameters
         ----------
         local_mat_index : int
-            Index for material tracked on this process that
-            exists in :attr:`local_mat_index` and fits within
-            the first axis in :attr:`results`
+            Index for specific burnable material. Effective
+            yields will be produced using
+            ``self.results[local_mat_index]``
 
         Returns
         -------
         library : dict
             Dictionary of ``{parent: {product: fyield}}``
         """
-        tally_results = self.results[local_mat_index]
+        rates = self.results[local_mat_index]
+        yields = self.constant_yields
+        # iterate over thermal then fast yields, prefer __mul__ to __rmul__
+        for therm_frac, nuc in zip(rates[0], self._tally_index):
+            yields[nuc.name] = self._thermal_yields[nuc.name] * therm_frac
 
-        # Dictionary {parent_nuclide : [product, yield_vector]}
-        initial_library = {}
-        for i_energy, energy in enumerate(self._energy_bounds[1:]):
-            for i_nuc, fiss_frac in enumerate(tally_results[i_energy]):
-                parent = self._tally_index[i_nuc]
-                yield_data = parent.yield_data.get(energy)
-                if yield_data is None:
-                    continue
-                if parent not in initial_library:
-                    initial_library[parent] = yield_data * fiss_frac
-                    continue
-                initial_library[parent] += yield_data * fiss_frac
+        for fast_frac, nuc in zip(rates[1], self._tally_index):
+            yields[nuc.name] += self._fast_yields[nuc.name] * fast_frac
+        return yields
 
-        # convert to dictionary that can be passed to Chain.form_matrix
-        library = {}
-        for k, yield_obj in initial_library.items():
-            library[k.name] = dict(zip(yield_obj.products, yield_obj.yields))
+    @property
+    def thermal_yields(self):
+        out = {}
+        for key, sub in self._thermal_yields.items():
+            out[key] = sub.copy()
+        return out
 
-        return library
+    @property
+    def fast_yields(self):
+        out = {}
+        for key, sub in self._fast_yields.items():
+            out[key] = sub.copy()
+        return out
diff --git a/tests/unit_tests/test_deplete_fission_yields.py b/tests/unit_tests/test_deplete_fission_yields.py
index b9eaf9d2f..f3911d80d 100644
--- a/tests/unit_tests/test_deplete_fission_yields.py
+++ b/tests/unit_tests/test_deplete_fission_yields.py
@@ -1,5 +1,148 @@
 """Test the FissionYieldHelpers"""
 
+from collections import namedtuple
+from unittest.mock import Mock
+
+import pytest
+import numpy
+from numpy.testing import assert_array_equal
+
+from openmc.capi import Material
+from openmc.deplete.nuclide import Nuclide, FissionYieldDistribution
+from openmc.deplete.helpers import (
+    FissionYieldCutoffHelper, ConstantFissionYieldHelper)
+
+
+@pytest.fixture(scope="module")
+def nuclide_bundle():
+    u5yield_dict = {
+        0.0253: {"Xe135": 7.85e-4, "Gd155": 4.08e-12, "Sm149": 1.71e-12},
+        5.0e5: {"Xe135": 7.85e-4, "Sm149": 1.71e-12},
+        1.40e7: {"Xe135": 4.54e-3, "Gd155": 5.83e-8}}
+    u235 = Nuclide()
+    u235.name = "U235"
+    u235.yield_energies = (0.0253, 5.0e5, 1.40e7)
+    u235.yield_data = FissionYieldDistribution.from_dict(u5yield_dict)
+
+    u8yield_dict = {5.00e5: {"Xe135": 1.12e-3, "Gd155": 1.32e-12}}
+    u238 = Nuclide()
+    u238.name = "U238"
+    u238.yield_energies = (5.00e5, )
+    u238.yield_data = FissionYieldDistribution.from_dict(u8yield_dict)
+
+    xe135 = Nuclide()
+    xe135.name = "Xe135"
+
+    NuclideBundle = namedtuple("NuclideBundle", "u235 u238 xe135")
+    return NuclideBundle(u235, u238, xe135)
+@pytest.mark.parametrize("input_energy, u5_yield_energy", (
+    (0.0253, 0.0253), (0.01, 0.0253), (4e5, 5e5)))
+def test_constant_helper(nuclide_bundle, input_energy, u5_yield_energy):
+    helper = ConstantFissionYieldHelper(
+        nuclide_bundle, energy=input_energy)
+    assert helper.energy == input_energy
+    assert helper.constant_yields == {
+        "U235": nuclide_bundle.u235.yield_data[u5_yield_energy],
+        "U238": nuclide_bundle.u238.yield_data[5.00e5]}  # only epithermal
+    assert helper.constant_yields == helper.weighted_yields(1)
+
+
+# ---------------------------------
+# Test the FissionYieldCutoffHelper
+# ---------------------------------
+
+
+def test_cutoff_construction(nuclide_bundle):
+    # defaults
+    helper = FissionYieldCutoffHelper(
+        nuclide_bundle, 1)
+    assert helper.constant_yields == {
+        "U238": nuclide_bundle.u238.yield_data[5.0e5]}
+    assert helper.thermal_yields == {
+        "U235": nuclide_bundle.u235.yield_data[0.0253]}
+    assert helper.fast_yields == {
+        "U235": nuclide_bundle.u235.yield_data[5e5]}
+    # use 14 MeV yields
+    helper = FissionYieldCutoffHelper(
+        nuclide_bundle, 1, fast_energy=14e6)
+    assert helper.constant_yields == {
+        "U238": nuclide_bundle.u238.yield_data[5.0e5]}
+    assert helper.thermal_yields == {
+        "U235": nuclide_bundle.u235.yield_data[0.0253]}
+    assert helper.fast_yields == {
+        "U235": nuclide_bundle.u235.yield_data[14e6]}
+    # specify missing thermal yields -> use 0.0253
+    helper = FissionYieldCutoffHelper(
+        nuclide_bundle, 1, thermal_energy=1)
+    assert helper.thermal_yields == {
+        "U235": nuclide_bundle.u235.yield_data[0.0253]}
+    assert helper.fast_yields == {
+        "U235": nuclide_bundle.u235.yield_data[5e5]}
+    # request missing fast yields -> use epithermal
+    helper = FissionYieldCutoffHelper(
+        nuclide_bundle, 1, fast_energy=1e4)
+    assert helper.thermal_yields == {
+        "U235": nuclide_bundle.u235.yield_data[0.0253]}
+    assert helper.fast_yields == {
+        "U235": nuclide_bundle.u235.yield_data[5e5]}
+    # test failures in cutoff: super low, super high
+    with pytest.raises(ValueError, match="thermal yields"):
+        FissionYieldCutoffHelper(
+            nuclide_bundle, 1, thermal_energy=0.001, cutoff=0.002)
+    with pytest.raises(ValueError, match="fast yields"):
+        FissionYieldCutoffHelper(
+            nuclide_bundle, 1, cutoff=15e6, fast_energy=17e6)
+
+
+@pytest.mark.parametrize("key", ("cutoff", "thermal_energy", "fast_energy"))
+def test_cutoff_failure(key):
+    with pytest.raises(TypeError, match=key):
+        FissionYieldCutoffHelper(None, None, **{key: None})
+    with pytest.raises(ValueError, match=key):
+        FissionYieldCutoffHelper(None, None, **{key: -1})
+
+
+class CutoffProxy(FissionYieldCutoffHelper):
+    """Proxy that supplies a set of tallies"""
+
+    def generate_tallies(self, materials, mat_indexes):
+        self._fission_rate_tally = Mock()
+        self._local_indexes = numpy.asarray(mat_indexes)
+
+
+# emulate some split between fast and thermal U235 fissions
+@pytest.mark.parametrize("therm_frac", (0.5, 0.2, 0.8))
+def test_cutoff_helper(nuclide_bundle, therm_frac):
+    materials = ["1", "2"]  # TODO Use real C API materials?
+    n_bmats = len(materials)
+    local_mats = [0]
+    proxy = CutoffProxy(nuclide_bundle, n_bmats)
+    proxy.generate_tallies(materials, local_mats)
+    non_zero_nucs = [n.name for n in nuclide_bundle]
+    tally_nucs = proxy.update_nuclides_from_operator(non_zero_nucs)
+    assert tally_nucs == ("U235", )
+    # Emulate building tallies
+    # material x energy, tallied_nuclides, 3
+    proxy_flux = 1e6
+    tally_data = numpy.empty((n_bmats * 2, 1, 3))
+    tally_data[0, 0, 1] = therm_frac * proxy_flux
+    tally_data[1, 0, 1] = (1  - therm_frac) * proxy_flux
+    proxy._fission_rate_tally.results = tally_data
+
+    proxy.unpack()
+    # expected results of shape (n_mats, 2, n_tnucs)
+    expected_results = numpy.empty((1, 2, 1))
+    expected_results[:, 0] = therm_frac
+    expected_results[:, 1] = (1 - therm_frac)
+    assert proxy.results == pytest.approx(expected_results)
+
+    actual_yields = proxy.weighted_yields(0)
+    assert actual_yields["U238"] == nuclide_bundle.u238.yield_data[5e5]
+    assert actual_yields["U235"] == (
+        proxy.thermal_yields["U235"] * therm_frac
+        + proxy.fast_yields["U235"] * (1 - therm_frac))
+"""Test the FissionYieldHelpers"""
+
 from collections import namedtuple
 
 import pytest