added test for combined tallies

examples/plot_rz_slice_point_source_combined_multiple_tallies.py

@@ -0,0 +1,98 @@
+# this example creates a simple CylindricalMesh tally and performs an openmc
+# simulation to populate the tally. Slices of the resulting tally is then
+# plotted using the openmc_cylindrical_plotter
+
+import openmc
+import numpy as np
+from math import pi
+import matplotlib.pyplot as plt
+from openmc_cylindrical_mesh_plotter import plot_mesh_tally_rz_slice
+
+mesh = openmc.CylindricalMesh(
+    phi_grid=np.linspace(0.0, 2 * pi, 3),
+    r_grid=np.linspace(0, 300, 20),
+    z_grid=np.linspace(0, 300, 10),
+    origin=(0,0,0)
+)
+
+neutron_filter = openmc.ParticleFilter("neutron")
+photon_filter = openmc.ParticleFilter("photon")
+
+tally1 = openmc.Tally(name="my_neutron_heating_tally")
+mesh_filter = openmc.MeshFilter(mesh)
+tally1.filters = [mesh_filter, neutron_filter]
+tally1.scores = ["heating"]
+
+tally2 = openmc.Tally(name="my_photon_heating_tally")
+mesh_filter = openmc.MeshFilter(mesh)
+tally2.filters = [mesh_filter, photon_filter]
+tally2.scores = ["heating"]
+
+tallies = openmc.Tallies([tally1, tally2])
+
+material = openmc.Material()
+material.add_element("Li", 1)
+material.set_density("g/cm3", 0.1)
+my_materials = openmc.Materials([material])
+
+outer_surface = openmc.Sphere(r=300, boundary_type="vacuum")
+cell = openmc.Cell(region=-outer_surface, fill=material)
+
+my_geometry = openmc.Geometry([cell])
+
+source_n = openmc.IndependentSource()
+source_n.space = openmc.stats.Point((200, 0, 0))
+source_n.angle = openmc.stats.Isotropic()
+source_n.energy = openmc.stats.Discrete([0.1e6], [1])
+source_n.strength = 1
+source_n.particle = "neutron"
+
+source_p = openmc.IndependentSource()
+source_p.space = openmc.stats.Point((0, 0, 200))
+source_p.angle = openmc.stats.Isotropic()
+source_p.energy = openmc.stats.Discrete([10e6], [1])
+source_p.strength = 1
+source_p.particle = "photon"
+
+
+my_settings = openmc.Settings()
+my_settings.run_mode = "fixed source"
+my_settings.batches = 10
+my_settings.particles = 100000
+my_settings.source = [source_p, source_n]
+my_settings.photon_transport = True
+
+model = openmc.model.Model(my_geometry, my_materials, my_settings, tallies)
+sp_filename = model.run()
+
+statepoint = openmc.StatePoint(sp_filename)
+
+my_tally1_result = statepoint.get_tally(name="my_neutron_heating_tally")
+my_tally2_result = statepoint.get_tally(name="my_photon_heating_tally")
+
+plot = plot_mesh_tally_rz_slice(
+    tally=[my_tally1_result, my_tally2_result],
+    outline=True,
+    geometry=my_geometry,
+    # norm=LogNorm(),
+    slice_index=1
+)
+plot.figure.savefig(f"rz_point_source_photon_and_neutron_heating.png")
+
+plot = plot_mesh_tally_rz_slice(
+    tally=[my_tally1_result],
+    outline=True,
+    geometry=my_geometry,
+    # norm=LogNorm(),
+    slice_index=1
+)
+plot.figure.savefig(f"rz_point_source_neutron_heating.png")
+
+plot = plot_mesh_tally_rz_slice(
+    tally=[my_tally2_result],
+    outline=True,
+    geometry=my_geometry,
+    # norm=LogNorm(),
+    slice_index=1
+)
+plot.figure.savefig(f"rz_point_source_photon_heating.png")

examples/plot_rz_slices_point_source.py

@@ -21,32 +21,22 @@
 tally.scores.append("flux")
 tallies = openmc.Tallies([tally])
 
-outer_surface = openmc.Sphere(r=100, boundary_type="vacuum")
-cell = openmc.Cell(region=-outer_surface)
-
 material = openmc.Material()
 material.add_nuclide("Fe56", 1)
 material.set_density("g/cm3", 0.1)
+
 my_materials = openmc.Materials([material])
+outer_surface = openmc.Sphere(r=100, boundary_type="vacuum")
+cell = openmc.Cell(region=-outer_surface, fill=material)
 
-universe = openmc.Universe(cells=[cell])
-my_geometry = openmc.Geometry(universe)
 
-my_source = openmc.Source()
+my_geometry = openmc.Geometry([cell])
 
