MNT: Cleanup PolyQuadMesh

greglucas · greglucas · commit d2bad4357305 · 2023-04-23T12:58:15.000-06:00
Better handle the subclassing in PolyQuadMesh. Most of the methods
should come from the PolyCollection as that is doing the primary
drawing. The QuadMesh is really contributing the set_array()/get_array()
methods.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5731,13 +5731,7 @@ def _pcolorargs(self, funcname, *args, shading='auto', **kwargs):
                     raise ValueError(
                         'x and y arguments to pcolormesh cannot have '
                         'non-finite values or be of type '
-                        'numpy.ma.core.MaskedArray with masked values')
-                # safe_masked_invalid() returns an ndarray for dtypes other
-                # than floating point.
-                if isinstance(X, np.ma.core.MaskedArray):
-                    X = X.data  # strip mask as downstream doesn't like it...
-                if isinstance(Y, np.ma.core.MaskedArray):
-                    Y = Y.data
+                        'numpy.ma.MaskedArray with masked values')
             nrows, ncols = C.shape[:2]
         else:
             raise _api.nargs_error(funcname, takes="1 or 3", given=len(args))
@@ -5787,10 +5781,8 @@ def _interp_grid(X):
                                 "This may lead to incorrectly calculated cell "
                                 "edges, in which case, please supply "
                                 f"explicit cell edges to {funcname}.")
-                        if isinstance(X, np.ma.core.MaskedArray):
-                            hstack = np.ma.hstack
-                        else:
-                            hstack = np.hstack
+
+                        hstack = np.ma.hstack if np.ma.isMA(X) else np.hstack
                         X = hstack((X[:, [0]] - dX[:, [0]],
                                     X[:, :-1] + dX,
                                     X[:, [-1]] + dX[:, [-1]]))
@@ -5983,8 +5975,7 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
 
         kwargs.setdefault('snap', False)
 
-        if (isinstance(X, np.ma.MaskedArray)
-                or isinstance(Y, np.ma.MaskedArray)):
+        if np.ma.isMaskedArray(X) or np.ma.isMaskedArray(Y):
             stack = np.ma.stack
             X = np.ma.asarray(X)
             Y = np.ma.asarray(Y)
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -2170,18 +2170,17 @@ def get_cursor_data(self, event):
 
 class PolyQuadMesh(PolyCollection, QuadMesh):
     def __init__(self, coordinates, **kwargs):
-        QuadMesh.__init__(self, coordinates=coordinates, **kwargs)
-        # Need sizes set because we are using Polycollection.draw()
-        self._sizes = np.ones(len(self._coordinates))
-        C = self.get_array()
-        X = self._coordinates[..., 0]
-        Y = self._coordinates[..., 1]
+        X = coordinates[..., 0]
+        Y = coordinates[..., 1]
 
         mask = np.ma.getmaskarray(X) | np.ma.getmaskarray(Y)
-        xymask = (mask[0:-1, 0:-1] + mask[1:, 1:] +
-                  mask[0:-1, 1:] + mask[1:, 0:-1])
-        # don't plot if C or any of the surrounding vertices are masked.
-        mask = np.ma.getmaskarray(C) | xymask
+        # We want the shape of the polygon, which is the corner of
+        # each X/Y array
+        mask = (mask[0:-1, 0:-1] | mask[1:, 1:] |
+                mask[0:-1, 1:] | mask[1:, 0:-1])
+        # Take account of the array data too
+        C = kwargs.get("array", np.empty(coordinates.shape[:2]))
+        mask |= np.ma.getmaskarray(C)
 
         unmask = ~mask
         X1 = np.ma.filled(X[:-1, :-1])[unmask]
@@ -2196,9 +2195,5 @@ def __init__(self, coordinates, **kwargs):
 
         xy = np.ma.stack([X1, Y1, X2, Y2, X3, Y3, X4, Y4, X1, Y1], axis=-1)
         verts = xy.reshape((npoly, 5, 2))
-        self.set_verts(verts)
-
-    @artist.allow_rasterization
-    def draw(self, renderer):
-        # We want to use the PolyCollection's draw
-        return PolyCollection.draw(self, renderer)
+        # We need both verts and coordinates here to go through the super chain
+        super().__init__(coordinates=coordinates, verts=verts, **kwargs)
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -1068,17 +1068,21 @@ def test_LineCollection_args():
     assert mcolors.same_color(lc.get_facecolor(), 'none')
 
 
-def test_array_wrong_dimensions():
+@pytest.mark.parametrize("func", (plt.pcolor, plt.pcolormesh))
+def test_array_dimensions(func):
+    # Make sure we can set the 1D, 2D, and 3D array shapes
     z = np.arange(12).reshape(3, 4)
-    pc = plt.pcolor(z)
-    pc.set_array(z)
+    pc = func(z)
+    # 1D
+    pc.set_array(z.ravel())
     pc.update_scalarmappable()
-    pc = plt.pcolormesh(z)
-    pc.set_array(z)  # 2D is OK for Quadmesh
+    # 2D
+    pc.set_array(z)
     pc.update_scalarmappable()
     # 3D RGB is OK as well
-    z = np.arange(36).reshape(3, 4, 3)
+    z = np.arange(36, dtype=np.uint8).reshape(3, 4, 3)
     pc.set_array(z)
+    pc.update_scalarmappable()
 
 
 def test_get_segments():
