diff --git a/openmc/plots.py b/openmc/plots.py index 6afe3ea33b8..d875960ab54 100644 --- a/openmc/plots.py +++ b/openmc/plots.py @@ -1,7 +1,9 @@ +import math from collections.abc import Iterable, Mapping from numbers import Integral, Real from pathlib import Path import lxml.etree as ET +from tempfile import TemporaryDirectory from typing import Optional import h5py @@ -166,6 +168,192 @@ 'yellowgreen': (154, 205, 50) } +_default_outline_kwargs = { + 'colors': 'black', + 'linestyles': 'solid', + 'linewidths': 1 +} + + +def plot_mesh_tally( + tally: 'openmc.Tally', + basis: str = 'xy', + slice_index: Optional[int] = None, + score: Optional[str] = None, + axes: Optional[str] = None, + axis_units: str = 'cm', + value: str = 'mean', + outline: bool = False, + outline_by: str = 'cell', + geometry: Optional['openmc.Geometry'] = None, + pixels: int = 40000, + colorbar: bool = True, + volume_normalization: bool = True, + scaling_factor: Optional[float] = None, + colorbar_kwargs: dict = {}, + outline_kwargs: dict = _default_outline_kwargs, + **kwargs +) -> 'matplotlib.image.AxesImage': + """Display a slice plot of the mesh tally score. + + Parameters + ---------- + tally : openmc.Tally + The openmc tally to plot. Tally must contain a MeshFilter that uses a RegularMesh. + basis : {'xy', 'xz', 'yz'} + The basis directions for the plot + slice_index : int + The mesh index to plot + score : str + Score to plot, e.g. 'flux' + axes : matplotlib.Axes + Axes to draw to + axis_units : {'km', 'm', 'cm', 'mm'} + Units used on the plot axis + value : str + A string for the type of value to return - 'mean' (default), + 'std_dev', 'rel_err', 'sum', or 'sum_sq' are accepted + outline : tuple + If set then an outline will be added to the plot. The outline can be + by cell or by material. + outline_by : {'cell', 'material'} + Indicate whether the plot should be colored by cell or by material + geometry : openmc.Geometry + The geometry to use for the outline. + pixels : int + This sets the total number of pixels in the plot and the number of + pixels in each basis direction is calculated from this total and + the image aspect ratio. + colorbar : bool + Whether or not to add a colorbar to the plot. + volume_normalization : bool, optional + Whether or not to normalize the data by the volume of the mesh elements. + scaling_factor : float + A optional multiplier to apply to the tally data prior to ploting. + colorbar_kwargs : dict + Keyword arguments passed to :func:`matplotlib.colorbar.Colorbar`. + outline_kwargs : dict + Keyword arguments passed to :func:`matplotlib.pyplot.contour`. + **kwargs + Keyword arguments passed to :func:`matplotlib.pyplot.imshow` + + Returns + ------- + matplotlib.image.AxesImage + Resulting image + + """ + + import matplotlib.pyplot as plt + + cv.check_value('basis', basis, _BASES) + cv.check_value('axis_units', axis_units, ['km', 'm', 'cm', 'mm']) + cv.check_type('volume_normalization', volume_normalization, bool) + cv.check_type('outline', outline, bool) + + mesh = tally.find_filter(filter_type=openmc.MeshFilter).mesh + if not isinstance(mesh, openmc.RegularMesh): + raise NotImplemented(f'Only RegularMesh are currently supported not {type(mesh)}') + + # 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() + + if slice_index is None: + basis_to_index = {'xy': 2, 'xz': 1, 'yz': 0}[basis] + slice_index = int(tally_data.shape[basis_to_index]/2) + + if basis == 'xz': + slice_data = tally_data[:, slice_index, :] + data = np.flip(np.rot90(slice_data, -1)) + xlabel, ylabel = f'x [{axis_units}]', f'z [{axis_units}]' + elif basis == 'yz': + slice_data = tally_data[slice_index, :, :] + data = np.flip(np.rot90(slice_data, -1)) + xlabel, ylabel = f'y [{axis_units}]', f'z [{axis_units}]' + else: # basis == 'xy' + slice_data = tally_data[:, :, slice_index] + data = np.rot90(slice_data, -3) + xlabel, ylabel = f'x [{axis_units}]', f'y [{axis_units}]' + + if volume_normalization: + # in a regular mesh all volumes are the same so we just divide by the first + data = data / mesh.volumes[0][0][0] + + if scaling_factor: + data = data * scaling_factor + + axis_scaling_factor = {'km': 0.00001, 'm': 0.01, 'cm': 1, 'mm': 10}[axis_units] + + x_min, x_max, y_min, y_max = [i * axis_scaling_factor for i in mesh.bounding_box.extent[basis]] + + if axes is None: + fig, axes = plt.subplots() + axes.set_xlabel(xlabel) + axes.set_ylabel(ylabel) + + im = axes.imshow(data, extent=(x_min, x_max, y_min, y_max), **kwargs) + + if colorbar: + fig.colorbar(im, **colorbar_kwargs) + + if outline and geometry is not None: + import matplotlib.image as mpimg + model = openmc.Model() + model.geometry = geometry + plot = openmc.Plot() + plot.origin = mesh.bounding_box.center + bb_width = mesh.bounding_box.extent[basis] + plot.width = (bb_width[0]-bb_width[1], bb_width[2]-bb_width[3]) + aspect_ratio = (bb_width[0]-bb_width[1]) / (bb_width[2]-bb_width[3]) + pixels_y = math.sqrt(pixels / aspect_ratio) + pixels = (int(pixels / pixels_y), int(pixels_y)) + plot.pixels = pixels + plot.basis = basis + plot.color_by = outline_by + model.plots.append(plot) + + with TemporaryDirectory() as tmpdir: + # Run OpenMC in geometry plotting mode + model.plot_geometry(False, cwd=tmpdir) + + # Read image from file + img_path = Path(tmpdir) / f'plot_{plot.id}.png' + if not img_path.is_file(): + img_path = img_path.with_suffix('.ppm') + img = mpimg.imread(str(img_path)) + + # Combine R, G, B values into a single int + rgb = (img * 256).astype(int) + image_value = (rgb[..., 0] << 16) + \ + (rgb[..., 1] << 8) + (rgb[..., 2]) + + if basis == 'xz': + image_value = np.rot90(image_value, 2) + elif basis == 'yz': + image_value = np.rot90(image_value, 2) + else: # basis == 'xy' + image_value = np.rot90(image_value, 2) + + # Plot image and return the axes + axes.contour( + image_value, + origin="upper", + levels=np.unique(image_value), + extent=(x_min, x_max, y_min, y_max), + **outline_kwargs + ) + + return axes + def _get_plot_image(plot, cwd): from IPython.display import Image diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py index 922e4d9dc95..8101e0eb254 100644 --- a/tests/unit_tests/test_plots.py +++ b/tests/unit_tests/test_plots.py @@ -221,3 +221,91 @@ def test_voxel_plot_roundtrip(): assert new_plot.origin == plot.origin assert new_plot.width == plot.width assert new_plot.color_by == plot.color_by + + +def test_plot_mesh_tally(run_in_tmpdir): + from matplotlib.colors import LogNorm + mat1 = openmc.Material() + mat1.add_nuclide('Li6', 1, percent_type='ao') + mats = openmc.Materials([mat1]) + + # this shape is chose to create axis with testable ranges + surface1 = openmc.model.RectangularParallelepiped( + -50, 25, -100, 125, -150, 175, boundary_type='vacuum' + ) + surface = openmc.model.RectangularParallelepiped( + -100, 50, -200, 250, -300, 350, boundary_type='vacuum' + ) + cell1 = openmc.Cell(region=-surface1) + cell2 = openmc.Cell(region=-surface & +surface1) + cell1.fill = mat1 + cell2.fill = mat1 + geom = openmc.Geometry([cell1, cell2]) + + source = openmc.IndependentSource() + source.angle = openmc.stats.Isotropic() + source.energy = openmc.stats.Discrete([14e6], [1]) + # puts the source in the center of the RectangularParallelepiped + source.space = openmc.stats.Point(cell2.bounding_box.center) + + sett = openmc.Settings() + sett.batches = 2 + sett.inactive = 0 + sett.particles = 5000 + sett.run_mode = 'fixed source' + sett.source = source + + mesh = openmc.RegularMesh().from_domain(geom, dimension=[10, 20, 30]) + mesh_filter = openmc.MeshFilter(mesh) + mesh_tally = openmc.Tally(name='mesh-tal') + mesh_tally.filters = [mesh_filter] + mesh_tally.scores = ['flux'] + tallies = openmc.Tallies([mesh_tally]) + + model = openmc.Model(geom, mats, sett, tallies) + sp_filename = model.run() + statepoint = openmc.StatePoint(sp_filename) + tally_result = statepoint.get_tally(name='mesh-tal') + + plot = openmc.plot_mesh_tally( + tally=tally_result, + basis='xy', + slice_index=29 # max value of slice selected + ) + # axis_units defaults to cm + assert plot.xaxis.get_label().get_text() == 'x [cm]' + assert plot.yaxis.get_label().get_text() == 'y [cm]' + assert plot.get_xlim() == (-100., 50) + assert plot.get_ylim() == (-200., 250.) + + plot = openmc.plot_mesh_tally( + tally=tally_result, + basis='yz', + axis_units='m', + slice_index=9, # max value of slice selected + value='std_dev' + ) + plot.figure.savefig('x.png') + assert plot.xaxis.get_label().get_text() == 'y [m]' + assert plot.yaxis.get_label().get_text() == 'z [m]' + assert plot.get_xlim() == (-2., 2.5) # note that units are in m + assert plot.get_ylim() == (-3., 3.5) + + plot = openmc.plot_mesh_tally( + tally=tally_result, + basis='xz', + slice_index=19, # max value of slice selected + axis_units='mm', + score='flux', + value='mean', + outline=True, + geometry=geom, + outline_by='material', + colorbar_kwargs={'label': 'neutron flux'}, + norm=LogNorm(vmin=1e-6, vmax=max(tally_result.mean.flatten())), + ) + assert plot.xaxis.get_label().get_text() == 'x [mm]' + assert plot.yaxis.get_label().get_text() == 'z [mm]' + assert plot.get_xlim() == (-1000., 500) # note that units are in mm + assert plot.get_ylim() == (-3000.0, 3500.0) + plot.figure.savefig('z.png')