Merge pull request #1882 from ojschumann/stats_branch

Implementation of openmc.stats.Mixture distribution

docs/source/io_formats/settings.rst

@@ -619,16 +619,17 @@ variable and whose sub-elements/attributes are as follows:
 
 :type:
   The type of the distribution. Valid options are "uniform", "discrete",
-  "tabular", "maxwell", and "watt". The "uniform" option produces variates
-  sampled from a uniform distribution over a finite interval. The "discrete"
-  option produces random variates that can assume a finite number of values
-  (i.e., a distribution characterized by a probability mass function). The
-  "tabular" option produces random variates sampled from a tabulated
+  "tabular", "maxwell", "watt", and "mixture". The "uniform" option produces
+  variates sampled from a uniform distribution over a finite interval. The 
+  "discrete" option produces random variates that can assume a finite number 
+  of values (i.e., a distribution characterized by a probability mass function). 
+  The "tabular" option produces random variates sampled from a tabulated
   distribution where the density function is either a histogram or
   linearly-interpolated between tabulated points. The "watt" option produces
   random variates is sampled from a Watt fission spectrum (only used for
   energies). The "maxwell" option produce variates sampled from a Maxwell
-  fission spectrum (only used for energies).
+  fission spectrum (only used for energies). The "mixture" option produces samples
+  from univariate sub-distributions with given probabilities.
 
   *Default*: None
 
@@ -651,12 +652,22 @@ variable and whose sub-elements/attributes are as follows:
 
   .. note:: The above format should be used even when using the multi-group
             :ref:`energy_mode`.
+
 :interpolation:
   For a "tabular" distribution, ``interpolation`` can be set to "histogram" or
   "linear-linear" thereby specifying how tabular points are to be interpolated.
 
   *Default*: histogram
 
+:pair:
+  For a "mixture" distribution, this element provides a distribution and its corresponding probability.
+
+  :probability:
+    An attribute or ``pair`` that provides the probability of a univariate distribution within a "mixture" distribution. 
+      
+  :dist:
+    This sub-element of a ``pair`` element provides information on the corresponding univariate distribution.
+
 -------------------------
 ``<state_point>`` Element
 -------------------------

include/openmc/distribution.h

@@ -24,6 +24,14 @@ public:
   virtual double sample(uint64_t* seed) const = 0;
 };
 
+using UPtrDist = unique_ptr<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);
+
+
 //==============================================================================
 //! A discrete distribution (probability mass function)
 //==============================================================================
@@ -222,12 +230,27 @@ private:
   vector<double> x_; //! Possible outcomes
 };
 
-using UPtrDist = unique_ptr<Distribution>;
+//==============================================================================
+//! 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;
+
+private:
+  // Storrage for probability + distribution
+  using DistPair = std::pair<double, UPtrDist>;
+
+  vector<DistPair>
+    distribution_; //!< sub-distributions + cummulative probabilities
+};
 
-//! 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
 

openmc/stats/univariate.py

@@ -812,8 +812,48 @@ 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):
-        raise NotImplementedError
+        """Generate mixture distribution from an XML element
+
+        Parameters
+        ----------
+        elem : xml.etree.ElementTree.Element
+            XML element
+
+        Returns
+        -------
+        openmc.stats.Mixture
+            Mixture distribution generated from XML element
+
+        """
+        probability = []
+        distribution = []
+        for pair in elem.findall('pair'):
+            probability.append(float(get_text(pair, 'probability')))
+            distribution.append(Univariate.from_xml_element(pair.find("dist")))
+
+        return cls(probability, distribution)

src/distribution.cpp

@@ -7,6 +7,8 @@
 #include <stdexcept> // for runtime_error
 #include <string>    // for string, stod
 
+#include <gsl/gsl-lite.hpp>
+
 #include "openmc/error.h"
 #include "openmc/math_functions.h"
 #include "openmc/random_dist.h"
@@ -287,6 +289,48 @@ double Equiprobable::sample(uint64_t* seed) const
   return xl + ((n - 1) * r - i) * (xr - xl);
 }
 