-# the distribution of radius is just a single value
-radius = openmc.stats.Discrete([4], [1])
-# the distribution of source z values is just a single value
-z_values = openmc.stats.Discrete([4], [1])
-# the distribution of source azimuthal angles values is a uniform distribution between 0 and 2 Pi
-angle = openmc.stats.Uniform(a=0.0, b=2 * 3.14159265359)
-# this makes the ring source using the three distributions and a radius
-# could do a point source instead with my_source.space = openmc.stats.Point((5,5., 5))
-my_source.space = openmc.stats.CylindricalIndependent(
-    r=radius, phi=angle, z=z_values, origin=(0.0, 0.0, 0.0)
-)
-# sets the direction to isotropic
+
+my_source = openmc.IndependentSource()
+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 = openmc.Settings()
@@ -65,6 +55,10 @@
 
 for slice_index in range(1, len(mesh.phi_grid)):
     plot = plot_mesh_tally_rz_slice(
-        tally=my_tally_result, outline=True, geometry=my_geometry, norm=LogNorm()
+        tally=my_tally_result,
+        outline=True,
+        geometry=my_geometry,
+        norm=LogNorm(),
+        slice_index=slice_index
     )
     plot.figure.savefig(f"rz_point_source_{slice_index}.png")

src/openmc_cylindrical_mesh_plotter/core.py

@@ -1,7 +1,7 @@
 import math
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Optional
+import typing
 import openmc
 import numpy as np
 import openmc
