diff --git a/openmc/universe.py b/openmc/universe.py index ab9858f68f..3984a0a61d 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -426,7 +426,8 @@ class Universe(UniverseBase): 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]) + image_value = (rgb[..., 0] << 16) + \ + (rgb[..., 1] << 8) + (rgb[..., 2]) axes.contour( image_value, @@ -462,7 +463,19 @@ class Universe(UniverseBase): # this works whether we're doing cells or materials label = key.name if key.name != '' else key.id - key_patch = mpatches.Patch(color=color, label=label) + + # matplotlib takes RGB on 0-1 scale rather than 0-255. at + # this point PlotBase has already checked that 3-tuple + # based colors are already valid, so if the length is three + # then we know it just needs to be converted to the 0-1 + # format. + if len(color) == 3 and not isinstance(color, str): + scaled_color = ( + color[0]/255, color[1]/255, color[2]/255) + else: + scaled_color = color + + key_patch = mpatches.Patch(color=scaled_color, label=label) patches.append(key_patch) axes.legend(handles=patches, **legend_kwargs) diff --git a/tests/unit_tests/test_universe.py b/tests/unit_tests/test_universe.py index 1bfbbf7ff8..cb4f2b7030 100644 --- a/tests/unit_tests/test_universe.py +++ b/tests/unit_tests/test_universe.py @@ -48,18 +48,25 @@ def test_bounding_box(): assert_unbounded(u) -def test_plot(run_in_tmpdir, sphere_model): - m = sphere_model.materials[0] - univ = sphere_model.geometry.root_universe +def test_plot(run_in_tmpdir): + + pincell = openmc.examples.pwr_pin_cell() + materials = pincell.materials + + mat_colors = { + materials[0]: (200, 1, 1), + materials[1]: "gray", + materials[2]: "limegreen" + } - colors = {m: 'limegreen'} for basis in ('xy', 'yz', 'xz'): - univ.plot( - basis=basis, + pincell.geometry.root_universe.plot( + colors=mat_colors, + color_by="material", + legend=True, pixels=(10, 10), - color_by='material', - colors=colors, - outline=True, + basis=basis, + outline=True )