Merge pull request #47 from fusion-energy/develop

adding tally finding utlis

src/regular_mesh_plotter/core.py

@@ -3,6 +3,41 @@
 import openmc
 
 
+def get_tallies_with_regular_mesh_filters(statepoint: openmc.StatePoint):
+    """scans the statepoint object to find all tallies and with regular mesh
+    filters, returns a list of tally indexes"""
+
+    matching_tally_ids = []
+    for tally_id, tally in statepoint.tallies.items():
+        print("tally id", tally_id)
+        try:
+            mf = tally.find_filter(filter_type=openmc.MeshFilter)
+            if isinstance(mf.mesh, openmc.RegularMesh):
+                matching_tally_ids.append(tally.id)
+                print("found regmeshfilter")
+        except ValueError:
+            mf = None
+
+    return sorted(matching_tally_ids)
+
+
+def get_regularmesh_tallies_and_scores(statepoint: openmc.StatePoint):
+    """scans the statepoint object to find all tallies and scores,
+    returns list of dictionaries. Each dictionary contains tally id,
+    score and tally name"""
+
+    tallies_of_interest = get_tallies_with_regular_mesh_filters(statepoint)
+
+    tally_score_info = []
+    for tally_id in tallies_of_interest:
+        tally = statepoint.tallies[tally_id]
+        for score in tally.scores:
+            entry = {"id": tally.id, "score": score, "name": tally.name}
+            tally_score_info.append(entry)
+
+    return tally_score_info
+
+
 def get_mpl_plot_extent(self, view_direction: str = "x"):
     """Returns the (x_min, x_max, y_min, y_max) of the bounding box. The
     view_direction is taken into account and can be set using

tests/create_statepoint_file_for_testing_using_csg.py

@@ -0,0 +1,49 @@
+import openmc
+
+
+mat1 = openmc.Material()
+mat1.add_nuclide("He4", 1)
+mat1.set_density("g/cm3", 1)
+my_materials = openmc.Materials([mat1])
+surf1 = openmc.Sphere(r=10, boundary_type="vacuum")
+region1 = -surf1
+cell1 = openmc.Cell(region=region1)
+my_geometry = openmc.Geometry([cell1])
+my_settings = openmc.Settings()
+my_settings.batches = 3
+my_settings.particles = 500
+my_settings.run_mode = "fixed source"
+my_source = openmc.Source()
+my_source.space = openmc.stats.Point((0, 0, 0))
+my_source.angle = openmc.stats.Isotropic()
+my_source.energy = openmc.stats.Discrete([14e6], [1])
+my_settings.source = my_source
+cymesh = openmc.CylindricalMesh().from_domain(my_geometry)
+regmesh = openmc.RegularMesh().from_domain(my_geometry)
+regmesh_filter = openmc.MeshFilter(regmesh)
+cmesh_filter = openmc.MeshFilter(cymesh)
+cell_filter = openmc.CellFilter(cell1)
+
+cell_tally_1 = openmc.Tally(name="tallies_on_cell_flux_and_heating")
+cell_tally_1.id = 1
+cell_tally_1.filters = [cell_filter]
+cell_tally_1.scores = ["flux", "heating"]
+
+mesh_tally_1 = openmc.Tally(name="tallies_on_regmesh_flux_and_heating")
+mesh_tally_1.id = 4
+mesh_tally_1.filters = [regmesh_filter]
+mesh_tally_1.scores = ["flux", "heating"]
+
+mesh_tally_2 = openmc.Tally(name="tallies_on_regmesh_absorption")
+mesh_tally_2.id = 3
+mesh_tally_2.filters = [regmesh_filter]
+mesh_tally_2.scores = ["absorption"]
+
+mesh_tally_3 = openmc.Tally(name="tallies_on_cymesh_absorption")
+mesh_tally_3.id = 5
+mesh_tally_3.filters = [cmesh_filter]
+mesh_tally_3.scores = ["absorption"]
+
+my_tallies = openmc.Tallies([cell_tally_1, mesh_tally_1, mesh_tally_2, mesh_tally_3])
+model = openmc.model.Model(my_geometry, my_materials, my_settings, my_tallies)
+output_filename = model.run()

tests/test_tally_finding.py

@@ -0,0 +1,29 @@
+import openmc
+from regular_mesh_plotter import (
+    get_tallies_with_regular_mesh_filters,
+    get_regularmesh_tallies_and_scores,
+)
+
+
+def test_tally_with_reg_mesh_finding():
+    # checks the correct number of tally score combinations are found
+
+    statepoint = openmc.StatePoint("statepoint.3.h5")
+
+    # one tally has a cell filter and another one has a cylinder mesh filter
+    assert len(statepoint.tallies) == 4
+
+    t_and_scores = get_tallies_with_regular_mesh_filters(statepoint)
+    assert t_and_scores == [3, 4]
+
+
+def test_get_regularmesh_tallies_and_scores():
+    statepoint = openmc.StatePoint("statepoint.3.h5")
+
+    t_and_c = get_regularmesh_tallies_and_scores(statepoint)
+
+    assert t_and_c == [
+        {"id": 3, "name": "tallies_on_regmesh_absorption", "score": "absorption"},
+        {"id": 4, "name": "tallies_on_regmesh_flux_and_heating", "score": "flux"},
+        {"id": 4, "name": "tallies_on_regmesh_flux_and_heating", "score": "heating"},
+    ]