Adding 2d support to quadmesh set_array

greglucas · greglucas · commit 5e5ac0178564 · 2020-05-25T07:03:52.000-06:00
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -779,7 +779,8 @@ def update_scalarmappable(self):
         """Update colors from the scalar mappable array, if it is not None."""
         if self._A is None:
             return
-        if self._A.ndim > 1:
+        # QuadMesh can map 2d arrays
+        if self._A.ndim > 1 and not isinstance(self, QuadMesh):
             raise ValueError('Collections can only map rank 1 arrays')
         if not self._check_update("array"):
             return
@@ -2044,8 +2045,10 @@ def draw(self, renderer):
         else:
             renderer.draw_quad_mesh(
                 gc, transform.frozen(), self._meshWidth, self._meshHeight,
-                coordinates, offsets, transOffset, self.get_facecolor(),
-                self._antialiased, self.get_edgecolors())
+                coordinates, offsets, transOffset,
+                # Backends expect flattened rgba arrays (n*m, 4) for fc and ec
+                self.get_facecolor().reshape((-1, 4)),
+                self._antialiased, self.get_edgecolors().reshape((-1, 4)))
         gc.restore()
         renderer.close_group(self.__class__.__name__)
         self.stale = False
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -637,3 +637,20 @@ def test_singleton_autolim():
     ax.scatter(0, 0)
     np.testing.assert_allclose(ax.get_ylim(), [-0.06, 0.06])
     np.testing.assert_allclose(ax.get_xlim(), [-0.06, 0.06])
+
+
+def test_quadmesh_set_array():
+    x = np.arange(4)
+    y = np.arange(4)
+    z = np.arange(9).reshape((3, 3))
+    fig, ax = plt.subplots()
+    coll = ax.pcolormesh(x, y, np.ones(z.shape))
+    # Test that the collection is able to update with a 2d array
+    coll.set_array(z)
+    fig.canvas.draw()
+    assert np.array_equal(coll.get_array(), z)
+
+    # Check that pre-flattened arrays work too
+    coll.set_array(np.ones(9))
+    fig.canvas.draw()
+    assert np.array_equal(coll.get_array(), np.ones(9))
