from __future__ import annotations

from abc import ABC, abstractmethod

from collections import defaultdict

from collections.abc import Iterable, Sequence

from functools import cache

from math import sqrt, pi, exp, log

from numbers import Real

from pathlib import Path

from warnings import warn

import lxml.etree as ET

import numpy as np

from scipy.integrate import trapezoid

from scipy.special import exprel, hyp1f1, lambertw

import scipy

import openmc.checkvalue as cv

from openmc.data import atomic_mass, NEUTRON_MASS

import openmc.data

from .._xml import get_elem_list, get_text

from ..mixin import EqualityMixin

_INTERPOLATION_SCHEMES = {

'histogram',

'linear-linear',

'linear-log',

'log-linear',

'log-log'

}

def exprel2(x):

"""Evaluate 2*(exp(x)-1-x)/x^2 without loss of precision near 0"""

return hyp1f1(1, 3, x)

def log1prel(x):

"""Evaluate log(1+x)/x without loss of precision near 0"""

return np.where(np.abs(x) < 1e-16, 1.0, np.log1p(x) / x)

class Univariate(EqualityMixin, ABC):

"""Probability distribution of a single random variable.

The Univariate class is an abstract class that can be derived to implement a

specific probability distribution.

Parameters

----------

bias : Iterable of float, optional

Distribution or discrete probabilities for biased sampling or discrete

probabilities for biased sampling.

"""

def __init__(self, bias: Univariate | Sequence[float] | None = None):

self.bias = bias

@property

def bias(self):

return self._bias

@bias.setter

def bias(self, bias):

check_bias_support(self, bias)

self._bias = bias

def _append_bias_to_xml(self, element: ET.Element) -> None:

"""Append bias distribution element to XML if present."""

if self.bias is not None:

if self.bias.bias is not None:

raise RuntimeError('Biasing distributions should not have their own bias.')

bias_elem = self.bias.to_xml_element("bias")

element.append(bias_elem)

@classmethod

def _read_bias_from_xml(cls, elem: ET.Element):

"""Read bias distribution from XML element if present."""

bias_elem = elem.find('bias')

if bias_elem is not None:

return Univariate.from_xml_element(bias_elem)

return None

def _append_array_bias_to_xml(self, element: ET.Element) -> None:

"""Append array-based bias probabilities to XML."""

if self.bias is not None:

bias_elem = ET.SubElement(element, "bias")

bias_elem.text = ' '.join(map(str, self.bias))

@classmethod

def _read_array_bias_from_xml(cls, elem: ET.Element):

"""Read array-based bias probabilities from XML."""

bias_elem = elem.find('bias')

if bias_elem is not None:

return get_elem_list(elem, "bias", float)

return None

@abstractmethod

def to_xml_element(self, element_name):

return ''

@abstractmethod

def __len__(self):

return 0

@classmethod

@abstractmethod

def from_xml_element(cls, elem):

distribution = get_text(elem, 'type')

if distribution == 'discrete':

return Discrete.from_xml_element(elem)

elif distribution == 'uniform':

return Uniform.from_xml_element(elem)

elif distribution == 'powerlaw':

return PowerLaw.from_xml_element(elem)

elif distribution == 'maxwell':

return Maxwell.from_xml_element(elem)

elif distribution == 'watt':

return Watt.from_xml_element(elem)

elif distribution == 'normal':

return Normal.from_xml_element(elem)

elif distribution == 'muir':

# Support older files where Muir had its own class

return muir(*get_elem_list(elem, "parameters", float))

elif distribution == 'tabular':

return Tabular.from_xml_element(elem)

elif distribution == 'legendre':

return Legendre.from_xml_element(elem)

elif distribution == 'mixture':

return Mixture.from_xml_element(elem)

elif distribution == 'decay_spectrum':

return DecaySpectrum.from_xml_element(elem)

@abstractmethod

def _sample_unbiased(self, n_samples: int = 1, seed: int | None = None):

"""Sample without bias handling.

Parameters

----------

n_samples : int

Number of sampled values to generate

seed : int or None

Initial random number seed.

Returns

-------

numpy.ndarray

The array of sampled values

"""

pass

def sample(self, n_samples: int = 1, seed: int | None = None):

"""Sample the univariate distribution, handling biasing automatically.

Parameters

----------

n_samples : int

Number of sampled values to generate

seed : int or None

Initial random number seed.

Returns

-------

tuple of numpy.ndarray

A tuple of (samples, weights)

"""

