From 2854b3a3a0e1806db23c04bdc89606f2c4af2b69 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 25 Feb 2020 09:57:40 -0500 Subject: [PATCH] Disallow numpy ufuncs being applied directly to FissionYields Set ``__array_ufunc__`` to ``None`` to ensure that numpy floats respect __rmul__ and __radd__ for fission yields. Allows weights to be left or right multiplied as numpy floats, which can be computed during fission yield weighting. A unit test has been added to ensure this behavior is preserved Closes #1492 --- openmc/deplete/nuclide.py | 8 ++++++++ tests/unit_tests/test_deplete_nuclide.py | 8 +++++++- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/openmc/deplete/nuclide.py b/openmc/deplete/nuclide.py index 7e352fbf07..0aae5adb72 100644 --- a/openmc/deplete/nuclide.py +++ b/openmc/deplete/nuclide.py @@ -533,6 +533,7 @@ class FissionYield(Mapping): return zip(self.products, self.yields) def __add__(self, other): + """Add one set of fission yields to this set, return new yields""" if not isinstance(other, FissionYield): return NotImplemented new = FissionYield(self.products, self.yields.copy()) @@ -550,12 +551,14 @@ class FissionYield(Mapping): return self + other def __imul__(self, scalar): + """Scale these fission yields by a real scalar""" if not isinstance(scalar, Real): return NotImplemented self.yields *= scalar return self def __mul__(self, scalar): + """Return a new set of yields scaled by a real scalar""" if not isinstance(scalar, Real): return NotImplemented new = FissionYield(self.products, self.yields.copy()) @@ -568,3 +571,8 @@ class FissionYield(Mapping): def __repr__(self): return "<{} containing {} products and yields>".format( self.__class__.__name__, len(self)) + + # Avoid greedy numpy operations like np.float64 * fission_yield + # converting this to an array on the fly. Force __rmul__ and + # __radd__. See issue #1492 + __array_ufunc__ = None diff --git a/tests/unit_tests/test_deplete_nuclide.py b/tests/unit_tests/test_deplete_nuclide.py index 7d68b3a3d8..b8a258df00 100644 --- a/tests/unit_tests/test_deplete_nuclide.py +++ b/tests/unit_tests/test_deplete_nuclide.py @@ -149,12 +149,18 @@ def test_fission_yield_distribution(): # __getitem__ return yields as a view into yield matrix assert orig_yields.yields.base is yield_dist.yield_matrix - # Fission yield feature uses scaled and incremented + # Scale and increment fission yields mod_yields = orig_yields * 2 assert numpy.array_equal(orig_yields.yields * 2, mod_yields.yields) mod_yields += orig_yields assert numpy.array_equal(orig_yields.yields * 3, mod_yields.yields) + mod_yields = 2.0 * orig_yields + assert numpy.array_equal(orig_yields.yields * 2, mod_yields.yields) + + mod_yields = numpy.float64(2.0) * orig_yields + assert numpy.array_equal(orig_yields.yields * 2, mod_yields.yields) + # Failure modes for adding, multiplying yields similar = numpy.empty_like(orig_yields.yields) with pytest.raises(TypeError):