From 0c3d30e6d1ebf373c7d79de1aeacfce88504609e Mon Sep 17 00:00:00 2001 From: Patrick Shriwise Date: Sun, 19 Jun 2022 09:31:38 -0500 Subject: [PATCH] Update openmc/stats/univariate.py Co-authored-by: Paul Romano --- openmc/stats/univariate.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py index d733b61935..e4c53c559a 100644 --- a/openmc/stats/univariate.py +++ b/openmc/stats/univariate.py @@ -139,14 +139,7 @@ class Discrete(Univariate): 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]) - - xi = np.random.rand(n_samples) - out = np.zeros_like(xi) - for i, c in enumerate(self.cdf()[:-1]): - out[xi >= c] = self.x[i] - return out + return np.random.choice(self.x, n_samples, p=self.p) def normalize(self): norm = sum(self.p)