From 7484618c7c93d3bb75fe6decc0104e3201129ad6 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Tue, 5 Apr 2022 13:01:03 -0500
Subject: [PATCH 01/22] Adding sample methods to univariate distribution
 classes

---
 openmc/stats/univariate.py | 123 +++++++++++++++++++++++++++++++++++++
 1 file changed, 123 insertions(+)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index e4a9b8f924..ea8c9f9432 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -60,6 +60,10 @@ class Univariate(EqualityMixin, ABC):
         elif distribution == 'mixture':
             return Mixture.from_xml_element(elem)
 
+    @abstractmethod
+    def sample(p, seed=None):
+        pass
+
 
 class Discrete(Univariate):
     """Distribution characterized by a probability mass function.
@@ -115,6 +119,26 @@ class Discrete(Univariate):
             cv.check_greater_than('discrete probability', pk, 0.0, True)
         self._p = p
 
+    def sample(self, n_samples=1, seed=None):
+        np.random.seed(seed)
+
+        if len(self.x) == 1:
+            return np.full((n_samples), self.x[0])
+
+        out = np.zeros((n_samples,))
+        for i, xi in enumerate(np.random.rand(n_samples)):
+            c = 0.0
+            for j, prob in enumerate(self.p):
+                c += prob
+                if (xi < c):
+                    out[i] = self.x[j]
+                    break
+        return out
+
+    def normalize(self):
+        norm = sum(self.p)
+        self.p = [val / norm for val in self.p]
+
     def to_xml_element(self, element_name):
         """Return XML representation of the discrete distribution
 
@@ -240,6 +264,10 @@ class Uniform(Univariate):
         t.c = [0., 1.]
         return t
 
+    def sample(self, n_samples=1, seed=None):
+        np.random.seed(seed)
+        return np.random.uniform(self.a, self.b, n_samples)
+
     def to_xml_element(self, element_name):
         """Return XML representation of the uniform distribution
 
@@ -341,6 +369,14 @@ class PowerLaw(Univariate):
         cv.check_type('power law exponent', n, Real)
         self._n = n
 
+    def sample(self, n_samples=1, seed=None):
+        np.random.seed(seed)
+        xi = np.random.rand(n_samples)
+        pwr = self.n + 1
+        offset = self.a**pwr
+        span = self.b**pwr - offset
+        return np.power(offset + xi * span, (1/pwr))
+
     def to_xml_element(self, element_name):
         """Return XML representation of the power law distribution
 
@@ -414,6 +450,16 @@ class Maxwell(Univariate):
         cv.check_greater_than('Maxwell temperature', theta, 0.0)
         self._theta = theta
 
+    def sample(self, n_samples=1, seed=None):
+        np.random.seed(seed)
+        self.sample_maxwell(self.theta, n_samples)
+
+    @staticmethod
+    def sample_maxwell(t, n_samples):
+        r1, r2, r3 = np.random.rand(3, n_samples)
+        c = np.power(np.cos(0.5 * np.pi * r3), 2)
+        return t * (np.log(r1) + np.log(r2) * c)
+
     def to_xml_element(self, element_name):
         """Return XML representation of the Maxwellian distribution
 
@@ -502,6 +548,13 @@ class Watt(Univariate):
         cv.check_greater_than('Watt b', b, 0.0)
         self._b = b
 
+    def sample(self, n_samples=1, seed=None):
+        np.random.seed(seed)
+        w = Maxwell.sample_maxwell(self.a, n_samples)
+        u = np.random.uniform(-1., 1., n_samples)
+        aab = self.a * self.a * self.b
+        return w + 0.25*aab + u * np.sqrt(aab*w)
+
     def to_xml_element(self, element_name):
         """Return XML representation of the Watt distribution
 
@@ -588,6 +641,10 @@ class Normal(Univariate):
         cv.check_greater_than('Normal std_dev', std_dev, 0.0)
         self._std_dev = std_dev
 
+    def sample(self, n_samples=1, seed=None):
+        np.random.seed(seed)
+        return np.random.normal(self.mean_value, self.std_dev, n_samples)
+
     def to_xml_element(self, element_name):
         """Return XML representation of the Normal distribution
 
@@ -694,6 +751,12 @@ class Muir(Univariate):
         cv.check_greater_than('Muir kt', kt, 0.0)
         self._kt = kt
 
+    def sample(self, n_samples=1, seed=None):
+        # https://permalink.lanl.gov/object/tr?what=info:lanl-repo/lareport/LA-05411-MS
+        np.random.seed(seed)
+        sigma = np.sqrt(4.*self.e0*self.kt/self.m_rat)
+        return np.random.normal(self.e0, sigma, n_samples)
+
     def to_xml_element(self, element_name):
         """Return XML representation of the Watt distribution
 
@@ -803,6 +866,63 @@ class Tabular(Univariate):
         cv.check_value('interpolation', interpolation, _INTERPOLATION_SCHEMES)
         self._interpolation = interpolation
 
+    def cdf(self):
+        if self.interpolation not in ('histogram', 'linear-linear'):
+            raise NotImplementedError('Can only generate CDFs for tabular '
+                                      'distributions using histogram or '
+                                      'linear-linear interpolation')
+        c = np.zeros_like((len(self.x),))
+        if self.interpolation == 'histogram':
+            c[1:] = self.p * np.diff(self.x)
+        elif self.interpolation == 'linear-linear':
+            c[1:] = 0.5 * np.diff(self.p) * np.diff(self.x)
+        return np.cumsum(c)
+
+    def sample(self, n_samples=1, seed=None):
+        np.random.seed(seed)
+        xi = np.random.rand(n_samples)
+        cdf = self.cdf()
+
+        # get CDF bins that are above the
+        # sampled values
+        c_i = np.full(n_samples, cdf[0])
+        cdf_idx = np.zeros_like(n_samples)
+        for i, val in enumerate(cdf):
+            mask = xi > val
+            c_i[mask] = val
+            cdf_idx[mask] = i
+
+        # get table values at each index where
+        # the random number is less than the next cdf
+        # entry
+        x_i = np.array([self.x[i] for i in cdf_idx])
+        p_i = np.array([self.p[i] for i in cdf_idx])
+
+        if self.interpolation == 'histogram':
+            # mask where probability
+            pos_mask = p_i > 0.0
+            p_i[pos_mask] = x_i[pos_mask] + (xi[pos_mask] - c_i[pos_mask]) \
+                           / p_i[pos_mask]
+            neg_mask = not pos_mask
+            p_i[neg_mask] = x_i[neg_mask]
+            return p_i
+        elif self.interpolation == 'linear-linear':
+            # get variable and probability values for the
+            # next entry
+            x_i1 = np.array([self.x[i+1] for i in cdf_idx])
+            p_i1 = np.array([self.p[i+1] for i in cdf_idx])
+            # compute slope between entries
+            m = (p_i1 - p_i) / (x_i1 - x_i)
+            # set values for zero slope
+            zero = m == 0.0
+            m[zero] = x_i[zero] + (xi[zero] - c_i[zero]) / p_i[zero]
+            # set values for non-zero slope
+            non_zero = not zero
+            quad = np.pow(p_i[non_zero]**2) + 2. * m[non_zero] * (xi[non_zero] - c_i[non_zero])
+            quad[quad < 0.0] = 0.0
+            m[non_zero] = x_i[non_zero] +  (np.sqrt(quad) - p_i) / m[non_zero]
+            return m
+
     def to_xml_element(self, element_name):
         """Return XML representation of the tabular distribution
 
@@ -890,6 +1010,9 @@ class Legendre(Univariate):
     def coefficients(self, coefficients):
         self._coefficients = np.asarray(coefficients)
 
+    def sample(self, n_samples=1, seed=None):
+        raise NotImplementedError
+
     def to_xml_element(self, element_name):
         raise NotImplementedError
 

From e856c6a27065eb9b2b64dd480f1200020937e276 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Tue, 5 Apr 2022 13:24:18 -0500
Subject: [PATCH 02/22] Correction to maxwell dist sampling

---
 openmc/stats/univariate.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index ea8c9f9432..72a5568aa3 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -452,13 +452,13 @@ class Maxwell(Univariate):
 
     def sample(self, n_samples=1, seed=None):
         np.random.seed(seed)
-        self.sample_maxwell(self.theta, n_samples)
+        return self.sample_maxwell(self.theta, n_samples)
 
     @staticmethod
     def sample_maxwell(t, n_samples):
         r1, r2, r3 = np.random.rand(3, n_samples)
         c = np.power(np.cos(0.5 * np.pi * r3), 2)
-        return t * (np.log(r1) + np.log(r2) * c)
+        return -t * (np.log(r1) + np.log(r2) * c)
 
     def to_xml_element(self, element_name):
         """Return XML representation of the Maxwellian distribution

