diff --git a/openmc/statepoint.py b/openmc/statepoint.py index 29c11921cb..a763db3971 100644 --- a/openmc/statepoint.py +++ b/openmc/statepoint.py @@ -536,7 +536,7 @@ class StatePoint: def get_tally(self, scores=[], filters=[], nuclides=[], name=None, id=None, estimator=None, exact_filters=False, exact_nuclides=False, exact_scores=False, - multiply_density=None, derivative=None): + multiply_density=None, derivative=None, filter_type=None): """Finds and returns a Tally object with certain properties. This routine searches the list of Tallies and returns the first Tally @@ -580,6 +580,9 @@ class StatePoint: to the same value as this parameter. derivative : openmc.TallyDerivative, optional TallyDerivative object to match. + filter_type : type, optional + If not None, the Tally must have at least one Filter that is an + instance of this type. For example `openmc.MeshFilter`. Returns ------- @@ -653,6 +656,10 @@ class StatePoint: if not contains_filters: continue + if filter_type is not None: + if not any(isinstance(f, filter_type) for f in test_tally.filters): + continue + # Determine if Tally has the queried Nuclide(s) if nuclides: if not all(nuclide in test_tally.nuclides for nuclide in nuclides): diff --git a/tests/unit_tests/test_statepoint.py b/tests/unit_tests/test_statepoint.py new file mode 100644 index 0000000000..7ffaf7ec2c --- /dev/null +++ b/tests/unit_tests/test_statepoint.py @@ -0,0 +1,65 @@ +import openmc + + +def test_get_tally_filter_type(run_in_tmpdir): + """Test various ways of retrieving tallies from a StatePoint object.""" + + mat = openmc.Material() + mat.add_nuclide("H1", 1.0) + mat.set_density("g/cm3", 10.0) + + sphere = openmc.Sphere(r=10.0, boundary_type="vacuum") + cell = openmc.Cell(fill=mat, region=-sphere) + geometry = openmc.Geometry([cell]) + + settings = openmc.Settings() + settings.particles = 10 + settings.batches = 2 + settings.run_mode = "fixed source" + + reg_mesh = openmc.RegularMesh().from_domain(cell) + tally1 = openmc.Tally(tally_id=1) + mesh_filter = openmc.MeshFilter(reg_mesh) + tally1.filters = [mesh_filter] + tally1.scores = ["flux"] + + tally2 = openmc.Tally(tally_id=2, name="heating tally") + cell_filter = openmc.CellFilter(cell) + tally2.filters = [cell_filter] + tally2.scores = ["heating"] + + tallies = openmc.Tallies([tally1, tally2]) + model = openmc.Model( + geometry=geometry, materials=[mat], settings=settings, tallies=tallies + ) + + sp_filename = model.run() + + sp = openmc.StatePoint(sp_filename) + + tally_found = sp.get_tally(filter_type=openmc.MeshFilter) + assert tally_found.id == 1 + + tally_found = sp.get_tally(filter_type=openmc.CellFilter) + assert tally_found.id == 2 + + tally_found = sp.get_tally(filters=[mesh_filter]) + assert tally_found.id == 1 + + tally_found = sp.get_tally(filters=[cell_filter]) + assert tally_found.id == 2 + + tally_found = sp.get_tally(scores=["heating"]) + assert tally_found.id == 2 + + tally_found = sp.get_tally(name="heating tally") + assert tally_found.id == 2 + + tally_found = sp.get_tally(name=None) + assert tally_found.id == 1 + + tally_found = sp.get_tally(id=1) + assert tally_found.id == 1 + + tally_found = sp.get_tally(id=2) + assert tally_found.id == 2