change p to prob and fix alias table

This commit is contained in:
Patrick A Myers 2023-01-13 16:50:50 -05:00
parent 8288b55822
commit 4122c02ad4
2 changed files with 23 additions and 23 deletions

View file

@ -47,13 +47,13 @@ public:
// Properties
const vector<double>& x() const { return x_; }
const vector<double>& p() const { return p_; }
const vector<double>& prob() const { return prob_; }
const vector<size_t>& alias() const { return alias_; }
private:
vector<double> x_; //!< Possible outcomes
vector<double>
p_; //!< Probability of each outcome, mapped to alias method table
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

View file

@ -1,6 +1,7 @@
#include "openmc/distribution.h"
#include <algorithm> // for copy
#include <array>
#include <cmath> // for sqrt, floor, max
#include <iterator> // for back_inserter
#include <numeric> // for accumulate
@ -26,34 +27,33 @@ Discrete::Discrete(pugi::xml_node node)
auto params = get_node_array<double>(node, "parameters");
std::size_t n = params.size();
std::copy(params.begin(), params.begin() + n / 2, std::back_inserter(x_));
std::copy(params.begin() + n / 2, params.end(), std::back_inserter(p_));
double x[n / 2], p[n / 2];
// Initialize alias tables
init_alias();
std::copy(params.begin(), params.begin() + n / 2, x);
std::copy(params.begin() + n / 2, params.end(), p);
new (this) Discrete(x, p, n / 2);
}
Discrete::Discrete(const double* x, const double* p, int n)
: x_ {x, x + n}, p_ {p, p + n}
: x_ {x, x + n}, prob_ {p, p + n}
{
// Initialize alias tables
init_alias();
}
void Discrete::init_alias() {
normalize();
// The initialization and sampling method is base on Vose
// (DOI: 10.1109/32.92917)
// Vectors for large and small probabilities based on 1/n
vector<size_t> large;
vector<size_t> small;
// Set and allocate memory
alias_.reserve(x_.size());
vector<size_t> alias(x_.size(), 0);
alias_ = alias;
// Fill large and small vectors based on 1/n
for (int i = 0; i < x_.size(); i++) {
p_[i] *= x_.size();
if (p_[i] > 1.0) {
for (size_t i = 0; i < x_.size(); i++) {
prob_[i] *= x_.size();
if (prob_[i] > 1.0) {
large.push_back(i);
} else {
small.push_back(i);
@ -68,11 +68,11 @@ void Discrete::init_alias() {
small.pop_back();
// Update probability and alias based on Vose's algorithm
p_[k] += p_[j] - 1.0;
prob_[k] += prob_[j] - 1.0;
alias_[j] = k;
// Move large index to small vector, if it is no longer large
if (p_[k] < 1.0) {
if (prob_[k] < 1.0) {
small.push_back(k);
large.pop_back();
}
@ -85,7 +85,7 @@ double Discrete::sample(uint64_t* seed) const
int n = x_.size();
if (n > 1) {
int u = prn(seed) * n;
if (prn(seed) < p_[u]) {
if (prn(seed) < prob_[u]) {
return x_[u];
} else {
return x_[alias_[u]];
@ -98,8 +98,8 @@ double Discrete::sample(uint64_t* seed) const
void Discrete::normalize()
{
// Renormalize density function so that it sums to unity
double norm = std::accumulate(p_.begin(), p_.end(), 0.0);
for (auto& p_i : p_) {
double norm = std::accumulate(prob_.begin(), prob_.end(), 0.0);
for (auto& p_i : prob_) {
p_i /= norm;
}
}