Criticality search method on the Model class (#3569)

Author: paulromano

openmc/model/model.py

@@ -1,18 +1,21 @@
 from __future__ import annotations
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 import copy
+from dataclasses import dataclass, field
 from functools import cache
 from pathlib import Path
 import math
 from numbers import Integral, Real
 import random
 import re
 from tempfile import NamedTemporaryFile, TemporaryDirectory
+from typing import Any, Protocol
 import warnings
 
 import h5py
 import lxml.etree as ET
 import numpy as np
+from scipy.optimize import curve_fit
 
 import openmc
 import openmc._xml as xml
@@ -24,6 +27,12 @@
 from openmc.utility_funcs import change_directory
 
 
+# Protocol for a function that is passed to search_keff
+class ModelModifier(Protocol):
+    def __call__(self, val: float, **kwargs: Any) -> None:
+        ...
+
+
 class Model:
     """Model container.
 
@@ -2196,3 +2205,262 @@ def _replace_infinity(value):
 
         # Take a wild guess as to how many rays are needed
         self.settings.particles = 2 * int(max_length)
+
+    def keff_search(
+        self,
+        func: ModelModifier,
+        x0: float,
+        x1: float,
+        target: float = 1.0,
+        k_tol: float = 1e-4,
+        sigma_final: float = 3e-4,
+        p: float = 0.5,
+        q: float = 0.95,
+        memory: int = 4,
+        x_min: float | None = None,
+        x_max: float | None = None,
+        b0: int | None = None,
+        b_min: int = 20,
+        b_max: int | None = None,
+        maxiter: int = 50,
+        output: bool = False,
+        func_kwargs: dict[str, Any] | None = None,
+        run_kwargs: dict[str, Any] | None = None,
+    ) -> SearchResult:
+        r"""Perform a keff search on a model parametrized by a single variable.
+
+        This method uses the GRsecant method described in a paper by `Price and
+        Roskoff <https://doi.org/10.1016/j.pnucene.2023.104731>`_. The GRsecant
+        method is a modification of the secant method that accounts for
+        uncertainties in the function evaluations. The method uses a weighted
+        linear fit of the most recent function evaluations to predict the next
+        point to evaluate. It also adaptively changes the number of batches to
+        meet the target uncertainty value at each iteration.
+
+        The target uncertainty for iteration :math:`n+1` is determined by the
+        following equation (following Eq. (8) in the paper):
+
+        .. math::
+            \sigma_{i+1} = q \sigma_\text{final} \left ( \frac{ \min \left \{
+            \left\lvert k_i - k_\text{target} \right\rvert : k=0,1,\dots,n
+            \right \} }{k_\text{tol}} \right )^p
+
+        where :math:`q` is a multiplicative factor less than 1, given as the
+        ``sigma_factor`` parameter below.
+
+        Parameters
+        ----------
+        func : ModelModifier
+            Function that takes the parameter to be searched and makes a
+            modification to the model.
+        x0 : float
+            First guess for the parameter passed to `func`
+        x1 : float
+            Second guess for the parameter passed to `func`
+        target : float, optional
+            keff value to search for
+        k_tol : float, optional
+            Stopping criterion on the function value; the absolute value must be
+            within ``k_tol`` of zero to be accepted.
+        sigma_final : float, optional
+            Maximum accepted k-effective uncertainty for the stopping criterion.
+        p : float, optional
+            Exponent used in the stopping criterion.
+        q : float, optional
+            Multiplicative factor used in the stopping criterion.
+        memory : int, optional
+            Number of most-recent points used in the weighted linear fit of
+            ``f(x) = a + b x`` to predict the next point.
+        x_min : float, optional
+            Minimum allowed value for the parameter ``x``.
+        x_max : float, optional
+            Maximum allowed value for the parameter ``x``.
+        b0 : int, optional
+            Number of active batches to use for the initial function
+            evaluations. If None, uses the model's current setting.
+        b_min : int, optional
+            Minimum number of active batches to use in a function evaluation.
+        b_max : int, optional
+            Maximum number of active batches to use in a function evaluation.
+        maxiter : int, optional
+            Maximum number of iterations to perform.
+        output : bool, optional
+            Whether or not to display output showing iteration progress.
+        func_kwargs : dict, optional
+            Keyword-based arguments to pass to the `func` function.
+        run_kwargs : dict, optional
+            Keyword arguments to pass to :meth:`openmc.Model.run` or
+            :meth:`openmc.lib.run`.
+
+        Returns
+        -------
+        SearchResult
+            Result object containing the estimated root (parameter value) and
+            evaluation history (parameters, means, standard deviations, and
+            batches), plus convergence status and termination reason.
+
+        """
+        import openmc.lib
+
+        check_type('model modifier', func, Callable)
+        check_type('target', target, Real)
+        if memory < 2:
+            raise ValueError("memory must be ≥ 2")
+        func_kwargs = {} if func_kwargs is None else dict(func_kwargs)
+        run_kwargs = {} if run_kwargs is None else dict(run_kwargs)
+        run_kwargs.setdefault('output', False)
+
+        # Create lists to store the history of evaluations
+        xs: list[float] = []
+        fs: list[float] = []
+        ss: list[float] = []
+        gs: list[int] = []
+        count = 0
+
+        # Helper function to evaluate f and store results
+        def eval_at(x: float, batches: int) -> tuple[float, float]:
+            # Modify the model with the current guess
+            func(x, **func_kwargs)
+
+            # Change the number of batches and run the model
+            batches += self.settings.inactive
+            if openmc.lib.is_initialized:
+                openmc.lib.settings.set_batches(batches)
+                openmc.lib.reset()
+                openmc.lib.run(**run_kwargs)
+                sp_filepath = f'statepoint.{batches}.h5'
+            else:
+                self.settings.batches = batches
+                sp_filepath = self.run(**run_kwargs)
+
+            # Extract keff and its uncertainty
+            with openmc.StatePoint(sp_filepath) as sp:
+                keff = sp.keff
+
+            if output:
+                nonlocal count
+                count += 1
+                print(f'Iteration {count}: {batches=}, {x=:.6g}, {keff=:.5f}')
+
+            xs.append(float(x))
+            fs.append(float(keff.n - target))
+            ss.append(float(keff.s))
+            gs.append(int(batches))
+            return fs[-1], ss[-1]
+
+        # Default b0 to current model settings if not explicitly provided
+        if b0 is None:
+            b0 = self.settings.batches - self.settings.inactive
+
+        # Perform the search (inlined GRsecant) in a temporary directory
+        with TemporaryDirectory() as tmpdir:
+            if not openmc.lib.is_initialized:
+                run_kwargs.setdefault('cwd', tmpdir)
+
+            # ---- Seed with two evaluations
+            f0, s0 = eval_at(x0, b0)
+            if abs(f0) <= k_tol and s0 <= sigma_final:
+                return SearchResult(x0, xs, fs, ss, gs, True, "converged")
+            f1, s1 = eval_at(x1, b0)
+            if abs(f1) <= k_tol and s1 <= sigma_final:
+                return SearchResult(x1, xs, fs, ss, gs, True, "converged")
+
+            for _ in range(maxiter - 2):
+                # ------ Step 1: propose next x via GRsecant
+                m = min(memory, len(xs))
+
+                # Perform a curve fit on f(x) = a + bx accounting for
+                # uncertainties. This is equivalent to minimizing the function
+                # in Equation (A.14)
+                (a, b), _ = curve_fit(
+                    lambda x, a, b: a + b*x,
+                    xs[-m:], fs[-m:], sigma=ss[-m:], absolute_sigma=True
+                )
+                x_new = float(-a / b)
+
+                # Clamp x_new to the bounds if provided
+                if x_min is not None:
+                    x_new = max(x_new, x_min)
+                if x_max is not None:
+                    x_new = min(x_new, x_max)
+
+                # ------ Step 2: choose target σ for next run (Eq. 8 + clamp)
+
+                min_abs_f = float(np.min(np.abs(fs)))
+                base = q * sigma_final
+                ratio = min_abs_f / k_tol if k_tol > 0 else 1.0
+                sig = base * (ratio ** p)
+                sig_target = max(sig, base)
+
+                # ------ Step 3: choose generations to hit σ_target (Appendix C)
+
+                # Use at least two past points for regression
+                if len(gs) >= 2 and np.var(np.log(gs)) > 0.0:
+                    # Perform a curve fit based on Eq. (C.3) to solve for ln(k).
+                    # Note that unlike in the paper, we do not leave r as an
+                    # undetermined parameter and choose r=0.5.
+                    (ln_k,), _ = curve_fit(
+                        lambda ln_b, ln_k: ln_k - 0.5*ln_b,
+                        np.log(gs[-4:]), np.log(ss[-4:]),
+                    )
+                    k = float(np.exp(ln_k))
+                else:
+                    k = float(ss[-1] * math.sqrt(gs[-1]))
+
+                b_new = (k / sig_target) ** 2
+
+                # Clamp and round up to integer
+                b_new = max(b_min, math.ceil(b_new))
+                if b_max is not None:
+                    b_new = min(b_new, b_max)
+
+                # Evaluate at proposed x with batches determined above
+                f_new, s_new = eval_at(x_new, b_new)
+
+                # Termination based on both criteria (|f| and σ)
+                if abs(f_new) <= k_tol and s_new <= sigma_final:
+                    return SearchResult(x_new, xs, fs, ss, gs, True, "converged")
+
+            return SearchResult(xs[-1], xs, fs, ss, gs, False, "maxiter")
+
+
+@dataclass
+class SearchResult:
+    """Result of a GRsecant keff search.
+
+    Attributes
+    ----------
+    root : float
+        Estimated parameter value where f(x) = 0 at termination.
+    parameters : list[float]
+        Parameter values (x) evaluated during the search, in order.
+    keffs : list[float]
+        Estimated keff values for each evaluation.
+    stdevs : list[float]
+        One-sigma uncertainties of keff for each evaluation.
+    batches : list[int]
+        Number of active batches used for each evaluation.
+    converged : bool
+        Whether both |f| <= k_tol and sigma <= sigma_final were met.
+    flag : str
+        Reason for termination (e.g., "converged", "maxiter").
+    """
+    root: float
+    parameters: list[float] = field(repr=False)
+    means: list[float] = field(repr=False)
+    stdevs: list[float] = field(repr=False)
+    batches: list[int] = field(repr=False)
+    converged: bool
+    flag: str
+
+    @property
+    def function_calls(self) -> int:
+        """Number of function evaluations performed."""
+        return len(self.parameters)
+
+    @property
+    def total_batches(self) -> int:
+        """Total number of active batches used across all evaluations."""
+        return sum(self.batches)
+
+

src/settings.cpp

@@ -1220,11 +1220,6 @@ extern "C" int openmc_set_n_batches(
     return OPENMC_E_INVALID_ARGUMENT;
   }
 
-  if (simulation::current_batch >= n_batches) {
-    set_errmsg("Number of batches must be greater than current batch.");
-    return OPENMC_E_INVALID_ARGUMENT;
-  }
-
   if (!settings::trigger_on) {
     // Set n_batches and n_max_batches to same value
     settings::n_batches = n_batches;

tests/unit_tests/test_lib.py

@@ -496,9 +496,6 @@ def test_set_n_batches(lib_run):
 
     for i in range(7):
         openmc.lib.next_batch()
-    # Setting n_batches less than current_batch should raise error
-    with pytest.raises(exc.InvalidArgumentError):
-        settings.set_batches(6)
     # n_batches should stay the same
     assert settings.get_batches() == 10
 

tests/unit_tests/test_model.py

@@ -901,6 +901,7 @@ def test_id_map_aligned_model():
     assert tr_instance == 3, f"Expected cell instance 3 at top-right corner, got {tr_instance}"
     assert tr_material == 5, f"Expected material ID 5 at top-right corner, got {tr_material}"
 
+
 def test_setter_from_list():
     mat = openmc.Material()
     model = openmc.Model(materials=[mat])
@@ -913,3 +914,59 @@ def test_setter_from_list():
     plot = openmc.Plot()
     model = openmc.Model(plots=[plot])
     assert isinstance(model.plots, openmc.Plots)
+
+
+def test_keff_search(run_in_tmpdir):
+    """Test the Model.keff_search method"""
+
+    # Create model of a sphere of U235
+    mat = openmc.Material()
+    mat.set_density('g/cm3', 18.9)
+    mat.add_nuclide('U235', 1.0)
+    sphere = openmc.Sphere(r=10.0, boundary_type='vacuum')
+    cell = openmc.Cell(fill=mat, region=-sphere)
+    geometry = openmc.Geometry([cell])
+    settings = openmc.Settings(particles=1000, inactive=10, batches=30)
+    model = openmc.Model(geometry=geometry, settings=settings)
+
+    # Define function to modify sphere radius
+    def modify_radius(radius):
+        sphere.r = radius
+
+    # Perform keff search
+    k_tol = 4e-3
+    sigma_final = 2e-3
+    result = model.keff_search(
+        func=modify_radius,
+        x0=6.0,
+        x1=9.0,
+        k_tol=k_tol,
+        sigma_final=sigma_final,
+        output=True,
+    )
+
+    final_keff = result.means[-1] + 1.0  # Add back target since means are (keff - target)
+    final_sigma = result.stdevs[-1]
+
+    # Check for convergence and that tolerances are met
+    assert result.converged, "keff_search did not converge"
+    assert abs(final_keff - 1.0) <= k_tol, \
+        f"Final keff {final_keff:.5f} not within k_tol {k_tol}"
+    assert final_sigma <= sigma_final, \
+        f"Final uncertainty {final_sigma:.5f} exceeds sigma_final {sigma_final}"
+
+    # Check type of result
+    assert isinstance(result, openmc.model.SearchResult)
+
+    # Check that we have function evaluation history
+    assert len(result.parameters) >= 2
+    assert len(result.means) == len(result.parameters)
+    assert len(result.stdevs) == len(result.parameters)
+    assert len(result.batches) == len(result.parameters)
+
+    # Check that function_calls property works
+    assert result.function_calls == len(result.parameters)
+
+    # Check that total_batches property works
+    assert result.total_batches == sum(result.batches)
+    assert result.total_batches > 0