if self.bias is None:

x = self._sample_unbiased(n_samples, seed)

return x, np.ones_like(x)

else:

if self.bias.bias is not None:

raise RuntimeError('Biasing distributions should not have their own bias.')

x, _ = self.bias.sample(n_samples=n_samples, seed=seed)

weight = self.evaluate(x) / self.bias.evaluate(x)

return x, weight

def integral(self):

"""Return integral of distribution

.. versionadded:: 0.13.1

Returns

-------

float

Integral of distribution

"""

return 1.0

@abstractmethod

def evaluate(self, x: float | Sequence[float]):

"""Evaluate the probability density at the provided value.

Parameters

----------

x : float or sequence of float

Location to evaluate p(x)

Returns

-------

float or numpy.ndarray

Value of p(x)

"""

pass

@property

@abstractmethod

def support(self):

"""Return the support of the probability distribution.

Returns

-------

set or tuple of float or dict

Returns the set of unique points assigned probability mass in a

discrete distribution, the sampling interval for a continuous

distribution, or a dictionary storing the discrete and continuous

parts of the support of a mixed random variable

"""

pass

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.

The Discrete distribution assigns probability values to discrete values of a

random variable, rather than expressing the distribution as a continuous

random variable.

Parameters

----------

x : Iterable of float

Values of the random variable

p : Iterable of float

Discrete probability for each value

bias : Iterable of float, optional

Alternative discrete probabilities for biased sampling. Defaults to

None for unbiased sampling.

Attributes

----------

x : numpy.ndarray

Values of the random variable

p : numpy.ndarray

Discrete probability for each value

support : set

Values of the random variable over which the distribution is

nonzero-valued

bias : numpy.ndarray or None

Discrete probabilities for biased sampling

"""

def __init__(self, x, p, bias=None):

self.x = x

self.p = p

super().__init__(bias)

def __len__(self):

return len(self.x)

@property

def x(self):

return self._x

@x.setter

def x(self, x):

if isinstance(x, Real):

x = [x]

cv.check_type('discrete values', x, Iterable, Real)

self._x = np.array(x, dtype=float)

@property

def p(self):

return self._p

@p.setter

def p(self, p):

if isinstance(p, Real):

p = [p]

cv.check_type('discrete probabilities', p, Iterable, Real)

for pk in p:

cv.check_greater_than('discrete probability', pk, 0.0, True)

self._p = np.array(p, dtype=float)

@property

def support(self):

return set(np.unique(self._x))

@Univariate.bias.setter

def bias(self, bias):

if bias is None:

self._bias = bias

else:

if isinstance(bias, Real):

bias = [bias]

cv.check_type('discrete bias probabilities', bias, Iterable, Real)

for bk in bias:

cv.check_greater_than('discrete probability', bk, 0.0, True)

if len(bias) != len(self.x):

raise RuntimeError("Discrete distribution has unequal number of "

"biased and unbiased probability entries.")

self._bias = np.array(bias, dtype=float)

def cdf(self):

return np.insert(np.cumsum(self.p), 0, 0.0)

def sample(self, n_samples=1, seed=None):

if self.bias is None:

samples = self._sample_unbiased(n_samples, seed)

return samples, np.ones_like(samples)

else:

rng = np.random.RandomState(seed)

p = self.p / self.p.sum()

b = self.bias / self.bias.sum()

indices = rng.choice(self.x.size, n_samples, p=b)

biased_sample = self.x[indices]

wgt = p[indices] / b[indices]

return biased_sample, wgt

def _sample_unbiased(self, n_samples=1, seed=None):

rng = np.random.RandomState(seed)

p = self.p / self.p.sum()

return rng.choice(self.x, n_samples, p=p)

def normalize(self):

"""Normalize the probabilities stored on the distribution"""

norm = sum(self.p)

self.p = [val / norm for val in self.p]

def evaluate(self, x):

raise NotImplementedError

def to_xml_element(self, element_name):

"""Return XML representation of the discrete distribution

Parameters

----------

element_name : str

XML element name

Returns

-------

element : lxml.etree._Element

XML element containing discrete distribution data

