Convert univariate probability distributions

CMakeLists.txt

@@ -384,6 +384,7 @@ add_library(libopenmc SHARED
   src/tallies/trigger.F90
   src/tallies/trigger_header.F90
   src/cell.cpp
+  src/distribution.cpp
   src/initialize.cpp
   src/finalize.cpp
   src/geometry_aux.cpp

src/constants.h

@@ -86,6 +86,11 @@ extern "C" double FP_PRECISION;
 constexpr double INFTY {std::numeric_limits<double>::max()};
 constexpr int C_NONE {-1};
 
+// Interpolation rules
+enum class Interpolation {
+  histogram, lin_lin, lin_log, log_lin, log_log
+};
+
 } // namespace openmc
 
 #endif // CONSTANTS_H

src/distribution.cpp

@@ -0,0 +1,137 @@
+#include "distribution.h"
+
+#include <algorithm>
+#include <cmath>
+#include <numeric>
+
+#include "error.h"
+#include "math_functions.h"
+#include "random_lcg.h"
+
+namespace openmc {
+
+Discrete::Discrete(const double* x, const double* p, int n)
+  : x_{x, x+n}, p_{p, p+n}
+{
+  // Renormalize density function so that it sums to unity
+  double norm = std::accumulate(p_.begin(), p_.end(), 0.0);
+  for (auto& p_i : p_)
+    p_i /= norm;
+}
+
+
+double Discrete::sample()
+{
+  int n = x_.size();
+  if (n > 1) {
+    double xi = prn();
+    double c = 0.0;
+    for (int i = 0; i < n; ++i) {
+      c += p_[i];
+      if (xi < c) return x_[i];
+    }
+    // throw exception?
+  } else {
+    return x_[0];
+  }
+}
+
+
+double Uniform::sample()
+{
+  return a_ + prn()*(b_ - a_);
+}
+
+
+double Maxwell::sample()
+{
+  return maxwell_spectrum_c(theta_);
+}
+
+
+double Watt::sample()
+{
+  return watt_spectrum_c(a_, b_);
+}
+
+
+Tabular::Tabular(const double* x, const double* p, int n, Interpolation interp)
+  : x_{x, x+n}, p_{p, p+n}, interp_{interp}, c_(n, 0.0)
+{
+  // Check interpolation parameter
+  if (interp_ != Interpolation::histogram &&
+      interp_ != Interpolation::lin_lin) {
+    openmc::fatal_error("Only histogram and linear-linear interpolation "
+                        "for tabular distribution is supported.");
+  }
+
+  // Calculate cumulative distribution function
+  for (int i = 1; i < n; ++i) {
+    if (interp_ == Interpolation::histogram) {
+      c_[i] = c_[i-1] + p_[i-1]*(x_[i] - x_[i-1]);
+    } else if (interp_ == Interpolation::lin_lin) {
+      c_[i] = c_[i-1] + 0.5*(p_[i-1] + p_[i]) * (x_[i] - x_[i-1]);
+    }
+  }
+
+  // Normalize density and distribution functions
+  for (int i = 0; i < n; ++i) {
+    p_[i] = p_[i]/c_[n-1];
+    c_[i] = c_[i]/c_[n-1];
+  }
+}
+
+
+double Tabular::sample()
+{
+  // Sample value of CDF
+  double c = prn();
+
+  // Find first CDF bin which is above the sampled value
+  double c_i = c_[0];
+  int i;
+  size_t n = c_.size();
+  for (i = 0; i < n - 1; ++i) {
+    if (c <= c_[i+1]) break;
+    c_i = c_[i+1];
+  }
+
+  // Determine bounding PDF values
+  double x_i = x_[i];
+  double p_i = p_[i];
+
+  if (interp_ == Interpolation::histogram) {
+    // Histogram interpolation
+    if (p_i > 0.0) {
+      return x_i + (c - c_i)/p_i;
+    } else {
+      return x_i;
+    }
+  } else {
+    // Linear-linear interpolation
+    double x_i1 = x_[i + 1];
+    double p_i1 = p_[i + 1];
+
+    double m = (p_i1 - p_i)/(x_i1 - x_i);
+    if (m == 0.0) {
+      return x_i + (c - c_i)/p_i;
+    } else {
+      return x_i + (std::sqrt(std::max(0.0, p_i*p_i + 2*m*(c - c_i))) - p_i)/m;
+    }
+  }
+}
+
+
+double Equiprobable::sample()
+{
+  size_t n = x_.size();
+
+  double r = prn();
+  int i = std::floor((n - 1)*r);
+
+  double xl = x_[i];
+  double xr = x_[i+i];
+  return xl + ((n - 1)*r - i) * (xr - xl);
+}
+
+} // namespace openmc

src/distribution.h

@@ -0,0 +1,78 @@
+#ifndef DISTRIBUTION_H
+#define DISTRIBUTION_H
+
+#include <vector>
+
+#include "constants.h"
+
+namespace openmc {
+
+class Distribution {
+public:
+  virtual double sample() = 0;
+  virtual ~Distribution() = default;
+};
+
+
+class Discrete : public Distribution {
+public:
+  Discrete(const double* x, const double* p, int n);
+  double sample();
+private:
+  std::vector<double> x_;
+  std::vector<double> p_;
+};
+
+
+class Uniform : public Distribution {
+public:
+  Uniform(double a, double b) : a_{a}, b_{b} {};
+  double sample();
+private:
+  double a_;
+  double b_;
+};
+
+
+class Maxwell : public Distribution {
+public:
+  Maxwell(double theta) : theta_{theta} { };
+  double sample();
+private:
+  double theta_;
+};
+
+
+class Watt : public Distribution {
+public:
+  Watt(double a, double b) : a_{a}, b_{b} { };
+  double sample();
+private:
+  double a_;
+  double b_;
+};
+
+
+class Tabular : public Distribution {
+public:
+  Tabular(const double* x, const double* p, int n, Interpolation interp);
+  double sample();
+private:
+  std::vector<double> x_;
+  std::vector<double> p_;
+  std::vector<double> c_;
+  Interpolation interp_;
+};
+
+
+class Equiprobable : public Distribution {
+public:
+  Equiprobable(const double* x, int n) : x_{x, x+n} { };
+  double sample();
+private:
+  std::vector<double> x_;
+};
+
+} // namespace openmc
+
+#endif // DISTRIBUTION_H