Merge pull request #2056 from paulromano/merge-discrete

Ability to merge discrete distributions

openmc/model/surface_composite.py

@@ -212,7 +212,7 @@ class CylinderSector(CompositeSurface):
 
 
 class IsogonalOctagon(CompositeSurface):
-    """Infinite isogonal octagon composite surface
+    r"""Infinite isogonal octagon composite surface
 
     An isogonal octagon is composed of eight planar surfaces. The prism is
     parallel to the x, y, or z axis. The remaining two axes (y and z, z and x,

openmc/stats/univariate.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from collections.abc import Iterable
 from numbers import Real
 from xml.etree import ElementTree as ET
@@ -156,6 +157,38 @@ class Discrete(Univariate):
         p = params[len(params)//2:]
         return cls(x, p)
 
+    @classmethod
+    def merge(cls, dists, probs):
+        """Merge multiple discrete distributions into a single distribution
+
+        Parameters
+        ----------
+        dists : iterable of openmc.stats.Discrete
+            Discrete distributions to combine
+        probs : iterable of float
+            Probability of each distribution
+
+        Returns
+        -------
+        openmc.stats.Discrete
+            Combined discrete distribution
+
+        """
+        if len(dists) != len(probs):
+            raise ValueError("Number of distributions and probabilities must match.")
+
+        # Combine distributions accounting for duplicate x values
+        x_merged = set()
+        p_merged = defaultdict(float)
+        for dist, p_dist in zip(dists, probs):
+            for x, p in zip(dist.x, dist.p):
+                x_merged.add(x)
+                p_merged[x] += p*p_dist
+
+        # Create values and probabilities as arrays
+        x_arr = np.array(sorted(x_merged))
+        p_arr = np.array([p_merged[x] for x in x_arr])
+        return cls(x_arr, p_arr)
 
 class Uniform(Univariate):
     """Distribution with constant probability over a finite interval [a,b]

tests/unit_tests/test_stats.py

@@ -27,6 +27,29 @@ def test_discrete():
     assert len(d2) == 1
 
 
+def test_merge_discrete():
+    x1 = [0.0, 1.0, 10.0]
+    p1 = [0.3, 0.2, 0.5]
+    d1 = openmc.stats.Discrete(x1, p1)
+
+    x2 = [0.5, 1.0, 5.0]
+    p2 = [0.4, 0.5, 0.1]
+    d2 = openmc.stats.Discrete(x2, p2)
+
+    # Merged distribution should have x values sorted and probabilities
+    # appropriately combined. Duplicate x values should appear once.
+    merged = openmc.stats.Discrete.merge([d1, d2], [0.6, 0.4])
+    assert merged.x == pytest.approx([0.0, 0.5, 1.0, 5.0, 10.0])
+    assert merged.p == pytest.approx(
+        [0.6*0.3, 0.4*0.4, 0.6*0.2 + 0.4*0.5, 0.4*0.1, 0.6*0.5])
+
+    # Probabilities add up but are not normalized
+    d1 = openmc.stats.Discrete([3.0], [1.0])
+    triple = openmc.stats.Discrete.merge([d1, d1, d1], [1.0, 2.0, 3.0])
+    assert triple.x == pytest.approx([3.0])
+    assert triple.p == pytest.approx([6.0])
+
+
 def test_uniform():
     a, b = 10.0, 20.0
     d = openmc.stats.Uniform(a, b)