From ca99195e8172d761042ac88e17efc28b303b29a5 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Thu, 11 Aug 2022 07:27:04 -0500
Subject: [PATCH] Fix Mixture.sample to account for intensities of underlying
 distributions

---
 openmc/stats/univariate.py     | 26 +++++++++++++++++++++++---
 tests/unit_tests/test_stats.py |  6 +++++-
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index 58b7c6172a..0c30526251 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -79,6 +79,18 @@ class Univariate(EqualityMixin, ABC):
         """
         pass
 
+    def integral(self):
+        """Return integral of distribution
+
+        .. versionadded:: 0.13.1
+
+        Returns
+        -------
+        float
+            Integral of distribution
+        """
+        return 1.0
+
 
 class Discrete(Univariate):
     """Distribution characterized by a probability mass function.
@@ -1168,9 +1180,17 @@ class Mixture(Univariate):
 
     def sample(self, n_samples=1, seed=None):
         np.random.seed(seed)
-        idx = np.random.choice(range(len(self.distribution)),
-                               n_samples, p=self.probability)
 
+        # Get probability of each distribution accounting for its intensity
+        p = np.array([prob*dist.integral() for prob, dist in
+                      zip(self.probability, self.distribution)])
+        p /= p.sum()
+
+        # Sample from the distributions
+        idx = np.random.choice(range(len(self.distribution)),
+                               n_samples, p=p)
+
+        # Draw samples from the distributions sampled above
         out = np.empty_like(idx, dtype=float)
         for i in np.unique(idx):
             n_dist_samples = np.count_nonzero(idx == i)
@@ -1294,7 +1314,7 @@ def combine_distributions(dists, probs):
 
     # Combine discrete and continuous if present
     if len(dist_list) > 1:
-        probs = [d.integral() for d in dist_list]
+        probs = [1.0]*len(dist_list)
         dist_list[:] = [Mixture(probs, dist_list.copy())]
 
     return dist_list[0]
diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py
index b8bb94f37c..dbf06d0b56 100644
--- a/tests/unit_tests/test_stats.py
+++ b/tests/unit_tests/test_stats.py
@@ -61,12 +61,14 @@ def test_merge_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])
+    assert merged.integral() == pytest.approx(1.0)
 
     # Probabilities add up but are not normalized
     d1 = openmc.stats.Discrete([3.0], [1.0])
     triple = openmc.stats.Discrete.merge([d1, d1, d1], [1.0, 2.0, 3.0])
     assert triple.x == pytest.approx([3.0])
     assert triple.p == pytest.approx([6.0])
+    assert triple.integral() == pytest.approx(6.0)
 
 
 def test_uniform():
@@ -186,6 +188,7 @@ def test_tabular():
     # test histogram sampling
     d = openmc.stats.Tabular(x, p, interpolation='histogram')
     d.normalize()
+    assert d.integral() == pytest.approx(1.0)
 
     samples = d.sample(n_samples, seed=100)
     diff = np.abs(samples - d.mean())
@@ -433,8 +436,9 @@ def test_combine_distributions():
     t2 = openmc.stats.Tabular([0., 1.], [0.0, 2.0])
     d1 = openmc.stats.Discrete([0.0], [1.0])
     combined = openmc.stats.combine_distributions([t1, t2, d1], [0.25, 0.25, 0.5])
+    assert combined.integral() == pytest.approx(1.0)
 
     # Sample the combined distribution and make sure the sample mean is within
     # uncertainty of the expected value
-    samples = combined.sample(1000)
+    samples = combined.sample(10_000)
     assert_sample_mean(samples, 0.25)