From cbe0c914ef17d8b820ed32399913d1da34e82571 Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Tue, 18 Jul 2023 18:07:57 +0100
Subject: [PATCH 01/17] draft mesh plot method

---
 openmc/plots.py | 75 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 75 insertions(+)

diff --git a/openmc/plots.py b/openmc/plots.py
index 5bc8a2e6e8a..b170b0714f5 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -167,6 +167,81 @@
 }
 
 
+def plot_mesh_tally(
+    tally: 'openmc.Tally',
+    basis: Optional[str] = 'xy',
+    slice_index: int = None,
+    score: Optional[str] = None,
+    axes: Optional[str] = None,
+    axis_units: Optional[str] = 'cm',
+    **kwargs
+):
+    """Display a slice plot of the mesh tally score.
+
+    Parameters
+    ----------
+    tally : openmc.Tally
+        todo
+    basis : str
+        todo
+    slice_index : int
+        todo
+    score : str
+        todo
+    axes : str
+        todo
+    axis_units : str
+        todo
+
+    Returns
+    -------
+    matplotlib.image.AxesImage
+        Resulting image
+
+    """
+
+    import matplotlib.pyplot as plt
+
+    cv.check_value('basis', basis, ['xy', 'xz', 'yz'])
+    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 and len(tally.scores) == 1:
+        score = tally.scores[0]
+
+    data = tally.get_reshaped_data(expand_dims=True).squeeze()
+
+    if slice_index is None:
+        basis_to_index = {'xy': 2, 'xz': 1, 'yz': 0}[basis]
+        slice_index = int(data.shape[basis_to_index]/2)
+
+    if basis == 'xz':
+        slice_data = data[:, slice_index, :]
+        oriented_data = np.flip(np.rot90(slice_data, -1))
+        xlabel, ylabel = f'x [{axis_units}]', f'z [{axis_units}]'
+    elif basis == 'yz':
+        slice_data = data[slice_index, :, :]
+        oriented_data = np.flip(np.rot90(slice_data, -1))
+        xlabel, ylabel = f'y [{axis_units}]', f'z [{axis_units}]'
+    else:  # basis == 'xy'
+        slice_data = data[:, :, slice_index]
+        oriented_data = np.rot90(slice_data, 1)
+        xlabel, ylabel = f'x [{axis_units}]', f'y [{axis_units}]'
+
+    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)
+
+    return axes.imshow(oriented_data, extent=(x_min, x_max, y_min, y_max), **kwargs)
+
+
 def _get_plot_image(plot, cwd):
     from IPython.display import Image
 

From 8bc87a734d2416d63dafd39d03d43d74544818ea Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <drshimwell@gmail.com>
Date: Wed, 19 Jul 2023 11:20:19 +0100
Subject: [PATCH 02/17] improved doc string

