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