From 27f2c6b15fe951104d48fda726e4ddc8764b2abb Mon Sep 17 00:00:00 2001
From: Andrew Johnson <drewej@protonmail.com>
Date: Mon, 5 Aug 2019 09:28:47 -0500
Subject: [PATCH] Guard against divide by zero in fission yield helper

Create an empty matrix to fit energy-dependent fission
rates for [materials, energies, nuclides].
Use the nonzero method on total fission rate to find
indicies along the first and last axis for non-zero total
fission rate. Populate corresponding fractional fission
rates by dividing group fission rates by total fission rate,
using the indices for materials and nuclides with non-zero
total fission rate.

Use numpy.where to find where total fission rate is zero,
and directly set these locations to zero in the final
result matrix.

Also clean up some lint in openmc.deplete.nuclide related to
old variable names and unused imports.
---
 openmc/deplete/helpers.py                | 20 +++++++++++++-------
 openmc/deplete/nuclide.py                |  7 ++++---
 tests/unit_tests/test_deplete_nuclide.py |  1 -
 3 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/openmc/deplete/helpers.py b/openmc/deplete/helpers.py
index 07c8e20242..43a835908c 100644
--- a/openmc/deplete/helpers.py
+++ b/openmc/deplete/helpers.py
@@ -1,10 +1,9 @@
 """
 Class for normalizing fission energy deposition
 """
-from collections import defaultdict
 from itertools import product
 
-from numpy import dot, zeros, newaxis, divide, asarray
+from numpy import dot, zeros, newaxis, asarray, empty_like, where
 
 from openmc.capi import (
     Tally, MaterialFilter, EnergyFilter)
@@ -290,12 +289,19 @@ class FissionYieldHelper(object):
             len(self._reaction_tally.nuclides))
 
         # Get results specific to this process
-        self.results = result_view[self.local_indexes, ...]
+        fission_rates = result_view[self.local_indexes]
+        self.results = empty_like(fission_rates)
 
-        # scale fission yields proportional to total fission rate
-        fsn_rate = self.results.sum(axis=1)
-        # TODO Guard against divide by zero
-        self.results /= fsn_rate[:, newaxis, :]
+        # scale group fission rates proportional to total fission rate
+        fission_total = fission_rates.sum(axis=1)
+        nz_mat, nz_nuc = fission_total.nonzero()
+        self.results[nz_mat, :, nz_nuc] = (
+            fission_rates[nz_mat, :, nz_nuc]
+            / fission_total[nz_mat, newaxis, nz_nuc])
+
+        # directly set values to zero where total fission rate is zero
+        z_mat, z_nuc = where(fission_total == 0.0)
+        self.results[z_mat, :, z_nuc] = 0.0
 
     def compute_yields(self, local_mat_index):
         """Compute single fission yields using :attr:`results`
diff --git a/openmc/deplete/nuclide.py b/openmc/deplete/nuclide.py
index 80ec648802..8a0a9dc642 100644
--- a/openmc/deplete/nuclide.py
+++ b/openmc/deplete/nuclide.py
@@ -12,9 +12,9 @@ try:
 except ImportError:
     import xml.etree.ElementTree as ET
 
-from numpy import asarray, fromiter, empty
+from numpy import asarray, empty
 
-from openmc.checkvalue import check_type, check_length
+from openmc.checkvalue import check_type
 
 
 DecayTuple = namedtuple('DecayTuple', 'type target branching_ratio')
@@ -260,7 +260,8 @@ class FissionYieldDistribution(Mapping):
             raise ValueError(
                 "Shape of yield matrix inconsistent. "
                 "Should be ({}, {}), is {}".format(
-                    len(energy_map), len(ordered_products), yield_matrix.shape))
+                    len(ordered_energies), len(ordered_products),
+                    yield_matrix.shape))
         self.yield_matrix = yield_matrix
 
     def __len__(self):
diff --git a/tests/unit_tests/test_deplete_nuclide.py b/tests/unit_tests/test_deplete_nuclide.py
index 07f2edb85c..2a3993921f 100644
--- a/tests/unit_tests/test_deplete_nuclide.py
+++ b/tests/unit_tests/test_deplete_nuclide.py
@@ -151,4 +151,3 @@ def test_fission_yield_distribution():
     assert numpy.array_equal(orig_yield_obj.yields * 2, mod_yields.yields)
     mod_yields += orig_yield_obj
     assert numpy.array_equal(orig_yield_obj.yields * 3, mod_yields.yields)
-