From ef3019a140d7e9b511f7e09bfb029c2b1e1e4015 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Tue, 5 Apr 2022 12:49:48 -0500
Subject: [PATCH] Add merge classmethod on openmc.stats.Discrete

---
 openmc/stats/univariate.py     | 31 +++++++++++++++++++++++++++++++
 tests/unit_tests/test_stats.py | 23 +++++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py
index f8e50044fc..53b19713bd 100644
--- a/openmc/stats/univariate.py
+++ b/openmc/stats/univariate.py
@@ -1,5 +1,7 @@
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from collections.abc import Iterable
+from functools import reduce
 from numbers import Real
 from xml.etree import ElementTree as ET
 
@@ -156,6 +158,35 @@ 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
+
+        """
+        # 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]
diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py
index bbcb12ff11..51a561bd36 100644
--- a/tests/unit_tests/test_stats.py
+++ 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)