diff --git a/openmc/universe.py b/openmc/universe.py index 7a429c019e..67296dda8c 100644 --- a/openmc/universe.py +++ b/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: diff --git a/tests/unit_tests/test_plots.py b/tests/unit_tests/test_plots.py index 183341c26e..0dddc827cc 100644 --- a/tests/unit_tests/test_plots.py +++ b/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 diff --git a/tests/unit_tests/test_universe.py b/tests/unit_tests/test_universe.py index a70bd46673..1bfbbf7ff8 100644 --- a/tests/unit_tests/test_universe.py +++ b/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, )