diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py index e0475bf78fb..556cc69e040 100644 --- a/openmc/stats/univariate.py +++ b/openmc/stats/univariate.py @@ -249,7 +249,7 @@ def from_xml_element(cls, elem: ET.Element): def merge( cls, dists: Sequence[Discrete], - probs: Sequence[int] + probs: Sequence[float] ): """Merge multiple discrete distributions into a single distribution @@ -830,6 +830,36 @@ def from_xml_element(cls, elem: ET.Element): params = get_text(elem, 'parameters').split() return cls(*map(float, params)) + @classmethod + def merge( + cls, + dists: Sequence[Normal], + probs: Sequence[float] + ): + """Merge multiple normal distributions into a single distribution + + Parameters + ---------- + dists : iterable of openmc.stats.Normal + Normal distributions to combine + probs : iterable of float + Probability of each distribution + + Returns + ------- + openmc.stats.Normal + Combined normal distribution + + """ + if len(dists) != len(probs): + raise ValueError("Number of distributions and probabilities must match.") + means = np.array([d.mean_value for d in dists]) + stds = np.array([d.std_dev for d in dists]) + probs = np.asarray(probs) + combined_mean = np.dot(probs,means) + combined_std = np.sqrt(np.dot(probs**2,stds**2)) + return cls(combined_mean, combined_std) + def muir(e0: float, m_rat: float, kt: float): """Generate a Muir energy spectrum @@ -1324,6 +1354,59 @@ def from_xml_element(cls, elem: ET.Element): return cls(probability, distribution) + @classmethod + def merge( + cls, + dists: Sequence[Univariate], + probs: Sequence[float] + ): + """Merge multiple distributions into a single distribution + + Parameters + ---------- + dists : iterable of openmc.stats.Univariate + Univariate distributions to combine + probs : iterable of float + Probability of each distribution + + Returns + ------- + openmc.stats.Univariate + Combined distribution + + """ + if len(dists) != len(probs): + raise ValueError("Number of distributions and probabilities must match.") + + discrete = [[],[]] + normal = [[],[]] + others = [[],[]] + for p,d in zip(probs,dists): + if isinstance(d, Mixture): + iterator = zip(p * m.probability, m.distribution) + else: + iterator = [(p,d)] + for prob,dist in iterator: + if isinstance(dist,Discrete): + discrete[0].append(dist) + discrete[1].append(prob) + elif isinstance(dist,Normal): + normal[0].append(dist) + normal[1].append(prob) + else: + others[0].append(dist) + others[1].append(prob) + if discrete[0]: + others[0].append(Discrete.merge(*discrete)) + others[1].append(1.0) + if normal[0]: + others[0].append(Normal.merge(*normal)) + others[1].append(1.0) + if others[1]==[1.0]: + return others[0][0] + else: + return cls(others[1], others[0]) + def integral(self): """Return integral of the distribution @@ -1406,8 +1489,8 @@ def combine_distributions( """Combine distributions with specified probabilities This function can be used to combine multiple instances of - :class:`~openmc.stats.Discrete` and `~openmc.stats.Tabular`. Multiple - discrete distributions are merged into a single distribution and the + :class:`~openmc.stats.Univariate`. Multiple + discrete or normal distributions are merged into a single distribution and the remainder of the distributions are put into a :class:`~openmc.stats.Mixture` distribution. @@ -1421,32 +1504,4 @@ def combine_distributions( Probability (or intensity) of each distribution """ - # Get copy of distribution list so as not to modify the argument - dist_list = deepcopy(dists) - - # Get list of discrete/continuous distribution indices - discrete_index = [i for i, d in enumerate(dist_list) if isinstance(d, Discrete)] - cont_index = [i for i, d in enumerate(dist_list) if isinstance(d, Tabular)] - - # Apply probabilites to continuous distributions - for i in cont_index: - dist = dist_list[i] - dist._p *= probs[i] - - if discrete_index: - # Create combined discrete distribution - dist_discrete = [dist_list[i] for i in discrete_index] - discrete_probs = [probs[i] for i in discrete_index] - combined_dist = Discrete.merge(dist_discrete, discrete_probs) - - # Replace multiple discrete distributions with merged - for idx in reversed(discrete_index): - dist_list.pop(idx) - dist_list.append(combined_dist) - - # Combine discrete and continuous if present - if len(dist_list) > 1: - probs = [1.0]*len(dist_list) - dist_list[:] = [Mixture(probs, dist_list.copy())] - - return dist_list[0] + return Mixture.merge(dists, probs) diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py index 386181f34d2..9bb2a0955f6 100644 --- a/tests/unit_tests/test_stats.py +++ b/tests/unit_tests/test_stats.py @@ -444,6 +444,15 @@ def test_normal(): assert_sample_mean(samples, mean) +def test_combine_normal(): + norm1=openmc.stats.Normal(mean_value=10, std_dev=1) + norm2=openmc.stats.Normal(mean_value=1, std_dev=1) + combined=openmc.stats.combine_distributions([norm1,norm2],probs=[0.1,0.9]) + assert isinstance(combined, openmc.stats.Normal) + assert combined.mean_value == pytest.approx(1.9) + assert combined.std_dev**2 == pytest.approx(0.82) + + def test_muir(): mean = 10.0 mass = 5.0