From 5e5ac0178564a2e429776451ba818f4e47f27147 Mon Sep 17 00:00:00 2001 From: Greg Lucas Date: Wed, 25 Mar 2020 15:55:34 -0600 Subject: [PATCH] Adding 2d support to quadmesh set_array --- lib/matplotlib/collections.py | 9 ++++++--- lib/matplotlib/tests/test_collections.py | 17 +++++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py index feef6c8d9e98..549cc350e4f4 100644 --- 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 index 2eaf703474e0..77f6e1884813 100644 --- 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))