Add truncated normal distribution support (#3761)

examples/parameterized_custom_source/parameterized_source_ring.cpp

@@ -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"

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();
 };
 
 //==============================================================================

include/openmc/math_functions.h

@@ -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

openmc/stats/univariate.py

@@ -1049,11 +1049,16 @@ 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
     ----------
@@ -1061,6 +1066,10 @@ class Normal(Univariate):
         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)
+        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,11 +1198,15 @@ 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")
+        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

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<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
 {
+  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)) /
+  // 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_;
 }
 
 //==============================================================================

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.;

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());
+}

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