diff --git a/src/openmc_regular_mesh_plotter/core.py b/src/openmc_regular_mesh_plotter/core.py index 5f8e69f..7247e91 100644 --- a/src/openmc_regular_mesh_plotter/core.py +++ b/src/openmc_regular_mesh_plotter/core.py @@ -145,8 +145,8 @@ def plot_mesh_tally( default_imshow_kwargs.update(kwargs) if isinstance(tally, typing.Sequence): - data = np.zeros(shape=(40,40)) - for one_tally in tally: + + for counter, one_tally in enumerate(tally): new_data = _get_tally_data( scaling_factor, mesh, @@ -157,6 +157,8 @@ def plot_mesh_tally( score, slice_index ) + if counter == 0: + data = np.zeros(shape=new_data.shape) data=data+new_data else: # single tally data = _get_tally_data( diff --git a/tests/test_units.py b/tests/test_units.py index 92c730d..e4fceb1 100644 --- a/tests/test_units.py +++ b/tests/test_units.py @@ -151,12 +151,15 @@ def test_plot_two_mesh_tallies(model): mesh = openmc.RegularMesh().from_domain(geometry, dimension=[1, 20, 30]) mesh_filter = openmc.MeshFilter(mesh) + mesh_tally_1 = openmc.Tally(name="mesh-tal-1") mesh_tally_1.filters = [mesh_filter] mesh_tally_1.scores = ["flux"] + mesh_tally_2 = openmc.Tally(name="mesh-tal-2") mesh_tally_2.filters = [mesh_filter] mesh_tally_2.scores = ["heating"] + tallies = openmc.Tallies([mesh_tally_1, mesh_tally_2]) model.tallies = tallies @@ -179,7 +182,7 @@ def test_plot_two_mesh_tallies(model): plot.figure.savefig("t.png") plot = plot_mesh_tally( - tally=tally_result, + tally=[tally_result_1, tally_result_2], basis="yz", axis_units="m", # slice_index=9, # max value of slice selected