Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve clipping of Mixture distributions #3154

Merged
merged 8 commits into from
Oct 11, 2024
Merged
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion openmc/material.py
Original file line number Diff line number Diff line change
Expand Up @@ -320,7 +320,7 @@ def get_decay_photon_energy(
probs = []
for nuc, atoms_per_bcm in self.get_nuclide_atom_densities().items():
source_per_atom = openmc.data.decay_photon_energy(nuc)
if source_per_atom is not None:
if source_per_atom is not None and atoms_per_bcm > 0.0:
dists.append(source_per_atom)
probs.append(1e24 * atoms_per_bcm * multiplier)

Expand All @@ -333,6 +333,11 @@ def get_decay_photon_energy(
if isinstance(combined, (Discrete, Mixture)):
combined.clip(clip_tolerance, inplace=True)

# If clipping resulted in a single distribution within a mixture, pick
# out that single distribution
if isinstance(combined, Mixture) and len(combined.distribution) == 1:
combined = combined.distribution[0]

return combined

@classmethod
Expand Down
105 changes: 74 additions & 31 deletions openmc/stats/univariate.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,45 @@ def integral(self):
return 1.0


def _intensity_clip(intensity: Sequence[float], tolerance: float = 1e-6) -> np.ndarray:
"""Clip low-importance points from an array of intensities.

Given an array of intensities, this function returns an array of indices for
points that contribute non-negligibly to the total sum of intensities.

Parameters
----------
intensity : sequence of float
Intensities in arbitrary units.
tolerance : float
Maximum fraction of intensities that will be discarded.

Returns
-------
Array of indices

"""
# Get indices of intensities from largest to smallest
index_sort = np.argsort(intensity)[::-1]

# Get intensities from largest to smallest
sorted_intensity = np.asarray(intensity)[index_sort]

# Determine cumulative sum of probabilities
cumsum = np.cumsum(sorted_intensity)
cumsum /= cumsum[-1]

# Find index that satisfies cutoff
index_cutoff = np.searchsorted(cumsum, 1.0 - tolerance)

# Now get indices up to cutoff
new_indices = index_sort[:index_cutoff + 1]

# Put back in the order of the original array and return
new_indices.sort()
return new_indices


class Discrete(Univariate):
"""Distribution characterized by a probability mass function.

Expand Down Expand Up @@ -283,32 +322,20 @@ def clip(self, tolerance: float = 1e-6, inplace: bool = False) -> Discrete:
cv.check_less_than("tolerance", tolerance, 1.0, equality=True)
cv.check_greater_than("tolerance", tolerance, 0.0, equality=True)

# Determine (reversed) sorted order of probabilities
# Compute intensities
intensity = self.p * self.x
index_sort = np.argsort(intensity)[::-1]

# Get probabilities in above order
sorted_intensity = intensity[index_sort]

# Determine cumulative sum of probabilities
cumsum = np.cumsum(sorted_intensity)
cumsum /= cumsum[-1]

# Find index which satisfies cutoff
index_cutoff = np.searchsorted(cumsum, 1.0 - tolerance)

# Now get indices up to cutoff
new_indices = index_sort[:index_cutoff + 1]
new_indices.sort()
# Get indices for intensities above threshold
indices = _intensity_clip(intensity, tolerance=tolerance)

# Create new discrete distribution
if inplace:
self.x = self.x[new_indices]
self.p = self.p[new_indices]
self.x = self.x[indices]
self.p = self.p[indices]
return self
else:
new_x = self.x[new_indices]
new_p = self.p[new_indices]
new_x = self.x[indices]
new_p = self.p[indices]
return type(self)(new_x, new_p)


Expand Down Expand Up @@ -1206,7 +1233,7 @@ def probability(self, probability):
for p in probability:
cv.check_greater_than('mixture distribution probabilities',
p, 0.0, True)
self._probability = probability
self._probability = np.array(probability, dtype=float)

@property
def distribution(self):
Expand Down Expand Up @@ -1312,40 +1339,56 @@ def integral(self):
])

def clip(self, tolerance: float = 1e-6, inplace: bool = False) -> Mixture:
r"""Remove low-importance points from contained discrete distributions.
r"""Remove low-importance points / distributions

Given a probability mass function :math:`p(x)` with :math:`\{x_1, x_2,
x_3, \dots\}` the possible values of the random variable with
corresponding probabilities :math:`\{p_1, p_2, p_3, \dots\}`, this
function will remove any low-importance points such that :math:`\sum_i
x_i p_i` is preserved to within some threshold.
Like :meth:`Discrete.clip`, this method will remove low-importance
points from discrete distributions contained within the mixture but it
will also clip any distributions that have negligible contributions to
the overall intensity.

.. versionadded:: 0.14.0

Parameters
----------
tolerance : float
Maximum fraction of :math:`\sum_i x_i p_i` that will be discarded
for any discrete distributions within the mixture distribution.
Maximum fraction of intensities that will be discarded.
inplace : bool
Whether to modify the current object in-place or return a new one.

Returns
-------
Discrete distribution with low-importance points removed
Distribution with low-importance points / distributions removed

"""
# Determine indices for any distributions that contribute non-negligibly
# to overall intensity
intensities = [prob*dist.integral() for prob, dist in
zip(self.probability, self.distribution)]
indices = _intensity_clip(intensities, tolerance=tolerance)

if inplace:
# Clip mixture of distributions
self.probability = self.probability[indices]
self.distribution = [self.distribution[i] for i in indices]

# Clip points from Discrete distributions
paulromano marked this conversation as resolved.
Show resolved Hide resolved
for dist in self.distribution:
if isinstance(dist, Discrete):
dist.clip(tolerance, inplace=True)

return self
else:
# Clip mixture of distributions
probability = self.probability[indices]
distribution = [self.distribution[i] for i in indices]

# Clip points from Discrete distributions
distribution = [
dist.clip(tolerance) if isinstance(dist, Discrete) else dist
for dist in self.distribution
for dist in distribution
]
paulromano marked this conversation as resolved.
Show resolved Hide resolved
return type(self)(self.probability, distribution)

return type(self)(probability, distribution)


def combine_distributions(
Expand Down
3 changes: 3 additions & 0 deletions src/distribution.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -394,6 +394,9 @@ Mixture::Mixture(pugi::xml_node node)
distribution_.push_back(std::make_pair(cumsum, std::move(dist)));
}

// Save integral of distribution
integral_ = cumsum;

// Normalize cummulative probabilities to 1
for (auto& pair : distribution_) {
pair.first /= cumsum;
Expand Down
18 changes: 9 additions & 9 deletions tests/regression_tests/source/inputs_true.dat
Original file line number Diff line number Diff line change
Expand Up @@ -85,13 +85,13 @@
</space>
<angle type="isotropic"/>
<energy type="mixture">
<pair probability="1">
<pair probability="1.0">
<dist parameters="1289500.0" type="maxwell"/>
</pair>
<pair probability="2">
<pair probability="2.0">
<dist parameters="988000.0 2.249e-06" type="watt"/>
</pair>
<pair probability="3">
<pair probability="3.0">
<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>
Expand All @@ -108,13 +108,13 @@
</space>
<angle type="isotropic"/>
<energy type="mixture">
<pair probability="1">
<pair probability="1.0">
<dist parameters="1289500.0" type="maxwell"/>
</pair>
<pair probability="2">
<pair probability="2.0">
<dist parameters="988000.0 2.249e-06" type="watt"/>
</pair>
<pair probability="3">
<pair probability="3.0">
<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>
Expand All @@ -132,13 +132,13 @@
</space>
<angle type="isotropic"/>
<energy type="mixture">
<pair probability="1">
<pair probability="1.0">
<dist parameters="1289500.0" type="maxwell"/>
</pair>
<pair probability="2">
<pair probability="2.0">
<dist parameters="988000.0 2.249e-06" type="watt"/>
</pair>
<pair probability="3">
<pair probability="3.0">
<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>
Expand Down
10 changes: 8 additions & 2 deletions tests/unit_tests/test_stats.py
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ def test_mixture():
d2 = openmc.stats.Uniform(3, 7)
p = [0.5, 0.5]
mix = openmc.stats.Mixture(p, [d1, d2])
assert mix.probability == p
np.testing.assert_allclose(mix.probability, p)
assert mix.distribution == [d1, d2]
assert len(mix) == 4

Expand All @@ -274,7 +274,7 @@ def test_mixture():
elem = mix.to_xml_element('distribution')

d = openmc.stats.Mixture.from_xml_element(elem)
assert d.probability == p
np.testing.assert_allclose(d.probability, p)
assert d.distribution == [d1, d2]
assert len(d) == 4

Expand All @@ -296,6 +296,12 @@ def test_mixture_clip():
mix_same = mix.clip(1e-6, inplace=True)
assert mix_same is mix

# Make sure clip removes low probability distributions
d_small = openmc.stats.Uniform(0., 1.)
d_large = openmc.stats.Uniform(2., 5.)
mix = openmc.stats.Mixture([1e-10, 1.0], [d_small, d_large])
mix_clip = mix.clip(1e-3)
assert mix_clip.distribution == [d_large]
Comment on lines +303 to +304
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

only question here is if inplace=False by default should we expect to get back the same Univariate objects in the resulting clipped Mixture? My initial reaction was that I would expect the underlying distributions to be different objects. I don't think any users will really care about this, but it could have side-effects if someone then modified d_large after mixing, not expecting it to effect the Mixture distribution, because there is a specific inplace kwarg to that effect in clip.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's an interesting thought. We could default to copying any distribution that is not a discrete distribution when inplace=False, although that could be a bit wasteful from a memory perspective. I don't have a strong feeling either way since I don't think this will be a common issue, so let me know which way you prefer.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

me neither we can leave as is


def test_polar_azimuthal():
# default polar-azimuthal should be uniform in mu and phi
Expand Down