Clip low-importance points from discrete distributions (#2715)

Co-authored-by: Ethan Peterson <ethan.peterson@mit.edu>

openmc/data/decay.py

@@ -599,7 +599,7 @@ def decay_photon_energy(nuclide: str) -> Optional[Univariate]:
     openmc.stats.Univariate or None
         Distribution of energies in [eV] of photons emitted from decay, or None
         if no photon source exists. Note that the probabilities represent
-        intensities, given as [decay/sec].
+        intensities, given as [Bq].
     """
     if not _DECAY_PHOTON_ENERGY:
         chain_file = openmc.config.get('chain_file')

openmc/geometry.py

@@ -764,8 +764,9 @@ class Geometry:
             Assigns colors to specific materials or cells. Keys are instances of
             :class:`Cell` or :class:`Material` and values are RGB 3-tuples, RGBA
             4-tuples, or strings indicating SVG color names. Red, green, blue,
-            and alpha should all be floats in the range [0.0, 1.0], for example:
-            .. code-block:: python
+            and alpha should all be floats in the range [0.0, 1.0], for
+            example::
+
                # Make water blue
                water = openmc.Cell(fill=h2o)
                universe.plot(..., colors={water: (0., 0., 1.))

openmc/material.py

@@ -19,7 +19,7 @@ import openmc.checkvalue as cv
 from ._xml import clean_indentation, reorder_attributes
 from .mixin import IDManagerMixin
 from openmc.checkvalue import PathLike
-from openmc.stats import Univariate
+from openmc.stats import Univariate, Discrete, Mixture
 
 
 # Units for density supported by OpenMC
@@ -93,13 +93,6 @@ class Material(IDManagerMixin):
     fissionable_mass : float
         Mass of fissionable nuclides in the material in [g]. Requires that the
         :attr:`volume` attribute is set.
-    decay_photon_energy : openmc.stats.Univariate or None
-        Energy distribution of photons emitted from decay of unstable nuclides
-        within the material, or None if no photon source exists. The integral of
-        this distribution is the total intensity of the photon source in
-        [decay/sec].
-
-        .. versionadded:: 0.13.2
     ncrystal_cfg : str
         NCrystal configuration string
 
@@ -282,15 +275,67 @@ class Material(IDManagerMixin):
 
     @property
     def decay_photon_energy(self) -> Optional[Univariate]:
-        atoms = self.get_nuclide_atoms()
+        warnings.warn(
+            "The 'decay_photon_energy' property has been replaced by the "
+            "get_decay_photon_energy() method and will be removed in a future "
+            "version.", FutureWarning)
+        return self.get_decay_photon_energy(0.0)
+
+    def get_decay_photon_energy(
+            self,
+            clip_tolerance: float = 1e-6,
+            units: str = 'Bq',
+            volume: Optional[float] = None
+        ) -> Optional[Univariate]:
+        r"""Return energy distribution of decay photons from unstable nuclides.
+
+        .. versionadded:: 0.13.4
+
+        Parameters
+        ----------
+        clip_tolerance : float
+            Maximum fraction of :math:`\sum_i x_i p_i` for discrete
+            distributions that will be discarded.
+        units : {'Bq', 'Bq/g', 'Bq/cm3'}
+            Specifies the units on the integral of the distribution.
+        volume : float, optional
+            Volume of the material. If not passed, defaults to using the
+            :attr:`Material.volume` attribute.
+
+        Returns
+        -------
+        Decay photon energy distribution. The integral of this distribution is
+        the total intensity of the photon source in the requested units.
+
+        """
+        cv.check_value('units', units, {'Bq', 'Bq/g', 'Bq/cm3'})
+        if units == 'Bq':
+            multiplier = volume if volume is not None else self.volume
+            if multiplier is None:
+                raise ValueError("volume must be specified if units='Bq'")
+        elif units == 'Bq/cm3':
+            multiplier = 1
+        elif units == 'Bq/g':
+            multiplier = 1.0 / self.get_mass_density()
+
         dists = []
         probs = []
-        for nuc, num_atoms in atoms.items():
+        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:
                 dists.append(source_per_atom)
-                probs.append(num_atoms)
-        return openmc.data.combine_distributions(dists, probs) if dists else None
+                probs.append(1e24 * atoms_per_bcm * multiplier)
+
+        # If no photon sources, exit early
+        if not dists:
+            return None
+
+        # Get combined distribution, clip low-intensity values in discrete spectra
+        combined = openmc.data.combine_distributions(dists, probs)
+        if isinstance(combined, (Discrete, Mixture)):
+            combined.clip(clip_tolerance, inplace=True)
+
+        return combined
 
     @classmethod
     def from_hdf5(cls, group: h5py.Group) -> Material:

openmc/model/model.py

@@ -609,6 +609,7 @@ class Model:
         this method creates the XML files and runs OpenMC via a system call. In
         both cases this method returns the path to the last statepoint file
         generated.
+
         .. versionchanged:: 0.12
             Instead of returning the final k-effective value, this function now
             returns the path to the final statepoint written.

openmc/stats/univariate.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import math
 import typing
 from abc import ABC, abstractmethod
@@ -208,7 +209,7 @@ class Discrete(Univariate):
     @classmethod
     def merge(
         cls,
-        dists: typing.Sequence['openmc.stats.Discrete'],
+        dists: typing.Sequence[Discrete],
         probs: typing.Sequence[int]
     ):
         """Merge multiple discrete distributions into a single distribution
@@ -256,6 +257,58 @@ class Discrete(Univariate):
         """
         return np.sum(self.p)
 
+    def clip(self, tolerance: float = 1e-6, inplace: bool = False) -> Discrete:
+        r"""Remove low-importance points from discrete distribution.
+
+        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.
+
+        .. versionadded:: 0.13.4
+
+        Parameters
+        ----------
+        tolerance : float
+            Maximum fraction of :math:`\sum_i x_i p_i` 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
+
+        """
+        # Determine (reversed) sorted order of probabilities
+        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()
+
+        # Create new discrete distribution
+        if inplace:
+            self.x = self.x[new_indices]
+            self.p = self.p[new_indices]
+            return self
+        else:
+            new_x = self.x[new_indices]
+            new_p = self.p[new_indices]
+            return type(self)(new_x, new_p)
+
+
 class Uniform(Univariate):
     """Distribution with constant probability over a finite interval [a,b]
 
@@ -1094,7 +1147,7 @@ class Mixture(Univariate):
     def __init__(
         self,
         probability: typing.Sequence[float],
-        distribution: typing.Sequence['openmc.Univariate']
+        distribution: typing.Sequence[Univariate]
     ):
         self.probability = probability
         self.distribution = distribution
@@ -1219,9 +1272,45 @@ class Mixture(Univariate):
             for p, dist in zip(self.probability, self.distribution)
         ])
 