+//==============================================================================
+// Mixture implementation
+//==============================================================================
+
+Mixture::Mixture(pugi::xml_node node)
+{
+  double cumsum = 0.0;
+  for (pugi::xml_node pair : node.children("pair")) {
+    // Check that required data exists
+    if (!pair.attribute("probability")) fatal_error("Mixture pair element does not have probability.");
+    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());
+
+    // Save cummulative probybility and distrubution
+    distribution_.push_back(
+      std::make_pair(cumsum, distribution_from_xml(pair.child("dist"))));
+  }
+
+  // Normalize cummulative probabilities to 1
+  for (auto& pair : distribution_) {
+    pair.first /= cumsum;
+  }
+}
+
+double Mixture::sample(uint64_t* seed) const
+{
+  // Sample value of CDF
+  const double p = prn(seed);
+
+  // find matching distribution
+  const auto it = std::lower_bound(distribution_.cbegin(), distribution_.cend(),
+    p, [](const DistPair& pair, double p) { return pair.first < p; });
+
+  // This should not happen. Catch it
+  Ensures(it != distribution_.cend());
+
+  // Sample the chosen distribution
+  return it->second->sample(seed);
+}
+
 //==============================================================================
 // Helper function
 //==============================================================================
@@ -315,6 +359,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);
   }

tests/regression_tests/source/inputs_true.dat

@@ -76,4 +76,27 @@
       <parameters>1.0 1.3894954943731377 1.93069772888325 2.6826957952797255 3.72759372031494 5.17947467923121 7.196856730011519 10.0 13.894954943731374 19.306977288832496 26.826957952797247 37.2759372031494 51.7947467923121 71.96856730011518 100.0 138.94954943731375 193.06977288832496 268.26957952797244 372.7593720314938 517.9474679231207 719.6856730011514 1000.0 1389.4954943731375 1930.6977288832495 2682.6957952797247 3727.593720314938 5179.474679231207 7196.856730011514 10000.0 13894.95494373136 19306.977288832495 26826.95795279722 37275.93720314938 51794.74679231213 71968.56730011514 100000.0 138949.5494373136 193069.77288832495 268269.5795279722 372759.3720314938 517947.4679231202 719685.6730011514 1000000.0 1389495.494373136 1930697.7288832497 2682695.7952797217 3727593.720314938 5179474.679231202 7196856.730011513 10000000.0 0.0 2.9086439299358713e-08 5.80533561806147e-08 8.67817193689187e-08 1.1515347785771536e-07 1.4305204600565115e-07 1.7036278261198208e-07 1.9697346200185813e-07 2.227747351856934e-07 2.4766057919761985e-07 2.715287327665956e-07 2.9428111652990295e-07 3.1582423606228735e-07 3.360695660646056e-07 3.549339141332686e-07 3.723397626156721e-07 3.882155871468592e-07 4.024961505584776e-07 4.151227709522976e-07 4.260435628367196e-07 4.3521365033538783e-07 4.4259535159179273e-07 4.4815833361210174e-07 4.5187973690993757e-07 4.5374426944091084e-07 4.5374426944091084e-07 4.5187973690993757e-07 4.4815833361210174e-07 4.4259535159179273e-07 4.352136503353879e-07 4.2604356283671966e-07 4.1512277095229767e-07 4.0249615055847764e-07 3.8821558714685926e-07 3.723397626156722e-07 3.5493391413326864e-07 3.360695660646057e-07 3.158242360622874e-07 2.942811165299031e-07 2.715287327665957e-07 2.4766057919762e-07 2.2277473518569352e-07 1.9697346200185819e-07 1.7036278261198226e-07 1.4305204600565126e-07 1.1515347785771556e-07 8.678171936891881e-08 5.805335618061493e-08 2.9086439299358858e-08 5.559621115282002e-23</parameters>
     </energy>
   </source>
