diff --git a/openmc/model/surface_composite.py b/openmc/model/surface_composite.py index 74f64c54ac..6271154fb8 100644 --- a/openmc/model/surface_composite.py +++ b/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, diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py index f8e50044fc..e4a9b8f924 100644 --- a/openmc/stats/univariate.py +++ b/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] 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)