From c93feed91f6f77714fab708fe16aeef7e0a526a4 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Thu, 31 Aug 2023 09:52:24 -0500
Subject: [PATCH] Account for integral of underlying distributions when
 sampling Mixture (#2658)

---
 include/openmc/distribution.h | 17 +++++++++++++++--
 src/distribution.cpp          | 29 ++++++++++++++++++-----------
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/include/openmc/distribution.h b/include/openmc/distribution.h
index 03b046bbbe1..65c68ad6457 100644
--- a/include/openmc/distribution.h
+++ b/include/openmc/distribution.h
@@ -22,6 +22,10 @@ class Distribution {
 public:
   virtual ~Distribution() = default;
   virtual double sample(uint64_t* seed) const = 0;
+
+  //! Return integral of distribution
+  //! \return Integral of distribution
+  virtual double integral() const { return 1.0; };
 };
 
 using UPtrDist = unique_ptr<Distribution>;
@@ -51,11 +55,13 @@ class DiscreteIndex {
   // Properties
   const vector<double>& prob() const { return prob_; }
   const vector<size_t>& alias() const { return alias_; }
+  double integral() const { return integral_; }
 
 private:
   vector<double> prob_; //!< Probability of accepting the uniformly sampled bin,
                         //!< mapped to alias method table
   vector<size_t> alias_; //!< Alias table
+  double integral_;      //!< Integral of distribution
 
   //! Normalize distribution so that probabilities sum to unity
   void normalize();
@@ -78,6 +84,8 @@ class Discrete : public Distribution {
   //! \return Sampled value
   double sample(uint64_t* seed) const override;
 
+  double integral() const override { return di_.integral(); };
+
   // Properties
   const vector<double>& x() const { return x_; }
   const vector<double>& prob() const { return di_.prob(); }
@@ -219,17 +227,19 @@ class Tabular : public Distribution {
   //! \return Sampled value
   double sample(uint64_t* seed) const override;
 
-  // x property
+  // properties
   vector<double>& x() { return x_; }
   const vector<double>& x() const { return x_; }
   const vector<double>& p() const { return p_; }
   Interpolation interp() const { return interp_; }
+  double integral() const override { return integral_; };
 
 private:
   vector<double> x_;     //!< tabulated independent variable
   vector<double> p_;     //!< tabulated probability density
   vector<double> c_;     //!< cumulative distribution at tabulated values
   Interpolation interp_; //!< interpolation rule
+  double integral_;      //!< Integral of distribution
 
   //! Initialize tabulated probability density function
   //! \param x Array of values for independent variable
@@ -272,12 +282,15 @@ class Mixture : public Distribution {
   //! \return Sampled value
   double sample(uint64_t* seed) const override;
 
+  double integral() const override { return integral_; }
+
 private:
   // Storrage for probability + distribution
   using DistPair = std::pair<double, UPtrDist>;
 
   vector<DistPair>
-    distribution_; //!< sub-distributions + cummulative probabilities
+    distribution_;  //!< sub-distributions + cummulative probabilities
+  double integral_; //!< integral of distribution
 };
 
 } // namespace openmc
diff --git a/src/distribution.cpp b/src/distribution.cpp
index 10d404c10f4..9c76147ee27 100644
--- a/src/distribution.cpp
+++ b/src/distribution.cpp
@@ -104,10 +104,12 @@ size_t DiscreteIndex::sample(uint64_t* seed) const
 
 void DiscreteIndex::normalize()
 {
-  // Renormalize density function so that it sums to unity
-  double norm = std::accumulate(prob_.begin(), prob_.end(), 0.0);
+  // Renormalize density function so that it sums to unity. Note that we save
+  // the integral of the distribution so that if it is used as part of another
+  // distribution (e.g., Mixture), we know its relative strength.
+  integral_ = std::accumulate(prob_.begin(), prob_.end(), 0.0);
   for (auto& p_i : prob_) {
-    p_i /= norm;
+    p_i /= integral_;
   }
 }
 
@@ -300,10 +302,13 @@ void Tabular::init(
     }
   }
 
-  // Normalize density and distribution functions
+  // Normalize density and distribution functions. Note that we save the
+  // integral of the distribution so that if it is used as part of another
+  // distribution (e.g., Mixture), we know its relative strength.
+  integral_ = c_[n - 1];
   for (int i = 0; i < n; ++i) {
-    p_[i] = p_[i] / c_[n - 1];
-    c_[i] = c_[i] / c_[n - 1];
+    p_[i] = p_[i] / integral_;
+    c_[i] = c_[i] / integral_;
   }
 }
 
@@ -379,12 +384,14 @@ Mixture::Mixture(pugi::xml_node node)
     if (!pair.child("dist"))
       fatal_error("Mixture pair element does not have a distribution.");
 
-    // cummulative sum of probybilities
-    cumsum += std::stod(pair.attribute("probability").value());
+    // cummulative sum of probabilities
+    double p = std::stod(pair.attribute("probability").value());
 
-    // Save cummulative probybility and distrubution
-    distribution_.push_back(
-      std::make_pair(cumsum, distribution_from_xml(pair.child("dist"))));
+    // Save cummulative probability and distribution
+    auto dist = distribution_from_xml(pair.child("dist"));
+    cumsum += p * dist->integral();
+
+    distribution_.push_back(std::make_pair(cumsum, std::move(dist)));
   }
 
   // Normalize cummulative probabilities to 1