Merge pull request matplotlib#24638 from greglucas/pcolormesh-no-flatten

MNT: Remove auto-flattening of input data to pcolormesh

Author: jklymak

lib/matplotlib/axes/_axes.py

@@ -6220,9 +6220,6 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
         X, Y, C, shading = self._pcolorargs('pcolormesh', *args,
                                             shading=shading, kwargs=kwargs)
         coords = np.stack([X, Y], axis=-1)
-        # convert to one dimensional array, except for 3D RGB(A) arrays
-        if C.ndim != 3:
-            C = C.ravel()
 
         kwargs.setdefault('snap', mpl.rcParams['pcolormesh.snap'])
 

lib/matplotlib/tests/test_collections.py

@@ -914,13 +914,13 @@ def test_quadmesh_vmin_vmax():
     norm = mpl.colors.Normalize(vmin=0, vmax=1)
     coll = ax.pcolormesh([[1]], cmap=cmap, norm=norm)
     fig.canvas.draw()
-    assert np.array_equal(coll.get_facecolors()[0, :], cmap(norm(1)))
+    assert np.array_equal(coll.get_facecolors()[0, 0, :], cmap(norm(1)))
 
     # Change the vmin/vmax of the norm so that the color is from
     # the bottom of the colormap now
     norm.vmin, norm.vmax = 1, 2
     fig.canvas.draw()
-    assert np.array_equal(coll.get_facecolors()[0, :], cmap(norm(1)))
+    assert np.array_equal(coll.get_facecolors()[0, 0, :], cmap(norm(1)))
 
 
 def test_quadmesh_alpha_array():
@@ -935,16 +935,16 @@ def test_quadmesh_alpha_array():
     coll2 = ax1.pcolormesh(x, y, z)
     coll2.set_alpha(alpha)
     plt.draw()
-    assert_array_equal(coll1.get_facecolors()[:, -1], alpha_flat)
-    assert_array_equal(coll2.get_facecolors()[:, -1], alpha_flat)
+    assert_array_equal(coll1.get_facecolors()[..., -1], alpha)
+    assert_array_equal(coll2.get_facecolors()[..., -1], alpha)
     # Or provide 1-D alpha:
     fig, (ax0, ax1) = plt.subplots(2)
-    coll1 = ax0.pcolormesh(x, y, z, alpha=alpha_flat)
+    coll1 = ax0.pcolormesh(x, y, z, alpha=alpha)
     coll2 = ax1.pcolormesh(x, y, z)
-    coll2.set_alpha(alpha_flat)
+    coll2.set_alpha(alpha)
     plt.draw()
-    assert_array_equal(coll1.get_facecolors()[:, -1], alpha_flat)
-    assert_array_equal(coll2.get_facecolors()[:, -1], alpha_flat)
+    assert_array_equal(coll1.get_facecolors()[..., -1], alpha)
+    assert_array_equal(coll2.get_facecolors()[..., -1], alpha)
 
 
 def test_alpha_validation():
@@ -992,7 +992,7 @@ def test_color_logic(pcfunc):
     pc.update_scalarmappable()  # This is called in draw().
     # Define 2 reference "colors" here for multiple use.
     face_default = mcolors.to_rgba_array(pc._get_default_facecolor())
-    mapped = pc.get_cmap()(pc.norm(z.ravel()))
+    mapped = pc.get_cmap()(pc.norm(z.ravel() if pcfunc == plt.pcolor else z))
     # GitHub issue #1302:
     assert mcolors.same_color(pc.get_edgecolor(), 'red')
     # Check setting attributes after initialization:
@@ -1011,30 +1011,30 @@ def test_color_logic(pcfunc):
     # Reset edgecolor to default.
     pc.set_edgecolor(None)
     pc.update_scalarmappable()
-    assert mcolors.same_color(pc.get_edgecolor(), mapped)
+    assert np.array_equal(pc.get_edgecolor(), mapped)
     pc.set_facecolor(None)  # restore default for facecolor
     pc.update_scalarmappable()
-    assert mcolors.same_color(pc.get_facecolor(), mapped)
+    assert np.array_equal(pc.get_facecolor(), mapped)
     assert mcolors.same_color(pc.get_edgecolor(), 'none')
     # Turn off colormapping entirely:
     pc.set_array(None)
     pc.update_scalarmappable()
     assert mcolors.same_color(pc.get_edgecolor(), 'none')
     assert mcolors.same_color(pc.get_facecolor(), face_default)  # not mapped
     # Turn it back on by restoring the array (must be 1D!):
-    pc.set_array(z.ravel())
+    pc.set_array(z.ravel() if pcfunc == plt.pcolor else z)
     pc.update_scalarmappable()
-    assert mcolors.same_color(pc.get_facecolor(), mapped)
+    assert np.array_equal(pc.get_facecolor(), mapped)
     assert mcolors.same_color(pc.get_edgecolor(), 'none')
     # Give color via tuple rather than string.
     pc = pcfunc(z, edgecolors=(1, 0, 0), facecolors=(0, 1, 0))
     pc.update_scalarmappable()
-    assert mcolors.same_color(pc.get_facecolor(), mapped)
+    assert np.array_equal(pc.get_facecolor(), mapped)
     assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
     # Provide an RGB array; mapping overrides it.
     pc = pcfunc(z, edgecolors=(1, 0, 0), facecolors=np.ones((12, 3)))
     pc.update_scalarmappable()
-    assert mcolors.same_color(pc.get_facecolor(), mapped)
+    assert np.array_equal(pc.get_facecolor(), mapped)
     assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
     # Turn off the mapping.
     pc.set_array(None)
@@ -1044,7 +1044,7 @@ def test_color_logic(pcfunc):
     # And an RGBA array.
     pc = pcfunc(z, edgecolors=(1, 0, 0), facecolors=np.ones((12, 4)))
     pc.update_scalarmappable()
-    assert mcolors.same_color(pc.get_facecolor(), mapped)
+    assert np.array_equal(pc.get_facecolor(), mapped)
     assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
     # Turn off the mapping.
     pc.set_array(None)