FIX: Handle masked arrays within PolyQuadMesh

greglucas · greglucas · commit 3c85ef1a144d · 2023-04-30T15:44:41.000-06:00
Update the PolyQuadMesh handling to factor out the mask calculation
and setting of vertices so we make sure we are updating properly
when new arrays get passed in that may have a different mask.
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -2242,24 +2242,38 @@ class PolyQuadMesh(_MeshData, PolyCollection):
 
     def __init__(self, coordinates, **kwargs):
         super().__init__(coordinates=coordinates)
-        X = coordinates[..., 0]
-        Y = coordinates[..., 1]
-
-        mask = np.ma.getmaskarray(X) | np.ma.getmaskarray(Y)
-        # 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]))
-        if C.ndim == 3:
-            # RGB(A) case
-            mask |= np.any(np.ma.getmaskarray(C), axis=-1)
-        else:
-            mask |= np.ma.getmaskarray(C)
+        PolyCollection.__init__(self, verts=[], **kwargs)
+        # Setting the verts updates the paths of the PolyCollection
+        # This is called after the initializers to make sure the kwargs
+        # have all been processed and available for the masking calculations
+        self._set_unmasked_verts()
+
+    def _get_unmasked_polys(self):
+        """Get the unmasked regions using the coordinates and array"""
+        # mask(X) | mask(Y)
+        mask = np.any(np.ma.getmaskarray(self._coordinates), axis=-1)
+
+        # 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, expand if it is None
+        arr = self.get_array()
+        if arr is not None:
+            arr = np.ma.getmaskarray(arr)
+            if arr.ndim == 3:
+                # RGB(A) case
+                mask |= np.any(arr, axis=-1)
+            elif arr.ndim == 2:
+                mask |= arr
+            else:
+                mask |= arr.reshape(self._coordinates[:-1, :-1, :].shape[:2])
+        return ~mask
+
+    def _set_unmasked_verts(self):
+        X = self._coordinates[..., 0]
+        Y = self._coordinates[..., 1]
 
-        unmask = ~mask
-        self._valid_polys = unmask.ravel()
+        unmask = self._get_unmasked_polys()
         X1 = np.ma.filled(X[:-1, :-1])[unmask]
         Y1 = np.ma.filled(Y[:-1, :-1])[unmask]
         X2 = np.ma.filled(X[1:, :-1])[unmask]
@@ -2272,25 +2286,34 @@ 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))
-        # Setting the verts updates the paths of the PolyCollection
-        PolyCollection.__init__(self, verts=verts, **kwargs)
+        self.set_verts(verts)
 
     def get_edgecolor(self):
         # docstring inherited
         # We only want to return the facecolors of the polygons
         # that were drawn.
         ec = super().get_edgecolor()
-        if len(ec) != len(self._valid_polys):
+        unmasked_polys = self._get_unmasked_polys().ravel()
+        if len(ec) != len(unmasked_polys):
             # Mapping is off
             return ec
-        return ec[self._valid_polys, :]
+        return ec[unmasked_polys, :]
 
     def get_facecolor(self):
         # docstring inherited
         # We only want to return the facecolors of the polygons
         # that were drawn.
         fc = super().get_facecolor()
-        if len(fc) != len(self._valid_polys):
+        unmasked_polys = self._get_unmasked_polys().ravel()
+        if len(fc) != len(unmasked_polys):
             # Mapping is off
             return fc
-        return fc[self._valid_polys, :]
+        return fc[unmasked_polys, :]
+
+    def set_array(self, A):
+        prev_mask = self._get_unmasked_polys()
+        super().set_array(A)
+        # If the mask has changed at all we need to update
+        # the set of Polys that we are drawing
+        if not np.array_equal(prev_mask, self._get_unmasked_polys()):
+            self._set_unmasked_verts()
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -871,7 +871,7 @@ def test_quadmesh_set_array_validation(pcfunc):
     coll = ax.pcolormesh(x, y, z, shading='gouraud')
 
 
-def test_polyquadmesh_masked_vertices():
+def test_polyquadmesh_masked_vertices_array():
     xx, yy = np.meshgrid([0, 1, 2], [0, 1, 2, 3])
     # 2 x 3 mesh data
     zz = (xx*yy)[:-1, :-1]
@@ -898,13 +898,26 @@ def test_polyquadmesh_masked_vertices():
 
     # Mask the origin cell data
     zz = np.ma.masked_where((xx[:-1, :-1] == 0) & (yy[:-1, :-1] == 0), zz)
-    polymesh = plt.pcolor(xx, yy, zz)
+    polymesh = plt.pcolor(zz)
     polymesh.update_scalarmappable()
     # One cell should be left out
     assert len(polymesh.get_paths()) == 5
     # Poly version should have the same facecolors as the end of the quadmesh
     assert_array_equal(quadmesh_fc, polymesh.get_facecolor())
 
+    # We should also be able to call set_array with a new mask and get
+    # updated polys
+    # Remove mask, should add all polys back
+    zz = np.arange(6).reshape((3, 2))
+    polymesh.set_array(zz)
+    polymesh.update_scalarmappable()
+    assert len(polymesh.get_paths()) == 6
+    # Add mask should remove polys
+    zz = np.ma.masked_less(zz, 2)
+    polymesh.set_array(zz)
+    polymesh.update_scalarmappable()
+    assert len(polymesh.get_paths()) == 4
+
 
 def test_quadmesh_get_coordinates(pcfunc):
     x = [0, 1, 2]