---
 openmc/plots.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index b170b0714f5..a8709883199 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -181,17 +181,17 @@ def plot_mesh_tally(
     Parameters
     ----------
     tally : openmc.Tally
-        todo
-    basis : str
-        todo
+        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
-        todo
+        The mesh index to plot
     score : str
-        todo
-    axes : str
-        todo
-    axis_units : str
-        todo
+        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
 
     Returns
     -------

From dd2e3561d36f0c7b8455f0565e20c73956eb01b7 Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Wed, 26 Jul 2023 19:56:02 +0100
Subject: [PATCH 03/17] added tests

---
 openmc/plots.py                | 20 ++++++---
 tests/unit_tests/test_plots.py | 80 ++++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 6 deletions(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index 349ab9287c5..7c2ef835dfe 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -169,13 +169,14 @@
 
 def plot_mesh_tally(
     tally: 'openmc.Tally',
-    basis: Optional[str] = 'xy',
-    slice_index: int = None,
+    basis: str = 'xy',
+    slice_index: Optional[int] = None,
     score: Optional[str] = None,
     axes: Optional[str] = None,
-    axis_units: Optional[str] = 'cm',
+    axis_units: str = 'cm',
+    value: str = 'mean',
     **kwargs
-):
+) -> 'matplotlib.image.AxesImage':
     """Display a slice plot of the mesh tally score.
 
     Parameters
@@ -192,6 +193,11 @@ def plot_mesh_tally(
         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
+    **kwargs
+        Keyword arguments passed to :func:`matplotlib.pyplot.imshow`
 
     Returns
     -------
@@ -211,7 +217,7 @@ def plot_mesh_tally(
     if score is None and len(tally.scores) == 1:
         score = tally.scores[0]
 
-    data = tally.get_reshaped_data(expand_dims=True).squeeze()
+    data = tally.get_reshaped_data(expand_dims=True, value=value).squeeze()
 
     if slice_index is None:
         basis_to_index = {'xy': 2, 'xz': 1, 'yz': 0}[basis]
@@ -238,8 +244,10 @@ def plot_mesh_tally(
         fig, axes = plt.subplots()
         axes.set_xlabel(xlabel)
         axes.set_ylabel(ylabel)
+    axes.imshow(oriented_data, extent=(x_min, x_max, y_min, y_max), **kwargs)
+    fig.colorbar()
 
-    return axes.imshow(oriented_data, extent=(x_min, x_max, y_min, y_max), **kwargs)
+    return axes
 
 
 def _get_plot_image(plot, cwd):
diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py
index 922e4d9dc95..6d3eba5d938 100644
--- a/tests/unit_tests/test_plots.py
+++ b/tests/unit_tests/test_plots.py
@@ -3,6 +3,7 @@
 import openmc
 import openmc.examples
 import pytest
+from matplotlib.colors import LogNorm
 
 
 @pytest.fixture(scope='module')
@@ -221,3 +222,82 @@ 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):
+
+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
+surface = openmc.model.RectangularParallelepiped(
+    -100, 50, -200, 250, -300, 350, boundary_type='vacuum'
+)
+cell = openmc.Cell(region=-surface)
+cell.fill = mat1
+geom = openmc.Geometry([cell])
+
+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(cell.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'
+)
+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',
+    norm=LogNorm(vmin=0.1, vmax=100)
+)
+plot.colorbar()
+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')
\ No newline at end of file

From a0f052bfc161281d49a614d126754455fc12f4d2 Mon Sep 17 00:00:00 2001
From: shimwell <mail@jshimwell.com>
Date: Wed, 26 Jul 2023 23:24:37 +0100
Subject: [PATCH 04/17] started outline

---
 openmc/plots.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/openmc/plots.py b/openmc/plots.py
index 7c2ef835dfe..4881c35ff2f 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -175,6 +175,8 @@ def plot_mesh_tally(
     axes: Optional[str] = None,
     axis_units: str = 'cm',
     value: str = 'mean',
+    pixels: int = 40000,
+    outline: Optional[str] = None,
     **kwargs
 ) -> 'matplotlib.image.AxesImage':
     """Display a slice plot of the mesh tally score.
@@ -196,6 +198,9 @@ def plot_mesh_tally(
     value : str
         A string for the type of value to return  - 'mean' (default),
         'std_dev', 'rel_err', 'sum', or 'sum_sq' are accepted
+    outline : {'material', 'cell'}
+        If set then an outline will be added to the plot. The outline can be
+        by cell or by material.
     **kwargs
         Keyword arguments passed to :func:`matplotlib.pyplot.imshow`
 
@@ -247,6 +252,15 @@ def plot_mesh_tally(
     axes.imshow(oriented_data, extent=(x_min, x_max, y_min, y_max), **kwargs)
     fig.colorbar()
 
+    if outline is not None:
+        plot = openmc.Plot()
+        plot.origin = mesh.center
+        plot.width = mesh.bounding_box.extent[basis]
+        plot.pixels = pixels
+        plot.basis = basis
+        plot.color_by = outline
+        model.plots.append(plot)
+
     return axes
 
 

From dd2d4e6bb8dfbc92c2cb10e2f22bc2177ec3b29b Mon Sep 17 00:00:00 2001
From: shimwell <mail@jshimwell.com>
Date: Wed, 26 Jul 2023 23:53:59 +0100
Subject: [PATCH 05/17] started outline

---
 openmc/plots.py                | 49 ++++++++++++++++++++++++++++++----
 tests/unit_tests/test_plots.py |  6 ++---
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index 4881c35ff2f..5bc6a38f3fa 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -177,6 +177,7 @@ def plot_mesh_tally(
     value: str = 'mean',
     pixels: int = 40000,
     outline: Optional[str] = None,
+    geometry: Optional['openmc.Geometry'] = None,
     **kwargs
 ) -> 'matplotlib.image.AxesImage':
     """Display a slice plot of the mesh tally score.
@@ -198,7 +199,7 @@ def plot_mesh_tally(
     value : str
         A string for the type of value to return  - 'mean' (default),
         'std_dev', 'rel_err', 'sum', or 'sum_sq' are accepted
-    outline : {'material', 'cell'}
+    outline : tuple
         If set then an outline will be added to the plot. The outline can be
         by cell or by material.
     **kwargs
@@ -250,17 +251,55 @@ def plot_mesh_tally(
         axes.set_xlabel(xlabel)
         axes.set_ylabel(ylabel)
     axes.imshow(oriented_data, extent=(x_min, x_max, y_min, y_max), **kwargs)
-    fig.colorbar()
+    # fig.colorbar()
 
     if outline is not None:
+        import matplotlib.image as mpimg
+        import math
+        from tempfile import TemporaryDirectory
+        model = openmc.Model()
+        model.geometry = outline[0]
         plot = openmc.Plot()
-        plot.origin = mesh.center
-        plot.width = mesh.bounding_box.extent[basis]
+        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])
+        print(plot.width )
+        print(bb_width)
+        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
+        plot.color_by = outline[1]
         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])
+
+        # Plot image and return the axes
+        axes.contour(
+            image_value,
+            origin="upper",
+            colors="k",
+            linestyles="solid",
+            linewidths=1,
+            levels=np.unique(image_value),
+            extent=(x_min, x_max, y_min, y_max),
+            **kwargs
+        )
+
     return axes
 
 
diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py
index 6d3eba5d938..58c62ddcebf 100644
--- a/tests/unit_tests/test_plots.py
+++ b/tests/unit_tests/test_plots.py
@@ -293,11 +293,11 @@ def test_voxel_plot_roundtrip():
     axis_units='mm',
     score='flux',
     value= 'mean',
-    norm=LogNorm(vmin=0.1, vmax=100)
+    norm=LogNorm(vmin=0.1, vmax=100),
+    outline=(geom, 'material')
 )
-plot.colorbar()
 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')
\ No newline at end of file
+plot.figure.savefig('z.png')

From 2e06415941bc47861714db905491a8be0d4a1f65 Mon Sep 17 00:00:00 2001
From: shimwell <mail@jshimwell.com>
Date: Thu, 27 Jul 2023 00:03:16 +0100
Subject: [PATCH 06/17] added 2nd cell to show outline

---
 openmc/plots.py                |  2 +-
 tests/unit_tests/test_plots.py | 40 +++++++++++++++++++---------------
 2 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index 5bc6a38f3fa..4b8cea66a36 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -292,7 +292,7 @@ def plot_mesh_tally(
         axes.contour(
             image_value,
             origin="upper",
-            colors="k",
+            colors="r",
             linestyles="solid",
             linewidths=1,
             levels=np.unique(image_value),
diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py
index 58c62ddcebf..5e5500769c9 100644
--- a/tests/unit_tests/test_plots.py
+++ b/tests/unit_tests/test_plots.py
@@ -231,18 +231,22 @@ def test_voxel_plot_roundtrip():
 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'
 )
-cell = openmc.Cell(region=-surface)
-cell.fill = mat1
-geom = openmc.Geometry([cell])
+cell1 = openmc.Cell(region=-surface1)
+cell2 = openmc.Cell(region=-surface&+surface1)
+cell1.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(cell.bounding_box.center)
+source.space = openmc.stats.Point(cell2.bounding_box.center)
 
 sett = openmc.Settings()
 sett.batches = 2
@@ -269,10 +273,10 @@ def test_voxel_plot_roundtrip():
     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.)
+# 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,
@@ -281,23 +285,23 @@ def test_voxel_plot_roundtrip():
     slice_index=9,  # max value of slice selected
     value= 'std_dev'
 )
-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)
+# 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',
+    # axis_units='mm',
     score='flux',
     value= 'mean',
-    norm=LogNorm(vmin=0.1, vmax=100),
+    # norm=LogNorm(vmin=0.1, vmax=100),
     outline=(geom, 'material')
 )
-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)
+# 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')

From e6a0989f1292455b835930c508542491af99e4dd Mon Sep 17 00:00:00 2001
From: shimwell <mail@jshimwell.com>
Date: Thu, 27 Jul 2023 08:43:38 +0100
Subject: [PATCH 07/17] allowing outline color

---
 openmc/plots.py                | 24 +++++++++++++++++-------
 tests/unit_tests/test_plots.py |  4 +++-
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index 4b8cea66a36..0b003fdd207 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -175,8 +175,10 @@ def plot_mesh_tally(
     axes: Optional[str] = None,
     axis_units: str = 'cm',
     value: str = 'mean',
+    outline: bool = False,
+    outline_by: str = 'cell',
+    outline_color: str = 'black',
     pixels: int = 40000,
-    outline: Optional[str] = None,
     geometry: Optional['openmc.Geometry'] = None,
     **kwargs
 ) -> 'matplotlib.image.AxesImage':
@@ -202,6 +204,16 @@ def plot_mesh_tally(
     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
+    outline_color : str
+        The matplotlib color to use for the plot.
+    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.
     **kwargs
         Keyword arguments passed to :func:`matplotlib.pyplot.imshow`
 
@@ -253,24 +265,22 @@ def plot_mesh_tally(
     axes.imshow(oriented_data, extent=(x_min, x_max, y_min, y_max), **kwargs)
     # fig.colorbar()
 
-    if outline is not None:
+    if outline and geometry is not None:
         import matplotlib.image as mpimg
         import math
         from tempfile import TemporaryDirectory
         model = openmc.Model()
-        model.geometry = outline[0]
+        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])
-        print(plot.width )
-        print(bb_width)
         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[1]
+        plot.color_by = outline_by
         model.plots.append(plot)
 
         with TemporaryDirectory() as tmpdir:
@@ -292,7 +302,7 @@ def plot_mesh_tally(
         axes.contour(
             image_value,
             origin="upper",
-            colors="r",
+            colors=outline_color,
             linestyles="solid",
             linewidths=1,
             levels=np.unique(image_value),
diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py
index 5e5500769c9..0ddf76a5a2d 100644
--- a/tests/unit_tests/test_plots.py
+++ b/tests/unit_tests/test_plots.py
@@ -298,7 +298,9 @@ def test_voxel_plot_roundtrip():
     score='flux',
     value= 'mean',
     # norm=LogNorm(vmin=0.1, vmax=100),
-    outline=(geom, 'material')
+    outline=True,
+    geometry=geom,
+    outline_by='material'
 )
 # assert plot.xaxis.get_label().get_text() == 'x [mm]'
 # assert plot.yaxis.get_label().get_text() == 'z [mm]'

From 0ed4548f8dd3fee92dbd4cc47e6393762eac3f30 Mon Sep 17 00:00:00 2001
From: shimwell <mail@jshimwell.com>
Date: Thu, 27 Jul 2023 09:12:57 +0100
Subject: [PATCH 08/17] cbar title added

---
 openmc/plots.py                | 18 ++++++++++++++++--
 tests/unit_tests/test_plots.py |  4 +++-
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index 0b003fdd207..578ff1a4e01 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -180,6 +180,9 @@ def plot_mesh_tally(
     outline_color: str = 'black',
     pixels: int = 40000,
     geometry: Optional['openmc.Geometry'] = None,
+    colorbar: bool = True,
+    color_bar_title : str = None,
+    volume_normalization: bool = True,
     **kwargs
 ) -> 'matplotlib.image.AxesImage':
     """Display a slice plot of the mesh tally score.
@@ -214,6 +217,12 @@ def plot_mesh_tally(
         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
+    color_bar_title : str
+        The title to assign the color bar.
+    volume_normalization : bool, optional
+        Whether or not to normalize the data by
+        the volume of the mesh elements.
     **kwargs
         Keyword arguments passed to :func:`matplotlib.pyplot.imshow`
 
@@ -262,8 +271,13 @@ def plot_mesh_tally(
         fig, axes = plt.subplots()
         axes.set_xlabel(xlabel)
         axes.set_ylabel(ylabel)
-    axes.imshow(oriented_data, extent=(x_min, x_max, y_min, y_max), **kwargs)
-    # fig.colorbar()
+    im = axes.imshow(oriented_data, extent=(x_min, x_max, y_min, y_max), **kwargs)
+    if colorbar:
+        cbar = fig.colorbar(im)
+        if color_bar_title is None:
+            cbar.set_label(score)
+        else:
+            cbar.set_label(color_bar_title)
 
     if outline and geometry is not None:
         import matplotlib.image as mpimg
diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py
index 0ddf76a5a2d..a1347e0e298 100644
--- a/tests/unit_tests/test_plots.py
+++ b/tests/unit_tests/test_plots.py
@@ -285,6 +285,7 @@ def test_voxel_plot_roundtrip():
     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
@@ -300,7 +301,8 @@ def test_voxel_plot_roundtrip():
     # norm=LogNorm(vmin=0.1, vmax=100),
     outline=True,
     geometry=geom,
-    outline_by='material'
+    outline_by='material',
+    color_bar_title='neutron flux'
 )
 # assert plot.xaxis.get_label().get_text() == 'x [mm]'
 # assert plot.yaxis.get_label().get_text() == 'z [mm]'

From 7c4f56d8db27635428843852fefacbe105f813f4 Mon Sep 17 00:00:00 2001
From: shimwell <mail@jshimwell.com>
Date: Fri, 28 Jul 2023 01:15:24 +0100
Subject: [PATCH 09/17] vol norm

---
 openmc/plots.py                | 28 +++++++++++++++++++++-------
 tests/unit_tests/test_plots.py |  7 ++++---
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index 578ff1a4e01..26279202107 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -166,6 +166,11 @@
     'yellowgreen': (154, 205, 50)
 }
 
+_default_outline_kwargs = {
+    'colors': 'black',
+    'linestyles': 'solid',
+    'linewidths': 1
+}
 
 def plot_mesh_tally(
     tally: 'openmc.Tally',
@@ -177,12 +182,13 @@ def plot_mesh_tally(
     value: str = 'mean',
     outline: bool = False,
     outline_by: str = 'cell',
-    outline_color: str = 'black',
     pixels: int = 40000,
     geometry: Optional['openmc.Geometry'] = None,
     colorbar: bool = True,
     color_bar_title : str = None,
     volume_normalization: bool = True,
+    scaling_factor: Optional[float] = None,
+    outline_kwargs: dict = _default_outline_kwargs,
     **kwargs
 ) -> 'matplotlib.image.AxesImage':
     """Display a slice plot of the mesh tally score.
@@ -221,8 +227,11 @@ def plot_mesh_tally(
     color_bar_title : str
         The title to assign the color bar.
     volume_normalization : bool, optional
-        Whether or not to normalize the data by
-        the volume of the mesh elements.
+        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.
+    outline_kwargs : dict
+        Keyword arguments passed to :func:`matplotlib.pyplot.contour`.
     **kwargs
         Keyword arguments passed to :func:`matplotlib.pyplot.imshow`
 
@@ -263,6 +272,12 @@ def plot_mesh_tally(
         oriented_data = np.rot90(slice_data, 1)
         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
+        oriented_data = oriented_data / mesh.volumes[0][0][0]
+    if scaling_factor:
+        oriented_data = oriented_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]]
@@ -271,7 +286,9 @@ def plot_mesh_tally(
         fig, axes = plt.subplots()
         axes.set_xlabel(xlabel)
         axes.set_ylabel(ylabel)
+
     im = axes.imshow(oriented_data, extent=(x_min, x_max, y_min, y_max), **kwargs)
+
     if colorbar:
         cbar = fig.colorbar(im)
         if color_bar_title is None:
@@ -316,12 +333,9 @@ def plot_mesh_tally(
         axes.contour(
             image_value,
             origin="upper",
-            colors=outline_color,
-            linestyles="solid",
-            linewidths=1,
             levels=np.unique(image_value),
             extent=(x_min, x_max, y_min, y_max),
-            **kwargs
+            **outline_kwargs
         )
 
     return axes
diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py
index a1347e0e298..b49283a9ee0 100644
--- a/tests/unit_tests/test_plots.py
+++ b/tests/unit_tests/test_plots.py
@@ -240,6 +240,7 @@ def test_voxel_plot_roundtrip():
 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()
@@ -270,7 +271,7 @@ def test_voxel_plot_roundtrip():
 plot = openmc.plot_mesh_tally(
     tally=tally_result,
     basis='xy',
-    slice_index=29  # max value of slice selected
+    # slice_index=29  # max value of slice selected
 )
 # axis_units defaults to cm
 # assert plot.xaxis.get_label().get_text() == 'x [cm]'
@@ -298,11 +299,11 @@ def test_voxel_plot_roundtrip():
     # axis_units='mm',
     score='flux',
     value= 'mean',
-    # norm=LogNorm(vmin=0.1, vmax=100),
     outline=True,
     geometry=geom,
     outline_by='material',
-    color_bar_title='neutron flux'
+    color_bar_title='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]'

From 4e74c35b1cac0cf369842b80c4cd803a562820f2 Mon Sep 17 00:00:00 2001
From: shimwell <mail@jshimwell.com>
Date: Sun, 30 Jul 2023 23:06:59 +0100
Subject: [PATCH 10/17] rotating contours;

---
 openmc/plots.py | 43 ++++++++++++++++++++++++++++++-------------
 1 file changed, 30 insertions(+), 13 deletions(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index 26279202107..1deabc02516 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -224,6 +224,7 @@ def plot_mesh_tally(
         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.
     color_bar_title : str
         The title to assign the color bar.
     volume_normalization : bool, optional
@@ -245,38 +246,47 @@ def plot_mesh_tally(
     import matplotlib.pyplot as plt
 
     cv.check_value('basis', basis, ['xy', 'xz', 'yz'])
+    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 and len(tally.scores) == 1:
-        score = tally.scores[0]
+    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)
 
-    data = tally.get_reshaped_data(expand_dims=True, value=value).squeeze()
+    tally_data = tally.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(data.shape[basis_to_index]/2)
+        slice_index = int(tally_data.shape[basis_to_index]/2)
 
     if basis == 'xz':
-        slice_data = data[:, slice_index, :]
-        oriented_data = np.flip(np.rot90(slice_data, -1))
+        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 = data[slice_index, :, :]
-        oriented_data = np.flip(np.rot90(slice_data, -1))
+        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 = data[:, :, slice_index]
-        oriented_data = np.rot90(slice_data, 1)
+        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
-        oriented_data = oriented_data / mesh.volumes[0][0][0]
+        data = data / mesh.volumes[0][0][0]
+
     if scaling_factor:
-        oriented_data = oriented_data * scaling_factor
+        data = data * scaling_factor
 
     axis_scaling_factor = {'km': 0.00001, 'm': 0.01, 'cm': 1, 'mm': 10}[axis_units]
 
@@ -287,7 +297,7 @@ def plot_mesh_tally(
         axes.set_xlabel(xlabel)
         axes.set_ylabel(ylabel)
 
-    im = axes.imshow(oriented_data, extent=(x_min, x_max, y_min, y_max), **kwargs)
+    im = axes.imshow(data, extent=(x_min, x_max, y_min, y_max), **kwargs)
 
     if colorbar:
         cbar = fig.colorbar(im)
@@ -329,6 +339,13 @@ def plot_mesh_tally(
         image_value = (rgb[..., 0] << 16) + \
             (rgb[..., 1] << 8) + (rgb[..., 2])
 
+        if basis == 'xz':
+            image_value = np.flip(np.rot90(image_value, -1))
+        elif basis == 'yz':
+            image_value = np.flip(np.rot90(image_value, -1))
+        else:  # basis == 'xy'
+            image_value = np.rot90(image_value, 2)
+
         # Plot image and return the axes
         axes.contour(
             image_value,

From f8e794705b13774dd12e24d8841352bac5206ec6 Mon Sep 17 00:00:00 2001
From: shimwell <mail@jshimwell.com>
Date: Sun, 30 Jul 2023 23:10:17 +0100
Subject: [PATCH 11/17] testing tally plot

---
 tests/unit_tests/test_plots.py | 172 ++++++++++++++++-----------------
 1 file changed, 86 insertions(+), 86 deletions(-)

diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py
index b49283a9ee0..e9b1e261682 100644
--- a/tests/unit_tests/test_plots.py
+++ b/tests/unit_tests/test_plots.py
@@ -224,89 +224,89 @@ def test_voxel_plot_roundtrip():
     assert new_plot.color_by == plot.color_by
 
 
-# def test_plot_mesh_tally(run_in_tmpdir):
-
-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',
-    color_bar_title='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')
+def test_plot_mesh_tally(run_in_tmpdir):
+
+    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',
+        color_bar_title='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')

From 49f56805d7eae11490bdec075933e7cc650ce24e Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Fri, 4 Aug 2023 16:45:26 +0100
Subject: [PATCH 12/17] tally plot working with multi score tallies

---
 openmc/plots.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index 1deabc02516..5605f06d4c4 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -172,6 +172,7 @@
     'linewidths': 1
 }
 
+
 def plot_mesh_tally(
     tally: 'openmc.Tally',
     basis: str = 'xy',
@@ -185,7 +186,7 @@ def plot_mesh_tally(
     pixels: int = 40000,
     geometry: Optional['openmc.Geometry'] = None,
     colorbar: bool = True,
-    color_bar_title : str = None,
+    color_bar_title: str = None,
     volume_normalization: bool = True,
     scaling_factor: Optional[float] = None,
     outline_kwargs: dict = _default_outline_kwargs,
@@ -245,7 +246,7 @@ def plot_mesh_tally(
 
     import matplotlib.pyplot as plt
 
-    cv.check_value('basis', basis, ['xy', 'xz', 'yz'])
+    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)
@@ -262,7 +263,9 @@ def plot_mesh_tally(
             msg = 'score was not specified and there are multiple scores in the tally.'
             raise ValueError(msg)
 
-    tally_data = tally.get_reshaped_data(expand_dims=True, value=value).squeeze()
+    tally_slice = tally.get_slice(scores=[score])
+    print(tally_slice)
+    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]
@@ -340,9 +343,9 @@ def plot_mesh_tally(
             (rgb[..., 1] << 8) + (rgb[..., 2])
 
         if basis == 'xz':
-            image_value = np.flip(np.rot90(image_value, -1))
+            image_value = np.rot90(image_value, 2)
         elif basis == 'yz':
-            image_value = np.flip(np.rot90(image_value, -1))
+            image_value = np.rot90(image_value, 2)
         else:  # basis == 'xy'
             image_value = np.rot90(image_value, 2)
 

From c6127d201e229cdfd034bfefdd298a77a799533a Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Fri, 4 Aug 2023 16:52:47 +0100
Subject: [PATCH 13/17] moved imports to top

---
 openmc/plots.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index 5605f06d4c4..96e1ab235cc 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
@@ -311,8 +313,6 @@ def plot_mesh_tally(
 
     if outline and geometry is not None:
         import matplotlib.image as mpimg
-        import math
-        from tempfile import TemporaryDirectory
         model = openmc.Model()
         model.geometry = geometry
         plot = openmc.Plot()

From 49998b1475673129f8680b24f187821da4a6eebb Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <drshimwell@gmail.com>
Date: Fri, 4 Aug 2023 18:20:21 +0100
Subject: [PATCH 14/17] moving matplotlib import into function for tests

---
 tests/unit_tests/test_plots.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py
index e9b1e261682..31663a98ead 100644
--- a/tests/unit_tests/test_plots.py
+++ b/tests/unit_tests/test_plots.py
@@ -3,7 +3,6 @@
 import openmc
 import openmc.examples
 import pytest
-from matplotlib.colors import LogNorm
 
 
 @pytest.fixture(scope='module')
@@ -225,7 +224,7 @@ def test_voxel_plot_roundtrip():
 
 
 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])

From fbb110898f10a1d7260e2e4c6c6496692089416c Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Wed, 9 Aug 2023 17:52:33 +0100
Subject: [PATCH 15/17] removed print statement

---
 openmc/plots.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index 96e1ab235cc..e49637caafe 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -266,7 +266,7 @@ def plot_mesh_tally(
             raise ValueError(msg)
 
     tally_slice = tally.get_slice(scores=[score])
-    print(tally_slice)
+
     tally_data = tally_slice.get_reshaped_data(expand_dims=True, value=value).squeeze()
 
     if slice_index is None:

From 513ef6a76abe3d08c7c861ac0dc07c473ee4a11e Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Wed, 9 Aug 2023 17:56:36 +0100
Subject: [PATCH 16/17] tidied up docstring

---
 openmc/plots.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index e49637caafe..e2e0293bded 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -185,8 +185,8 @@ def plot_mesh_tally(
     value: str = 'mean',
     outline: bool = False,
     outline_by: str = 'cell',
-    pixels: int = 40000,
     geometry: Optional['openmc.Geometry'] = None,
+    pixels: int = 40000,
     colorbar: bool = True,
     color_bar_title: str = None,
     volume_normalization: bool = True,
@@ -218,8 +218,6 @@ def plot_mesh_tally(
         by cell or by material.
     outline_by : {'cell', 'material'}
         Indicate whether the plot should be colored by cell or by material
-    outline_color : str
-        The matplotlib color to use for the plot.
     geometry : openmc.Geometry
         The geometry to use for the outline.
     pixels : int

From 640016c8a167bb928c07bc88d27bb4b02b1a5f2a Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Tue, 22 Aug 2023 12:02:14 +0100
Subject: [PATCH 17/17] flexible input for cbar kwargs

---
 openmc/plots.py                | 12 ++++--------
 tests/unit_tests/test_plots.py |  8 ++++----
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/openmc/plots.py b/openmc/plots.py
index e2e0293bded..d875960ab54 100644
--- a/openmc/plots.py
+++ b/openmc/plots.py
@@ -188,9 +188,9 @@ def plot_mesh_tally(
     geometry: Optional['openmc.Geometry'] = None,
     pixels: int = 40000,
     colorbar: bool = True,
-    color_bar_title: str = None,
     volume_normalization: bool = True,
     scaling_factor: Optional[float] = None,
+    colorbar_kwargs: dict = {},
     outline_kwargs: dict = _default_outline_kwargs,
     **kwargs
 ) -> 'matplotlib.image.AxesImage':
@@ -226,12 +226,12 @@ def plot_mesh_tally(
         the image aspect ratio.
     colorbar : bool
         Whether or not to add a colorbar to the plot.
-    color_bar_title : str
-        The title to assign the color bar.
     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
@@ -303,11 +303,7 @@ def plot_mesh_tally(
     im = axes.imshow(data, extent=(x_min, x_max, y_min, y_max), **kwargs)
 
     if colorbar:
-        cbar = fig.colorbar(im)
-        if color_bar_title is None:
-            cbar.set_label(score)
-        else:
-            cbar.set_label(color_bar_title)
+        fig.colorbar(im, **colorbar_kwargs)
 
     if outline and geometry is not None:
         import matplotlib.image as mpimg
diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py
index 31663a98ead..8101e0eb254 100644
--- a/tests/unit_tests/test_plots.py
+++ b/tests/unit_tests/test_plots.py
@@ -237,7 +237,7 @@ def test_plot_mesh_tally(run_in_tmpdir):
         -100, 50, -200, 250, -300, 350, boundary_type='vacuum'
     )
     cell1 = openmc.Cell(region=-surface1)
-    cell2 = openmc.Cell(region=-surface&+surface1)
+    cell2 = openmc.Cell(region=-surface & +surface1)
     cell1.fill = mat1
     cell2.fill = mat1
     geom = openmc.Geometry([cell1, cell2])
@@ -283,7 +283,7 @@ def test_plot_mesh_tally(run_in_tmpdir):
         basis='yz',
         axis_units='m',
         slice_index=9,  # max value of slice selected
-        value= 'std_dev'
+        value='std_dev'
     )
     plot.figure.savefig('x.png')
     assert plot.xaxis.get_label().get_text() == 'y [m]'
@@ -297,11 +297,11 @@ def test_plot_mesh_tally(run_in_tmpdir):
         slice_index=19,  # max value of slice selected
         axis_units='mm',
         score='flux',
-        value= 'mean',
+        value='mean',
         outline=True,
         geometry=geom,
         outline_by='material',
-        color_bar_title='neutron flux',
+        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]'