diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py index c6ce217fb..a70737a34 100644 --- a/openmc/stats/univariate.py +++ b/openmc/stats/univariate.py @@ -2065,31 +2065,48 @@ class Mixture(Univariate): def combine_distributions( - dists: Sequence[Discrete | Tabular], + dists: Sequence[Discrete | Tabular | Mixture], probs: Sequence[float] ): """Combine distributions with specified probabilities This function can be used to combine multiple instances of - :class:`~openmc.stats.Discrete` and `~openmc.stats.Tabular`. Multiple - discrete distributions are merged into a single distribution and the - remainder of the distributions are put into a :class:`~openmc.stats.Mixture` - distribution. + :class:`~openmc.stats.Discrete`, :class:`~openmc.stats.Tabular` and + :class:`~openmc.stats.Mixture` of them. Multiple discrete distributions are + merged into a single distribution and the remainder of the distributions are + put into a :class:`~openmc.stats.Mixture` distribution. .. versionadded:: 0.13.1 Parameters ---------- - dists : sequence of openmc.stats.Discrete or openmc.stats.Tabular + dists : sequence of openmc.stats.Discrete, openmc.stats.Tabular, or openmc.stats.Mixture Distributions to combine probs : sequence of float Probability (or intensity) of each distribution """ + new_probs = [] + new_dists = [] for i, dist in enumerate(dists): - cv.check_type(f'dists[{i}]', dist, (Discrete, Tabular)) + cv.check_type(f'dists[{i}]', dist, (Discrete, Tabular, Mixture)) cv.check_type(f'probs[{i}]', probs[i], Real) cv.check_greater_than(f'probs[{i}]', probs[i], 0.0) + if isinstance(dist, Mixture): + if dist.bias is not None: + warn("A Mixture distribution with a bias specified was passed " + "to combine_distributions. The bias will be discarded " + "during flattening.") + for j, d in enumerate(dist.distribution): + cv.check_type(f'dists[{i}].distribution[{j}]', d, (Discrete, Tabular)) + new_probs.append(probs[i]*dist.probability[j]) + new_dists.append(d) + else: + new_probs.append(probs[i]) + new_dists.append(dist) + + probs = new_probs + dists = new_dists # Get list of discrete/continuous distribution indices discrete_index = [i for i, d in enumerate(dists) if isinstance(d, Discrete)] diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py index 71679024f..669d5b74c 100644 --- a/tests/unit_tests/test_stats.py +++ b/tests/unit_tests/test_stats.py @@ -835,6 +835,23 @@ def test_combine_distributions(): assert isinstance(mixed, openmc.stats.Mixture) assert len(mixed.distribution) == 2 assert len(mixed.probability) == 2 + assert mixed == openmc.stats.combine_distributions([mixed], [1.0]) + + # Mixture combined with another distribution: probabilities should be + # correctly scaled when the Mixture is flattened + d_a = openmc.stats.delta_function(1.0) + d_b = openmc.stats.delta_function(2.0) + m = openmc.stats.Mixture([0.3, 0.7], [d_a, d_b]) + extra = openmc.stats.delta_function(3.0) + result = openmc.stats.combine_distributions([m, extra], [0.5, 0.5]) + assert isinstance(result, openmc.stats.Discrete) + assert result.x == pytest.approx([1.0, 2.0, 3.0]) + assert result.p == pytest.approx([0.5*0.3, 0.5*0.7, 0.5]) + + # Passing a Mixture with a bias should warn that the bias is dropped + biased_m = openmc.stats.Mixture([0.5, 0.5], [d_a, d_b], bias=[0.8, 0.2]) + with pytest.warns(UserWarning, match='bias'): + openmc.stats.combine_distributions([biased_m], [1.0]) # Single tabular returns a tabular distribution with scaled probabilities t_single = openmc.stats.Tabular([0.0, 1.0], [2.0, 0.0])