Merge pull request #2482 from shimwell/adding_outlines_to_plot

adding optional outline to plots

openmc/universe.py

@@ -1,3 +1,4 @@
+import typing
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 from collections.abc import Iterable
@@ -302,7 +303,8 @@ class Universe(UniverseBase):
     def plot(self, origin=(0., 0., 0.), width=(1., 1.), pixels=(200, 200),
              basis='xy', color_by='cell', colors=None, seed=None,
              openmc_exec='openmc', axes=None, legend=False,
-             legend_kwargs=_default_legend_kwargs, **kwargs):
+             legend_kwargs=_default_legend_kwargs, outline=False,
+             **kwargs):
         """Display a slice plot of the universe.
 
         Parameters
@@ -343,11 +345,14 @@ class Universe(UniverseBase):
         legend_kwargs : dict
             Keyword arguments passed to :func:`matplotlib.pyplot.legend`.
 
+            .. versionadded:: 0.13.4
+        outline : bool
+            Whether outlines between color boundaries should be drawn
+
             .. versionadded:: 0.13.4
         **kwargs
             Keyword arguments passed to :func:`matplotlib.pyplot.imshow`
 
-            .. versionadded:: 0.13.4
         Returns
         -------
         matplotlib.image.AxesImage
@@ -418,6 +423,21 @@ class Universe(UniverseBase):
                 height = pixels[0]*px/(params.top - params.bottom)
                 fig.set_size_inches(width, height)
 
+            if outline:
+                # Combine R, G, B values into a single int
+                rgb = (img * 256).astype(int)
+                image_value = (rgb[..., 0] << 16) + (rgb[..., 1] << 8) + (rgb[..., 2])
+
+                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),
+                )
+
             # add legend showing which colors represent which material
             # or cell if that was requested
             if legend:

tests/unit_tests/test_plots.py

@@ -59,26 +59,31 @@ def myprojectionplot():
 
     plot.overlap_color = (255, 211, 0)
     plot.overlap_color = 'yellow'
-    
+
     plot.wireframe_thickness = 2
 
     plot.level = 1
     return plot
 
+
 def test_attributes(myplot):
     assert myplot.name == 'myplot'
 
+
 def test_attributes_proj(myprojectionplot):
     assert myprojectionplot.name == 'myprojectionplot'
 
+
 def test_repr(myplot):
     r = repr(myplot)
     assert isinstance(r, str)
 
+
 def test_repr_proj(myprojectionplot):
     r = repr(myprojectionplot)
     assert isinstance(r, str)
 
+
 def test_from_geometry():
     width = 25.
     s = openmc.Sphere(r=width/2, boundary_type='vacuum')
@@ -155,8 +160,8 @@ def test_plots(run_in_tmpdir):
 
     # from_xml
     new_plots = openmc.Plots.from_xml()
-    assert len(plots)
-    assert plots[0].origin == p1.origin
-    assert plots[0].colors == p1.colors
-    assert plots[0].mask_components == p1.mask_components
-    assert plots[1].origin == p2.origin
+    assert len(new_plots)
+    assert new_plots[0].origin == p1.origin
+    assert new_plots[0].colors == p1.colors
+    assert new_plots[0].mask_components == p1.mask_components
+    assert new_plots[1].origin == p2.origin

tests/unit_tests/test_universe.py

@@ -59,6 +59,7 @@ def test_plot(run_in_tmpdir, sphere_model):
             pixels=(10, 10),
             color_by='material',
             colors=colors,
+            outline=True,
         )