ENH: Add a PolyQuadMesh for 2d meshes as Polygons

greglucas · greglucas · commit 083689dc5e1f · 2023-04-23T12:58:15.000-06:00
This adds a new collection PolyQuadMesh that is a combination of
a PolyCollection and QuadMesh, storing all of the coordinate data
and arrays through QuadMesh, but drawing each quad individually
through the PolyCollection.draw() method.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5787,9 +5787,13 @@ def _interp_grid(X):
                                 "This may lead to incorrectly calculated cell "
                                 "edges, in which case, please supply "
                                 f"explicit cell edges to {funcname}.")
-                        X = np.hstack((X[:, [0]] - dX[:, [0]],
-                                       X[:, :-1] + dX,
-                                       X[:, [-1]] + dX[:, [-1]]))
+                        if isinstance(X, np.ma.core.MaskedArray):
+                            hstack = np.ma.hstack
+                        else:
+                            hstack = np.hstack
+                        X = hstack((X[:, [0]] - dX[:, [0]],
+                                    X[:, :-1] + dX,
+                                    X[:, [-1]] + dX[:, [-1]]))
                     else:
                         # This is just degenerate, but we can't reliably guess
                         # a dX if there is just one value.
@@ -5907,7 +5911,7 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
 
         Returns
         -------
-        `matplotlib.collections.Collection`
+        `matplotlib.collections.PolyQuadMesh`
 
         Other Parameters
         ----------------
@@ -5925,7 +5929,7 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
 
         **kwargs
             Additionally, the following arguments are allowed. They are passed
-            along to the `~matplotlib.collections.PolyCollection` constructor:
+            along to the `~matplotlib.collections.PolyQuadMesh` constructor:
 
         %(PolyCollection:kwdoc)s
 
@@ -5959,35 +5963,6 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
         shading = shading.lower()
         X, Y, C, shading = self._pcolorargs('pcolor', *args, shading=shading,
                                             kwargs=kwargs)
-        Ny, Nx = X.shape
-
-        # convert to MA, if necessary.
-        C = ma.asarray(C)
-        X = ma.asarray(X)
-        Y = ma.asarray(Y)
-
-        mask = ma.getmaskarray(X) + 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 = ma.getmaskarray(C) + xymask
-
-        unmask = ~mask
-        X1 = ma.filled(X[:-1, :-1])[unmask]
-        Y1 = ma.filled(Y[:-1, :-1])[unmask]
-        X2 = ma.filled(X[1:, :-1])[unmask]
-        Y2 = ma.filled(Y[1:, :-1])[unmask]
-        X3 = ma.filled(X[1:, 1:])[unmask]
-        Y3 = ma.filled(Y[1:, 1:])[unmask]
-        X4 = ma.filled(X[:-1, 1:])[unmask]
-        Y4 = ma.filled(Y[:-1, 1:])[unmask]
-        npoly = len(X1)
-
-        xy = np.stack([X1, Y1, X2, Y2, X3, Y3, X4, Y4, X1, Y1], axis=-1)
-        verts = xy.reshape((npoly, 5, 2))
-
-        C = ma.filled(C[:Ny - 1, :Nx - 1])[unmask]
-
         linewidths = (0.25,)
         if 'linewidth' in kwargs:
             kwargs['linewidths'] = kwargs.pop('linewidth')
@@ -6001,19 +5976,30 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
         # unless the boundary is not stroked, in which case the
         # default will be False; with unstroked boundaries, aa
         # makes artifacts that are often disturbing.
-        if 'antialiased' in kwargs:
-            kwargs['antialiaseds'] = kwargs.pop('antialiased')
-        if 'antialiaseds' not in kwargs and cbook._str_lower_equal(ec, "none"):
-            kwargs['antialiaseds'] = False
+        if 'antialiaseds' in kwargs:
+            kwargs['antialiased'] = kwargs.pop('antialiaseds')
+        if 'antialiased' not in kwargs and cbook._str_lower_equal(ec, "none"):
+            kwargs['antialiased'] = False
 
         kwargs.setdefault('snap', False)
 
-        collection = mcoll.PolyCollection(
-            verts, array=C, cmap=cmap, norm=norm, alpha=alpha, **kwargs)
-        collection._scale_norm(norm, vmin, vmax)
+        if (isinstance(X, np.ma.MaskedArray)
+                or isinstance(Y, np.ma.MaskedArray)):
+            stack = np.ma.stack
+            X = np.ma.asarray(X)
+            Y = np.ma.asarray(Y)
+            # For bounds collections later
+            x = X.compressed()
+            y = Y.compressed()
+        else:
+            stack = np.stack
+            x = X
+            y = Y
+        coords = stack([X, Y], axis=-1)
 
-        x = X.compressed()
-        y = Y.compressed()
+        collection = mcoll.PolyQuadMesh(
+            coords, array=C, cmap=cmap, norm=norm, alpha=alpha, **kwargs)
+        collection._scale_norm(norm, vmin, vmax)
 
         # Transform from native to data coordinates?
         t = collection._transform
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -2166,3 +2166,39 @@ def get_cursor_data(self, event):
         if contained and self.get_array() is not None:
             return self.get_array().ravel()[info["ind"]]
         return None
+
+
+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]
+
+        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
+
+        unmask = ~mask
+        X1 = np.ma.filled(X[:-1, :-1])[unmask]
+        Y1 = np.ma.filled(Y[:-1, :-1])[unmask]
+        X2 = np.ma.filled(X[1:, :-1])[unmask]
+        Y2 = np.ma.filled(Y[1:, :-1])[unmask]
+        X3 = np.ma.filled(X[1:, 1:])[unmask]
+        Y3 = np.ma.filled(Y[1:, 1:])[unmask]
+        X4 = np.ma.filled(X[:-1, 1:])[unmask]
+        Y4 = np.ma.filled(Y[:-1, 1:])[unmask]
+        npoly = len(X1)
+
+        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)
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -993,7 +993,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() if pcfunc == plt.pcolor else z))
+    mapped = pc.get_cmap()(pc.norm(z))
     # GitHub issue #1302:
     assert mcolors.same_color(pc.get_edgecolor(), 'red')
     # Check setting attributes after initialization:
@@ -1023,7 +1023,7 @@ def test_color_logic(pcfunc):
     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() if pcfunc == plt.pcolor else z)
+    pc.set_array(z)
     pc.update_scalarmappable()
     assert np.array_equal(pc.get_facecolor(), mapped)
     assert mcolors.same_color(pc.get_edgecolor(), 'none')
@@ -1071,9 +1071,8 @@ def test_LineCollection_args():
 def test_array_wrong_dimensions():
     z = np.arange(12).reshape(3, 4)
     pc = plt.pcolor(z)
-    with pytest.raises(ValueError, match="^Collections can only map"):
-        pc.set_array(z)
-        pc.update_scalarmappable()
+    pc.set_array(z)
+    pc.update_scalarmappable()
     pc = plt.pcolormesh(z)
     pc.set_array(z)  # 2D is OK for Quadmesh
     pc.update_scalarmappable()
