Add truncated normal distribution support (#3761)

This commit is contained in:
Paul Romano 2026-02-12 10:25:56 -06:00 committed by GitHub
parent a3426cf833
commit 8198e6021d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 385 additions and 28 deletions

View file

@ -1,5 +1,5 @@
#include <cmath> // for M_PI
#include <memory> // for unique_ptr
#include <cmath>
#include <memory>
#include <unordered_map>
#include "openmc/particle.h"

View file

@ -265,23 +265,29 @@ private:
};
//==============================================================================
//! Normal distributions with form 1/2*std_dev*sqrt(pi) exp
//! (-(e-E0)/2*std_dev)^2
//! Normal distribution with optional truncation bounds.
//!
//! The standard normal PDF is 1/(sqrt(2*pi)*sigma) *
//! exp(-(x-mu)^2/(2*sigma^2)). When truncated to [lower, upper], the PDF is
//! renormalized so that it integrates to 1 over the truncation interval.
//==============================================================================
class Normal : public Distribution {
public:
explicit Normal(pugi::xml_node node);
Normal(double mean_value, double std_dev)
: mean_value_ {mean_value}, std_dev_ {std_dev} {};
Normal(double mean_value, double std_dev, double lower = -INFTY,
double upper = INFTY);
//! Evaluate probability density, f(x), at a point
//! \param x Point to evaluate f(x)
//! \return f(x)
//! \return f(x), accounting for truncation normalization
double evaluate(double x) const override;
double mean_value() const { return mean_value_; }
double std_dev() const { return std_dev_; }
double lower() const { return lower_; }
double upper() const { return upper_; }
bool is_truncated() const { return is_truncated_; }
protected:
//! Sample a value (unbiased) from the distribution
@ -290,8 +296,15 @@ protected:
double sample_unbiased(uint64_t* seed) const override;
private:
double mean_value_; //!< middle of distribution [eV]
double std_dev_; //!< standard deviation [eV]
double mean_value_; //!< Mean of distribution
double std_dev_; //!< Standard deviation
double lower_; //!< Lower truncation bound (default: -INFTY)
double upper_; //!< Upper truncation bound (default: +INFTY)
bool is_truncated_; //!< True if bounds are finite
double norm_factor_; //!< Normalization factor for truncated distribution
//! Compute normalization factor for truncated distribution
void compute_normalization();
};
//==============================================================================

View file

@ -223,5 +223,15 @@ double log1prel(double x);
void get_energy_index(
const vector<double>& energies, double E, int& i, double& f);
//==============================================================================
//! Calculate the cumulative distribution function of the standard normal
//! distribution at a given value.
//!
//! \param z The value at which to evaluate the CDF
//! \return Phi(z) = P(X <= z) for X ~ N(0,1)
//==============================================================================
double standard_normal_cdf(double z);
} // namespace openmc
#endif // OPENMC_MATH_FUNCTIONS_H

View file

