diff --git a/include/openmc/distribution.h b/include/openmc/distribution.h index 44fd985d2f..dcd31f4293 100644 --- a/include/openmc/distribution.h +++ b/include/openmc/distribution.h @@ -24,6 +24,14 @@ public: virtual double sample(uint64_t* seed) const = 0; }; +using UPtrDist = unique_ptr; + +//! Return univariate probability distribution specified in XML file +//! \param[in] node XML node representing distribution +//! \return Unique pointer to distribution +UPtrDist distribution_from_xml(pugi::xml_node node); + + //============================================================================== //! A discrete distribution (probability mass function) //============================================================================== @@ -222,12 +230,28 @@ private: vector x_; //! Possible outcomes }; -using UPtrDist = unique_ptr; +//============================================================================== +//! Mixture distribution +//============================================================================== + +class Mixture : public Distribution { +public: + explicit Mixture(pugi::xml_node node); + + //! Sample a value from the distribution + //! \param seed Pseudorandom number seed pointer + //! \return Sampled value + double sample(uint64_t* seed) const; + + // d and p property + const vector& d() const { return d_; } + const vector& p() const { return p_; } +private: + vector d_; //!< Pointer to sub-distributions + vector p_; //!< tabulated probability density + vector c_; //!< cumulative distribution +}; -//! Return univariate probability distribution specified in XML file -//! \param[in] node XML node representing distribution -//! \return Unique pointer to distribution -UPtrDist distribution_from_xml(pugi::xml_node node); } // namespace openmc diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py index 4bbc509277..c7b6dadbe7 100644 --- a/openmc/stats/univariate.py +++ b/openmc/stats/univariate.py @@ -812,7 +812,28 @@ class Mixture(Univariate): self._distribution = distribution def to_xml_element(self, element_name): - raise NotImplementedError + """Return XML representation of the mixture distribution + + Parameters + ---------- + element_name : str + XML element name + + Returns + ------- + element : xml.etree.ElementTree.Element + XML element containing mixture distribution data + + """ + element = ET.Element(element_name) + element.set("type", "mixture") + + for p,d in zip(self.probability, self.distribution): + data = ET.SubElement(element, "pair") + data.set("probability", str(p)) + data.append(d.to_xml_element("dist")) + + return element @classmethod def from_xml_element(cls, elem): diff --git a/src/distribution.cpp b/src/distribution.cpp index 4713110ab4..f473513ba7 100644 --- a/src/distribution.cpp +++ b/src/distribution.cpp @@ -287,6 +287,42 @@ double Equiprobable::sample(uint64_t* seed) const return xl + ((n - 1) * r - i) * (xr - xl); } +//============================================================================== +// Mixture implementation +//============================================================================== + +Mixture::Mixture(pugi::xml_node node) +{ + // Read and initialize tabular distribution + //auto params = get_node_array(node, "parameters"); + + double cumsum = 0.0; + for (pugi::xml_node pair = node.child("pair"); pair; pair = pair.next_sibling("pair")) { + if (!pair.attribute("probability")) openmc::fatal_error("Mixture pair element does not have probability."); + if (!pair.child("dist")) openmc::fatal_error("Mixture pair element does not have a distribution."); + double p = std::stod(pair.attribute("probability").value()); + p_.push_back(p); + cumsum += p; + c_.push_back(cumsum); + d_.push_back(distribution_from_xml(pair.child("dist"))); + } + std::transform(c_.begin(), c_.end(), c_.begin(), [cumsum](double c){ return c/cumsum;}); +} + +double Mixture::sample(uint64_t* seed) const { + // + // Sample value of CDF + double c = prn(seed); + + // Find first CDF bin which is above the sampled value and sample corresponding sub-distribution + for (size_t i = 0; i < c_.size(); i++) { + if (c<=c_[i]) return d_[i]->sample(seed); + } + + // This should not happen. Catch it + openmc::fatal_error("Bad Sampling in Mixture Distribution."); +} + //============================================================================== // Helper function //============================================================================== @@ -315,6 +351,8 @@ UPtrDist distribution_from_xml(pugi::xml_node node) dist = UPtrDist {new Discrete(node)}; } else if (type == "tabular") { dist = UPtrDist {new Tabular(node)}; + } else if (type == "mixture") { + dist = UPtrDist{new Mixture(node)}; } else { openmc::fatal_error("Invalid distribution type: " + type); }