ENH: Adding a get_coordinates() method to Quadmesh

greglucas · greglucas · commit b5c5b9be612d · 2021-05-15T19:44:23.000-06:00
This returns the coordinates used to create the Quadmesh collection.
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -2019,6 +2019,15 @@ def set_paths(self):
     def get_datalim(self, transData):
         return (self.get_transform() - transData).transform_bbox(self._bbox)
 
+    def get_coordinates(self):
+        """
+        Return the x and y vertices of the mesh.
+
+        The shape of the returned coordinates is
+        (*meshWidth* + 1, *meshHeight* + 1, 2).
+        """
+        return self._coordinates
+
     @staticmethod
     def convert_mesh_to_paths(meshWidth, meshHeight, coordinates):
         """
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -714,6 +714,18 @@ def test_singleton_autolim():
     np.testing.assert_allclose(ax.get_xlim(), [-0.06, 0.06])
 
 
+def test_quadmesh_get_xydata():
+    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_xydata(), coords)
+
+
 def test_quadmesh_set_array():
     x = np.arange(4)
     y = np.arange(4)
