diff --git a/openmc/deplete/abc.py b/openmc/deplete/abc.py index 2c6d177c4..b5cea425c 100644 --- a/openmc/deplete/abc.py +++ b/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 @@ class Integrator(ABC): 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 @@ class Integrator(ABC): 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 @@ 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: - # 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 @@ class Integrator(ABC): 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 @@ class Integrator(ABC): 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,23 +1154,22 @@ class Integrator(ABC): ... 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 ... ) @@ -1186,12 +1177,7 @@ class Integrator(ABC): """ 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): diff --git a/openmc/deplete/keff_search_control.py b/openmc/deplete/keff_search_control.py index fc8652154..49f7cc4df 100644 --- a/openmc/deplete/keff_search_control.py +++ b/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 @@ class _KeffSearchControl: 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 \ No newline at end of file + 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/model/model.py b/openmc/model/model.py index 6e4c1c585..c2d4d6915 100644 --- a/openmc/model/model.py +++ b/openmc/model/model.py @@ -1179,8 +1179,8 @@ class Model: # Convert ID map to RGB image img = id_map_to_rgb( - id_map=id_map, - color_by=color_by, + id_map=id_map, + color_by=color_by, colors=colors, overlap_color=overlap_color ) @@ -1217,7 +1217,7 @@ class Model: extent=(x_min, x_max, y_min, y_max), **contour_kwargs ) - + # If only showing outline, set the axis limits and aspect explicitly if outline == 'only': axes.set_xlim(x_min, x_max) diff --git a/tests/regression_tests/deplete_with_keff_search_control/test.py b/tests/regression_tests/deplete_with_keff_search_control/test.py index 0af6acd53..82e601280 100644 --- a/tests/regression_tests/deplete_with_keff_search_control/test.py +++ b/tests/regression_tests/deplete_with_keff_search_control/test.py @@ -79,32 +79,35 @@ def model(): 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 - densities = openmc.lib.materials[f.id].densities + +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') - new_density = densities[nuc_idx] * density_factor - densities[nuc_idx] = new_density - openmc.lib.materials[f.id].set_densities(nuclides, densities) + 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'), - (adjust_material_density, 0.5, 2, [0.3, 3.0], 'depletion_with_refuel') - ]) - + (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( @@ -114,10 +117,10 @@ def test_keff_search_control(run_in_tmpdir, model, function, x0, x1, bracket, re x0=x0, x1=x1, bracket=bracket, - output=True, - k_tol=1e-1, + output=True, + k_tol=0.1, sigma_final=5e-2) - + integrator.integrate() # Get path to test and reference results diff --git a/tests/unit_tests/test_deplete_keff_search_control.py b/tests/unit_tests/test_deplete_keff_search_control.py index 12b4020f0..425b8f840 100644 --- a/tests/unit_tests/test_deplete_keff_search_control.py +++ b/tests/unit_tests/test_deplete_keff_search_control.py @@ -11,8 +11,8 @@ from openmc.deplete import CoupledOperator CHAIN_PATH = Path(__file__).parents[1] / "chain_simple.xml" -@pytest.fixture -def model(): + +def make_model(): f = openmc.Material(name="fuel") f.add_element("U", 1, percent_type="ao", enrichment=4.25) f.add_element("O", 2) @@ -60,14 +60,6 @@ def model(): 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 translate_cell(position): """Helper function to translate a cell""" @@ -75,47 +67,50 @@ def translate_cell(position): 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""" + +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 - 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 + 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, 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) + +@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, model, operator, integrator, - function, x0, x1, bracket, test_value): +def test_integrator_add_keff_search_control(run_in_tmpdir, function, x0, x1, bracket): """Test adding add_keff_search_control to integrator""" - model.export_to_xml() - openmc.lib.init() - test_function = function(test_value) - # Test that add_reactivity_control method exists and works + 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=test_function, + 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 == test_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 integrator.keff_search_control.search_kwargs['output'] == False - openmc.lib.finalize() + + 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']