diff --git a/include/openmc/distribution.h b/include/openmc/distribution.h index 815ffa21a5..03b046bbbe 100644 --- a/include/openmc/distribution.h +++ b/include/openmc/distribution.h @@ -31,6 +31,39 @@ using UPtrDist = unique_ptr; //! \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& prob() const { return prob_; } + const vector& alias() const { return alias_; } + +private: + vector prob_; //!< Probability of accepting the uniformly sampled bin, + //!< mapped to alias method table + vector 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& x() const { return x_; } - const vector& prob() const { return prob_; } - const vector& alias() const { return alias_; } + const vector& prob() const { return di_.prob(); } + const vector& alias() const { return di_.alias(); } private: - vector x_; //!< Possible outcomes - vector prob_; //!< Probability of accepting the uniformly sampled bin, - //!< mapped to alias method table - vector alias_; //!< Alias table - - //! Normalize distribution so that probabilities sum to unity - void normalize(); - - //! Initialize alias tables for distribution - void init_alias(vector& x, vector& p); + vector x_; //!< Possible outcomes + DiscreteIndex di_; //!< discrete probability distribution of + //!< outcome indices }; //============================================================================== diff --git a/include/openmc/distribution_spatial.h b/include/openmc/distribution_spatial.h index a8cc7c58e0..9fba950602 100644 --- a/include/openmc/distribution_spatial.h +++ b/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 mesh_CDF_; - std::vector mesh_strengths_; + DiscreteIndex elem_idx_dist_; //!< Distribution of + //!< mesh element indices }; //============================================================================== diff --git a/src/distribution.cpp b/src/distribution.cpp index 4b4ebd69d7..10d404c10f 100644 --- a/src/distribution.cpp +++ b/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(node, "parameters"); + std::size_t n = params.size() / 2; - std::size_t n = params.size(); - std::vector x_vec(params.begin(), params.begin() + n / 2); - std::vector 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 x_vec(x, x + n); - std::vector p_vec(p, p + n); - - this->init_alias(x_vec, p_vec); + assign(p, n); } -void Discrete::init_alias(vector& x, vector& 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& x, vector& p) vector large; vector small; + size_t n = prob_.size(); + // Set and allocate memory - vector 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& x, vector& 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(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 //============================================================================== diff --git a/src/distribution_spatial.cpp b/src/distribution_spatial.cpp index fb1fbf099e..e1a4c20c21 100644 --- a/src/distribution_spatial.cpp +++ b/src/distribution_spatial.cpp @@ -210,16 +210,11 @@ MeshSpatial::MeshSpatial(pugi::xml_node node) } int32_t n_bins = this->n_sources(); - std::vector strengths(n_bins, 0.0); - - mesh_CDF_.resize(n_bins + 1); - mesh_CDF_[0] = {0.0}; - total_strength_ = 0.0; + std::vector 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(n_bins, 1.0); if (check_for_node(node, "strengths")) { strengths = get_node_array(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); }