Merge pull request #18423 from danuo/2d_array_color_support

QuLogic · web-flow · commit 8676c2782838 · 2020-09-16T17:31:38.000-04:00
2-D array RGB and RGBA values not understood in plt.plot()
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -265,6 +265,10 @@ def _to_rgba_no_colorcycle(c, alpha=None):
                     f"Value must be within 0-1 range")
             return c, c, c, alpha if alpha is not None else 1.
         raise ValueError(f"Invalid RGBA argument: {orig_c!r}")
+    # turn 2-D array into 1-D array
+    if isinstance(c, np.ndarray):
+        if c.ndim == 2 and c.shape[0] == 1:
+            c = c.reshape(-1)
     # tuple color.
     if not np.iterable(c):
         raise ValueError(f"Invalid RGBA argument: {orig_c!r}")
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -6590,3 +6590,22 @@ def test_sharing_does_not_link_positions():
     init_pos = ax1.get_position()
     fig.subplots_adjust(left=0)
     assert (ax1.get_position().get_points() == init_pos.get_points()).all()
+
+
+@check_figures_equal(extensions=["pdf"])
+def test_2dcolor_plot(fig_test, fig_ref):
+    color = np.array([0.1, 0.2, 0.3])
+    # plot with 1D-color:
+    axs = fig_test.subplots(5)
+    axs[0].plot([1, 2], [1, 2], c=color.reshape(-1))
+    axs[1].scatter([1, 2], [1, 2], c=color.reshape(-1))
+    axs[2].step([1, 2], [1, 2], c=color.reshape(-1))
+    axs[3].hist(np.arange(10), color=color.reshape(-1))
+    axs[4].bar(np.arange(10), np.arange(10), color=color.reshape(-1))
+    # plot with 2D-color:
+    axs = fig_ref.subplots(5)
+    axs[0].plot([1, 2], [1, 2], c=color.reshape((1, -1)))
+    axs[1].scatter([1, 2], [1, 2], c=color.reshape((1, -1)))
+    axs[2].step([1, 2], [1, 2], c=color.reshape((1, -1)))
+    axs[3].hist(np.arange(10), color=color.reshape((1, -1)))
+    axs[4].bar(np.arange(10), np.arange(10), color=color.reshape((1, -1)))
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -1211,3 +1211,10 @@ def test_colormap_bad_data_with_alpha():
     assert_array_equal(c[0, 0], (0, 0, 0, 0))
     c = cmap([[np.nan, 0.5], [0, 0]], alpha=np.full((2, 2), 0.5))
     assert_array_equal(c[0, 0], (0, 0, 0, 0))
+
+
+def test_2d_to_rgba():
+    color = np.array([0.1, 0.2, 0.3])
+    rgba_1d = mcolors.to_rgba(color.reshape(-1))
+    rgba_2d = mcolors.to_rgba(color.reshape((1, -1)))
+    assert rgba_1d == rgba_2d