From 136ff45cc56a4a552eea93b35aeaf792e9fad0e7 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Tue, 5 Apr 2022 13:52:50 -0500
Subject: [PATCH 03/22] Adding docstrings. Corrections to Tabular sampling

---
 openmc/stats/univariate.py | 34 +++++++++++++++++++++++++---------
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index 72a5568aa3..3d99f8b8ce 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -61,7 +61,21 @@ class Univariate(EqualityMixin, ABC):
             return Mixture.from_xml_element(elem)
 
     @abstractmethod
-    def sample(p, seed=None):
+    def sample(n_samples=1, seed=None):
+        """Sample the univariate distribution
+
+        Parameters
+        ----------
+        n_samples : int
+            Number of sampled values to generate
+        seed : int or None
+            Initial random number seed.
+
+        Returns
+        -------
+        numpy.ndarray
+            A 1-D array of sampled values
+        """
         pass
 
 
@@ -867,15 +881,17 @@ class Tabular(Univariate):
         self._interpolation = interpolation
 
     def cdf(self):
-        if self.interpolation not in ('histogram', 'linear-linear'):
+        if not self.interpolation in ('histogram', 'linear-linear'):
             raise NotImplementedError('Can only generate CDFs for tabular '
                                       'distributions using histogram or '
                                       'linear-linear interpolation')
-        c = np.zeros_like((len(self.x),))
+        c = np.zeros_like(self.x)
+        x = np.asarray(self.x)
+        p = np.asarray(self.p)
         if self.interpolation == 'histogram':
-            c[1:] = self.p * np.diff(self.x)
+            c[1:] = p * np.diff(x)
         elif self.interpolation == 'linear-linear':
-            c[1:] = 0.5 * np.diff(self.p) * np.diff(self.x)
+            c[1:] = 0.5 * (p[0:-1] + p[1:]) * np.diff(x)
         return np.cumsum(c)
 
     def sample(self, n_samples=1, seed=None):
@@ -886,7 +902,7 @@ class Tabular(Univariate):
         # get CDF bins that are above the
         # sampled values
         c_i = np.full(n_samples, cdf[0])
-        cdf_idx = np.zeros_like(n_samples)
+        cdf_idx = np.zeros((n_samples,), dtype=int)
         for i, val in enumerate(cdf):
             mask = xi > val
             c_i[mask] = val
@@ -903,7 +919,7 @@ class Tabular(Univariate):
             pos_mask = p_i > 0.0
             p_i[pos_mask] = x_i[pos_mask] + (xi[pos_mask] - c_i[pos_mask]) \
                            / p_i[pos_mask]
-            neg_mask = not pos_mask
+            neg_mask = np.invert(pos_mask)
             p_i[neg_mask] = x_i[neg_mask]
             return p_i
         elif self.interpolation == 'linear-linear':
@@ -917,8 +933,8 @@ class Tabular(Univariate):
             zero = m == 0.0
             m[zero] = x_i[zero] + (xi[zero] - c_i[zero]) / p_i[zero]
             # set values for non-zero slope
-            non_zero = not zero
-            quad = np.pow(p_i[non_zero]**2) + 2. * m[non_zero] * (xi[non_zero] - c_i[non_zero])
+            non_zero = np.invert(zero)
+            quad = np.power(p_i[non_zero], 2) + 2. * m[non_zero] * (xi[non_zero] - c_i[non_zero])
             quad[quad < 0.0] = 0.0
             m[non_zero] = x_i[non_zero] +  (np.sqrt(quad) - p_i) / m[non_zero]
             return m

From 17be0257a0ec05994ff6d2a2fea8edcd15b5dfc7 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Tue, 5 Apr 2022 14:05:35 -0500
Subject: [PATCH 04/22] Correction to histogram cdf method

---
 openmc/stats/univariate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index 3d99f8b8ce..7fc479e630 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -889,7 +889,7 @@ class Tabular(Univariate):
         x = np.asarray(self.x)
         p = np.asarray(self.p)
         if self.interpolation == 'histogram':
-            c[1:] = p * np.diff(x)
+            c[1:] = p[:-1] * np.diff(x)
         elif self.interpolation == 'linear-linear':
             c[1:] = 0.5 * (p[0:-1] + p[1:]) * np.diff(x)
         return np.cumsum(c)

From 8070ccaf52abb93eca2e3ee99784284d5b730a26 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Tue, 5 Apr 2022 15:32:38 -0500
Subject: [PATCH 05/22] Adding sampling for Mixture. Improving performance of
 Discrete sampling

---
 openmc/stats/univariate.py | 36 ++++++++++++++++++++++++++++--------
 1 file changed, 28 insertions(+), 8 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index 7fc479e630..0e6b4ae61d 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -133,20 +133,19 @@ class Discrete(Univariate):
             cv.check_greater_than('discrete probability', pk, 0.0, True)
         self._p = p
 
+    def cdf(self):
+        return np.insert(np.cumsum(self.p), 0, 0.0)
+
     def sample(self, n_samples=1, seed=None):
         np.random.seed(seed)
 
         if len(self.x) == 1:
             return np.full((n_samples), self.x[0])
 
