diff --git a/openmc/deplete/stepresult.py b/openmc/deplete/stepresult.py index 27420246f..cd21df07b 100644 --- a/openmc/deplete/stepresult.py +++ b/openmc/deplete/stepresult.py @@ -12,8 +12,9 @@ import h5py import numpy as np import openmc -from openmc.mpi import comm, MPI from openmc.checkvalue import PathLike +from openmc.mpi import MPI, comm + from .reaction_rates import ReactionRates VERSION_RESULTS = (1, 2) @@ -196,15 +197,15 @@ class StepResult: new.rates = self.rates[ranges] return new - def get_material(self, mat_id): + def get_material(self, mat_id: str | int) -> openmc.Material: """Return material object for given depleted composition .. versionadded:: 0.13.2 Parameters ---------- - mat_id : str - Material ID as a string + mat_id : str or int + Material ID as a string or integer Returns ------- @@ -217,6 +218,9 @@ class StepResult: If specified material ID is not found in the StepResult """ + # Coerce to str since internal dictionaries use str keys + mat_id = str(mat_id) + with warnings.catch_warnings(): warnings.simplefilter('ignore', openmc.IDWarning) material = openmc.Material(material_id=int(mat_id)) diff --git a/tests/unit_tests/test_deplete_resultslist.py b/tests/unit_tests/test_deplete_resultslist.py index 9a4699a4f..39c532c54 100644 --- a/tests/unit_tests/test_deplete_resultslist.py +++ b/tests/unit_tests/test_deplete_resultslist.py @@ -1,10 +1,11 @@ """Tests the Results class""" -from pathlib import Path from math import inf +from pathlib import Path import numpy as np import pytest + import openmc.deplete @@ -221,3 +222,11 @@ def test_stepresult_get_material(res): densities = mat1.get_nuclide_atom_densities() assert densities['Xe135'] == pytest.approx(1e-14) assert densities['U234'] == pytest.approx(1.00506e-05) + + +def test_stepresult_get_material_mat_id_as_int(res): + # Get material at first timestep using int mat_id + step_result = res[0] + mat1 = step_result.get_material(1) + assert mat1.id == 1 + assert mat1.volume == step_result.volume["1"]