Merge pull request #2487 from gridley/legend_color_fix

pass 0-1 RGB value to matplotlib in universe plot legend

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)

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
         )