-        out = np.zeros((n_samples,))
-        for i, xi in enumerate(np.random.rand(n_samples)):
-            c = 0.0
-            for j, prob in enumerate(self.p):
-                c += prob
-                if (xi < c):
-                    out[i] = self.x[j]
-                    break
+        xi = np.random.rand(n_samples)
+        out = np.zeros_like(xi)
+        for i, c in enumerate(self.cdf()[:-1]):
+            out[xi >= c] = self.x[i]
         return out
 
     def normalize(self):
@@ -1086,6 +1085,27 @@ class Mixture(Univariate):
                       Iterable, Univariate)
         self._distribution = distribution
 
+    def cdf(self):
+        return np.insert(np.cumsum(self.probability), 0, 0.0)
+
+    def sample(self, n_samples=1, seed=None):
+        np.random.seed(seed)
+        xi = np.random.rand(n_samples)
+        idx = np.zeros_like(xi, dtype=int)
+        for i, c in enumerate(self.cdf()[:-1]):
+            idx[xi >= c] = i
+
+        out = np.zeros_like(xi)
+        for i in np.unique(idx):
+            n_dist_samples = np.count_nonzero(idx == i)
+            samples = self.distribution[i].sample(n_dist_samples)
+            out[idx == i] = samples
+        return out
+
+    def normalize(self):
+        norm = sum(self.probability)
+        self.probability = [val / norm for val in self.probability]
+
     def to_xml_element(self, element_name):
         """Return XML representation of the mixture distribution
 

From 2f22474aea252d2bdff2fecd5ef7e17aa31bba8f Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Wed, 6 Apr 2022 10:32:39 -0500
Subject: [PATCH 06/22] Adding simple test for normal distribution

---
 tests/unit_tests/test_stats.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py
index 51a561bd36..7bf67c1fa2 100644
--- a/tests/unit_tests/test_stats.py
+++ b/tests/unit_tests/test_stats.py
@@ -65,6 +65,7 @@ def test_uniform():
     assert t.p == [1/(b-a), 1/(b-a)]
     assert t.interpolation == 'histogram'
 
+
 def test_powerlaw():
     a, b, n = 10.0, 20.0, 2.0
     d = openmc.stats.PowerLaw(a, b, n)
@@ -76,6 +77,7 @@ def test_powerlaw():
     assert d.n == n
     assert len(d) == 3
 
+
 def test_maxwell():
     theta = 1.2895e6
     d = openmc.stats.Maxwell(theta)
@@ -264,6 +266,18 @@ def test_normal():
     assert d.std_dev == pytest.approx(std_dev)
     assert len(d) == 2
 
+    # sample normal distribution
+    n_samples = 10000
+    samples = d.sample(n_samples, seed=100)
+    samples = np.abs(samples - mean)
+    within_1_sigma = np.count_nonzero(samples < std_dev)
+    assert within_1_sigma / n_samples >= 0.68
+    within_2_sigma = np.count_nonzero(samples < 2*std_dev)
+    assert within_2_sigma / n_samples >= 0.95
+    within_3_sigma = np.count_nonzero(samples < 3*std_dev)
+    assert within_3_sigma / n_samples >= 0.99
+
+
 def test_muir():
     mean = 10.0
     mass = 5.0

From 2c74dcf336f135a11649391616bb9741aa139ee3 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Wed, 6 Apr 2022 10:35:46 -0500
Subject: [PATCH 07/22] Adding similar test for Muir distribution

---
 openmc/stats/univariate.py     |  7 +++++--
 tests/unit_tests/test_stats.py | 11 +++++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index 0e6b4ae61d..4f2fbf66c2 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -764,11 +764,14 @@ class Muir(Univariate):
         cv.check_greater_than('Muir kt', kt, 0.0)
         self._kt = kt
 
+    @property
+    def std_dev(self):
+        return np.sqrt(4.*self.e0*self.kt/self.m_rat)
+
     def sample(self, n_samples=1, seed=None):
         # https://permalink.lanl.gov/object/tr?what=info:lanl-repo/lareport/LA-05411-MS
         np.random.seed(seed)
-        sigma = np.sqrt(4.*self.e0*self.kt/self.m_rat)
-        return np.random.normal(self.e0, sigma, n_samples)
+        return np.random.normal(self.e0, self.std_dev, n_samples)
 
     def to_xml_element(self, element_name):
         """Return XML representation of the Watt distribution
diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py
index 7bf67c1fa2..532936e422 100644
--- a/tests/unit_tests/test_stats.py
+++ b/tests/unit_tests/test_stats.py
@@ -292,3 +292,14 @@ def test_muir():
     assert d.m_rat == pytest.approx(mass)
     assert d.kt == pytest.approx(temp)
     assert len(d) == 3
+
+    # sample normal distribution
+    n_samples = 10000
+    samples = d.sample(n_samples, seed=100)
+    samples = np.abs(samples - mean)
+    within_1_sigma = np.count_nonzero(samples < d.std_dev)
+    assert within_1_sigma / n_samples >= 0.68
+    within_2_sigma = np.count_nonzero(samples < 2*d.std_dev)
+    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

From 12b8a0802f7579ca1da24a7d2a0ba6bebdad6b3f Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Thu, 2 Jun 2022 23:42:06 -0500
Subject: [PATCH 08/22] Updating masks and adding some comments for clarity

---
 openmc/stats/univariate.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index 4f2fbf66c2..120dc62b6d 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -897,6 +897,10 @@ class Tabular(Univariate):
         return np.cumsum(c)
 
     def sample(self, n_samples=1, seed=None):
+        if not self.interpolation in ('histogram', 'linear-linear'):
+            raise NotImplementedError('Can only sample tabular distributions '
+                                      'using histogram or '
+                                      'linear-linear interpolation')
         np.random.seed(seed)
         xi = np.random.rand(n_samples)
         cdf = self.cdf()
@@ -905,7 +909,7 @@ class Tabular(Univariate):
         # sampled values
         c_i = np.full(n_samples, cdf[0])
         cdf_idx = np.zeros((n_samples,), dtype=int)
-        for i, val in enumerate(cdf):
+        for i, val in enumerate(cdf[:-1]):
             mask = xi > val
             c_i[mask] = val
             cdf_idx[mask] = i
@@ -916,13 +920,17 @@ class Tabular(Univariate):
         x_i = np.array([self.x[i] for i in cdf_idx])
         p_i = np.array([self.p[i] for i in cdf_idx])
 
+        # TODO: check that probability doesn't exceed the last value
+
         if self.interpolation == 'histogram':
-            # mask where probability
+            # mask where probability is greater than zero
             pos_mask = p_i > 0.0
+            # probabilities greater than zero are set proportional to the
+            # position of the random numebers in relation to the cdf value
             p_i[pos_mask] = x_i[pos_mask] + (xi[pos_mask] - c_i[pos_mask]) \
                            / p_i[pos_mask]
-            neg_mask = np.invert(pos_mask)
-            p_i[neg_mask] = x_i[neg_mask]
+            # probabilities smaller than zero are set to the random number value
+            p_i[~pos_mask] = x_i[~pos_mask]
             return p_i
         elif self.interpolation == 'linear-linear':
             # get variable and probability values for the
