From 4072fc430b763e7026e0cefff11fbda6dc019eb8 Mon Sep 17 00:00:00 2001
From: myerspat <myerspat@umich.edu>
Date: Sat, 12 Nov 2022 13:31:27 -0500
Subject: [PATCH] Removing excess probability vector

---
 include/openmc/distribution.h |  9 ++++-----
 src/distribution.cpp          | 16 ++++++++--------
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/include/openmc/distribution.h b/include/openmc/distribution.h
index de16642753..e2c710ae9a 100644
--- a/include/openmc/distribution.h
+++ b/include/openmc/distribution.h
@@ -48,14 +48,13 @@ public:
   // Properties
   const vector<double>& x() const { return x_; }
   const vector<double>& p() const { return p_; }
+  const vector<size_t>& alias() const { return alias_; }
 
 private:
   vector<double> x_; //!< Possible outcomes
-  vector<double> p_; //!< Probability of each outcome
-  
-  //! Alies method tables
-  vector<double> prob_; 
-  vector<double> alias_;
+  vector<double>
+    p_; //!< Probability of each outcome, mapped to alias method table
+  vector<size_t> alias_; //!< Alias table
 
   //! Normalize distribution so that probabilities sum to unity
   void normalize();
diff --git a/src/distribution.cpp b/src/distribution.cpp
index 55b6c53b73..436e34116a 100644
--- a/src/distribution.cpp
+++ b/src/distribution.cpp
@@ -38,17 +38,16 @@ Discrete::Discrete(const double* x, const double* p, int n)
   normalize();
 
   // Vectors for large and small probabilities based on 1/n
-  vector<int> large;
-  vector<int> small;
+  vector<size_t> large;
+  vector<size_t> 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) {
+    p_[i] *= n;
+    if (p_[i] > 1.0) {
       large.push_back(i);
     } else {
       small.push_back(i);
@@ -63,11 +62,11 @@ Discrete::Discrete(const double* x, const double* p, int n)
     small.pop_back();
 
     // Update probability and alias based on Vose's algorithm
-    prob_[k] += prob_[j] - 1;
+    p_[k] += p_[j] - 1.0;
     alias_[j] = k;
 
     // Remove last vector element and or move large index to small vector
-    if (prob_[k] < 1.0) {
+    if (p_[k] < 1.0) {
       small.push_back(k);
       large.pop_back();
     }
@@ -76,10 +75,11 @@ Discrete::Discrete(const double* x, const double* p, int n)
 
 double Discrete::sample(uint64_t* seed) const
 {
+  // Alias sampling of discrete distribution
   int n = x_.size();
   if (n > 1) {
     int u = prn(seed) * n;
-    if (prn(seed) < prob_[u]) {
+    if (prn(seed) < p_[u]) {
       return x_[u];
     } else {
       return x_[alias_[u]];