Merge pull request #55 from paulromano/keff-search-fixes

Fixes for PR 2693

Author: church89

openmc/deplete/abc.py

@@ -618,11 +618,6 @@ 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
 
     """)
 
@@ -736,15 +731,6 @@ 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(
@@ -853,12 +839,36 @@ 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 _restore_keff_search_control(self, res):
+        """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(res)
+            else:
+                keff_search_root = None
+
+        return bos_conc, res, keff_search_root
 
     def integrate(
             self,
@@ -898,21 +908,9 @@ 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:
-                    # 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
+                # 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)
 
@@ -940,7 +938,7 @@ def integrate(
             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)
+                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)
@@ -1091,40 +1089,34 @@ def add_keff_search_control(
         function: Callable,
         x0: float,
         x1: float,
-        bracket: list[float],
+        bracket: Sequence[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.
+        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
+            ``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
+            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
+        x1 : float
             Initial upper bound for the keff search.
-        bracket : list[float]
+        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 
+            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
@@ -1162,36 +1154,30 @@ def add_keff_search_control(
         ...     k_tol=1e-4
         ... )
 
-        Add keff search that adjusts material density:
+        Add keff search that adjusts the U235 density:
 
-        >>> def adjust_material_density(density_factor):
+        >>> 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
-        ...     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)
+        ...     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(
-        ...     adjust_material_density,
-        ...     x0=0.8,
-        ...     x1=1.2,
-        ...     bracket=[0.2, 1.5],
+        ...     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)
+            self.operator, function, x0, x1, bracket, **search_kwargs)
 
 @add_params
 class SIIntegrator(Integrator):

openmc/deplete/keff_search_control.py

@@ -1,14 +1,15 @@
-import openmc.lib
-from openmc.mpi import comm
 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.
-    
+
+    This class performs keff searches to maintain a target keff by adjusting a
+    model parameter through a provided function.
+
     Parameters
     ----------
     operator : openmc.deplete.Operator
@@ -20,11 +21,12 @@ class _KeffSearchControl:
     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`
+        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:
@@ -39,108 +41,88 @@ def __init__(self, operator, function: Callable, x0: float, x1: float, bracket:
         self.search_kwargs['x_min'] = bracket[0]
         self.search_kwargs['x_max'] = bracket[1]
 
-    def search_for_keff(self, x, step_index):
+    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)
-        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
-        
+        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(model=self.operator.model) as session:
+        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
+                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
-            )
-
+            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 bound "
-                f"({self.search_kwargs['x_max']:.6f}). Clamping to bracket upper bound.",
-                UserWarning
-            )
-        
-        #restore the number of initial batches
+            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.
-        
-        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.
-        
+
+        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)
-            
-        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
+        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