@@ -935,10 +943,10 @@ class Tabular(Univariate):
             zero = m == 0.0
             m[zero] = x_i[zero] + (xi[zero] - c_i[zero]) / p_i[zero]
             # set values for non-zero slope
-            non_zero = np.invert(zero)
+            non_zero = ~zero
             quad = np.power(p_i[non_zero], 2) + 2. * m[non_zero] * (xi[non_zero] - c_i[non_zero])
             quad[quad < 0.0] = 0.0
-            m[non_zero] = x_i[non_zero] +  (np.sqrt(quad) - p_i) / m[non_zero]
+            m[non_zero] = x_i[non_zero] +  (np.sqrt(quad) - p_i[non_zero]) / m[non_zero]
             return m
 
     def to_xml_element(self, element_name):

From 94c91cc6593fa6f61816e06b169dd6c82995bb19 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Thu, 2 Jun 2022 23:42:24 -0500
Subject: [PATCH 09/22] Adding more mean value tests for sampling distributions

---
 tests/unit_tests/test_stats.py | 84 +++++++++++++++++++++++++++++++++-
 1 file changed, 83 insertions(+), 1 deletion(-)

diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py
index 532936e422..a3b5cfafb5 100644
--- a/tests/unit_tests/test_stats.py
+++ b/tests/unit_tests/test_stats.py
@@ -253,6 +253,88 @@ def test_point():
     d = openmc.stats.Point.from_xml_element(elem)
     assert d.xyz == pytest.approx(p)
 
+
+def test_discrete():
+
+    vals = np.array([1.0, 2.0, 3.0])
+    probs = np.array([0.1, 0.7, 0.2])
+
+    exp_mean = (vals * probs).sum()
+
+    d = openmc.stats.Discrete(vals, probs)
+
+    # sample discrete distribution
+    n_samples = 1_000_000
+    samples = d.sample(n_samples, seed=100)
+    # check that the mean of the samples is close to the true mean
+    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+
+
+def test_uniform():
+
+    lower = 1.1
+    upper = 23.3
+
+    exp_mean = 0.5 * (lower + upper)
+
+    d = openmc.stats.Uniform(lower, upper)
+
+    # sample the uniform distribution
+    n_samples = 1_000_000
+    samples = d.sample(n_samples, seed=100)
+    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+
+
+def test_power_law():
+
+    lower = 0.0
+    upper = 1.0
+    exponent = 2.0
+
+    d = openmc.stats.PowerLaw(lower, upper, exponent)
+
+    exp_mean = (exponent + 1) / (exponent + 2)
+
+    # sample power law distribution
+    n_samples = 1_000_000
+    samples = d.sample(n_samples, seed=100)
+    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+
+
+def test_maxwell():
+    theta = 1000
+
+    exp_mean = 0
+
+    d = openmc.stats.Maxwell(theta)
+
+    exp_mean = 3/2 * theta
+
+    # sample maxwell distribution
+    n_samples = 1_000_000
+    samples = d.sample(n_samples, seed=100)
+    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+
+def test_watt():
+
+    a = 10
+    b = 20
+
+    d = openmc.stats.Watt(a, b)
+
+    # mean value form adapted from
+    # "Prompt-fission-neutron average energy for 238U(n, f ) from
+    # threshold to 200 MeV" Ethvignot et. al.
+    # https://doi.org/10.1016/j.physletb.2003.09.048
+    exp_mean = 3/2 * a + a**2 * b / 4
+
+    # sample Watt distribution
+    n_samples = 100_000
+    samples = d.sample(n_samples, seed=100)
+    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+
+
+
 def test_normal():
     mean = 10.0
     std_dev = 2.0
@@ -293,7 +375,7 @@ def test_muir():
     assert d.kt == pytest.approx(temp)
     assert len(d) == 3
 
-    # sample normal distribution
+    # sample muir distribution
     n_samples = 10000
     samples = d.sample(n_samples, seed=100)
     samples = np.abs(samples - mean)

From 9dc84fdb1859c73886b1f7596f30d3f6b9fa2cc9 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Thu, 9 Jun 2022 01:17:46 -0500
Subject: [PATCH 10/22] Always normalize distributions when sampling. Adding
 normalization method for tabular.

---
 openmc/stats/univariate.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index 120dc62b6d..df1a721502 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -890,12 +890,17 @@ class Tabular(Univariate):
         c = np.zeros_like(self.x)
         x = np.asarray(self.x)
         p = np.asarray(self.p)
+
         if self.interpolation == 'histogram':
             c[1:] = p[:-1] * np.diff(x)
         elif self.interpolation == 'linear-linear':
             c[1:] = 0.5 * (p[0:-1] + p[1:]) * np.diff(x)
+
         return np.cumsum(c)
 
