diff --git a/docs/source/io_formats/depletion_results.rst b/docs/source/io_formats/depletion_results.rst index 7fb088268..856993db6 100644 --- a/docs/source/io_formats/depletion_results.rst +++ b/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//** diff --git a/openmc/deplete/abc.py b/openmc/deplete/abc.py index 7736b8c13..80d8474ec 100644 --- a/openmc/deplete/abc.py +++ b/openmc/deplete/abc.py @@ -31,6 +31,7 @@ from .results import Results, _SECONDS_PER_MINUTE, _SECONDS_PER_HOUR, \ 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 @@ class TransportOperator(ABC): 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 @@ class Integrator(ABC): 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 @@ class Integrator(ABC): 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 @@ class Integrator(ABC): 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 @@ class Integrator(ABC): self._i_res + i, proc_time, write_rates=write_rates, + keff_search_root=keff_search_root, path=path ) @@ -908,6 +939,10 @@ class Integrator(ABC): # 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 @@ class Integrator(ABC): 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 @@ class Integrator(ABC): 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 diff --git a/openmc/deplete/coupled_operator.py b/openmc/deplete/coupled_operator.py index 34bb28b49..a21d57d46 100644 --- a/openmc/deplete/coupled_operator.py +++ b/openmc/deplete/coupled_operator.py @@ -399,7 +399,7 @@ class CoupledOperator(OpenMCOperator): 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: diff --git a/openmc/deplete/independent_operator.py b/openmc/deplete/independent_operator.py index c192907cf..c12863956 100644 --- a/openmc/deplete/independent_operator.py +++ b/openmc/deplete/independent_operator.py @@ -384,7 +384,7 @@ class IndependentOperator(OpenMCOperator): # 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 diff --git a/openmc/deplete/keff_search_control.py b/openmc/deplete/keff_search_control.py new file mode 100644 index 000000000..49f7cc4df --- /dev/null +++ b/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 diff --git a/openmc/deplete/stepresult.py b/openmc/deplete/stepresult.py index 6da7f7adc..9f638a058 100644 --- a/openmc/deplete/stepresult.py +++ b/openmc/deplete/stepresult.py @@ -17,7 +17,7 @@ from openmc.mpi import MPI, comm from .reaction_rates import ReactionRates -VERSION_RESULTS = (1, 2) +VERSION_RESULTS = (1, 3) __all__ = ["StepResult"] @@ -58,6 +58,8 @@ class StepResult: proc_time : int Average time spent depleting a material across all materials and processes + keff_search_root : float + The root returned by the keff search control. """ def __init__(self): @@ -74,6 +76,7 @@ class StepResult: self.name_list = None self.data = None + self.keff_search_root = None def __repr__(self): t = self.time[0] @@ -368,6 +371,10 @@ class StepResult: "depletion time", (1,), maxshape=(None,), dtype="float64") + handle.create_dataset( + "keff_search_root", (1,), maxshape=(None,), + dtype="float64") + def _to_hdf5(self, handle, index, parallel=False, write_rates: bool = False): """Converts results object into an hdf5 object. @@ -400,6 +407,7 @@ class StepResult: time_dset = handle["/time"] source_rate_dset = handle["/source_rate"] proc_time_dset = handle["/depletion time"] + keff_search_root_dset = handle["/keff_search_root"] # Get number of results stored number_shape = list(number_dset.shape) @@ -433,6 +441,10 @@ class StepResult: proc_shape[0] = new_shape proc_time_dset.resize(proc_shape) + keff_search_root_shape = list(keff_search_root_dset.shape) + keff_search_root_shape[0] = new_shape + keff_search_root_dset.resize(keff_search_root_shape) + # If nothing to write, just return if len(self.index_mat) == 0: return @@ -452,6 +464,7 @@ class StepResult: proc_time_dset[index] = ( self.proc_time / (comm.size * self.n_hdf5_mats) ) + keff_search_root_dset[index] = self.keff_search_root @classmethod def from_hdf5(cls, handle, step): @@ -500,6 +513,10 @@ class StepResult: if step < proc_time_dset.shape[0]: results.proc_time = proc_time_dset[step] + if "keff_search_root" in handle: + keff_search_root_dset = handle["/keff_search_root"] + results.keff_search_root = keff_search_root_dset[step] + if results.proc_time is None: results.proc_time = np.array([np.nan]) @@ -554,6 +571,7 @@ class StepResult: step_ind, proc_time=None, write_rates: bool = False, + keff_search_root=None, path: PathLike = "depletion_results.h5" ): """Creates and writes depletion results to disk @@ -578,6 +596,8 @@ class StepResult: processes. write_rates : bool, optional Whether reaction rates should be written to the results file. + keff_search_root : float + The root returned by the keff search control. path : PathLike Path to file to write. Defaults to 'depletion_results.h5'. @@ -605,6 +625,7 @@ class StepResult: results.proc_time = proc_time if results.proc_time is not None: results.proc_time = comm.reduce(proc_time, op=MPI.SUM) + results.keff_search_root = keff_search_root if not Path(path).is_file(): Path(path).parent.mkdir(parents=True, exist_ok=True) 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 000000000..e69de29bb diff --git a/tests/regression_tests/deplete_with_keff_search_control/ref_depletion_with_refuel.h5 b/tests/regression_tests/deplete_with_keff_search_control/ref_depletion_with_refuel.h5 new file mode 100644 index 000000000..a335e3b47 Binary files /dev/null and b/tests/regression_tests/deplete_with_keff_search_control/ref_depletion_with_refuel.h5 differ diff --git a/tests/regression_tests/deplete_with_keff_search_control/ref_depletion_with_rotation.h5 b/tests/regression_tests/deplete_with_keff_search_control/ref_depletion_with_rotation.h5 new file mode 100644 index 000000000..c8b4f67ff Binary files /dev/null and b/tests/regression_tests/deplete_with_keff_search_control/ref_depletion_with_rotation.h5 differ diff --git a/tests/regression_tests/deplete_with_keff_search_control/ref_depletion_with_translation.h5 b/tests/regression_tests/deplete_with_keff_search_control/ref_depletion_with_translation.h5 new file mode 100644 index 000000000..35285321b Binary files /dev/null and b/tests/regression_tests/deplete_with_keff_search_control/ref_depletion_with_translation.h5 differ 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 000000000..82e601280 --- /dev/null +++ b/tests/regression_tests/deplete_with_keff_search_control/test.py @@ -0,0 +1,140 @@ +""" Tests for KeffSearchControl class """ + +from hmac import new +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 set_u235_density(u235_density): + fuel = [material for material in openmc.lib.materials.values() + if material.name == 'f'][0] + nuclides = openmc.lib.materials[fuel.id].nuclides + densities = openmc.lib.materials[fuel.id].densities + nuc_idx = nuclides.index('U235') + densities[nuc_idx] = u235_density + openmc.lib.materials[fuel.id].set_densities(nuclides, 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'), + (set_u235_density, 2e-4, 1e-3, (1e-4, 2e-3), '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' + model.settings.verbosity = 1 + op = CoupledOperator(model, chain_file) + + integrator = openmc.deplete.PredictorIntegrator( + op, [1], 174., timestep_units = 'd') + integrator.add_keff_search_control( + function=function, + x0=x0, + x1=x1, + bracket=bracket, + output=True, + k_tol=0.1, + sigma_final=5e-2) + + 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 000000000..425b8f840 --- /dev/null +++ b/tests/unit_tests/test_deplete_keff_search_control.py @@ -0,0 +1,116 @@ +""" Tests for KeffSearchControl class """ + +from pathlib import Path + +import pytest +import numpy as np + +import openmc +import openmc.lib +from openmc.deplete import CoupledOperator + +CHAIN_PATH = Path(__file__).parents[1] / "chain_simple.xml" + + +def make_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) + + +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 set_u235_density(u235_density): + """Helper function to set the U235 density directly""" + fuel = [m for m in openmc.lib.materials.values() if m.name == 'fuel'][0] + nuclides = openmc.lib.materials[fuel.id].nuclides + densities = openmc.lib.materials[fuel.id].densities + u235_idx = nuclides.index('U235') + densities[u235_idx] = u235_density + openmc.lib.materials[fuel.id].set_densities(nuclides, densities) + return u235_density + + +@pytest.mark.parametrize("function, x0, x1, bracket", [ + (translate_cell, -1.0, 1.0, (-5.0, 5.0)), + (rotate_cell, -45.0, 45.0, (-90.0, 90.0)), + (set_u235_density, 0.8, 1.2, (0.5, 1.5)) +]) +def test_integrator_add_keff_search_control(run_in_tmpdir, function, x0, x1, bracket): + """Test adding add_keff_search_control to integrator""" + model = make_model() + operator = CoupledOperator(model, CHAIN_PATH) + integrator = openmc.deplete.PredictorIntegrator( + operator, [1, 1], 0.0, timestep_units='d') + + integrator.add_keff_search_control( + function=function, + x0=x0, + x1=x1, + bracket=bracket, + k_tol=0.1, + output=False, + ) + + assert integrator._keff_search_control.x0 == x0 + assert integrator._keff_search_control.x1 == x1 + assert integrator._keff_search_control.function == function + assert integrator._keff_search_control.search_kwargs['x_min'] == bracket[0] + assert integrator._keff_search_control.search_kwargs['x_max'] == bracket[1] + assert integrator._keff_search_control.search_kwargs['k_tol'] == 0.1 + assert not integrator._keff_search_control.search_kwargs['output']