From b469b3a16981269ca7202c070058587f78c2a542 Mon Sep 17 00:00:00 2001
From: myerspat <myerspat@umich.edu>
Date: Sat, 5 Nov 2022 13:50:27 -0400
Subject: [PATCH] Added alias sampling to discrete distribution

---
 include/openmc/distribution.h |  4 +++
 src/distribution.cpp          | 58 ++++++++++++++++++++++++++++++-----
 2 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/include/openmc/distribution.h b/include/openmc/distribution.h
index d3fe958caa..de16642753 100644
--- a/include/openmc/distribution.h
+++ b/include/openmc/distribution.h
@@ -52,6 +52,10 @@ public:
 private:
   vector<double> x_; //!< Possible outcomes
   vector<double> p_; //!< Probability of each outcome
+  
+  //! Alies method tables
+  vector<double> prob_; 
+  vector<double> alias_;
 
   //! Normalize distribution so that probabilities sum to unity
   void normalize();
diff --git a/src/distribution.cpp b/src/distribution.cpp
index 6bf51df6a6..29cf4014a2 100644
--- a/src/distribution.cpp
+++ b/src/distribution.cpp
@@ -36,20 +36,62 @@ Discrete::Discrete(const double* x, const double* p, int n)
   : x_ {x, x + n}, p_ {p, p + n}
 {
   normalize();
+
+  // Vectors for large and small probabilities based on 1/n
+  vector<int> large;
+  vector<int> small;
+
+  // Set and allocate memory
+  prob_ = p_;
+  alias_.reserve(n);
+
+  // Fill large and small vectors based on 1/n
+  for (int i = 0; i < n; i++) {
+    prob_[i] *= n;
+    if (prob_[i] > 1.0) {
+      large.push_back(i);
+    } else {
+      small.push_back(i);
+    }
+  }
+
+  while (!large.empty() && !small.empty()) {
+    int j = small.back();
+    int k = large.back();
+
+    // Update probability and alias based on Vose's algorithm
+    prob_[k] += prob_[j] - 1;
+    alias_[j] = k;
+
+    // Remove last vector element and or move large index to small vector
+    if (prob_[k] < 1) {
+      small.push_back(k);
+      large.pop_back();
+    } else {
+      small.pop_back();
+    }
+  }
 }
 
 double Discrete::sample(uint64_t* seed) const
 {
-  int n = x_.size();
+  int n = prob_.size();
   if (n > 1) {
-    double xi = prn(seed);
-    double c = 0.0;
-    for (int i = 0; i < n; ++i) {
-      c += p_[i];
-      if (xi < c)
-        return x_[i];
+    int u = prn(seed) * n;
+    if (prn(seed) < prob_[u]) {
+      return x_[u];
+    } else {
+      return x_[alias_[u]];
     }
-    throw std::runtime_error {"Error when sampling probability mass function."};
+
+    // double xi = prn(seed);
+    // double c = 0.0;
+    // for (int i = 0; i < n; ++i) {
+    //   c += p_[i];
+    //   if (xi < c)
+    //     return x_[i];
+    // }
+    // throw std::runtime_error {"Error when sampling probability mass function."};
   } else {
     return x_[0];
   }