OpenMC/tests/unit_tests/test_deplete_integrator.py
Perry 2d5c50080c
Allow the use of substeps for CRAM (#3908)
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
2026-04-22 18:24:34 -05:00

375 lines
12 KiB
Python

"""Tests for saving results
It is worth noting that openmc.deplete.integrate is extremely complex, to the
point I am unsure if it can be reasonably unit-tested. For the time being, it
will be left unimplemented and testing will be done via regression.
"""
import copy
from random import uniform
from unittest.mock import MagicMock
import h5py
import numpy as np
from uncertainties import ufloat
import pytest
from openmc.mpi import comm
from openmc.deplete import (
ReactionRates, StepResult, Results, OperatorResult, PredictorIntegrator,
CECMIntegrator, CF4Integrator, CELIIntegrator, EPCRK4Integrator,
LEQIIntegrator, SICELIIntegrator, SILEQIIntegrator, cram, pool)
from tests import dummy_operator
INTEGRATORS = [
PredictorIntegrator,
CECMIntegrator,
CF4Integrator,
CELIIntegrator,
EPCRK4Integrator,
LEQIIntegrator,
SICELIIntegrator,
SILEQIIntegrator
]
def test_results_save(run_in_tmpdir):
"""Test data save module"""
rng = np.random.RandomState(comm.rank)
# Mock geometry
op = MagicMock()
# Avoid DummyOperator thinking it's doing a restart calculation
op.prev_res = None
vol_dict = {}
full_burn_list = []
for i in range(comm.size):
vol_dict[str(2*i)] = 1.2
vol_dict[str(2*i + 1)] = 1.2
full_burn_list.append(str(2*i))
full_burn_list.append(str(2*i + 1))
burn_list = full_burn_list[2*comm.rank: 2*comm.rank + 2]
nuc_list = ["na", "nb"]
name_list = {mat: "" for mat in full_burn_list}
op.get_results_info.return_value = (
vol_dict, nuc_list, burn_list, full_burn_list, name_list)
# Construct end-of-step concentrations
x1 = [rng.random(2), rng.random(2)]
x2 = [rng.random(2), rng.random(2)]
# Construct reaction rates
r1 = ReactionRates(burn_list, ["na", "nb"], ["ra", "rb"])
r1[:] = rng.random((2, 2, 2))
rate1 = copy.deepcopy(r1)
r2 = ReactionRates(burn_list, ["na", "nb"], ["ra", "rb"])
r2[:] = rng.random((2, 2, 2))
rate2 = copy.deepcopy(r2)
# Create global terms (eigenvalue and uncertainty)
eigvl1 = rng.random(2)
eigvl2 = rng.random(2)
eigvl1 = comm.bcast(eigvl1, root=0)
eigvl2 = comm.bcast(eigvl2, root=0)
t1 = [0.0, 1.0]
t2 = [1.0, 2.0]
op_result1 = OperatorResult(ufloat(*eigvl1), rate1)
op_result2 = OperatorResult(ufloat(*eigvl2), rate2)
# saves within a subdirectory
StepResult.save(
op,
x1,
op_result1,
t1,
0,
0,
write_rates=True,
path='out/put/depletion.h5'
)
res = Results('out/put/depletion.h5')
# saves with default filename
StepResult.save(op, x1, op_result1, t1, 0, 0, write_rates=True)
StepResult.save(op, x2, op_result2, t2, 0, 1, write_rates=True)
# Load the files
res = Results("depletion_results.h5")
for mat_i, mat in enumerate(burn_list):
for nuc_i, nuc in enumerate(nuc_list):
assert res[0][mat, nuc] == x1[mat_i][nuc_i]
assert res[1][mat, nuc] == x2[mat_i][nuc_i]
np.testing.assert_array_equal(res[0].rates, rate1)
np.testing.assert_array_equal(res[1].rates, rate2)
np.testing.assert_array_equal(res[0].k, eigvl1)
np.testing.assert_array_equal(res[0].time, t1)
np.testing.assert_array_equal(res[1].k, eigvl2)
np.testing.assert_array_equal(res[1].time, t2)
def test_results_save_without_rates(run_in_tmpdir):
"""StepResult.save skips reaction-rate datasets by default"""
op = MagicMock()
op.prev_res = None
vol_dict = {"0": 1.0}
nuc_list = ["na"]
burn_list = ["0"]
name_list = {mat: "" for mat in burn_list}
op.get_results_info.return_value = (vol_dict, nuc_list, burn_list, burn_list, name_list)
x = [np.array([1.0])]
rates = ReactionRates(burn_list, nuc_list, ["ra"])
rates[:] = np.array([[[2.0]]])
op_result = OperatorResult(ufloat(1.0, 0.1), rates)
StepResult.save(op, x, op_result, [0.0, 1.0], 0.0, 0)
with h5py.File('depletion_results.h5', 'r') as handle:
assert 'reaction rates' not in handle
assert 'reactions' not in handle
res = Results('depletion_results.h5')
assert res[0].rates.size == 0
def test_bad_integrator_inputs():
"""Test failure modes for Integrator inputs"""
op = MagicMock()
op.prev_res = None
op.chain = None
op.heavy_metal = 1.0
timesteps = [1]
# No power nor power density given
with pytest.raises(ValueError, match="Either power"):
PredictorIntegrator(op, timesteps)
# Length of power != length time
with pytest.raises(ValueError, match="number of powers"):
PredictorIntegrator(op, timesteps, power=[1, 2])
# Length of power density != length time
with pytest.raises(ValueError, match="number of powers"):
PredictorIntegrator(op, timesteps, power_density=[1, 2])
# SI integrator with bad steps
with pytest.raises(TypeError, match="n_steps"):
SICELIIntegrator(op, timesteps, [1], n_steps=2.5)
with pytest.raises(ValueError, match="n_steps"):
SICELIIntegrator(op, timesteps, [1], n_steps=0)
with pytest.raises(ValueError, match="Solver failure"):
PredictorIntegrator(op, timesteps, power=1, solver="failure")
with pytest.raises(TypeError, match=".*callable.*NoneType"):
PredictorIntegrator(op, timesteps, power=1, solver=None)
with pytest.raises(ValueError, match="four arguments"):
PredictorIntegrator(op, timesteps, power=1, solver=mock_bad_solver_nargs)
with pytest.raises(ValueError, match="default to 1"):
PredictorIntegrator(op, timesteps, power=1,
solver=mock_bad_solver_fourth_required)
with pytest.raises(ValueError, match="substeps"):
PredictorIntegrator(op, timesteps, power=1, substeps=0)
with pytest.raises(ValueError, match="substeps"):
PredictorIntegrator(op, timesteps, power=1, substeps=-1)
def mock_good_solver(A, n, t, substeps=1):
return n.copy()
def mock_good_solver_substeps(A, n, t, substeps=1):
return n + substeps
def mock_unsupported_substeps_solver(A, n, t, substeps=1):
if substeps > 1:
raise NotImplementedError("substeps > 1 not supported")
return n.copy()
def mock_bad_solver_nargs(A, n):
pass
def mock_bad_solver_fourth_required(A, n, t, substeps):
pass
@pytest.mark.parametrize("scheme", dummy_operator.SCHEMES)
def test_integrator(run_in_tmpdir, scheme):
"""Test the integrators against their expected values"""
bundle = dummy_operator.SCHEMES[scheme]
operator = dummy_operator.DummyOperator()
bundle.solver(operator, [0.75, 0.75], 1.0).integrate()
# get expected results
res = Results(operator.output_dir / "depletion_results.h5")
t1, y1 = res.get_atoms("1", "1")
t2, y2 = res.get_atoms("1", "2")
assert (t1 == [0.0, 0.75, 1.5]).all()
assert y1 == pytest.approx(bundle.atoms_1)
assert (t2 == [0.0, 0.75, 1.5]).all()
assert y2 == pytest.approx(bundle.atoms_2)
# test structure of depletion time dataset
dep_time = res.get_depletion_time()
assert dep_time.shape == (2, )
assert all(dep_time > 0)
integrator = bundle.solver(operator, [0.75], 1, solver=cram.CRAM48)
assert integrator.solver is cram.CRAM48
integrator = bundle.solver(operator, [0.75], 1, solver="cram16")
assert integrator.solver is cram.CRAM16
integrator = bundle.solver(operator, [0.75], 1, solver=cram.Cram48Solver,
substeps=2)
assert integrator.solver is cram.Cram48Solver
assert integrator.substeps == 2
integrator.solver = mock_good_solver
assert integrator.solver is mock_good_solver
lfunc = lambda A, n, t, substeps=1: mock_good_solver(A, n, t, substeps)
integrator.solver = lfunc
assert integrator.solver is lfunc
integrator.solver = mock_good_solver_substeps
assert integrator.solver is mock_good_solver_substeps
def test_custom_solver_with_default_substeps(monkeypatch):
operator = dummy_operator.DummyOperator()
n = operator.initial_condition()
rates = operator(n, 1.0).rates
integrator = PredictorIntegrator(
operator, [0.75], power=1.0, solver=mock_good_solver)
monkeypatch.setattr(pool, "USE_MULTIPROCESSING", False)
_, result = integrator._timed_deplete(n, rates, 0.75)
np.testing.assert_array_equal(result[0], n[0])
def test_substep_aware_custom_solver_receives_substeps(monkeypatch):
operator = dummy_operator.DummyOperator()
n = operator.initial_condition()
rates = operator(n, 1.0).rates
integrator = PredictorIntegrator(
operator, [0.75], power=1.0, solver=mock_good_solver_substeps,
substeps=3)
monkeypatch.setattr(pool, "USE_MULTIPROCESSING", False)
_, result = integrator._timed_deplete(n, rates, 0.75)
np.testing.assert_array_equal(result[0], n[0] + 3)
def test_custom_solver_propagates_substeps_error(monkeypatch):
operator = dummy_operator.DummyOperator()
n = operator.initial_condition()
rates = operator(n, 1.0).rates
integrator = PredictorIntegrator(
operator, [0.75], power=1.0,
solver=mock_unsupported_substeps_solver, substeps=2)
monkeypatch.setattr(pool, "USE_MULTIPROCESSING", False)
with pytest.raises(NotImplementedError, match="not supported"):
integrator._timed_deplete(n, rates, 0.75)
def test_custom_solver_requires_four_args():
op = MagicMock()
op.prev_res = None
op.chain = None
op.heavy_metal = 1.0
with pytest.raises(ValueError, match="four arguments"):
PredictorIntegrator(op, [1], power=1, solver=mock_bad_solver_nargs)
@pytest.mark.parametrize("integrator", INTEGRATORS)
def test_timesteps(integrator):
# Crate fake operator
op = MagicMock()
op.prev_res = None
op.chain = None
# Set heavy metal mass and power randomly
op.heavy_metal = uniform(0, 10000)
power = uniform(0, 1e6)
# Reference timesteps in seconds
day = 86400.0
ref_timesteps = [1*day, 2*day, 5*day, 10*day]
# Case 1, timesteps in seconds
timesteps = ref_timesteps
x = integrator(op, timesteps, power, timestep_units='s')
assert np.allclose(x.timesteps, ref_timesteps)
# Case 2, timesteps in minutes
minute = 60
timesteps = [t / minute for t in ref_timesteps]
x = integrator(op, timesteps, power, timestep_units='min')
assert np.allclose(x.timesteps, ref_timesteps)
# Case 3, timesteps in hours
hour = 60*60
timesteps = [t / hour for t in ref_timesteps]
x = integrator(op, timesteps, power, timestep_units='h')
assert np.allclose(x.timesteps, ref_timesteps)
# Case 4, timesteps in days
timesteps = [t / day for t in ref_timesteps]
x = integrator(op, timesteps, power, timestep_units='d')
assert np.allclose(x.timesteps, ref_timesteps)
# Case 5, timesteps in MWd/kg
kilograms = op.heavy_metal / 1000.0
days = [t/day for t in ref_timesteps]
megawatts = power / 1000000.0
burnup = [t * megawatts / kilograms for t in days]
x = integrator(op, burnup, power, timestep_units='MWd/kg')
assert np.allclose(x.timesteps, ref_timesteps)
# Case 6, mixed units
burnup_per_day = (1e-6*power) / kilograms
timesteps = [(burnup_per_day, 'MWd/kg'), (2*day, 's'), (5, 'd'),
(10*burnup_per_day, 'MWd/kg')]
x = integrator(op, timesteps, power)
assert np.allclose(x.timesteps, ref_timesteps)
# Bad units should raise an exception
with pytest.raises(ValueError, match="unit"):
integrator(op, ref_timesteps, power, timestep_units='🐨')
with pytest.raises(ValueError, match="unit"):
integrator(op, [(800.0, 'gorillas')], power)