Merge pull request #2428 from gonuke/mesh_spatial_Discrete

Use Discrete distribution for sampling a mesh source

include/openmc/distribution.h

@@ -31,6 +31,39 @@ using UPtrDist = unique_ptr<Distribution>;
 //! \return Unique pointer to distribution
 UPtrDist distribution_from_xml(pugi::xml_node node);
 
+//==============================================================================
+//! A discrete distribution index (probability mass function)
+//==============================================================================
+
+class DiscreteIndex {
+public:
+  DiscreteIndex() {};
+  DiscreteIndex(pugi::xml_node node);
+  DiscreteIndex(const double* p, int n);
+
+  void assign(const double* p, int n);
+
+  //! Sample a value from the distribution
+  //! \param seed Pseudorandom number seed pointer
+  //! \return Sampled value
+  size_t sample(uint64_t* seed) const;
+
+  // Properties
+  const vector<double>& prob() const { return prob_; }
+  const vector<size_t>& alias() const { return alias_; }
+
+private:
+  vector<double> prob_; //!< Probability of accepting the uniformly sampled bin,
+                        //!< mapped to alias method table
+  vector<size_t> alias_; //!< Alias table
+
+  //! Normalize distribution so that probabilities sum to unity
+  void normalize();
+
+  //! Initialize alias tables for distribution
+  void init_alias();
+};
+
 //==============================================================================
 //! A discrete distribution (probability mass function)
 //==============================================================================
@@ -47,20 +80,13 @@ public:
 
   // Properties
   const vector<double>& x() const { return x_; }
-  const vector<double>& prob() const { return prob_; }
-  const vector<size_t>& alias() const { return alias_; }
+  const vector<double>& prob() const { return di_.prob(); }
+  const vector<size_t>& alias() const { return di_.alias(); }
 
 private:
-  vector<double> x_;    //!< Possible outcomes
-  vector<double> prob_; //!< Probability of accepting the uniformly sampled bin,
-                        //!< mapped to alias method table
-  vector<size_t> alias_; //!< Alias table
-
-  //! Normalize distribution so that probabilities sum to unity
-  void normalize();
-
-  //! Initialize alias tables for distribution
-  void init_alias(vector<double>& x, vector<double>& p);
+  vector<double> x_; //!< Possible outcomes
+  DiscreteIndex di_; //!< discrete probability distribution of
+                     //!< outcome indices
 };
 
 //==============================================================================

include/openmc/distribution_spatial.h

@@ -114,11 +114,8 @@ 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_;
+  DiscreteIndex elem_idx_dist_; //!< Distribution of
+                                //!< mesh element indices
 };
 
 //==============================================================================

src/distribution.cpp

@@ -19,32 +19,31 @@
 namespace openmc {
 
 //==============================================================================
-// Discrete implementation
+// DiscreteIndex implementation
 //==============================================================================
 
-Discrete::Discrete(pugi::xml_node node)
+DiscreteIndex::DiscreteIndex(pugi::xml_node node)
 {
   auto params = get_node_array<double>(node, "parameters");
+  std::size_t n = params.size() / 2;
 
-  std::size_t n = params.size();
-  std::vector<double> x_vec(params.begin(), params.begin() + n / 2);
-  std::vector<double> p_vec(params.begin() + n / 2, params.end());
-
-  this->init_alias(x_vec, p_vec);
+  assign(params.data() + n, n);
 }
 
-Discrete::Discrete(const double* x, const double* p, int n)
+DiscreteIndex::DiscreteIndex(const double* p, int n)
 {
-  std::vector<double> x_vec(x, x + n);
-  std::vector<double> p_vec(p, p + n);
-
-  this->init_alias(x_vec, p_vec);
+  assign(p, n);
 }
 
-void Discrete::init_alias(vector<double>& x, vector<double>& p)
+void DiscreteIndex::assign(const double* p, int n)
+{
+  prob_.assign(p, p + n);
+
+  this->init_alias();
+}
+
+void DiscreteIndex::init_alias()
 {
-  x_ = x;
-  prob_ = p;
   normalize();
 
   // The initialization and sampling method is based on Vose
@@ -53,13 +52,14 @@ void Discrete::init_alias(vector<double>& x, vector<double>& p)
   vector<size_t> large;
   vector<size_t> small;
 
+  size_t n = prob_.size();
+
   // Set and allocate memory
-  vector<size_t> alias(x_.size(), 0);
-  alias_ = alias;
+  alias_.assign(n, 0);
 
   // Fill large and small vectors based on 1/n
-  for (size_t i = 0; i < x_.size(); i++) {
-    prob_[i] *= x_.size();
+  for (size_t i = 0; i < n; i++) {
+    prob_[i] *= n;
     if (prob_[i] > 1.0) {
       large.push_back(i);
     } else {
@@ -86,23 +86,23 @@ void Discrete::init_alias(vector<double>& x, vector<double>& p)
   }
 }
 
-double Discrete::sample(uint64_t* seed) const
+size_t DiscreteIndex::sample(uint64_t* seed) const
 {
   // Alias sampling of discrete distribution
-  int n = x_.size();
+  size_t n = prob_.size();
   if (n > 1) {
-    int u = prn(seed) * n;
+    size_t u = prn(seed) * n;
     if (prn(seed) < prob_[u]) {
-      return x_[u];
+      return u;
     } else {
-      return x_[alias_[u]];
+      return alias_[u];
     }
   } else {
-    return x_[0];
+    return 0;
   }
 }
 
-void Discrete::normalize()
+void DiscreteIndex::normalize()
 {
   // Renormalize density function so that it sums to unity
   double norm = std::accumulate(prob_.begin(), prob_.end(), 0.0);
@@ -111,6 +111,30 @@ void Discrete::normalize()
   }
 }
 
+//==============================================================================
+// Discrete implementation
+//==============================================================================
+
+Discrete::Discrete(pugi::xml_node node) : di_(node)
+{
+  auto params = get_node_array<double>(node, "parameters");
+
+  std::size_t n = params.size() / 2;
+
+  x_.assign(params.begin(), params.begin() + n);
+}
+
+Discrete::Discrete(const double* x, const double* p, int n) : di_(p, n)
+{
+
+  x_.assign(x, x + n);
+}
+
+double Discrete::sample(uint64_t* seed) const
+{
+  return x_[di_.sample(seed)];
+}
+
 //==============================================================================
 // Uniform implementation
 //==============================================================================

src/distribution_spatial.cpp

@@ -210,16 +210,11 @@ 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;
+  std::vector<double> strengths(n_bins, 1.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
-  mesh_strengths_ = std::vector<double>(n_bins, 1.0);
   if (check_for_node(node, "strengths")) {
     strengths = get_node_array<double>(node, "strengths");
     if (strengths.size() != n_bins) {
@@ -228,36 +223,21 @@ 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);
   }
 
   if (get_node_value_bool(node, "volume_normalized")) {
     for (int i = 0; i < n_bins; i++) {
-      mesh_strengths_[i] *= mesh()->volume(i);
+      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;
+  elem_idx_dist_.assign(strengths.data(), n_bins);
 }
 
 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_dist_.sample(seed);
   return mesh()->sample(seed, elem_idx);
 }