diff --git a/openmc/model/model.py b/openmc/model/model.py index c1ffafafd3f..7963751d567 100644 --- a/openmc/model/model.py +++ b/openmc/model/model.py @@ -1,6 +1,7 @@ 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 @@ -8,11 +9,13 @@ 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 `_. 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) + + diff --git a/src/settings.cpp b/src/settings.cpp index 325256cdc50..13b91b0e4f8 100644 --- a/src/settings.cpp +++ b/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; diff --git a/tests/unit_tests/test_lib.py b/tests/unit_tests/test_lib.py index 43bc5a8f6f0..eb4dc3dce66 100644 --- a/tests/unit_tests/test_lib.py +++ b/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 diff --git a/tests/unit_tests/test_model.py b/tests/unit_tests/test_model.py index 60f8b1a25a3..f4f94a47cac 100644 --- a/tests/unit_tests/test_model.py +++ b/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