diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py index f571377bf8..e11c9b50ee 100644 --- a/openmc/stats/univariate.py +++ b/openmc/stats/univariate.py @@ -96,9 +96,9 @@ class Discrete(Univariate): Attributes ---------- - x : Iterable of float + x : numpy.ndarray Values of the random variable - p : Iterable of float + p : numpy.ndarray Discrete probability for each value """ @@ -123,7 +123,7 @@ class Discrete(Univariate): if isinstance(x, Real): x = [x] cv.check_type('discrete values', x, Iterable, Real) - self._x = x + self._x = np.array(x, dtype=float) @p.setter def p(self, p): @@ -132,7 +132,7 @@ class Discrete(Univariate): cv.check_type('discrete probabilities', p, Iterable, Real) for pk in p: cv.check_greater_than('discrete probability', pk, 0.0, True) - self._p = p + self._p = np.array(p, dtype=float) def cdf(self): return np.insert(np.cumsum(self.p), 0, 0.0) @@ -836,9 +836,9 @@ class Tabular(Univariate): Attributes ---------- - x : Iterable of float + x : numpy.ndarray Tabulated values of the random variable - p : Iterable of float + p : numpy.ndarray Tabulated probabilities interpolation : {'histogram', 'linear-linear', 'linear-log', 'log-linear', 'log-log'}, optional Indicate whether the density function is constant between tabulated @@ -871,7 +871,7 @@ class Tabular(Univariate): @x.setter def x(self, x): cv.check_type('tabulated values', x, Iterable, Real) - self._x = x + self._x = np.array(x, dtype=float) @p.setter def p(self, p): @@ -879,7 +879,7 @@ class Tabular(Univariate): if not self._ignore_negative: for pk in p: cv.check_greater_than('tabulated probability', pk, 0.0, True) - self._p = p + self._p = np.array(p, dtype=float) @interpolation.setter def interpolation(self, interpolation): @@ -892,8 +892,8 @@ class Tabular(Univariate): 'distributions using histogram or ' 'linear-linear interpolation') c = np.zeros_like(self.x) - x = np.asarray(self.x) - p = np.asarray(self.p) + x = self.x + p = self.p if self.interpolation == 'histogram': c[1:] = p[:-1] * np.diff(x) @@ -933,7 +933,7 @@ class Tabular(Univariate): def normalize(self): """Normalize the probabilities stored on the distribution""" - self.p = np.asarray(self.p) / self.cdf().max() + self.p /= self.cdf().max() def sample(self, n_samples=1, seed=None): if not self.interpolation in ('histogram', 'linear-linear'): diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py index b57f9578a1..404cecd139 100644 --- a/tests/unit_tests/test_stats.py +++ b/tests/unit_tests/test_stats.py @@ -13,8 +13,8 @@ def test_discrete(): elem = d.to_xml_element('distribution') d = openmc.stats.Discrete.from_xml_element(elem) - assert d.x == x - assert d.p == p + np.testing.assert_array_equal(d.x, x) + np.testing.assert_array_equal(d.p, p) assert len(d) == len(x) d = openmc.stats.Univariate.from_xml_element(elem) @@ -76,8 +76,8 @@ def test_uniform(): assert len(d) == 2 t = d.to_tabular() - assert t.x == [a, b] - assert t.p == [1/(b-a), 1/(b-a)] + np.testing.assert_array_equal(t.x, [a, b]) + np.testing.assert_array_equal(t.p, [1/(b-a), 1/(b-a)]) assert t.interpolation == 'histogram' exp_mean = 0.5 * (a + b) @@ -396,3 +396,34 @@ def test_muir(): assert within_2_sigma / n_samples >= 0.95 within_3_sigma = np.count_nonzero(samples < 3*d.std_dev) assert within_3_sigma / n_samples >= 0.99 + + +def test_combine_distributions(): + # Combine two discrete (same data as in test_merge_discrete) + x1 = [0.0, 1.0, 10.0] + p1 = [0.3, 0.2, 0.5] + d1 = openmc.stats.Discrete(x1, p1) + x2 = [0.5, 1.0, 5.0] + p2 = [0.4, 0.5, 0.1] + d2 = openmc.stats.Discrete(x2, p2) + + # Merged distribution should have x values sorted and probabilities + # appropriately combined. Duplicate x values should appear once. + merged = openmc.stats.combine_distributions([d1, d2], [0.6, 0.4]) + assert isinstance(merged, openmc.stats.Discrete) + assert merged.x == pytest.approx([0.0, 0.5, 1.0, 5.0, 10.0]) + assert merged.p == pytest.approx( + [0.6*0.3, 0.4*0.4, 0.6*0.2 + 0.4*0.5, 0.4*0.1, 0.6*0.5]) + + # Probabilities add up but are not normalized + d1 = openmc.stats.Discrete([3.0], [1.0]) + triple = openmc.stats.combine_distributions([d1, d1, d1], [1.0, 2.0, 3.0]) + assert triple.x == pytest.approx([3.0]) + assert triple.p == pytest.approx([6.0]) + + # Combine discrete and tabular + t1 = openmc.stats.Tabular(x2, p2) + mixed = openmc.stats.combine_distributions([d1, t1], [0.5, 0.5]) + assert isinstance(mixed, openmc.stats.Mixture) + assert len(mixed.distribution) == 2 + assert len(mixed.probability) == 2