"""

element = ET.Element(element_name)

element.set("type", "discrete")

params = ET.SubElement(element, "parameters")

params.text = ' '.join(map(str, self.x)) + ' ' + ' '.join(map(str, self.p))

self._append_array_bias_to_xml(element)

return element

@classmethod

def from_xml_element(cls, elem: ET.Element):

"""Generate discrete distribution from an XML element

Parameters

----------

elem : lxml.etree._Element

XML element

Returns

-------

openmc.stats.Discrete

Discrete distribution generated from XML element

"""

params = get_elem_list(elem, "parameters", float)

x = params[:len(params)//2]

p = params[len(params)//2:]

bias_dist = cls._read_array_bias_from_xml(elem)

return cls(x, p, bias=bias_dist)

@classmethod

def merge(

cls,

dists: Sequence[Discrete],

probs: Sequence[float]

):

"""Merge multiple discrete distributions into a single distribution

.. versionadded:: 0.13.1

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.")

biasing = False

for d in dists:

if d.bias is not None:

# If we find that at least one distribution is biased, all

# distributions which are not biased will be assigned their

# default probability vector as a "bias" so that biased

# sampling can occur on the merged distribution.

biasing = True

break

# Combine distributions accounting for duplicate x values

x_merged = set()

p_merged = defaultdict(float)

new_bias = None

if biasing:

b_merged = defaultdict(float)

# Generate any missing bias distributions

dists = dists.copy()

for i, d in enumerate(dists):

if d.bias is None:

dists[i] = Discrete(d.x, d.p, bias=d.p)

for dist, p_dist in zip(dists, probs):

for x, p, b in zip(dist.x, dist.p, dist.bias):

x_merged.add(x)

p_merged[x] += p*p_dist

b_merged[x] += b*p_dist

# Create values and bias probabilities as arrays

x_arr = np.array(sorted(x_merged))

new_bias = np.array([b_merged[x] for x in x_arr])

else:

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 as array

x_arr = np.array(sorted(x_merged))

# Create probabilities as array

p_arr = np.array([p_merged[x] for x in x_arr])

return cls(x_arr, p_arr, new_bias)

def integral(self):

"""Return integral of distribution

.. versionadded:: 0.13.1

Returns

-------

float

Integral of discrete distribution

"""

return np.sum(self.p)

def mean(self) -> float:

"""Return mean of the discrete distribution

The mean is the weighted average of the discrete values.

.. versionadded:: 0.15.3

Returns

-------

float

Mean of discrete distribution

"""

return np.sum(self.x * self.p) / 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.

For biased distributions, clipping should be performed before the bias

probabilities are added.

.. versionadded:: 0.14.0

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

"""

if self.bias is not None:

raise RuntimeError("Biased discrete distributions should be clipped "

"before applying bias.")

cv.check_less_than("tolerance", tolerance, 1.0, equality=True)

cv.check_greater_than("tolerance", tolerance, 0.0, equality=True)

# Compute intensities

intensity = self.p * self.x

# Get indices for intensities above threshold

indices = _intensity_clip(intensity, tolerance=tolerance)

# Create new discrete distribution

if inplace:

self.x = self.x[indices]

self.p = self.p[indices]

return self

else:

new_x = self.x[indices]

new_p = self.p[indices]

return type(self)(new_x, new_p)

def delta_function(value: float, intensity: float = 1.0) -> Discrete:

"""Return a discrete distribution with a single point.

.. versionadded:: 0.15.1

Parameters

----------

value : float

Value of the random variable.

intensity : float, optional

When used for an energy distribution, this can be used to assign an

intensity.

Returns

-------

Discrete distribution with a single point

"""

return Discrete([value], [intensity])

class Uniform(Univariate):

"""Distribution with constant probability over a finite interval [a,b]

Parameters

----------

a : float, optional

Lower bound of the sampling interval. Defaults to zero.

b : float, optional

Upper bound of the sampling interval. Defaults to unity.

bias : openmc.stats.Univariate, optional

Distribution for biased sampling.

Attributes

----------

a : float

Lower bound of the sampling interval

b : float

Upper bound of the sampling interval

support : tuple of float

A 2-tuple (lower, upper) defining the interval over which the

distribution is nonzero-valued

bias : openmc.stats.Univariate or None

Distribution for biased sampling