+  <source strength="0.1">
+    <space origin="1.0 1.0 0.0" type="cylindrical">
+      <r parameters="2.0 3.0" type="uniform" />
+      <phi parameters="0.0 6.283185307179586" type="uniform" />
+      <z interpolation="linear-linear" type="tabular">
+        <parameters>-2.0 0.0 2.0 0.2 0.3 0.2</parameters>
+      </z>
+    </space>
+    <angle type="isotropic" />
+    <energy type="mixture">
+      <pair probability="1">
+        <dist parameters="1289500.0" type="maxwell" />
+      </pair>
+      <pair probability="2">
+        <dist parameters="988000.0 2.249e-06" type="watt" />
+      </pair>
+      <pair probability="3">
+        <dist interpolation="histogram" type="tabular">
+          <parameters>1.0 1.3894954943731377 1.93069772888325 2.6826957952797255 3.72759372031494 5.17947467923121 7.196856730011519 10.0 13.894954943731374 19.306977288832496 26.826957952797247 37.2759372031494 51.7947467923121 71.96856730011518 100.0 138.94954943731375 193.06977288832496 268.26957952797244 372.7593720314938 517.9474679231207 719.6856730011514 1000.0 1389.4954943731375 1930.6977288832495 2682.6957952797247 3727.593720314938 5179.474679231207 7196.856730011514 10000.0 13894.95494373136 19306.977288832495 26826.95795279722 37275.93720314938 51794.74679231213 71968.56730011514 100000.0 138949.5494373136 193069.77288832495 268269.5795279722 372759.3720314938 517947.4679231202 719685.6730011514 1000000.0 1389495.494373136 1930697.7288832497 2682695.7952797217 3727593.720314938 5179474.679231202 7196856.730011513 10000000.0 0.0 2.9086439299358713e-08 5.80533561806147e-08 8.67817193689187e-08 1.1515347785771536e-07 1.4305204600565115e-07 1.7036278261198208e-07 1.9697346200185813e-07 2.227747351856934e-07 2.4766057919761985e-07 2.715287327665956e-07 2.9428111652990295e-07 3.1582423606228735e-07 3.360695660646056e-07 3.549339141332686e-07 3.723397626156721e-07 3.882155871468592e-07 4.024961505584776e-07 4.151227709522976e-07 4.260435628367196e-07 4.3521365033538783e-07 4.4259535159179273e-07 4.4815833361210174e-07 4.5187973690993757e-07 4.5374426944091084e-07 4.5374426944091084e-07 4.5187973690993757e-07 4.4815833361210174e-07 4.4259535159179273e-07 4.352136503353879e-07 4.2604356283671966e-07 4.1512277095229767e-07 4.0249615055847764e-07 3.8821558714685926e-07 3.723397626156722e-07 3.5493391413326864e-07 3.360695660646057e-07 3.158242360622874e-07 2.942811165299031e-07 2.715287327665957e-07 2.4766057919762e-07 2.2277473518569352e-07 1.9697346200185819e-07 1.7036278261198226e-07 1.4305204600565126e-07 1.1515347785771556e-07 8.678171936891881e-08 5.805335618061493e-08 2.9086439299358858e-08 5.559621115282002e-23</parameters>
+        </dist>
+      </pair>
+    </energy>
+  </source>
 </settings>

tests/regression_tests/source/results_true.dat

@@ -1,2 +1,2 @@
 k-combined:
-3.033600E-01 1.676108E-03
+2.980096E-01 9.632798E-04

tests/regression_tests/source/test.py

@@ -55,18 +55,20 @@ class SourceTestHarness(PyAPITestHarness):
         energy1 = openmc.stats.Maxwell(1.2895e6)
         energy2 = openmc.stats.Watt(0.988e6, 2.249e-6)
         energy3 = openmc.stats.Tabular(E, p, interpolation='histogram')
+        energy4 = openmc.stats.Mixture([1, 2, 3], [energy1, energy2, energy3])
 
         source1 = openmc.Source(spatial1, angle1, energy1, strength=0.5)
         source2 = openmc.Source(spatial2, angle2, energy2, strength=0.3)
         source3 = openmc.Source(spatial3, angle3, energy3, strength=0.1)
         source4 = openmc.Source(spatial4, angle3, energy3, strength=0.1)
         source5 = openmc.Source(spatial5, angle3, energy3, strength=0.1)
+        source6 = openmc.Source(spatial5, angle3, energy4, strength=0.1)
 
         settings = openmc.Settings()
         settings.batches = 10
         settings.inactive = 5
         settings.particles = 1000
-        settings.source = [source1, source2, source3, source4, source5]
+        settings.source = [source1, source2, source3, source4, source5, source6]
         settings.export_to_xml()
 
 

tests/unit_tests/test_stats.py

@@ -101,8 +101,12 @@ def test_mixture():
     assert mix.distribution == [d1, d2]
     assert len(mix) == 4
 
-    with pytest.raises(NotImplementedError):
-        mix.to_xml_element('distribution')
+    elem = mix.to_xml_element('distribution')
+
+    d = openmc.stats.Mixture.from_xml_element(elem)
+    assert d.probability == p
+    assert d.distribution == [d1, d2]
+    assert len(d) == 4
 
 
 def test_polar_azimuthal():