diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py index e4a9b8f924..5528bdb303 100644 --- a/openmc/stats/univariate.py +++ b/openmc/stats/univariate.py @@ -60,6 +60,24 @@ class Univariate(EqualityMixin, ABC): elif distribution == 'mixture': return Mixture.from_xml_element(elem) + @abstractmethod + 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 + class Discrete(Univariate): """Distribution characterized by a probability mass function. @@ -115,6 +133,18 @@ 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) + 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] + def to_xml_element(self, element_name): """Return XML representation of the discrete distribution @@ -240,6 +270,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 +375,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 +456,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) + 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.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 @@ -502,6 +554,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 +647,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 +757,15 @@ 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): + # Based on LANL report LA-05411-MS + np.random.seed(seed) + 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 @@ -803,6 +875,114 @@ class Tabular(Univariate): cv.check_value('interpolation', interpolation, _INTERPOLATION_SCHEMES) self._interpolation = interpolation + def cdf(self): + 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(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[:-1] + p[1:]) * np.diff(x) + + 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() + + 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() + cdf /= cdf.max() + # always use normalized probabilities when sampling + p = self.p / cdf.max() + + # 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) + for i, val in enumerate(cdf[:-1]): + 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 = self.x[cdf_idx] + p_i = self.p[cdf_idx] + + if self.interpolation == 'histogram': + # 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] + # probabilities smaller than zero are set to the random number value + p_i[~pos_mask] = x_i[~pos_mask] + + samples_out = p_i + + elif self.interpolation == 'linear-linear': + # get variable and probability values for the + # next entry + 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 + 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 = ~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] + 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 @@ -890,6 +1070,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 @@ -947,6 +1130,25 @@ 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) + idx = np.random.choice(self.distribution, n_samples, p=self.probability) + + 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) + out[idx == i] = samples + 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] + def to_xml_element(self, element_name): """Return XML representation of the mixture distribution diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py index 51a561bd36..b57f9578a1 100644 --- a/tests/unit_tests/test_stats.py +++ b/tests/unit_tests/test_stats.py @@ -26,6 +26,21 @@ 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 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(): x1 = [0.0, 1.0, 10.0] @@ -65,8 +80,17 @@ 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) + # 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(): - 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') @@ -76,6 +100,17 @@ 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) + # 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(): theta = 1.2895e6 d = openmc.stats.Maxwell(theta) @@ -85,6 +120,25 @@ 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) + # 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) + # 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(): a, b = 0.965e6, 2.29e-6 @@ -96,19 +150,59 @@ def test_watt(): assert d.b == b assert len(d) == 2 + # 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) + # 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(): - 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-linear sampling + d = openmc.stats.Tabular(x, p) + + n_samples = 100_000 + samples = d.sample(n_samples, seed=100) + 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) + 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(): # Pu239 elastic scattering at 100 keV @@ -251,6 +345,7 @@ def test_point(): d = openmc.stats.Point.from_xml_element(elem) assert d.xyz == pytest.approx(p) + def test_normal(): mean = 10.0 std_dev = 2.0 @@ -264,6 +359,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 @@ -278,3 +385,14 @@ def test_muir(): assert d.m_rat == pytest.approx(mass) assert d.kt == pytest.approx(temp) assert len(d) == 3 + + # sample muir 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