Add reactivity control to coupled transport-depletion analyses (#2693)

Co-authored-by: Andrew Johnson <drewejohnson@users.noreply.github.com>
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

Author: 3 people

docs/source/io_formats/depletion_results.rst

@@ -4,7 +4,7 @@
 Depletion Results File Format
 =============================
 
-The current version of the depletion results file format is 1.2.
+The current version of the depletion results file format is 1.3.
 
 **/**
 
@@ -29,6 +29,8 @@ The current version of the depletion results file format is 1.2.
            - **depletion time** (*double[]*) -- Average process time in [s]
              spent depleting a material across all burnable materials and,
              if applicable, MPI processes.
+           - **keff_search_root** (*double[]*) -- Root of the keff search at the
+             end of the timestep, if applicable.
 
 **/materials/<id>/**
 

openmc/deplete/abc.py

@@ -31,6 +31,7 @@
 from .pool import deplete
 from .reaction_rates import ReactionRates
 from .transfer_rates import TransferRates, ExternalSourceRates
+from .keff_search_control import _KeffSearchControl
 
 
 __all__ = [
@@ -159,7 +160,7 @@ def __init__(self, chain_file=None, fission_q=None, prev_results=None):
             self.prev_res = prev_results
 
     @abstractmethod
-    def __call__(self, vec, source_rate):
+    def __call__(self, vec, source_rate) -> OperatorResult:
         """Runs a simulation.
 
         Parameters
@@ -686,6 +687,7 @@ def __init__(
 
         self.transfer_rates = None
         self.external_source_rates = None
+        self._keff_search_control = None
 
         if isinstance(solver, str):
             # Delay importing of cram module, which requires this file
@@ -839,6 +841,37 @@ def _get_start_data(self) -> tuple[float, int]:
         return (self.operator.prev_res[-1].time[0],
                 len(self.operator.prev_res) - 1)
 
+    def _restore_keff_search_control(self, res: StepResult):
+        """Restore keff search control from restart results."""
+        keff_search_root = res.keff_search_root
+        if keff_search_root is None:
+            raise ValueError(
+                "Cannot restore keff search control from restart "
+                "results because no stored keff_search_root is "
+                "available."
+            )
+        self._keff_search_control.function(keff_search_root)
+        return keff_search_root
+
+    def _get_bos_data(self, step_index, source_rate, bos_conc):
+        """Get beginning-of-step concentrations, rates, and control state."""
+        if step_index > 0 or self.operator.prev_res is None:
+            if self._keff_search_control is not None and source_rate != 0.0:
+                keff_search_root = self._keff_search_control.run(bos_conc)
+            else:
+                keff_search_root = None
+            bos_conc, res = self._get_bos_data_from_operator(
+                step_index, source_rate, bos_conc)
+        else:
+            bos_conc, res = self._get_bos_data_from_restart(
+                source_rate, bos_conc)
+            if self._keff_search_control is not None and source_rate != 0.0:
+                keff_search_root = self._restore_keff_search_control(self.operator.prev_res[-1])
+            else:
+                keff_search_root = None
+
+        return bos_conc, res, keff_search_root
+
     def integrate(
             self,
             final_step: bool = True,
@@ -877,11 +910,8 @@ def integrate(
                 if output and comm.rank == 0:
                     print(f"[openmc.deplete] t={t} s, dt={dt} s, source={source_rate}")
 
-                # Solve transport equation (or obtain result from restart)
-                if i > 0 or self.operator.prev_res is None:
-                    n, res = self._get_bos_data_from_operator(i, source_rate, n)
-                else:
-                    n, res = self._get_bos_data_from_restart(source_rate, n)
+                # Get beginning-of-step data from operator or restart results
+                n, res, keff_search_root = self._get_bos_data(i, source_rate, n)
 
                 # Solve Bateman equations over time interval
                 proc_time, n_end = self(n, res.rates, dt, source_rate, i)
@@ -895,6 +925,7 @@ def integrate(
                     self._i_res + i,
                     proc_time,
                     write_rates=write_rates,
+                    keff_search_root=keff_search_root,
                     path=path
                 )
 
@@ -908,6 +939,10 @@ def integrate(
             # solve)
             if output and final_step and comm.rank == 0:
                 print(f"[openmc.deplete] t={t} (final operator evaluation)")
+            if self._keff_search_control is not None and source_rate != 0.0:
+                keff_search_root = self._keff_search_control.run(n)
+            else:
+                keff_search_root = None
             res_final = self.operator(n, source_rate if final_step else 0.0)
             StepResult.save(
                 self.operator,
@@ -918,6 +953,7 @@ def integrate(
                 self._i_res + len(self),
                 proc_time,
                 write_rates=write_rates,
+                keff_search_root=keff_search_root,
                 path=path
             )
             self.operator.write_bos_data(len(self) + self._i_res)
@@ -1050,6 +1086,101 @@ def add_redox(self, material, buffer, oxidation_states, timesteps=None):
 
         self.transfer_rates.set_redox(material, buffer, oxidation_states, timesteps)
 
+    def add_keff_search_control(
+        self,
+        function: Callable,
+        x0: float,
+        x1: float,
+        bracket: Sequence[float],
+        **search_kwargs
+    ):
+        """Add keff search to the integrator scheme.
+
+        This method causes OpenMC to perform a keff search during depletion to
+        maintain a target keff by adjusting a model parameter through the
+        provided function.
+
+        .. important::
+            The function **must** modify the model through ``openmc.lib`` (e.g.,
+            ``openmc.lib.cells``, ``openmc.lib.materials``) and **NOT** through
+            ``openmc.Model``. The function is called within a
+            :class:`openmc.lib.TemporarySession` context where only the C API
+            (``openmc.lib``) is available for modifications.
+
+        Parameters
+        ----------
+        function : Callable
+            Function that takes a single float argument and modifies the model
+            through :mod:`openmc.lib`.
+        x0 : float
+            Initial lower bound for the keff search.
+        x1 : float
+            Initial upper bound for the keff search.
+        bracket : sequence of float
+            Bracket interval [x_min, x_max] that constrains the allowed parameter
+            values during the keff search. This is a required parameter
+            that defines the absolute bounds for the search. The bracket must contain
+            exactly 2 elements with bracket[0] < bracket[1]. These values are passed
+            directly to the ``x_min`` and ``x_max`` optional arguments in
+            :meth:`openmc.Model.keff_search`, which enforce hard limits on the
+            parameter range. If the keff search converges to a value outside this
+            bracket, it will be clamped to the nearest bracket bound with a warning.
+        **search_kwargs
+            Additional keyword arguments passed to
+            :meth:`openmc.Model.keff_search`. Common options include:
+
+            * ``target`` : float, optional
+              Target keff value to search for. Defaults to 1.0.
+            * ``k_tol`` : float, optional
+              Stopping criterion on the function value. Defaults to 1e-4.
+            * ``sigma_final`` : float, optional
+              Maximum accepted k-effective uncertainty. Defaults to 3e-4.
+            * ``maxiter`` : int, optional
+              Maximum number of iterations. Defaults to 50.
+
+            See :meth:`openmc.Model.keff_search` for a complete list of
+            available options.
+
+        Examples
+        --------
+        Add keff search that adjusts a control rod position:
+
+        >>> def adjust_rod_position(position):
+        ...     openmc.lib.cells[rod_cell.id].translation = [0, 0, position]
+        >>> integrator.add_keff_search_control(
+        ...     adjust_rod_position,
+        ...     x0=0.0,
+        ...     x1=5.0,
+        ...     bracket=[-10,10],
+        ...     target=1.0,
+        ...     k_tol=1e-4
+        ... )
+
+        Add keff search that adjusts the U235 density:
+
+        >>> def set_u235_density(u235_density):
+        ...     # Get the material from openmc.lib
+        ...     lib_mat = openmc.lib.materials[material_id]
+        ...     # Get current nuclides and densities
+        ...     nuclides = lib_mat.nuclides
+        ...     densities = lib_mat.densities
+        ...     u235_idx = nuclides.index('U235')
+        ...     densities[u235_idx] = u235_density
+        ...     lib_mat.set_densities(nuclides, densities)
+        >>> integrator.add_keff_search_control(
+        ...     set_u235_density,
+        ...     x0=5.0e-4,
+        ...     x1=1.0e-3,
+        ...     bracket=[1.0e-4, 2.0e-3],
+        ...     target=1.0
+        ... )
+
+        .. versionadded:: 0.15.4
+
+        """
+        self._keff_search_control = _KeffSearchControl(
+            self.operator, function, x0, x1, bracket, **search_kwargs)
+
 @add_params
 class SIIntegrator(Integrator):
     r"""Abstract class for the Stochastic Implicit Euler integrators

openmc/deplete/coupled_operator.py

@@ -399,7 +399,7 @@ def _generate_materials_xml(self):
 
         self.materials.export_to_xml(nuclides_to_ignore=self._decay_nucs)
 
-    def __call__(self, vec, source_rate):
+    def __call__(self, vec, source_rate) -> OperatorResult:
         """Runs a simulation.
 
         Simulation will abort under the following circumstances:

openmc/deplete/independent_operator.py

@@ -384,7 +384,7 @@ def initial_condition(self):
         # Return number density vector
         return super().initial_condition(self.materials)
 
-    def __call__(self, vec, source_rate):
+    def __call__(self, vec, source_rate) -> OperatorResult:
         """Obtain the reaction rates
 
         Parameters

openmc/deplete/keff_search_control.py

@@ -0,0 +1,128 @@
+from typing import Callable
+from warnings import warn
+
+import openmc.lib
+
+
+class _KeffSearchControl:
+    """Controller for keff search during depletion calculations.
+
+    This class performs keff searches to maintain a target keff by adjusting a
+    model parameter through a provided function.
+
+    Parameters
+    ----------
+    operator : openmc.deplete.Operator
+        Depletion operator instance
+    function : Callable
+        Function that modifies the model based on a parameter value
+    x0 : float
+        Initial lower bound for the keff search
+    x1 : float
+        Initial upper bound for the keff search
+    bracket : list[float]
+        Absolute bracketing interval lower and upper. If the keff search
+        solution lies off these limits the closest limit will be set as new
+        result.
+    **search_kwargs : dict, optional
+        Additional keyword arguments to pass to :meth:`openmc.Model.keff_search`
+
+    """
+    def __init__(self, operator, function: Callable, x0: float, x1: float, bracket: list[float], **search_kwargs):
+        if len(bracket) != 2:
+            raise ValueError(f"bracket must have exactly 2 elements, got {len(bracket)}")
+        if bracket[0] >= bracket[1]:
+            raise ValueError(f"bracket[0] must be < bracket[1], got {bracket}")
+        self.x0 = x0
+        self.x1 = x1
+        self.operator = operator
+        self.function = function
+        self.search_kwargs = search_kwargs
+        self.search_kwargs['x_min'] = bracket[0]
+        self.search_kwargs['x_max'] = bracket[1]
+
+    def run(self, x):
+        """Perform keff search and update the atom density vector.
+
+        Parameters
+        ----------
+        x : list of numpy.ndarray
+            Current atom density vector (atoms per material)
+
+        Returns
+        -------
+        root : float
+            Parameter value that achieves target keff
+        """
+        root = self._search_for_keff()
+        self._update_vec(x)
+        return root
+
+    def _search_for_keff(self) -> float:
+        """Perform the keff search using the model's keff_search method.
+
+        Returns
+        -------
+        float
+            Parameter value that achieves target keff
+
+        Raises
+        ------
+        ValueError
+            If the keff search fails to converge
+        """
+        with openmc.lib.TemporarySession(self.operator.model):
+            # Only pass the first 3 required args plus explicitly provided kwargs
+            result = self.operator.model.keff_search(
+                self.function, self.x0, self.x1, **self.search_kwargs
+            )
+        if not result.converged:
+            raise ValueError(
+                f"Search for keff failed to converge. "
+                f"Termination reason: {result.flag}"
+            )
+
+        root = result.root
+
+        # Check if root is outside the bracket bounds and give a warning
+        if root < self.search_kwargs['x_min']:
+            warn(f"keff search result ({root:.6f}) is below the lower bracket "
+                 f"bound ({self.search_kwargs['x_min']:.6f}).", UserWarning)
+        elif root > self.search_kwargs['x_max']:
+            warn(f"keff search result ({root:.6f}) is above the upper bracket "
+                 f"bound ({self.search_kwargs['x_max']:.6f}).", UserWarning)
+
+        # Restore the number of initial batches
+        openmc.lib.settings.set_batches(self.operator.model.settings.batches)
+
+        return root
+
+    def _update_vec(self, x):
+        """Update the atom density vector from openmc.lib.materials and AtomNumber object.
+
+        The depletion vector ``x`` is rank-local, matching the materials owned
+        by ``self.operator.number`` on the current MPI rank. We therefore only
+        update entries for locally owned materials using the compositions
+        currently stored in ``openmc.lib.materials``.
+
+        Parameters
+        ----------
+        x : list of numpy.ndarray
+            Atom density vector to update (atoms per material)
+
+        """
+        number = self.operator.number
+
+        for mat_idx, mat in enumerate(number.materials):
+            lib_material = openmc.lib.materials[int(mat)]
+            nuclides = lib_material.nuclides
+            densities = 1e24 * lib_material.densities
+            volume = number.get_mat_volume(mat)
+
+            for nuc_idx, nuc in enumerate(number.burnable_nuclides):
+                if nuc in nuclides:
+                    lib_nuc_idx = nuclides.index(nuc)
+                    atom_density = densities[lib_nuc_idx]
+                else:
+                    atom_density = number.get_atom_density(mat, nuc)
+                x[mat_idx][nuc_idx] = atom_density * volume