add renamed files

openmc/deplete/keff_search_control.py

@@ -0,0 +1,141 @@
+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.termination_reason}"
+            )
+        
+        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 the operator's number 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 and volumes.
+        
+        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)
+                        atom_density = number_i.get_atom_density(mat, nuc)
+                        volume = number_i.get_mat_volume(mat)
+                        x[mat_idx][nuc_idx] = atom_density * volume
+        
+            x = comm.bcast(x, root=rank)
+        return x

tests/regression_tests/deplete_with_keff_search_control/test.py

@@ -0,0 +1,130 @@
+""" Tests for KeffSearchControl class """
+
+from pathlib import Path
+import shutil
+import sys
+
+import pytest
+import numpy as np
+
+import openmc
+import openmc.lib
+from openmc.deplete import CoupledOperator
+
+from tests.regression_tests import config
+
+@pytest.fixture
+def model():
+    f = openmc.Material(name='f')
+    f.set_density('g/cm3', 10.29769)
+    f.add_element('U', 1., enrichment=2.4)
+    f.add_element('O', 2.)
+
+    h = openmc.Material(name='h')
+    h.set_density('g/cm3', 0.001598)
+    h.add_element('He', 2.4044e-4)
+
+    w = openmc.Material(name='w')
+    w.set_density('g/cm3', 0.740582)
+    w.add_element('H', 2)
+    w.add_element('O', 1)
+
+    # Define overall material
+    materials = openmc.Materials([f, h, w])
+
+    # Define surfaces
+    radii = [0.5, 0.8, 1]
+    height = 80
+    surf_in = openmc.ZCylinder(r=radii[0])
+    surf_mid = openmc.ZCylinder(r=radii[1])
+    surf_out = openmc.ZCylinder(r=radii[2], boundary_type='reflective')
+    surf_top = openmc.ZPlane(z0=height/2, boundary_type='vacuum')
+    surf_bot = openmc.ZPlane(z0=-height/2, boundary_type='vacuum')
+
+    surf_trans = openmc.ZPlane(z0=0)
+    surf_rot1 = openmc.XPlane(x0=0)
+    surf_rot2 = openmc.YPlane(y0=0)
+
+    # Define cells
+    cell_f = openmc.Cell(name='fuel_cell', fill=f,
+            region=-surf_in  & -surf_top & +surf_bot)
+    cell_g = openmc.Cell(fill=h,
+            region = +surf_in & -surf_mid & -surf_top & +surf_bot & +surf_rot2)
+
+    # Define unbounded cells for rotation universe
+    cell_w = openmc.Cell(fill=w, region = -surf_rot1)
+    cell_h = openmc.Cell(fill=h, region = +surf_rot1)
+    universe_rot = openmc.Universe(cells=(cell_w, cell_h))
+    cell_rot = openmc.Cell(name="rot_cell", fill=universe_rot,
+            region = +surf_in & -surf_mid & -surf_top & +surf_bot & -surf_rot2)
+
+    # Define unbounded cells for translation universe
+    cell_w = openmc.Cell(fill=w, region=+surf_in  & -surf_trans )
+    cell_h = openmc.Cell(fill=h,  region=+surf_in & +surf_trans)
+    universe_trans = openmc.Universe(cells=(cell_w, cell_h))
+    cell_trans = openmc.Cell(name="trans_cell", fill=universe_trans,
+                          region=+surf_mid & -surf_out & -surf_top & +surf_bot)
+
+    # Define overall geometry
+    geometry = openmc.Geometry([cell_f, cell_g, cell_rot, cell_trans])
+
+    # Set material volume for depletion fuel.
+    f.volume = np.pi * radii[0]**2 * height
+
+    settings = openmc.Settings()
+    settings.particles = 1000
+    settings.inactive = 10
+    settings.batches = 50
+
+    return openmc.Model(geometry, materials, settings)
+
+def translate_cell(position):
+    cell_trans = [cell for cell in openmc.lib.cells.values() if cell.name == 'trans_cell'][0]
+    openmc.lib.cells[cell_trans.id].translation = [0, 0, position]
+
+def rotate_cell(angle):
+    cell_rot = [cell for cell in openmc.lib.cells.values() if cell.name == 'rot_cell'][0]
+    openmc.lib.cells[cell_rot.id].rotation = [0, 0, angle]
+
+def adjust_material_density(density_factor):
+    f = [material for material in openmc.lib.materials.values() if material.name == 'f'][0]
+    nuclides = openmc.lib.materials[f.id].nuclides
+    current_densities = openmc.lib.materials[f.id].densities
+    new_densities = [d * density_factor for d in current_densities]
+    openmc.lib.materials[f.id].set_densities(nuclides, new_densities)
+
+@pytest.mark.parametrize("function, x0, x1, bracket, ref_result", [
+    (translate_cell, -11, -5, [-15,0], 'depletion_with_translation'),
+    (rotate_cell, -80, -50, [-90,0], 'depletion_with_rotation'),
+    (adjust_material_density, 0.5, 1.2, [0.2, 1.5], 'depletion_with_refuel')
+    ])
+
+def test_keff_search_control(run_in_tmpdir, model, function, x0, x1, bracket, ref_result):
+
+    chain_file = Path(__file__).parents[2] / 'chain_simple.xml'
+    op = CoupledOperator(model, chain_file)
+
+    integrator = openmc.deplete.PredictorIntegrator(
+        op, [1], 174., timestep_units = 'd')
+    integrator.add_keff_search_control(function, x0, x1, bracket, 
+        output=True, 
+        k_tol=1e-2, 
+        sigma_final=5e-3)
+    
+    integrator.integrate()
+
+    # Get path to test and reference results
+    path_test = op.output_dir / 'depletion_results.h5'
+    path_reference = Path(__file__).with_name(f'ref_{ref_result}.h5')
+
+    # If updating results, do so and return
+    if config['update']:
+        shutil.copyfile(str(path_test), str(path_reference))
+        return
+
+    # Load the reference/test results
+    res_test = openmc.deplete.Results(path_test)
+    res_ref = openmc.deplete.Results(path_reference)
+
+    # Use high tolerance here
+    assert res_test[0].keff_search_root == pytest.approx(res_ref[0].keff_search_root, rel=2)

