From 8542e72838e306328c0780ca0f835d39b33848fd Mon Sep 17 00:00:00 2001 From: Jonathan Shimwell Date: Wed, 6 Dec 2023 18:34:54 +0000 Subject: [PATCH] added test for combined tallies --- ..._point_source_combined_multiple_tallies.py | 98 ++++++++++++ examples/plot_rz_slices_point_source.py | 32 ++-- src/openmc_cylindrical_mesh_plotter/core.py | 146 ++++++++++++------ tests/test_slice_of_data.py | 11 ++ 4 files changed, 219 insertions(+), 68 deletions(-) create mode 100644 examples/plot_rz_slice_point_source_combined_multiple_tallies.py diff --git a/examples/plot_rz_slice_point_source_combined_multiple_tallies.py b/examples/plot_rz_slice_point_source_combined_multiple_tallies.py new file mode 100644 index 0000000..a906396 --- /dev/null +++ b/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") diff --git a/examples/plot_rz_slices_point_source.py b/examples/plot_rz_slices_point_source.py index 9842a1f..8d44325 100644 --- a/examples/plot_rz_slices_point_source.py +++ b/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") diff --git a/src/openmc_cylindrical_mesh_plotter/core.py b/src/openmc_cylindrical_mesh_plotter/core.py index c5a2f3a..da9c5e0 100644 --- a/src/openmc_cylindrical_mesh_plotter/core.py +++ b/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, diff --git a/tests/test_slice_of_data.py b/tests/test_slice_of_data.py index 6aa0dac..3027ab1 100644 --- a/tests/test_slice_of_data.py +++ b/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, + )