Adding a 'mean' method to tabular dist

This commit is contained in:
Patrick Shriwise 2022-06-19 10:00:27 -05:00
parent 862189985b
commit 967fc9f88b
2 changed files with 31 additions and 26 deletions

View file

@ -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()

View file

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