"""

def __init__(self, a: float = 0.0, b: float = 1.0,

bias: Univariate | None = None):

self.a = a

self.b = b

super().__init__(bias)

def __len__(self):

return 2

@property

def a(self):

return self._a

@a.setter

def a(self, a):

cv.check_type('Uniform a', a, Real)

self._a = a

@property

def b(self):

return self._b

@b.setter

def b(self, b):

cv.check_type('Uniform b', b, Real)

self._b = b

@property

def support(self):

return (self._a, self._b)

def to_tabular(self):

if self.bias is not None:

raise RuntimeError("to_tabular() is not permitted for biased distributions.")

prob = 1./(self.b - self.a)

t = Tabular([self.a, self.b], [prob, prob], 'histogram')

t.c = [0., 1.]

return t

def _sample_unbiased(self, n_samples=1, seed=None):

rng = np.random.RandomState(seed)

return rng.uniform(self.a, self.b, n_samples)

def evaluate(self, x):

return np.where((self.a <= x) & (x <= self.b), 1/(self.b - self.a), 0.0)

def mean(self) -> float:

"""Return mean of the uniform distribution

.. versionadded:: 0.15.3

Returns

-------

float

Mean of uniform distribution

"""

return 0.5 * (self.a + self.b)

def to_xml_element(self, element_name: str):

"""Return XML representation of the uniform distribution

Parameters

----------

element_name : str

XML element name

Returns

-------

element : lxml.etree._Element

XML element containing uniform distribution data

"""

element = ET.Element(element_name)

element.set("type", "uniform")

element.set("parameters", f'{self.a} {self.b}')

self._append_bias_to_xml(element)

return element

@classmethod

def from_xml_element(cls, elem: ET.Element):

"""Generate uniform distribution from an XML element

Parameters

----------

elem : lxml.etree._Element

XML element

Returns

-------

openmc.stats.Uniform

Uniform distribution generated from XML element

"""

params = get_elem_list(elem, "parameters", float)

bias_dist = cls._read_bias_from_xml(elem)

return cls(*params, bias=bias_dist)

class PowerLaw(Univariate):

"""Distribution with power law probability over a finite interval [a,b]

The power law distribution has density function :math:`p(x) dx = c x^n dx`.

.. versionadded:: 0.13.0

Parameters

----------

a : float, optional

Lower bound of the sampling interval. Defaults to zero.

b : float, optional

Upper bound of the sampling interval. Defaults to unity.

n : float, optional

Power law exponent. Defaults to zero, which is equivalent to a uniform

distribution.

bias : openmc.stats.Univariate, optional

Distribution for biased sampling.

Attributes

----------

a : float

Lower bound of the sampling interval

b : float

Upper bound of the sampling interval

n : float

Power law exponent

support : tuple of float

A 2-tuple (lower, upper) defining the interval over which the

distribution is nonzero-valued

bias : openmc.stats.Univariate or None

Distribution for biased sampling

"""

def __init__(self, a: float = 0.0, b: float = 1.0, n: float = 0.,

bias: Univariate | None = None):

if a >= b:

raise ValueError(

"Lower bound of sampling interval must be less than upper bound.")

self.a = a

self.b = b

self.n = n

super().__init__(bias)

def __len__(self):

return 3

@property

def a(self):

return self._a

@a.setter

def a(self, a):

cv.check_type('interval lower bound', a, Real)

if a < 0:

raise ValueError(

"PowerLaw sampling is restricted to positive-valued intervals.")

self._a = a

@property

def b(self):

return self._b

@b.setter

def b(self, b):

cv.check_type('interval upper bound', b, Real)

if b < 0:

raise ValueError(

"PowerLaw sampling is restricted to positive-valued intervals.")

self._b = b

@property

def n(self):

return self._n

@n.setter

def n(self, n):

cv.check_type('power law exponent', n, Real)

self._n = n

@property

def support(self):

return (self._a, self._b)

def _sample_unbiased(self, n_samples=1, seed=None):

rng = np.random.RandomState(seed)

xi = rng.random(n_samples)

pwr = self.n + 1

offset = self.a**pwr

span = self.b**pwr - offset

return np.power(offset + xi * span, 1/pwr)

def evaluate(self, x):

c = (self.n + 1)/(self.b**(self.n + 1) - self.a**(self.n + 1))

return np.where((self.a <= x) & (x <= self.b), c * np.abs(x)**self.n, 0.0)

def to_xml_element(self, element_name: str):

"""Return XML representation of the power law distribution

Parameters

----------

element_name : str

XML element name

Returns