+    def clip(self, tolerance: float = 1e-6, inplace: bool = False) -> Mixture:
+        r"""Remove low-importance points from contained discrete 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.
+
+        .. versionadded:: 0.13.4
+
+        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.
+        inplace : bool
+            Whether to modify the current object in-place or return a new one.
+
+        Returns
+        -------
+        Discrete distribution with low-importance points removed
+
+        """
+        if inplace:
+            for dist in self.distribution:
+                if isinstance(dist, Discrete):
+                    dist.clip(tolerance, inplace=True)
+            return self
+        else:
+            distribution = [
+                dist.clip(tolerance) if isinstance(dist, Discrete) else dist
+                for dist in self.distribution
+            ]
+            return type(self)(self.probability, distribution)
+
 
 def combine_distributions(
-    dists: typing.Sequence['openmc.Univariate'],
+    dists: typing.Sequence[Univariate],
     probs: typing.Sequence[float]
 ):
     """Combine distributions with specified probabilities

tests/unit_tests/test_material.py

@@ -626,13 +626,19 @@ def test_decay_photon_energy():
     m.volume = 1.0
 
     # Get decay photon source and make sure it's the right type
-    src = m.decay_photon_energy
+    src = m.get_decay_photon_energy()
     assert isinstance(src, openmc.stats.Discrete)
 
+    # Make sure units/volume work as expected
+    src_v2 = m.get_decay_photon_energy(volume=2.0)
+    assert src.p * 2.0 == pytest.approx(src_v2.p)
+    src_per_cm3 = m.get_decay_photon_energy(units='Bq/cm3', volume=100.0)
+    assert (src.p == src_per_cm3.p).all()
+
     # If we add Xe135 (which has a tabular distribution), the photon source
     # should be a mixture distribution
     m.add_nuclide('Xe135', 1.0e-24)
-    src = m.decay_photon_energy
+    src = m.get_decay_photon_energy()
     assert isinstance(src, openmc.stats.Mixture)
 
     # With a single atom of each, the intensity of the photon source should be
@@ -640,6 +646,12 @@ def test_decay_photon_energy():
     def intensity(src):
         return src.integral() if src is not None else 0.0
 
+    assert src.integral() == pytest.approx(sum(
+        intensity(decay_photon_energy(nuc)) for nuc in m.get_nuclides()
+    ), rel=1e-3)
+
+    # When the clipping threshold is zero, the intensities should match exactly
+    src = m.get_decay_photon_energy(0.0)
     assert src.integral() == pytest.approx(sum(
         intensity(decay_photon_energy(nuc)) for nuc in m.get_nuclides()
     ))
@@ -648,4 +660,4 @@ def test_decay_photon_energy():
     stable = openmc.Material()
     stable.add_nuclide('Gd156', 1.0)
     stable.volume = 1.0
-    assert stable.decay_photon_energy is None
+    assert stable.get_decay_photon_energy() is None

tests/unit_tests/test_stats.py

@@ -74,6 +74,22 @@ def test_merge_discrete():
     assert triple.integral() == pytest.approx(6.0)
 
 
+def test_clip_discrete():
+    # Create discrete distribution with two points that are not important, one
+    # because the x value is very small, and one because the p value is very
+    # small
+    d = openmc.stats.Discrete([1e-8, 1.0, 2.0, 1000.0], [3.0, 2.0, 5.0, 1e-12])
+
+    # Clipping the distribution should result in two points
+    d_clip = d.clip(1e-6)
+    assert d_clip.x.size == 2
+    assert d_clip.p.size == 2
+
+    # Make sure inplace returns same object
+    d_same = d.clip(1e-6, inplace=True)
+    assert d_same is d
+
+
 def test_uniform():
     a, b = 10.0, 20.0
     d = openmc.stats.Uniform(a, b)
@@ -232,6 +248,24 @@ def test_mixture():
     assert len(d) == 4
 
 
+def test_mixture_clip():
+    # Create mixture distribution containing a discrete distribution with two
+    # points that are not important, one because the x value is very small, and
+    # one because the p value is very small
+    d1 = openmc.stats.Discrete([1e-8, 1.0, 2.0, 1000.0], [3.0, 2.0, 5.0, 1e-12])
+    d2 = openmc.stats.Uniform(0, 5)
+    mix = openmc.stats.Mixture([0.5, 0.5], [d1, d2])
+
+    # Clipping should reduce the contained discrete distribution to 2 points
+    mix_clip = mix.clip(1e-6)
+    assert mix_clip.distribution[0].x.size == 2
+    assert mix_clip.distribution[0].p.size == 2
+
+    # Make sure inplace returns same object
+    mix_same = mix.clip(1e-6, inplace=True)
+    assert mix_same is mix
+
+
 def test_polar_azimuthal():
     # default polar-azimuthal should be uniform in mu and phi
     d = openmc.stats.PolarAzimuthal()