@@ -23,20 +23,20 @@
 
 
 def plot_mesh_tally_rz_slice(
-    tally: "openmc.Tally",
-    slice_index: Optional[int] = None,
-    score: Optional[str] = None,
-    axes: Optional[str] = None,
+    tally: typing.Union["openmc.Tally", typing.Sequence["openmc.Tally"]],
+    slice_index: typing.Optional[int] = None,
+    score: typing.Optional[str] = None,
+    axes: typing.Optional[str] = None,
     axis_units: str = "cm",
     value: str = "mean",
     outline: bool = False,
     outline_by: str = "cell",
-    geometry: Optional["openmc.Geometry"] = None,
+    geometry: typing.Optional["openmc.Geometry"] = None,
     geometry_basis: str = "xz",
     pixels: int = 40000,
     colorbar: bool = True,
     volume_normalization: bool = True,
-    scaling_factor: Optional[float] = None,
+    scaling_factor: typing.Optional[float] = None,
     colorbar_kwargs: dict = {},
     outline_kwargs: dict = _default_outline_kwargs,
     **kwargs,
@@ -45,7 +45,9 @@ def plot_mesh_tally_rz_slice(
     Parameters
     ----------
     tally : openmc.Tally
-        The openmc tally to plot. Tally must contain a MeshFilter that uses a CylindricalMesh.
+        The openmc tally/tallies to plot. Tally must contain a MeshFilter that
+        uses a CylindricalMesh. If a sequence of multiple tallies are provided
+        the score will be added.
     slice_index : int
         The mesh index to plot
     score : str
@@ -93,50 +95,31 @@ def plot_mesh_tally_rz_slice(
     cv.check_type("volume_normalization", volume_normalization, bool)
     cv.check_type("outline", outline, bool)
 
-    mesh = tally.find_filter(filter_type=openmc.MeshFilter).mesh
+    # if tally is multiple tallies
+    if isinstance(tally, typing.Sequence):
+        mesh_ids = []
+        for one_tally in tally:
+            mesh = one_tally.find_filter(filter_type=openmc.MeshFilter).mesh
+            # TODO check the tallies use the same mesh
+            mesh_ids.append(mesh.id)
+        if not all(i == mesh_ids[0] for i in mesh_ids):
+            raise ValueError(
+                f"mesh ids {mesh_ids} are different, please use same mesh when combining tallies"
+            )
+    # if tally is single tally
+    else:
+
+        mesh = tally.find_filter(filter_type=openmc.MeshFilter).mesh
+
     if not isinstance(mesh, openmc.CylindricalMesh):
         raise NotImplemented(
             f"Only CylindricalMesh are currently supported not {type(mesh)}"
         )
+
     if mesh.n_dimension != 3:
         msg = "Your mesh has {mesh.n_dimension} dimension and currently only CylindricalMesh with 3 dimensions are supported"
         raise NotImplementedError(msg)
 
-    # if score is not specified and tally has a single score then we know which score to use
-    if score is None:
-        if len(tally.scores) == 1:
-            score = tally.scores[0]
-        else:
-            msg = "score was not specified and there are multiple scores in the tally."
-            raise ValueError(msg)
-
-    tally_slice = tally.get_slice(scores=[score])
-
-    tally_data = tally_slice.get_reshaped_data(expand_dims=True, value=value).squeeze()
-
-    # get the middle phi value
-    if slice_index is None:
-        slice_index = int(tally_data.shape[1] / 2)  # index 1 is the phi value
-
-    if len(tally_data.shape) == 3:
-        data = tally_data[:, slice_index, :]
-    elif len(tally_data.shape) == 2:
-        data = tally_data[:, :]
-    else:
-        raise NotImplementedError("Mesh is not 3d or 2d, can't plot")
-
-    if volume_normalization:
-        if len(tally_data.shape) == 3:
-            slice_volumes = mesh.volumes[:, slice_index, :].squeeze()
-        elif len(tally_data.shape) == 2:
-            slice_volumes = mesh.volumes[:, :].squeeze()
-        data = data / slice_volumes
-
-    if scaling_factor:
-        data = data * scaling_factor
-
-    data = np.rot90(data, 1)
-
     xlabel, ylabel = f"r [{axis_units}]", f"z [{axis_units}]"
     axis_scaling_factor = {"km": 0.00001, "m": 0.01, "cm": 1, "mm": 10}[axis_units]
 
@@ -163,6 +146,31 @@ def plot_mesh_tally_rz_slice(
     default_imshow_kwargs = {"interpolation": "none"}
     default_imshow_kwargs.update(kwargs)
 
+    if isinstance(tally, typing.Sequence):
+        for counter, one_tally in enumerate(tally):
+            new_data = _get_tally_data(
+                scaling_factor,
+                mesh,
+                one_tally,
+                value,
+                volume_normalization,
+                score,
+                slice_index,
+            )
+            if counter == 0:
+                data = np.zeros(shape=new_data.shape)
+            data = data + new_data
+    else:  # single tally
+        data = _get_tally_data(
+            scaling_factor,
+            mesh,
+            tally,
+            value,
+            volume_normalization,
+            score,
+            slice_index,
+        )
+
     im = axes.imshow(data, extent=(x_min, x_max, y_min, y_max), **default_imshow_kwargs)
 
     if colorbar:
@@ -224,21 +232,61 @@ def plot_mesh_tally_rz_slice(
 
     return axes
 
+def _get_tally_data(
+    scaling_factor, mesh, tally, value, volume_normalization, score, slice_index
+):
+    # if score is not specified and tally has a single score then we know which score to use
+    if score is None:
+        if len(tally.scores) == 1:
+            score = tally.scores[0]
+        else:
+            msg = "score was not specified and there are multiple scores in the tally."
+            raise ValueError(msg)
+    tally_slice = tally.get_slice(scores=[score])
+
+
+    tally_slice = tally.get_slice(scores=[score])
+
+    tally_data = tally_slice.get_reshaped_data(expand_dims=True, value=value).squeeze()
+
+    # get the middle phi value
+    if slice_index is None:
+        slice_index = int(tally_data.shape[1] / 2)  # index 1 is the phi value
+
+    if len(tally_data.shape) == 3:
+        data = tally_data[:, slice_index, :]
+    elif len(tally_data.shape) == 2:
+        data = tally_data[:, :]
+    else:
+        raise NotImplementedError("Mesh is not 3d or 2d, can't plot")
+
+    if volume_normalization:
+        if len(tally_data.shape) == 3:
+            slice_volumes = mesh.volumes[:, slice_index, :].squeeze()
+        elif len(tally_data.shape) == 2:
+            slice_volumes = mesh.volumes[:, :].squeeze()
+        data = data / slice_volumes
+
+    if scaling_factor:
+        data = data * scaling_factor
+
+    data = np.rot90(data, 1)
+    return data
 
 def plot_mesh_tally_phir_slice(
     tally: "openmc.Tally",
-    slice_index: Optional[int] = None,
-    score: Optional[str] = None,
-    axes: Optional[str] = None,
+    slice_index: typing.Optional[int] = None,
+    score: typing.Optional[str] = None,
+    axes: typing.Optional[str] = None,
     axis_units: str = "cm",
     value: str = "mean",
     outline: bool = False,
     outline_by: str = "cell",
-    geometry: Optional["openmc.Geometry"] = None,
+    geometry: typing.Optional["openmc.Geometry"] = None,
     pixels: int = 40000,
     colorbar: bool = True,
     volume_normalization: bool = True,
-    scaling_factor: Optional[float] = None,
+    scaling_factor: typing.Optional[float] = None,
     colorbar_kwargs: dict = {},
     outline_kwargs: dict = _default_outline_kwargs,
     **kwargs,

tests/test_slice_of_data.py

@@ -180,3 +180,14 @@ def test_phir_slice_of_data_circular_simulation_normalization(
         )
     for slice_index in range(0, len(mesh.phi_grid) - 1):
         plot_mesh_tally_rz_slice(tally=tally, slice_index=slice_index)
+
+def test_rz_slice_of_data_point_simulation_combined(point_source_simulation):
+    tally = point_source_simulation
+    mesh = tally.find_filter(openmc.MeshFilter).mesh
+    for slice_index in range(0, len(mesh.z_grid) - 1):
+        plot_mesh_tally_phir_slice(
+            tally=[tally, tally],
+            slice_index=slice_index,
+            colorbar=True,
+            volume_normalization=True,
+        )