Add truncated normal distribution support (#3761)

Author: paulromano

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 @@ class Watt : public Distribution {
 };
 
 //==============================================================================
-//! 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 @@ class Normal : public Distribution {
   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,18 +1049,27 @@ def from_xml_element(cls, elem: ET.Element):
 
 
 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,19 +1079,29 @@ 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
     bias : openmc.stats.Univariate or None
         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 @@ def std_dev(self, std_dev):
         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 @@ def to_xml_element(self, element_name: str):
         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 @@ def from_xml_element(cls, elem: ET.Element):
         """
         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
 {
-  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_;
 }
 
 //==============================================================================

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