Add reactivity control to coupled transport-depletion analyses

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__ = [
@@ -617,6 +618,11 @@ class Integrator(ABC):
         External source rates for the depletion system.
 
         .. versionadded:: 0.15.3
+    keff_search_control : openmc.deplete._KeffSearchControl
+        Instance of _KeffSearchControl class to perform keff search during
+        transport-depletion simulation.
+
+        .. versionadded:: 0.15.4
 
     """)
 
@@ -684,6 +690,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
@@ -729,6 +736,15 @@ def solver(self, func):
 
         self._solver = func
 
+    @property
+    def keff_search_control(self):
+        return self._keff_search_control
+
+    @keff_search_control.setter
+    def keff_search_control(self, keff_search_control):
+        check_type('keff search control', keff_search_control, _KeffSearchControl)
+        self._keff_search_control = keff_search_control
+
     def _timed_deplete(self, n, rates, dt, i=None, matrix_func=None):
         start = time.time()
         results = deplete(
@@ -837,6 +853,13 @@ def _get_start_data(self) -> tuple[float, int]:
         return (self.operator.prev_res[-1].time[0],
                 len(self.operator.prev_res) - 1)
 
+    def _get_bos_from_keff_search_control(self, step_index, bos_conc):
+        """Get BOS from keff search control."""
+        x = deepcopy(bos_conc)
+        # Get new vector after keff criticality control
+        x, keff_search_root = self._keff_search_control.search_for_keff(x, step_index)
+        return x, keff_search_root
+
     def integrate(
             self,
             final_step: bool = True,
@@ -877,10 +900,19 @@ def integrate(
 
                 # Solve transport equation (or obtain result from restart)
                 if i > 0 or self.operator.prev_res is None:
+                    # Update geometry/material according to keff search control
+                    if self._keff_search_control is not None and source_rate != 0.0:
+                        n, keff_search_root = self._get_bos_from_keff_search_control(i, n)
+                    else:
+                        keff_search_root = 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 keff search root from keff search control
+                    if self._keff_search_control:
+                        n, keff_search_root = self._get_bos_from_keff_search_control(i, n)
+                    else:
+                        keff_search_root = None
                 # Solve Bateman equations over time interval
                 proc_time, n_end = self(n, res.rates, dt, source_rate, i)
 
@@ -893,6 +925,7 @@ def integrate(
                     self._i_res + i,
                     proc_time,
                     write_rates=write_rates,
+                    keff_search_root=keff_search_root,
                     path=path
                 )
 
@@ -906,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:
+                n, keff_search_root = self._get_bos_from_keff_search_control(i+1, n)
+            else:
+                keff_search_root = None
             res_final = self.operator(n, source_rate if final_step else 0.0)
             StepResult.save(
                 self.operator,
@@ -916,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)
@@ -1048,6 +1086,113 @@ 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: list[float],
+        **search_kwargs
+    ):
+        """Add keff search to the integrator scheme.
+
+        This method creates a :class:`openmc.deplete._KeffSearchControl` that
+        performs keff searches 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 modifies the model through ``openmc.lib`` based on a
+            parameter value. The function should take a single float parameter
+            and modify ``openmc.lib`` objects accordingly (e.g., adjust control
+            rod position via ``openmc.lib.cells[...].translation``, material
+            density via ``openmc.lib.materials[...].set_densities(...)``, etc.).
+
+            **Important**: The function must modify ``openmc.lib`` objects, not
+            ``openmc.model`` objects.
+        x0: float
+            Initial lower bound for the keff search.
+        x1: float
+            Initial upper bound for the keff search.
+        bracket : list[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 material density:
+
+        >>> def adjust_material_density(density_factor):
+        ...     # Get the material from openmc.lib
+        ...     lib_mat = openmc.lib.materials[material_id]
+        ...     # Get current nuclides and densities
+        ...     nuclides = lib_mat.nuclides
+        ...     current_densities = lib_mat.densities
+        ...     # Scale all densities by the factor
+        ...     new_densities = densities * density_factor
+        ...     # Update the material densities
+        ...     lib_mat.set_densities(nuclides, new_densities)
+        >>> integrator.add_keff_search_control(
+        ...     adjust_material_density,
+        ...     x0=0.8,
+        ...     x1=1.2,
+        ...     bracket=[0.2, 1.5],
+        ...     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/keff_search_control.py

@@ -0,0 +1,146 @@
+import openmc.lib
+from openmc.mpi import comm
+from typing import Callable
+from warnings import warn
+
+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 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 `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 search_for_keff(self, x, step_index):
+        """Perform keff search and update the atom density vector.
+
+        Parameters
+        ----------
+        x : list of numpy.ndarray
+            Current atom density vector (atoms per material)
+        step_index : int
+            Current depletion step index
+
+        Returns
+        -------
+        x : list of numpy.ndarray
+            Updated atom density vector
+        root : float
+            Parameter value that achieves target keff
+        """
+        root = self._search_for_keff()
+        x = self._update_vec(x)
+        return x, 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(model=self.operator.model) as session:
+            # 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 bound "
+                f"({self.search_kwargs['x_min']:.6f}). Clamping to bracket lower bound.",
+                UserWarning
+            )
+
+        elif root > self.search_kwargs['x_max']:
+            warn(
+                f"keff search result ({root:.6f}) is above the upper bracket bound "
+                f"({self.search_kwargs['x_max']:.6f}). Clamping to bracket upper bound.",
+                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.
+
+        This method synchronizes the atom densities across all MPI ranks by
+        broadcasting the number object from each rank and updating the x vector
+        with the current atom densities from openmc.lib.materials or AtomNumber object 
+        if the nuclide is not in openmc.lib.materials.
+
+        Parameters
+        ----------
+        x : list of numpy.ndarray
+            Atom density vector to update (atoms per material)
+
+        Returns
+        -------
+        list of numpy.ndarray
+            Updated atom density vector synchronized across all ranks
+        """
+        for rank in range(comm.size):
+            number_i = comm.bcast(self.operator.number, root=rank)
+
+            for mat_idx, mat in enumerate(number_i.materials):
+                for nuc in number_i.nuclides:
+                    if nuc in number_i.burnable_nuclides:
+                        nuc_idx = number_i.burnable_nuclides.index(nuc)
+                        volume = number_i.get_mat_volume(mat) # cm^3
+                        if nuc in openmc.lib.materials[int(mat)].nuclides:
+                            _nuc_idx = openmc.lib.materials[int(mat)].nuclides.index(nuc)
+                            val = 1.0e24 * openmc.lib.materials[int(mat)].densities[_nuc_idx] # atom/cm^3
+                        else:
+                            val = number_i.get_atom_density(mat, nuc) # atom/cm^3
+                        x[mat_idx][nuc_idx] = val * volume
+
+            x = comm.bcast(x, root=rank)
+        return x