Support Mixture distributions in combine_distributions (#3784)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

Author: GuySten

openmc/stats/univariate.py

@@ -2065,31 +2065,48 @@ def clip(self, tolerance: float = 1e-6, inplace: bool = False) -> Mixture:
 
 
 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)]

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])