From ee6975767d98f84133a0e4df3e443800a068fc32 Mon Sep 17 00:00:00 2001
From: "Paul P.H. Wilson" <paul.wilson@wisc.edu>
Date: Thu, 16 Mar 2023 17:34:44 -0500
Subject: [PATCH] Use Discrete distribution for mesh sampling

This version builds the independent variable by
converting int to double and then converts back
when sampling.
---
 include/openmc/distribution_spatial.h |  6 +-----
 src/distribution_spatial.cpp          | 27 ++-------------------------
 2 files changed, 3 insertions(+), 30 deletions(-)

diff --git a/include/openmc/distribution_spatial.h b/include/openmc/distribution_spatial.h
index a8cc7c58e..63cbbb956 100644
--- a/include/openmc/distribution_spatial.h
+++ b/include/openmc/distribution_spatial.h
@@ -114,11 +114,7 @@ public:
 
 private:
   int32_t mesh_idx_ {C_NONE};
-  double total_strength_ {0.0};
-  // TODO: move to an independent class in the future that's similar
-  // to a discrete distribution without outcomes
-  std::vector<double> mesh_CDF_;
-  std::vector<double> mesh_strengths_;
+  UPtrDist elem_idx_; //!< Distribution of mesh element indices
 };
 
 //==============================================================================
diff --git a/src/distribution_spatial.cpp b/src/distribution_spatial.cpp
index fb1fbf099..206f0b597 100644
--- a/src/distribution_spatial.cpp
+++ b/src/distribution_spatial.cpp
@@ -212,10 +212,6 @@ MeshSpatial::MeshSpatial(pugi::xml_node node)
   int32_t n_bins = this->n_sources();
   std::vector<double> strengths(n_bins, 0.0);
 
-  mesh_CDF_.resize(n_bins + 1);
-  mesh_CDF_[0] = {0.0};
-  total_strength_ = 0.0;
-
   // Create cdfs for sampling for an element over a mesh
   // Volume scheme is weighted by the volume of each tet
   // File scheme is weighted by an array given in the xml file
@@ -228,28 +224,9 @@ MeshSpatial::MeshSpatial(pugi::xml_node node)
                     "not match the number of entities in mesh {} ({}).",
           strengths.size(), mesh_id, n_bins));
     }
-    mesh_strengths_ = std::move(strengths);
+    elem_idx_ = UPtrDist {new Discrete {strengths, n_bins}};
   }
 
-  if (get_node_value_bool(node, "volume_normalized")) {
-    for (int i = 0; i < n_bins; i++) {
-      mesh_strengths_[i] *= mesh()->volume(i);
-    }
-  }
-
-  total_strength_ =
-    std::accumulate(mesh_strengths_.begin(), mesh_strengths_.end(), 0.0);
-
-  for (int i = 0; i < n_bins; i++) {
-    mesh_CDF_[i + 1] = mesh_CDF_[i] + mesh_strengths_[i] / total_strength_;
-  }
-
-  if (fabs(mesh_CDF_.back() - 1.0) > FP_COINCIDENT) {
-    fatal_error(
-      fmt::format("Mesh sampling CDF is incorrectly formed. Final value is: {}",
-        mesh_CDF_.back()));
-  }
-  mesh_CDF_.back() = 1.0;
 }
 
 Position MeshSpatial::sample(uint64_t* seed) const
@@ -257,7 +234,7 @@ Position MeshSpatial::sample(uint64_t* seed) const
   // Create random variable for sampling element from mesh
   double eta = prn(seed);
   // Sample over the CDF defined in initialization above
-  int32_t elem_idx = lower_bound_index(mesh_CDF_.begin(), mesh_CDF_.end(), eta);
+  int32_t elem_idx = elem_idx_->sample(eta);
   return mesh()->sample(seed, elem_idx);
 }