Merge pull request #18472 from greglucas/quadmesh_get_xydata

tacaswell · web-flow · commit 9091ba90f67a · 2021-05-26T13:23:45.000-04:00
ENH: Adding a get_coordinates() method to Quadmesh collections
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -2045,6 +2045,16 @@ def set_paths(self):
     def get_datalim(self, transData):
         return (self.get_transform() - transData).transform_bbox(self._bbox)
 
+    def get_coordinates(self):
+        """
+        Return the vertices of the mesh as an (M+1, N+1, 2) array.
+
+        M, N are the number of quadrilaterals in the rows / columns of the
+        mesh, corresponding to (M+1, N+1) vertices.
+        The last dimension specifies the components (x, y).
+        """
+        return self._coordinates
+
     @staticmethod
     @_api.deprecated("3.5", alternative="QuadMesh(coordinates).get_paths()")
     def convert_mesh_to_paths(meshWidth, meshHeight, coordinates):
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -786,6 +786,18 @@ def test_quadmesh_deprecated_positional(fig_test, fig_ref):
     ax.add_collection(qmesh)
 
 
+def test_quadmesh_get_coordinates():
+    x = [0, 1, 2]
+    y = [2, 4, 6]
+    z = np.ones(shape=(2, 2))
+    xx, yy = np.meshgrid(x, y)
+    coll = plt.pcolormesh(xx, yy, z)
+
+    # shape (3, 3, 2)
+    coords = np.stack([xx.T, yy.T]).T
+    assert_array_equal(coll.get_coordinates(), coords)
+
+
 def test_quadmesh_set_array():
     x = np.arange(4)
     y = np.arange(4)