@ -1049,18 +1049,27 @@ class Watt(Univariate):
class Normal(Univariate):
r"""Normally distributed sampling.
r"""Normally distributed sampling with optional truncation.
The Normal Distribution is characterized by two parameters
:math:`\mu` and :math:`\sigma` and has density function
:math:`p(X) dX = 1/(\sqrt{2\pi}\sigma) e^{-(X-\mu)^2/(2\sigma^2)}`
The normal distribution is characterized by parameters :math:`\mu` and
:math:`\sigma` and has density function :math:`p(X) = 1/(\sqrt{2\pi}\sigma)
e^{-(X-\mu)^2/(2\sigma^2)}`. When truncated to the interval [lower, upper],
the distribution is renormalized so that the PDF integrates to 1 over the
truncation interval.
.. versionchanged:: 0.15.4
Added optional truncation bounds via `lower` and `upper` parameters.
Parameters
----------
mean_value : float
Mean value of the distribution
Mean value of the distribution
std_dev : float
Standard deviation of the Normal distribution
lower : float, optional
Lower truncation bound. Defaults to -infinity (no lower bound).
upper : float, optional
Upper truncation bound. Defaults to +infinity (no upper bound).
bias : openmc.stats.Univariate, optional
Distribution for biased sampling.
@ -1070,6 +1079,10 @@ class Normal(Univariate):
Mean of the Normal distribution
std_dev : float
Standard deviation of the Normal distribution
lower : float
Lower truncation bound
upper : float
Upper truncation bound
support : tuple of float
A 2-tuple (lower, upper) defining the interval over which the
distribution is nonzero-valued
@ -1077,12 +1090,18 @@ class Normal(Univariate):
Distribution for biased sampling
"""
def __init__(self, mean_value, std_dev, bias: Univariate | None = None):
def __init__(self, mean_value, std_dev, lower=-np.inf, upper=np.inf,
bias: Univariate | None = None):
self.mean_value = mean_value
self.std_dev = std_dev
self.lower = lower
self.upper = upper
self._compute_normalization()
super().__init__(bias)
def __len__(self):
if self._is_truncated:
return 4
return 2
@property
@ -1104,16 +1123,69 @@ class Normal(Univariate):
cv.check_greater_than('Normal std_dev', std_dev, 0.0)
self._std_dev = std_dev
@property
def lower(self):
return self._lower
@lower.setter
def lower(self, lower):
cv.check_type('Normal lower bound', lower, Real)
self._lower = lower
@property
def upper(self):
return self._upper
@upper.setter
def upper(self, upper):
cv.check_type('Normal upper bound', upper, Real)
self._upper = upper
def _compute_normalization(self):
"""Compute normalization factor for truncated distribution."""
# Check if truncation bounds are finite
self._is_truncated = (self._lower > -np.inf or self._upper < np.inf)
if self._lower >= self._upper:
raise ValueError("Normal distribution lower bound must be less "
"than upper bound.")
if self._is_truncated:
alpha = (self._lower - self._mean_value) / self._std_dev
beta = (self._upper - self._mean_value) / self._std_dev
cdf_diff = scipy.stats.norm.cdf(beta) - scipy.stats.norm.cdf(alpha)
if cdf_diff <= 0:
raise ValueError("Truncation bounds exclude entire distribution")
self._norm_factor = 1.0 / cdf_diff
else:
self._norm_factor = 1.0
@property
def support(self):
return (-np.inf, np.inf)
return (self._lower, self._upper)
def _sample_unbiased(self, n_samples=1, seed=None):
rng = np.random.RandomState(seed)
return rng.normal(self.mean_value, self.std_dev, n_samples)
if not self._is_truncated:
return rng.normal(self.mean_value, self.std_dev, n_samples)
else:
# Use scipy's truncated normal for efficient direct sampling
a = (self._lower - self._mean_value) / self._std_dev
b = (self._upper - self._mean_value) / self._std_dev
return scipy.stats.truncnorm.rvs(
a, b, loc=self._mean_value, scale=self._std_dev,
size=n_samples, random_state=rng
)
def evaluate(self, x):
return scipy.stats.norm.pdf(x, self.mean_value, self.std_dev)
"""Evaluate PDF at x, returning normalized value for truncated dist."""
x = np.asarray(x)
f = scipy.stats.norm.pdf(x, self.mean_value, self.std_dev)
if self._is_truncated:
# PDF is zero outside bounds
in_bounds = (x >= self._lower) & (x <= self._upper)
f = np.where(in_bounds, f * self._norm_factor, 0.0)
return f
def to_xml_element(self, element_name: str):
"""Return XML representation of the Normal distribution
@ -1126,12 +1198,16 @@ class Normal(Univariate):
Returns
-------
element : lxml.etree._Element
XML element containing Watt distribution data
XML element containing Normal distribution data
"""
element = ET.Element(element_name)
element.set("type", "normal")
element.set("parameters", f'{self.mean_value} {self.std_dev}')
if self._is_truncated:
element.set("parameters",
f'{self.mean_value} {self.std_dev} {self.lower} {self.upper}')
else:
element.set("parameters", f'{self.mean_value} {self.std_dev}')
self._append_bias_to_xml(element)
return element
@ -1152,7 +1228,10 @@ class Normal(Univariate):
"""
params = get_elem_list(elem, "parameters", float)
bias_dist = cls._read_bias_from_xml(elem)
return cls(*map(float, params), bias=bias_dist)
if len(params) == 4:
return cls(params[0], params[1], params[2], params[3], bias=bias_dist)
else:
return cls(params[0], params[1], bias=bias_dist)
def muir(e0: float, m_rat: float, kt: float, bias: Univariate | None = None):
@ -1184,7 +1263,7 @@ def muir(e0: float, m_rat: float, kt: float, bias: Univariate | None = None):
"""
# https://permalink.lanl.gov/object/tr?what=info:lanl-repo/lareport/LA-05411-MS
std_dev = sqrt(2 * e0 * kt / m_rat)
return Normal(e0, std_dev, bias)
return Normal(e0, std_dev, bias=bias)
# Retain deprecated name for the time being

View file

@ -363,30 +363,92 @@ double Watt::evaluate(double x) const
//==============================================================================
// Normal implementation
//==============================================================================
Normal::Normal(double mean_value, double std_dev, double lower, double upper)
: mean_value_ {mean_value}, std_dev_ {std_dev}, lower_ {lower}, upper_ {upper}
{
compute_normalization();
}
Normal::Normal(pugi::xml_node node)
{
auto params = get_node_array<double>(node, "parameters");
if (params.size() != 2) {
if (params.size() != 2 && params.size() != 4) {
openmc::fatal_error("Normal energy distribution must have two "
"parameters specified.");
"parameters (mean, std_dev) or four parameters "
"(mean, std_dev, lower, upper) specified.");
}
mean_value_ = params.at(0);
std_dev_ = params.at(1);
// Optional truncation bounds
if (params.size() == 4) {
lower_ = params.at(2);
upper_ = params.at(3);
} else {
lower_ = -INFTY;
upper_ = INFTY;
}
compute_normalization();
read_bias_from_xml(node);
}
void Normal::compute_normalization()
{
// Validate bounds
if (lower_ >= upper_) {
openmc::fatal_error(
"Normal distribution lower bound must be less than upper bound.");
}
// Check if truncation bounds are finite
is_truncated_ = (lower_ > -INFTY || upper_ < INFTY);
if (is_truncated_) {
double alpha = (lower_ - mean_value_) / std_dev_;
double beta = (upper_ - mean_value_) / std_dev_;
double cdf_diff = standard_normal_cdf(beta) - standard_normal_cdf(alpha);
if (cdf_diff <= 0.0) {
openmc::fatal_error(
"Normal distribution truncation bounds exclude entire distribution.");
}
norm_factor_ = 1.0 / cdf_diff;
} else {
norm_factor_ = 1.0;
}
}
double Normal::sample_unbiased(uint64_t* seed) const
{
return normal_variate(mean_value_, std_dev_, seed);
if (!is_truncated_) {
return normal_variate(mean_value_, std_dev_, seed);
}
// Rejection sampling for truncated normal
double x;
do {
x = normal_variate(mean_value_, std_dev_, seed);
} while (x < lower_ || x > upper_);
return x;
}
double Normal::evaluate(double x) const
{
return (1.0 / (std::sqrt(2.0 / PI) * std_dev_)) *
std::exp(-(std::pow((x - mean_value_), 2.0)) /
(2.0 * std::pow(std_dev_, 2.0)));
// Return 0 outside truncation bounds
if (x < lower_ || x > upper_) {
return 0.0;
}
// Standard normal PDF value
double pdf = (1.0 / (std::sqrt(2.0 * PI) * std_dev_)) *
std::exp(-std::pow((x - mean_value_), 2.0) /
(2.0 * std::pow(std_dev_, 2.0)));
// Apply normalization for truncation
return pdf * norm_factor_;
}
//==============================================================================

View file

@ -95,6 +95,13 @@ double t_percentile(double p, int df)
return t;
}
double standard_normal_cdf(double z)
{
// Use the complementary error function to compute the standard normal CDF
// Phi(z) = 0.5 * (1 + erf(z / sqrt(2))) = 0.5 * erfc(-z / sqrt(2))
return 0.5 * std::erfc(-z / std::sqrt(2.0));
}
void calc_pn_c(int n, double x, double pnx[])
{
pnx[0] = 1.;

View file

@ -92,3 +92,103 @@ TEST_CASE("Test construction of SpatialBox with parameters")
REQUIRE(box.upper_right() == openmc::Position {30, 15, 5});
REQUIRE_FALSE(box.only_fissionable());
}
TEST_CASE("Test Normal distribution")
{
// Test untruncated normal distribution
openmc::Normal normal_unbounded(0.0, 1.0);
// Check PDF at mean (should be 1/sqrt(2*pi) ≈ 0.3989)
REQUIRE_THAT(
normal_unbounded.evaluate(0.0), Catch::Matchers::WithinRel(0.3989, 0.001));
// Check that it's not truncated
REQUIRE_FALSE(normal_unbounded.is_truncated());
// Check accessors
REQUIRE(normal_unbounded.mean_value() == 0.0);
REQUIRE(normal_unbounded.std_dev() == 1.0);
REQUIRE(normal_unbounded.lower() == -openmc::INFTY);
REQUIRE(normal_unbounded.upper() == openmc::INFTY);
}
TEST_CASE("Test truncated Normal distribution")
{
// Create a truncated normal: mean=0, std=1, bounds=[-1, 1]
openmc::Normal normal_truncated(0.0, 1.0, -1.0, 1.0);
// Check that it's truncated
REQUIRE(normal_truncated.is_truncated());
// Check accessors
REQUIRE(normal_truncated.lower() == -1.0);
REQUIRE(normal_truncated.upper() == 1.0);
// PDF should be zero outside bounds
REQUIRE(normal_truncated.evaluate(-2.0) == 0.0);
REQUIRE(normal_truncated.evaluate(2.0) == 0.0);
// PDF inside bounds should be higher than untruncated (due to
// renormalization)
openmc::Normal normal_unbounded(0.0, 1.0);
REQUIRE(normal_truncated.evaluate(0.0) > normal_unbounded.evaluate(0.0));
// The truncated PDF at mean should be approximately 0.3989 / 0.6827 ≈ 0.584
// (0.6827 is the probability mass of N(0,1) in [-1,1])
REQUIRE_THAT(
normal_truncated.evaluate(0.0), Catch::Matchers::WithinRel(0.584, 0.01));
}
TEST_CASE("Test truncated Normal sampling")
{
constexpr int n_samples = 10000;
openmc::Normal normal_truncated(0.0, 1.0, -1.0, 1.0);
uint64_t seed = openmc::init_seed(0, 0);
// Sample and verify all samples are within bounds
for (int i = 0; i < n_samples; ++i) {
auto [x, w] = normal_truncated.sample(&seed);
REQUIRE(x >= -1.0);
REQUIRE(x <= 1.0);
REQUIRE(w == 1.0); // Unbiased sampling should have weight 1
}
}
TEST_CASE("Test one-sided truncated Normal")
{
// Test lower-bounded only (positive half-normal)
openmc::Normal lower_bounded(0.0, 1.0, 0.0, openmc::INFTY);
REQUIRE(lower_bounded.is_truncated());
REQUIRE(lower_bounded.evaluate(-1.0) == 0.0);
REQUIRE(lower_bounded.evaluate(1.0) > 0.0);
// PDF at 0 should be approximately 2 * 0.3989 ≈ 0.798 (half-normal)
REQUIRE_THAT(
lower_bounded.evaluate(0.0), Catch::Matchers::WithinRel(0.798, 0.01));
// Test upper-bounded only
openmc::Normal upper_bounded(0.0, 1.0, -openmc::INFTY, 0.0);
REQUIRE(upper_bounded.is_truncated());
REQUIRE(upper_bounded.evaluate(1.0) == 0.0);
REQUIRE(upper_bounded.evaluate(-1.0) > 0.0);
}
TEST_CASE("Test Normal XML constructor with truncation")
{
// XML doc node for truncated Normal
pugi::xml_document doc;
pugi::xml_node energy = doc.append_child("energy");
energy.append_child("type")
.append_child(pugi::node_pcdata)
.set_value("normal");
energy.append_child("parameters")
.append_child(pugi::node_pcdata)
.set_value("1.0e6 1.0e5 0.8e6 1.2e6");
openmc::Normal dist(energy);
REQUIRE(dist.mean_value() == 1.0e6);
REQUIRE(dist.std_dev() == 1.0e5);
REQUIRE(dist.lower() == 0.8e6);
REQUIRE(dist.upper() == 1.2e6);
REQUIRE(dist.is_truncated());
}

View file

@ -591,6 +591,92 @@ def test_normal():
assert np.all(weights != 1.0)
@pytest.mark.flaky(reruns=1)
def test_normal_truncated():
mean = 10.0
std_dev = 2.0
lower = 6.0
upper = 14.0
d = openmc.stats.Normal(mean, std_dev, lower, upper)
# Check attributes
assert d.mean_value == pytest.approx(mean)
assert d.std_dev == pytest.approx(std_dev)
assert d.lower == pytest.approx(lower)
assert d.upper == pytest.approx(upper)
assert len(d) == 4
assert d.support == (lower, upper)
# Test XML round-trip
elem = d.to_xml_element('distribution')
assert elem.attrib['type'] == 'normal'
params = elem.attrib['parameters'].split()
assert len(params) == 4
d2 = openmc.stats.Normal.from_xml_element(elem)
assert d2.mean_value == pytest.approx(mean)
assert d2.std_dev == pytest.approx(std_dev)
assert d2.lower == pytest.approx(lower)
assert d2.upper == pytest.approx(upper)
# Test PDF evaluation
# PDF should be zero outside bounds
assert d.evaluate(lower - 1.0) == 0.0
assert d.evaluate(upper + 1.0) == 0.0
# PDF should be positive inside bounds
assert d.evaluate(mean) > 0.0
# PDF should be higher than untruncated at the mean (due to renormalization)
d_unbounded = openmc.stats.Normal(mean, std_dev)
assert d.evaluate(mean) > d_unbounded.evaluate(mean)
# Verify that PDF integrates to approximately 1
x = np.linspace(lower, upper, 1000)
integral = trapezoid(d.evaluate(x), x)
assert integral == pytest.approx(1.0, rel=0.01)
# Sample truncated distribution
n_samples = 10_000
samples, weights = d.sample(n_samples)
# All samples should be within bounds
assert np.all(samples >= lower)
assert np.all(samples <= upper)
# Weights should all be 1 (no biasing)
assert np.all(weights == 1.0)
def test_normal_truncated_one_sided():
# Test lower-bounded only (positive half-normal centered at 0)
d_lower = openmc.stats.Normal(0.0, 1.0, lower=0.0)
assert d_lower.lower == 0.0
assert d_lower.upper == np.inf
assert d_lower.evaluate(-1.0) == 0.0
assert d_lower.evaluate(1.0) > 0.0
# PDF at 0 should be approximately 2 * 0.3989 ≈ 0.798 (half-normal)
assert d_lower.evaluate(0.0) == pytest.approx(0.798, rel=0.01)
# Test upper-bounded only
d_upper = openmc.stats.Normal(0.0, 1.0, upper=0.0)
assert d_upper.lower == -np.inf
assert d_upper.upper == 0.0
assert d_upper.evaluate(1.0) == 0.0
assert d_upper.evaluate(-1.0) > 0.0
def test_normal_truncated_errors():
# Invalid bounds (lower >= upper)
with pytest.raises(ValueError):
openmc.stats.Normal(0.0, 1.0, lower=1.0, upper=0.0)
with pytest.raises(ValueError):
openmc.stats.Normal(0.0, 1.0, lower=1.0, upper=1.0)
@pytest.mark.flaky(reruns=1)
def test_muir():
mean = 10.0