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