From 8198e6021dae83e4e9519108ca1483ad656f06a4 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Thu, 12 Feb 2026 10:25:56 -0600 Subject: [PATCH] Add truncated normal distribution support (#3761) --- .../parameterized_source_ring.cpp | 4 +- include/openmc/distribution.h | 27 +++-- include/openmc/math_functions.h | 10 ++ openmc/stats/univariate.py | 105 +++++++++++++++--- src/distribution.cpp | 74 +++++++++++- src/math_functions.cpp | 7 ++ tests/cpp_unit_tests/test_distribution.cpp | 100 +++++++++++++++++ tests/unit_tests/test_stats.py | 86 ++++++++++++++ 8 files changed, 385 insertions(+), 28 deletions(-) diff --git a/examples/parameterized_custom_source/parameterized_source_ring.cpp b/examples/parameterized_custom_source/parameterized_source_ring.cpp index a983414f1..9756e1b97 100644 --- a/examples/parameterized_custom_source/parameterized_source_ring.cpp +++ b/examples/parameterized_custom_source/parameterized_source_ring.cpp @@ -1,5 +1,5 @@ -#include // for M_PI -#include // for unique_ptr +#include +#include #include #include "openmc/particle.h" diff --git a/include/openmc/distribution.h b/include/openmc/distribution.h index 80fe70bae..cd81d8801 100644 --- a/include/openmc/distribution.h +++ b/include/openmc/distribution.h @@ -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(); }; //============================================================================== diff --git a/include/openmc/math_functions.h b/include/openmc/math_functions.h index d67b8305e..dcc0e21fe 100644 --- a/include/openmc/math_functions.h +++ b/include/openmc/math_functions.h @@ -223,5 +223,15 @@ double log1prel(double x); void get_energy_index( const vector& 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 diff --git a/openmc/stats/univariate.py b/openmc/stats/univariate.py index 3d46024a4..c6ce217fb 100644 --- a/openmc/stats/univariate.py +++ b/openmc/stats/univariate.py @@ -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 diff --git a/src/distribution.cpp b/src/distribution.cpp index 756b098ed..c722d0366 100644 --- a/src/distribution.cpp +++ b/src/distribution.cpp @@ -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(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_; } //============================================================================== diff --git a/src/math_functions.cpp b/src/math_functions.cpp index f45165414..81f08fa04 100644 --- a/src/math_functions.cpp +++ b/src/math_functions.cpp @@ -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.; diff --git a/tests/cpp_unit_tests/test_distribution.cpp b/tests/cpp_unit_tests/test_distribution.cpp index f7c480c8d..e46088a63 100644 --- a/tests/cpp_unit_tests/test_distribution.cpp +++ b/tests/cpp_unit_tests/test_distribution.cpp @@ -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()); +} diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py index bdbdaa0f2..b9d9c4e1f 100644 --- a/tests/unit_tests/test_stats.py +++ b/tests/unit_tests/test_stats.py @@ -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