Merge pull request #2080 from pshriwise/sample-dist

Sample distributions from Python API

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
 

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