Fix restart with saved keff search

This commit is contained in:
Paul Romano 2026-03-11 23:52:52 -05:00
parent 5af848f6fd
commit 821bb5ab65
5 changed files with 63 additions and 47 deletions

View file

@ -903,9 +903,16 @@ class Integrator(ABC):
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)
if self._keff_search_control is not None and source_rate != 0.0:
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.apply_stored_root(
keff_search_root)
else:
keff_search_root = None
# Solve Bateman equations over time interval
@ -1100,7 +1107,9 @@ class Integrator(ABC):
``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.
(``openmc.lib``) is available for modifications. The function
should behave like a setter for the controlled parameter so that a
stored keff-search root can be safely reapplied during restart.
Parameters
----------
@ -1112,7 +1121,9 @@ class Integrator(ABC):
density via ``openmc.lib.materials[...].set_densities(...)``, etc.).
**Important**: The function must modify ``openmc.lib`` objects, not
``openmc.model`` objects.
``openmc.model`` objects, and should set the controlled parameter
directly rather than modifying the current state relative to its
existing value.
x0: float
Initial lower bound for the keff search.
x1: float
@ -1157,23 +1168,22 @@ class Integrator(ABC):
... k_tol=1e-4
... )
Add keff search that adjusts material density:
Add keff search that sets the U235 density directly:
>>> 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
... )

View file

@ -62,6 +62,16 @@ class _KeffSearchControl:
x = self._update_vec(x)
return x, root
def apply_stored_root(self, root: float):
"""Reapply a previously stored control parameter value.
Parameters
----------
root : float
Parameter value from a prior keff search result.
"""
self.function(root)
def _search_for_keff(self) -> float:
"""Perform the keff search using the model's keff_search method.

View file

@ -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)

View file

@ -90,24 +90,24 @@ def rotate_cell(angle):
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')
(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(
@ -118,7 +118,7 @@ def test_keff_search_control(run_in_tmpdir, model, function, x0, x1, bracket, re
x1=x1,
bracket=bracket,
output=True,
k_tol=1e-1,
k_tol=0.1,
sigma_final=5e-2)
integrator.integrate()

View file

@ -84,30 +84,27 @@ def rotate_cell(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):
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
integrator.add_keff_search_control(
function=test_function,
function=function,
x0=x0,
x1=x1,
bracket=bracket,
@ -117,9 +114,8 @@ def test_integrator_add_keff_search_control(run_in_tmpdir, model, operator, inte
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.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']
openmc.lib.finalize()