-------

element : lxml.etree._Element

XML element containing distribution data

"""

element = ET.Element(element_name)

element.set("type", "powerlaw")

element.set("parameters", f'{self.a} {self.b} {self.n}')

self._append_bias_to_xml(element)

return element

@classmethod

def from_xml_element(cls, elem: ET.Element):

"""Generate power law distribution from an XML element

Parameters

----------

elem : lxml.etree._Element

XML element

Returns

-------

openmc.stats.PowerLaw

Distribution generated from XML element

"""

params = get_elem_list(elem, "parameters", float)

bias_dist = cls._read_bias_from_xml(elem)

return cls(*map(float, params), bias=bias_dist)

class Maxwell(Univariate):

r"""Maxwellian distribution in energy.

The Maxwellian distribution in energy is characterized by a single parameter

:math:`\theta` and has a density function :math:`p(E) dE = c \sqrt{E}

e^{-E/\theta} dE`.

Parameters

----------

theta : float

Effective temperature for distribution in eV

bias : openmc.stats.Univariate, optional

Distribution for biased sampling.

Attributes

----------

theta : float

Effective temperature for distribution in eV

support : tuple of float

A 2-tuple (lower, upper) defining the interval over which the

distribution is nonzero-valued

bias : openmc.stats.Univariate or None

Distribution for biased sampling

"""

def __init__(self, theta, bias: Univariate | None = None):

self.theta = theta

super().__init__(bias)

def __len__(self):

return 1

@property

def theta(self):

return self._theta

@theta.setter

def theta(self, theta):

cv.check_type('Maxwell temperature', theta, Real)

cv.check_greater_than('Maxwell temperature', theta, 0.0)

self._theta = theta

@property

def support(self):

return (0.0, np.inf)

def _sample_unbiased(self, n_samples=1, seed=None):

rng = np.random.RandomState(seed)

return self.sample_maxwell(self.theta, n_samples, rng=rng)

@staticmethod

def sample_maxwell(t, n_samples: int, rng=None):

if rng is None:

rng = np.random.default_rng()

return rng.gamma(1.5, t, n_samples)

def evaluate(self, E):

return scipy.stats.gamma.pdf(E, 1.5, scale=self.theta)

def to_xml_element(self, element_name: str):

"""Return XML representation of the Maxwellian distribution

Parameters

----------

element_name : str

XML element name

Returns

-------

element : lxml.etree._Element

XML element containing Maxwellian distribution data

"""

element = ET.Element(element_name)

element.set("type", "maxwell")

element.set("parameters", str(self.theta))

self._append_bias_to_xml(element)

return element

@classmethod

def from_xml_element(cls, elem: ET.Element):

"""Generate Maxwellian distribution from an XML element

Parameters

----------

elem : lxml.etree._Element

XML element

Returns

-------

openmc.stats.Maxwell

Maxwellian distribution generated from XML element

"""

theta = float(get_text(elem, 'parameters'))

bias_dist = cls._read_bias_from_xml(elem)

return cls(theta, bias=bias_dist)

class Watt(Univariate):

r"""Watt fission energy spectrum.

The Watt fission energy spectrum is characterized by two parameters

:math:`a` and :math:`b` and has density function :math:`p(E) dE = c e^{-E/a}

\sinh \sqrt{b \, E} dE`.

Parameters

----------

a : float

First parameter of distribution in units of eV

b : float

Second parameter of distribution in units of 1/eV

bias : openmc.stats.Univariate, optional

Distribution for biased sampling.

Attributes

----------

a : float

First parameter of distribution in units of eV

b : float

Second parameter of distribution in units of 1/eV

support : tuple of float

A 2-tuple (lower, upper) defining the interval over which the

distribution is nonzero-valued

bias : openmc.stats.Univariate or None

Distribution for biased sampling

"""

def __init__(self, a=0.988e6, b=2.249e-6, bias: Univariate | None = None):

self.a = a

self.b = b

super().__init__(bias)

def __len__(self):

return 2

@property

def a(self):

return self._a

@a.setter

def a(self, a):

cv.check_type('Watt a', a, Real)

cv.check_greater_than('Watt a', a, 0.0)

self._a = a

@property

def b(self):

return self._b

@b.setter

def b(self, b):

cv.check_type('Watt b', b, Real)

cv.check_greater_than('Watt b', b, 0.0)

self._b = b