diff --git a/openmc/deplete/keff_search_control.py b/openmc/deplete/keff_search_control.py new file mode 100644 index 0000000000..7e87df459b --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/regression_tests/deplete_with_keff_search_control/__init__.py b/tests/regression_tests/deplete_with_keff_search_control/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/regression_tests/deplete_with_keff_search_control/test.py b/tests/regression_tests/deplete_with_keff_search_control/test.py new file mode 100644 index 0000000000..aa1ec09429 --- /dev/null +++ b/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) diff --git a/tests/unit_tests/test_deplete_keff_search_control.py b/tests/unit_tests/test_deplete_keff_search_control.py new file mode 100644 index 0000000000..b55870194b --- /dev/null +++ b/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 +