Adding regularmesh tally plot

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

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')