+    def normalize(self):
+        self.p = self.p / self.cdf().max()
+
     def sample(self, n_samples=1, seed=None):
         if not self.interpolation in ('histogram', 'linear-linear'):
             raise NotImplementedError('Can only sample tabular distributions '
@@ -904,6 +909,9 @@ class Tabular(Univariate):
         np.random.seed(seed)
         xi = np.random.rand(n_samples)
         cdf = self.cdf()
+        cdf /= cdf.max()
+        # always use normalized probabilities when sampling
+        p = self.p / cdf.max()
 
         # get CDF bins that are above the
         # sampled values
@@ -918,7 +926,7 @@ class Tabular(Univariate):
         # the random number is less than the next cdf
         # entry
         x_i = np.array([self.x[i] for i in cdf_idx])
-        p_i = np.array([self.p[i] for i in cdf_idx])
+        p_i = np.array([p[i] for i in cdf_idx])
 
         # TODO: check that probability doesn't exceed the last value
 
@@ -936,7 +944,7 @@ class Tabular(Univariate):
             # get variable and probability values for the
             # next entry
             x_i1 = np.array([self.x[i+1] for i in cdf_idx])
-            p_i1 = np.array([self.p[i+1] for i in cdf_idx])
+            p_i1 = np.array([p[i+1] for i in cdf_idx])
             # compute slope between entries
             m = (p_i1 - p_i) / (x_i1 - x_i)
             # set values for zero slope
@@ -944,9 +952,9 @@ class Tabular(Univariate):
             m[zero] = x_i[zero] + (xi[zero] - c_i[zero]) / p_i[zero]
             # set values for non-zero slope
             non_zero = ~zero
-            quad = np.power(p_i[non_zero], 2) + 2. * m[non_zero] * (xi[non_zero] - c_i[non_zero])
+            quad = np.power(p_i[non_zero], 2) + 2.0 * m[non_zero] * (xi[non_zero] - c_i[non_zero])
             quad[quad < 0.0] = 0.0
-            m[non_zero] = x_i[non_zero] +  (np.sqrt(quad) - p_i[non_zero]) / m[non_zero]
+            m[non_zero] = x_i[non_zero] + (np.sqrt(quad) - p_i[non_zero]) / m[non_zero]
             return m
 
     def to_xml_element(self, element_name):

From 85b6566fd74c9c082ee943dd8d64a8138a119259 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Thu, 9 Jun 2022 01:18:17 -0500
Subject: [PATCH 11/22] Updating tabular means. Removing re-defined functions

---
 tests/unit_tests/test_stats.py | 174 +++++++++++++++++----------------
 1 file changed, 88 insertions(+), 86 deletions(-)

diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py
index a3b5cfafb5..f847970c01 100644
--- a/tests/unit_tests/test_stats.py
+++ b/tests/unit_tests/test_stats.py
@@ -26,6 +26,20 @@ def test_discrete():
     assert d2.p == [1.0]
     assert len(d2) == 1
 
+    vals = np.array([1.0, 2.0, 3.0])
+    probs = np.array([0.1, 0.7, 0.2])
+
+    exp_mean = (vals * probs).sum()
+
+    d3 = openmc.stats.Discrete(vals, probs)
+
+    # sample discrete distribution
+    n_samples = 1_000_000
+    samples = d3.sample(n_samples, seed=100)
+    # check that the mean of the samples is close to the true mean
+    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+
+
 
 def test_merge_discrete():
     x1 = [0.0, 1.0, 10.0]
@@ -65,9 +79,14 @@ def test_uniform():
     assert t.p == [1/(b-a), 1/(b-a)]
     assert t.interpolation == 'histogram'
 
+    exp_mean = 0.5 * (a + b)
+    n_samples = 1_000_000
+    samples = d.sample(n_samples, seed=100)
+    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+
 
 def test_powerlaw():
-    a, b, n = 10.0, 20.0, 2.0
+    a, b, n = 10.0, 100.0, 2.0
     d = openmc.stats.PowerLaw(a, b, n)
     elem = d.to_xml_element('distribution')
 
@@ -77,6 +96,13 @@ def test_powerlaw():
     assert d.n == n
     assert len(d) == 3
 
+    exp_mean = 100.0 * (n+1) / (n+2)
+
+    # sample power law distribution
+    n_samples = 1_000_000
+    samples = d.sample(n_samples, seed=100)
+    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+
 
 def test_maxwell():
     theta = 1.2895e6
@@ -87,6 +113,14 @@ def test_maxwell():
     assert d.theta == theta
     assert len(d) == 1
 
+    exp_mean = 3/2 * theta
+
+    # sample maxwell distribution
+    n_samples = 1_000_000
+    samples = d.sample(n_samples, seed=100)
+    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+
+
 
 def test_watt():
     a, b = 0.965e6, 2.29e-6
@@ -98,19 +132,68 @@ def test_watt():
     assert d.b == b
     assert len(d) == 2
 
+    d = openmc.stats.Watt(a, b)
+
+    # mean value form adapted from
+    # "Prompt-fission-neutron average energy for 238U(n, f ) from
+    # threshold to 200 MeV" Ethvignot et. al.
+    # https://doi.org/10.1016/j.physletb.2003.09.048
+    exp_mean = 3/2 * a + a**2 * b / 4
+
+    # sample Watt distribution
+    n_samples = 1_000_000
+    samples = d.sample(n_samples, seed=100)
+    assert samples.mean() == pytest.approx(exp_mean, rel=1e-03)
+
 
 def test_tabular():
-    x = [0.0, 5.0, 7.0]
-    p = [0.1, 0.2, 0.05]
+    x = np.array([0.0, 5.0, 7.0])
+    p = np.array([0.1, 0.2, 0.05])
     d = openmc.stats.Tabular(x, p, 'linear-linear')
     elem = d.to_xml_element('distribution')
 
     d = openmc.stats.Tabular.from_xml_element(elem)
-    assert d.x == x
-    assert d.p == p
+    assert all(d.x == x)
+    assert all(d.p == p)
     assert d.interpolation == 'linear-linear'
     assert len(d) == len(x)
 
+    # test linear-lienar sampling
+    d = openmc.stats.Tabular(x, p)
+
+    # compute the expected value (mean) of the
+    # piecewise pdf
+    mean = 0.0
+    d.normalize()
+    for i in range(1, len(x)):
+        y_min = d.p[i-1]
+        y_max = d.p[i]
+        x_min = d.x[i-1]
+        x_max = d.x[i]
+
+        m = (y_max - y_min) / (x_max - x_min)
+
+        exp_val = (1./3.) * m * (x_max**3 - x_min**3)
+        exp_val += 0.5 * m * x_min * (x_min**2 - x_max**2)
+        exp_val += 0.5 * y_min * (x_max**2 - x_min**2)
+
+        mean += exp_val
+
+    n_samples = 100_000
+    samples = d.sample(n_samples, seed=100)
+    assert samples.mean() == pytest.approx(mean, rel=1e-03)
+
+    # test histogram sampling
+    d = openmc.stats.Tabular(x, p, interpolation='histogram')
+    d.normalize()
+
+    mean = 0.5 * (x[:-1] + x[1:])
+    mean *= np.diff(d.cdf())
+    mean = sum(mean)
+
+    samples = d.sample(n_samples, seed=100)
+    assert samples.mean() == pytest.approx(mean, rel=1e-03)
+
 
 def test_legendre():
     # Pu239 elastic scattering at 100 keV
@@ -254,87 +337,6 @@ def test_point():
     assert d.xyz == pytest.approx(p)
 
 
-def test_discrete():
-
-    vals = np.array([1.0, 2.0, 3.0])
-    probs = np.array([0.1, 0.7, 0.2])
-
-    exp_mean = (vals * probs).sum()
-
-    d = openmc.stats.Discrete(vals, probs)
-
-    # sample discrete distribution
-    n_samples = 1_000_000
-    samples = d.sample(n_samples, seed=100)
-    # check that the mean of the samples is close to the true mean
-    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
-
-
-def test_uniform():
-
-    lower = 1.1
-    upper = 23.3
-
-    exp_mean = 0.5 * (lower + upper)
-
-    d = openmc.stats.Uniform(lower, upper)
-
-    # sample the uniform distribution
-    n_samples = 1_000_000
-    samples = d.sample(n_samples, seed=100)
-    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
-
-
-def test_power_law():
-
-    lower = 0.0
-    upper = 1.0
-    exponent = 2.0
-
-    d = openmc.stats.PowerLaw(lower, upper, exponent)
-
-    exp_mean = (exponent + 1) / (exponent + 2)
-
-    # sample power law distribution
-    n_samples = 1_000_000
-    samples = d.sample(n_samples, seed=100)
-    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
-
-
-def test_maxwell():
-    theta = 1000
-
-    exp_mean = 0
-
-    d = openmc.stats.Maxwell(theta)
-
-    exp_mean = 3/2 * theta
-
-    # sample maxwell distribution
-    n_samples = 1_000_000
-    samples = d.sample(n_samples, seed=100)
-    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
-
-def test_watt():
-
-    a = 10
-    b = 20
-
-    d = openmc.stats.Watt(a, b)
-
-    # mean value form adapted from
-    # "Prompt-fission-neutron average energy for 238U(n, f ) from
-    # threshold to 200 MeV" Ethvignot et. al.
-    # https://doi.org/10.1016/j.physletb.2003.09.048
-    exp_mean = 3/2 * a + a**2 * b / 4
-
-    # sample Watt distribution
-    n_samples = 100_000
-    samples = d.sample(n_samples, seed=100)
-    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
-
-
-
 def test_normal():
     mean = 10.0
     std_dev = 2.0

From 1e83df9dc1d6709cfcb181fb818f6e7e28bb85d0 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Thu, 9 Jun 2022 01:20:30 -0500
Subject: [PATCH 12/22] Removing re-defined distribution

---
 tests/unit_tests/test_stats.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py
index f847970c01..8e38a7b487 100644
--- a/tests/unit_tests/test_stats.py
+++ b/tests/unit_tests/test_stats.py
@@ -132,8 +132,6 @@ def test_watt():
     assert d.b == b
     assert len(d) == 2
 
-    d = openmc.stats.Watt(a, b)
-
     # mean value form adapted from
     # "Prompt-fission-neutron average energy for 238U(n, f ) from
     # threshold to 200 MeV" Ethvignot et. al.

From e1fd9959bc3aeb9d69005d5e5e971d11b441850f Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Thu, 9 Jun 2022 08:00:17 -0500
Subject: [PATCH 13/22] A little cleanup

---
 openmc/stats/univariate.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index df1a721502..d733b61935 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -566,7 +566,7 @@ class Watt(Univariate):
         w = Maxwell.sample_maxwell(self.a, n_samples)
         u = np.random.uniform(-1., 1., n_samples)
         aab = self.a * self.a * self.b
-        return w + 0.25*aab + u * np.sqrt(aab*w)
+        return w + 0.25*aab + u*np.sqrt(aab*w)
 
     def to_xml_element(self, element_name):
         """Return XML representation of the Watt distribution
@@ -769,7 +769,7 @@ class Muir(Univariate):
         return np.sqrt(4.*self.e0*self.kt/self.m_rat)
 
     def sample(self, n_samples=1, seed=None):
-        # https://permalink.lanl.gov/object/tr?what=info:lanl-repo/lareport/LA-05411-MS
+        # Based on LANL report LA-05411-MS
         np.random.seed(seed)
         return np.random.normal(self.e0, self.std_dev, n_samples)
 
@@ -899,7 +899,7 @@ class Tabular(Univariate):
         return np.cumsum(c)
 
     def normalize(self):
-        self.p = self.p / self.cdf().max()
+        self.p = np.asarray(self.p) / self.cdf().max()
 
     def sample(self, n_samples=1, seed=None):
         if not self.interpolation in ('histogram', 'linear-linear'):

From aba235e0b46a33cd6dbbb388d52a27f56e149cba Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Fri, 10 Jun 2022 08:50:52 -0500
Subject: [PATCH 14/22] Update tests/unit_tests/test_stats.py

Co-authored-by: Jonathan Shimwell <drshimwell@gmail.com>
---
 tests/unit_tests/test_stats.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py
index 8e38a7b487..c2578ea8b2 100644
--- a/tests/unit_tests/test_stats.py
+++ b/tests/unit_tests/test_stats.py
@@ -120,6 +120,11 @@ def test_maxwell():
     samples = d.sample(n_samples, seed=100)
     assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
 
+    # A second sample with a different seed
+    samples_2 = d.sample(n_samples, seed=200)
+    assert samples_2.mean() == pytest.approx(exp_mean, rel=1e-02)
+    assert samples_2.mean() != samples.mean()
+
 
 
 def test_watt():

From 0c3d30e6d1ebf373c7d79de1aeacfce88504609e Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Sun, 19 Jun 2022 09:31:38 -0500
Subject: [PATCH 15/22] Update openmc/stats/univariate.py

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
---
 openmc/stats/univariate.py | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index d733b61935..e4c53c559a 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -139,14 +139,7 @@ class Discrete(Univariate):
     def sample(self, n_samples=1, seed=None):
         np.random.seed(seed)
 
-        if len(self.x) == 1:
-            return np.full((n_samples), self.x[0])
-
-        xi = np.random.rand(n_samples)
-        out = np.zeros_like(xi)
-        for i, c in enumerate(self.cdf()[:-1]):
-            out[xi >= c] = self.x[i]
-        return out
+        return np.random.choice(self.x, n_samples, p=self.p)
 
     def normalize(self):
         norm = sum(self.p)

From f52865854ba41746d58b6621215bb5c223c896f2 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Sun, 19 Jun 2022 09:39:53 -0500
Subject: [PATCH 16/22] Apply suggestions from @paulromano

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
---
 openmc/stats/univariate.py     | 18 +++++++++---------
 tests/unit_tests/test_stats.py |  2 +-
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index e4c53c559a..6875f4e4f0 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -381,7 +381,7 @@ class PowerLaw(Univariate):
         pwr = self.n + 1
         offset = self.a**pwr
         span = self.b**pwr - offset
-        return np.power(offset + xi * span, (1/pwr))
+        return np.power(offset + xi * span, 1/pwr)
 
     def to_xml_element(self, element_name):
         """Return XML representation of the power law distribution
@@ -463,8 +463,8 @@ class Maxwell(Univariate):
     @staticmethod
     def sample_maxwell(t, n_samples):
         r1, r2, r3 = np.random.rand(3, n_samples)
-        c = np.power(np.cos(0.5 * np.pi * r3), 2)
-        return -t * (np.log(r1) + np.log(r2) * c)
+        c = np.cos(0.5 * np.pi * r3)
+        return -t * (np.log(r1) + np.log(r2) * c * c)
 
     def to_xml_element(self, element_name):
         """Return XML representation of the Maxwellian distribution
@@ -887,7 +887,7 @@ class Tabular(Univariate):
         if self.interpolation == 'histogram':
             c[1:] = p[:-1] * np.diff(x)
         elif self.interpolation == 'linear-linear':
-            c[1:] = 0.5 * (p[0:-1] + p[1:]) * np.diff(x)
+            c[1:] = 0.5 * (p[:-1] + p[1:]) * np.diff(x)
 
         return np.cumsum(c)
 
@@ -909,7 +909,7 @@ class Tabular(Univariate):
         # get CDF bins that are above the
         # sampled values
         c_i = np.full(n_samples, cdf[0])
-        cdf_idx = np.zeros((n_samples,), dtype=int)
+        cdf_idx = np.zeros(n_samples, dtype=int)
         for i, val in enumerate(cdf[:-1]):
             mask = xi > val
             c_i[mask] = val
@@ -918,8 +918,8 @@ class Tabular(Univariate):
         # get table values at each index where
         # the random number is less than the next cdf
         # entry
-        x_i = np.array([self.x[i] for i in cdf_idx])
-        p_i = np.array([p[i] for i in cdf_idx])
+        x_i = self.x[cdf_idx]
+        p_i = self.p[cdf_idx]
 
         # TODO: check that probability doesn't exceed the last value
 
@@ -936,8 +936,8 @@ class Tabular(Univariate):
         elif self.interpolation == 'linear-linear':
             # get variable and probability values for the
             # next entry
-            x_i1 = np.array([self.x[i+1] for i in cdf_idx])
-            p_i1 = np.array([p[i+1] for i in cdf_idx])
+            x_i1 = self.x[cdf_idx + 1]
+            p_i1 = self.p[cdf_idx + 1]
             # compute slope between entries
             m = (p_i1 - p_i) / (x_i1 - x_i)
             # set values for zero slope
diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py
index c2578ea8b2..f62f2157b2 100644
--- a/tests/unit_tests/test_stats.py
+++ b/tests/unit_tests/test_stats.py
@@ -161,7 +161,7 @@ def test_tabular():
     assert d.interpolation == 'linear-linear'
     assert len(d) == len(x)
 
-    # test linear-lienar sampling
+    # test linear-linear sampling
     d = openmc.stats.Tabular(x, p)
 
     # compute the expected value (mean) of the

From f46402fd87f33daf2da4050010973bd00b6c5bae Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Sun, 19 Jun 2022 09:36:31 -0500
Subject: [PATCH 17/22] Adding docstrings to normalize methods

---
 openmc/stats/univariate.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index 6875f4e4f0..bbac4c030f 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -142,6 +142,7 @@ class Discrete(Univariate):
         return np.random.choice(self.x, n_samples, p=self.p)
 
     def normalize(self):
+        """Normalize the probabilities stored on the distribution"""
         norm = sum(self.p)
         self.p = [val / norm for val in self.p]
 
@@ -892,6 +893,7 @@ class Tabular(Univariate):
         return np.cumsum(c)
 
     def normalize(self):
+        """Normalize the probabilities stored on the distribution"""
         self.p = np.asarray(self.p) / self.cdf().max()
 
     def sample(self, n_samples=1, seed=None):
@@ -1115,6 +1117,7 @@ class Mixture(Univariate):
         return out
 
     def normalize(self):
+        """Normalize the probabilities stored on the distribution"""
         norm = sum(self.probability)
         self.probability = [val / norm for val in self.probability]
 

From 862189985b42253cebe12bcd6fc2888c954f6c09 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Sun, 19 Jun 2022 09:51:45 -0500
Subject: [PATCH 18/22] Check that output values do not exceed values in
 tabulated dist

---
 openmc/stats/univariate.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index bbac4c030f..630bb665b8 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -923,8 +923,6 @@ class Tabular(Univariate):
         x_i = self.x[cdf_idx]
         p_i = self.p[cdf_idx]
 
-        # TODO: check that probability doesn't exceed the last value
-
         if self.interpolation == 'histogram':
             # mask where probability is greater than zero
             pos_mask = p_i > 0.0
@@ -934,7 +932,9 @@ class Tabular(Univariate):
                            / p_i[pos_mask]
             # probabilities smaller than zero are set to the random number value
             p_i[~pos_mask] = x_i[~pos_mask]
-            return p_i
+
+            samples_out = p_i
+
         elif self.interpolation == 'linear-linear':
             # get variable and probability values for the
             # next entry
@@ -950,7 +950,10 @@ class Tabular(Univariate):
             quad = np.power(p_i[non_zero], 2) + 2.0 * m[non_zero] * (xi[non_zero] - c_i[non_zero])
             quad[quad < 0.0] = 0.0
             m[non_zero] = x_i[non_zero] + (np.sqrt(quad) - p_i[non_zero]) / m[non_zero]
-            return m
+            samples_out = m
+
+        assert all(samples_out < self.x[-1])
+        return samples_out
 
     def to_xml_element(self, element_name):
         """Return XML representation of the tabular distribution

From 967fc9f88bd7b22373505253eec7d5d6af3f7a4a Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Sun, 19 Jun 2022 10:00:27 -0500
Subject: [PATCH 19/22] Adding a 'mean' method to tabular dist

---
 openmc/stats/univariate.py     | 30 +++++++++++++++++++++++++++++-
 tests/unit_tests/test_stats.py | 27 ++-------------------------
 2 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index 630bb665b8..ae89cb79fa 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -138,7 +138,6 @@ class Discrete(Univariate):
 
     def sample(self, n_samples=1, seed=None):
         np.random.seed(seed)
-
         return np.random.choice(self.x, n_samples, p=self.p)
 
     def normalize(self):
@@ -892,6 +891,35 @@ class Tabular(Univariate):
 
         return np.cumsum(c)
 
+    def mean(self):
+        """Compute the mean of the tabular distribution"""
+        if not self.interpolation in ('histogram', 'linear-linear'):
+            raise NotImplementedError('Can only compute mean for tabular '
+                                      'distributions using histogram '
+                                      'or linear-linear interpolation.')
+        if self.interpolation == 'linear-linear':
+            mean = 0.0
+            self.normalize()
+            for i in range(1, len(self.x)):
+                y_min = self.p[i-1]
+                y_max = self.p[i]
+                x_min = self.x[i-1]
+                x_max = self.x[i]
+
+                m = (y_max - y_min) / (x_max - x_min)
+
+                exp_val = (1./3.) * m * (x_max**3 - x_min**3)
+                exp_val += 0.5 * m * x_min * (x_min**2 - x_max**2)
+                exp_val += 0.5 * y_min * (x_max**2 - x_min**2)
+                mean += exp_val
+
+        elif self.interpolation == 'histogram':
+            mean = 0.5 * (self.x[:-1] + self.x[1:])
+            mean *= np.diff(self.cdf())
+            mean = sum(mean)
+
+        return mean
+
     def normalize(self):
         """Normalize the probabilities stored on the distribution"""
         self.p = np.asarray(self.p) / self.cdf().max()
diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py
index f62f2157b2..f789a0dca7 100644
--- a/tests/unit_tests/test_stats.py
+++ b/tests/unit_tests/test_stats.py
@@ -40,7 +40,6 @@ def test_discrete():
     assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
 
 
-
 def test_merge_discrete():
     x1 = [0.0, 1.0, 10.0]
     p1 = [0.3, 0.2, 0.5]
@@ -164,38 +163,16 @@ def test_tabular():
     # test linear-linear sampling
     d = openmc.stats.Tabular(x, p)
 
-    # compute the expected value (mean) of the
-    # piecewise pdf
-    mean = 0.0
-    d.normalize()
-    for i in range(1, len(x)):
-        y_min = d.p[i-1]
-        y_max = d.p[i]
-        x_min = d.x[i-1]
-        x_max = d.x[i]
-
-        m = (y_max - y_min) / (x_max - x_min)
-
-        exp_val = (1./3.) * m * (x_max**3 - x_min**3)
-        exp_val += 0.5 * m * x_min * (x_min**2 - x_max**2)
-        exp_val += 0.5 * y_min * (x_max**2 - x_min**2)
-
-        mean += exp_val
-
     n_samples = 100_000
     samples = d.sample(n_samples, seed=100)
-    assert samples.mean() == pytest.approx(mean, rel=1e-03)
+    assert samples.mean() == pytest.approx(d.mean(), rel=1e-03)
 
     # test histogram sampling
     d = openmc.stats.Tabular(x, p, interpolation='histogram')
     d.normalize()
 
-    mean = 0.5 * (x[:-1] + x[1:])
-    mean *= np.diff(d.cdf())
-    mean = sum(mean)
-
     samples = d.sample(n_samples, seed=100)
-    assert samples.mean() == pytest.approx(mean, rel=1e-03)
+    assert samples.mean() == pytest.approx(d.mean(), rel=1e-03)
 
 
 def test_legendre():

From f92ea393f23a602ef56227e057091fd4a6f811d3 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Mon, 20 Jun 2022 06:39:55 -0500
Subject: [PATCH 20/22] Using np.random.choice in Mixture class

---
 openmc/stats/univariate.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index ae89cb79fa..5528bdb303 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -1135,12 +1135,9 @@ class Mixture(Univariate):
 
     def sample(self, n_samples=1, seed=None):
         np.random.seed(seed)
-        xi = np.random.rand(n_samples)
-        idx = np.zeros_like(xi, dtype=int)
-        for i, c in enumerate(self.cdf()[:-1]):
-            idx[xi >= c] = i
+        idx = np.random.choice(self.distribution, n_samples, p=self.probability)
 
-        out = np.zeros_like(xi)
+        out = np.zeros_like(idx)
         for i in np.unique(idx):
             n_dist_samples = np.count_nonzero(idx == i)
             samples = self.distribution[i].sample(n_dist_samples)

From 3e61076715b9a15285b2d731127b968c600d7812 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Sat, 25 Jun 2022 14:38:19 +0200
Subject: [PATCH 21/22] Changing check for Tabular to use std. dev.

---
 tests/unit_tests/test_stats.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py
index f789a0dca7..87d8118812 100644
--- a/tests/unit_tests/test_stats.py
+++ b/tests/unit_tests/test_stats.py
@@ -165,14 +165,26 @@ def test_tabular():
 
     n_samples = 100_000
     samples = d.sample(n_samples, seed=100)
-    assert samples.mean() == pytest.approx(d.mean(), rel=1e-03)
+    diff = np.abs(samples - d.mean())
+    # within_1_sigma = np.count_nonzero(diff < samples.std())
+    # assert within_1_sigma / n_samples >= 0.68
+    within_2_sigma = np.count_nonzero(diff < 2*samples.std())
+    assert within_2_sigma / n_samples >= 0.95
+    within_3_sigma = np.count_nonzero(diff < 3*samples.std())
+    assert within_3_sigma / n_samples >= 0.99
 
     # test histogram sampling
     d = openmc.stats.Tabular(x, p, interpolation='histogram')
     d.normalize()
 
     samples = d.sample(n_samples, seed=100)
-    assert samples.mean() == pytest.approx(d.mean(), rel=1e-03)
+    diff = np.abs(samples - d.mean())
+    # within_1_sigma = np.count_nonzero(diff < samples.std())
+    # assert within_1_sigma / n_samples >= 0.68
+    within_2_sigma = np.count_nonzero(diff < 2*samples.std())
+    assert within_2_sigma / n_samples >= 0.95
+    within_3_sigma = np.count_nonzero(diff < 3*samples.std())
+    assert within_3_sigma / n_samples >= 0.99
 
 
 def test_legendre():

From e9eaf2edef09fe4cec82810d8669e6437d512a4d Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Sun, 26 Jun 2022 01:10:53 -0400
Subject: [PATCH 22/22] Using std. dev. of sampled mean for comparison

---
 tests/unit_tests/test_stats.py | 33 +++++++++++++++++++++++++--------
 1 file changed, 25 insertions(+), 8 deletions(-)

diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py
index 87d8118812..b57f9578a1 100644
--- a/tests/unit_tests/test_stats.py
+++ b/tests/unit_tests/test_stats.py
@@ -36,8 +36,10 @@ def test_discrete():
     # sample discrete distribution
     n_samples = 1_000_000
     samples = d3.sample(n_samples, seed=100)
-    # check that the mean of the samples is close to the true mean
-    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+    # check that the mean of the samples is within 3 std. dev.
+    # of the expected mean
+    std_dev = samples.std() / np.sqrt(n_samples)
+    assert np.abs(exp_mean - samples.mean()) < 3*std_dev
 
 
 def test_merge_discrete():
@@ -81,7 +83,10 @@ def test_uniform():
     exp_mean = 0.5 * (a + b)
     n_samples = 1_000_000
     samples = d.sample(n_samples, seed=100)
-    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+    # check that the mean of the samples is within 3 std. dev.
+    # of the expected mean
+    std_dev = samples.std() / np.sqrt(n_samples)
+    assert np.abs(exp_mean - samples.mean()) < 3*std_dev
 
 
 def test_powerlaw():
@@ -100,7 +105,10 @@ def test_powerlaw():
     # sample power law distribution
     n_samples = 1_000_000
     samples = d.sample(n_samples, seed=100)
-    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+    # check that the mean of the samples is within 3 std. dev.
+    # of the expected mean
+    std_dev = samples.std() / np.sqrt(n_samples)
+    assert np.abs(exp_mean - samples.mean()) < 3*std_dev
 
 
 def test_maxwell():
@@ -117,13 +125,19 @@ def test_maxwell():
     # sample maxwell distribution
     n_samples = 1_000_000
     samples = d.sample(n_samples, seed=100)
-    assert samples.mean() == pytest.approx(exp_mean, rel=1e-02)
+    # check that the mean of the samples is within 3 std. dev.
+    # of the expected mean
+    std_dev = samples.std() / np.sqrt(n_samples)
+    assert np.abs(exp_mean - samples.mean()) < 3*std_dev
 
     # A second sample with a different seed
     samples_2 = d.sample(n_samples, seed=200)
-    assert samples_2.mean() == pytest.approx(exp_mean, rel=1e-02)
-    assert samples_2.mean() != samples.mean()
+    # check that the mean of the samples is within 3 std. dev.
+    # of the expected mean
+    std_dev = samples_2.std() / np.sqrt(n_samples)
+    assert np.abs(exp_mean - samples_2.mean()) < 3*std_dev
 
+    assert samples_2.mean() != samples.mean()
 
 
 def test_watt():
@@ -145,7 +159,10 @@ def test_watt():
     # sample Watt distribution
     n_samples = 1_000_000
     samples = d.sample(n_samples, seed=100)
-    assert samples.mean() == pytest.approx(exp_mean, rel=1e-03)
+    # check that the mean of the samples is within 3 std. dev.
+    # of the expected mean
+    std_dev = samples.std() / np.sqrt(n_samples)
+    assert np.abs(exp_mean - samples.mean()) < 3*std_dev
 
 
 def test_tabular():