tests/unit_tests/test_deplete_keff_search_control.py

@@ -0,0 +1,200 @@
+""" Tests for KeffSearchControl class """
+
+from pathlib import Path
+
+import pytest
+import numpy as np
+
+import openmc
+import openmc.lib
+from openmc.deplete import CoupledOperator
+from openmc.deplete import KeffSearchControl
+
+CHAIN_PATH = Path(__file__).parents[1] / "chain_simple.xml"
+
+@pytest.fixture
+def model():
+    f = openmc.Material(name="fuel")
+    f.add_element("U", 1, percent_type="ao", enrichment=4.25)
+    f.add_element("O", 2)
+    f.set_density("g/cc", 10.4)
+    f.temperature = 293.15
+
+    w = openmc.Material(name="water")
+    w.add_element("O", 1)
+    w.add_element("H", 2)
+    w.set_density("g/cc", 1.0)
+    w.temperature = 293.15
+    w.depletable = True
+
+    h = openmc.Material(name='helium')
+    h.add_element('He', 1)
+    h.set_density('g/cm3', 0.001598)
+
+    radii = [0.42, 0.45]
+    height = 0.5
+
+    f.volume = np.pi * radii[0] ** 2 * height
+    w.volume = np.pi * (radii[1]**2 - radii[0]**2) * height/2
+
+    materials = openmc.Materials([f, w, h])
+
+    surf_interface = openmc.ZPlane(z0=0)
+    surf_top = openmc.ZPlane(z0=height/2)
+    surf_bot = openmc.ZPlane(z0=-height/2)
+    surf_in = openmc.Sphere(r=radii[0])
+    surf_out = openmc.Sphere(r=radii[1], boundary_type='vacuum')
+
+    cell_water = openmc.Cell(fill=w, region=-surf_interface)
+    cell_helium = openmc.Cell(fill=h, region=+surf_interface)
+    universe = openmc.Universe(cells=(cell_water, cell_helium))
+    cell_fuel = openmc.Cell(name='fuel_cell', fill=f,
+                            region=-surf_in & -surf_top & +surf_bot)
+    cell_universe = openmc.Cell(name='universe_cell',fill=universe,
+                            region=+surf_in & -surf_out & -surf_top & +surf_bot)
+    geometry = openmc.Geometry([cell_fuel, cell_universe])
+
+    settings = openmc.Settings()
+    settings.particles = 1000
+    settings.inactive = 10
+    settings.batches = 50
+
+    return openmc.Model(geometry, materials, settings)
+
+@pytest.fixture
+def operator(model):
+    return CoupledOperator(model, CHAIN_PATH)
+
+@pytest.fixture
+def integrator(operator):
+    return openmc.deplete.PredictorIntegrator(
+            operator, [1,1], 0.0, timestep_units = 'd')
+
+def test_keff_search_control_init(operator):
+    """Test KeffSearchControl initialization"""
+    def dummy_function(x):
+        return x
+    
+    # Valid initialization
+    controller = KeffSearchControl(
+        operator=operator,
+        function=dummy_function,
+        x0=0.0,
+        x1=1.0,
+        bracket=[0.0, 2.0]
+    )
+    
+    assert controller.operator == operator
+    assert controller.function == dummy_function
+    assert controller.x0 == 0.0
+    assert controller.x1 == 1.0
+    assert controller.search_kwargs['x_min'] == 0.0
+    assert controller.search_kwargs['x_max'] == 2.0
+    
+    # Test invalid bracket with wrong length
+    with pytest.raises(ValueError, match="bracket must have exactly 2 elements"):
+        KeffSearchControl(operator, dummy_function, 0.0, 1.0, [0.0])
+    
+    # Test invalid bracket with wrong order
+    with pytest.raises(ValueError, match="bracket\\[0\\] must be < bracket\\[1\\]"):
+        KeffSearchControl(operator, dummy_function, 0.0, 1.0, [2.0, 1.0])
+
+def translate_cell(position):
+    """Helper function to translate a cell"""
+    cell = [c for c in openmc.lib.cells.values() if c.name == 'universe_cell'][0]
+    openmc.lib.cells[cell.id].translation = [0, 0, position]
+    return position
+
+def rotate_cell(angle):
+    """Helper function to rotate a cell"""
+    cell = [c for c in openmc.lib.cells.values() if c.name == 'universe_cell'][0]
+    openmc.lib.cells[cell.id].rotation = [0, 0, angle]
+    return angle
+
+def adjust_fuel_density(density_factor):
+    """Helper function to adjust fuel density"""
+    fuel = [m for m in openmc.lib.materials.values() if m.name == 'fuel'][0]
+    nuclides = openmc.lib.materials[fuel.id].nuclides
+    current_densities = openmc.lib.materials[fuel.id].densities
+    new_densities = [d * density_factor for d in current_densities]
+    openmc.lib.materials[fuel.id].set_densities(nuclides, new_densities)
+    return density_factor
+
+@pytest.mark.parametrize("function, x0, x1, bracket, test_value", [
+    (translate_cell, -1.0, 1.0, [-5.0, 5.0], 0.5),
+    (rotate_cell, -45.0, 45.0, [-90.0, 90.0], 10.0),
+    (adjust_fuel_density, 0.8, 1.2, [0.5, 1.5], 1.0)
+])
+def test_keff_search_control_with_openmc_functions(
+    run_in_tmpdir, model, operator, function, x0, x1, bracket, test_value
+):
+    """Test _KeffSearchControl with actual OpenMC geometry modification functions"""
+    
+    controller = KeffSearchControl(
+        operator=operator,
+        function=function,
+        x0=x0,
+        x1=x1,
+        bracket=bracket,
+        output=False, 
+        k_tol=0.1,
+    )
+    
+    # Export model and initialize OpenMC library
+    model.export_to_xml()
+    openmc.lib.init()
+    openmc.lib.finalize()
+
+def test_controller_kwargs_handling(operator):
+    """Test that additional kwargs are properly stored and accessible"""
+    def dummy_func(x):
+        return x
+    
+    custom_kwargs = {
+        'target': 1.0,
+        'k_tol': 1e-4,
+        'output': True
+    }
+    
+    controller = KeffSearchControl(
+        operator=operator,
+        function=dummy_func,
+        x0=0.0,
+        x1=1.0,
+        bracket=[0.0, 2.0],
+        **custom_kwargs
+    )
+    
+    # Check that bracket limits were added to kwargs
+    assert controller.search_kwargs['x_min'] == 0.0
+    assert controller.search_kwargs['x_max'] == 2.0
+    
+    # Check that custom kwargs were preserved
+    assert controller.search_kwargs['target'] == 1.0
+    assert controller.search_kwargs['k_tol'] == 1e-4
+    assert controller.search_kwargs['output'] == True
+
+def test_integrator_add_keff_search_control(run_in_tmpdir, model, operator, integrator):
+    """Test adding keff search control to integrator using the add_keff_search_control method"""
+    
+    def adjust_density(factor):
+        # This would modify material density in practice
+        return factor
+    
+    # Test that add_reactivity_control method exists and works
+    integrator.add_keff_search_control(
+        function=adjust_density,
+        x0=0.8,
+        x1=1.2,
+        bracket=[0.5, 1.5],
+        k_tol=0.1,
+        output=False,
+    )
+    
+    assert integrator.keff_search_control is not None
+    assert isinstance(integrator.keff_search_control, KeffSearchControl)
+    assert integrator.keff_search_control.x0 == 0.8
+    assert integrator.keff_search_control.x1 == 1.2
+    assert integrator.keff_search_control.search_kwargs['x_min'] == 0.5
+    assert integrator.keff_search_control.search_kwargs['x_max'] == 1.5   
+