Adding sample methods to univariate distribution classes

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