ENH: Adding a get_coordinates() method to Quadmesh

greglucas · greglucas · commit f52bbc16105d · 2021-05-25T20:23:06.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
@@ -2048,6 +2048,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 vertices of the mesh as an (M, N, 2) array.
+
+        M, N are the number of rows / columns in the mesh.
+        The last dimension specifies the components (x, y).
+        """
+        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
@@ -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)
