add tests for set_xydata

Author: nvaytet

src/matplotgl/mesh.py

@@ -38,11 +38,9 @@ def _maybe_centers_to_edges(coord: np.ndarray, M: int, N: int, axis: str) -> np.
         expected_edges = expected_centers + 1
 
         if len(coord) == expected_edges:
-            # Already edges
             return coord
 
         if len(coord) == expected_centers:
-            # Assume centers, convert to edges
             edges = np.zeros(expected_edges, dtype=coord.dtype)
             # Midpoints between centers, with extrapolation at ends
             edges[1:-1] = 0.5 * (coord[:-1] + coord[1:])
@@ -58,11 +56,9 @@ def _maybe_centers_to_edges(coord: np.ndarray, M: int, N: int, axis: str) -> np.
     elif coord.ndim == 2:
         m, n = coord.shape
         if (m, n) == (M + 1, N + 1):
-            # Already edges
             return coord
 
         if (m, n) == (M, N):
-            # Assume centers, convert to edges
             # Strategy: First pad the coordinate array by extrapolating one
             # layer of "ghost" cells around the perimeter, then compute all
             # edges as average of 4 surrounding padded centers
@@ -73,13 +69,13 @@ def _maybe_centers_to_edges(coord: np.ndarray, M: int, N: int, axis: str) -> np.
             padded[1:-1, 1:-1] = coord
 
             # Extrapolate edges (sides)
-            # Left edge: coord[:, 0] - (coord[:, 1] - coord[:, 0])
+            # Left edge
             padded[1:-1, 0] = coord[:, 0] - (coord[:, 1] - coord[:, 0])
-            # Right edge: coord[:, -1] + (coord[:, -1] - coord[:, -2])
+            # Right edge
             padded[1:-1, -1] = coord[:, -1] + (coord[:, -1] - coord[:, -2])
-            # Bottom edge: coord[0, :] - (coord[1, :] - coord[0, :])
+            # Bottom edge
             padded[0, 1:-1] = coord[0, :] - (coord[1, :] - coord[0, :])
-            # Top edge: coord[-1, :] + (coord[-1, :] - coord[-2, :])
+            # Top edge
             padded[-1, 1:-1] = coord[-1, :] + (coord[-1, :] - coord[-2, :])
 
             # Extrapolate corners
@@ -264,19 +260,15 @@ def get_xdata(self) -> np.ndarray:
 
     def set_xdata(self, x: np.ndarray):
         M, N = self._c.shape
-        self._x = np.asarray(x)
-        # Convert centers to edges if necessary
-        self._x = _maybe_centers_to_edges(self._x, M, N, axis='x')
+        self._x = _maybe_centers_to_edges(np.asarray(x), M, N, axis='x')
         self._update_positions()
 
     def get_ydata(self) -> np.ndarray:
         return self._y
 
     def set_ydata(self, y: np.ndarray):
         M, N = self._c.shape
-        self._y = np.asarray(y)
-        # Convert centers to edges if necessary
-        self._y = _maybe_centers_to_edges(self._y, M, N, axis='y')
+        self._y = _maybe_centers_to_edges(np.asarray(y), M, N, axis='y')
         self._update_positions()
 
     def set_array(self, c: np.ndarray):

tests/pcolormesh_test.py

@@ -99,3 +99,85 @@ def test_pcolormesh_2d_coords_midpoints():
 
     assert len(ax.collections) == 1
     assert m._c.shape == (M, N)
+
+
+def test_pcolormesh_1d_set_xdata():
+    _, ax = plt.subplots()
+
+    M, N = 80, 100
+    xx = np.linspace(1.0, 5.0, N + 1)
+    yy = np.linspace(1.0, 5.0, M + 1)
+    x, y = np.meshgrid(xx[:-1], yy[:-1])
+    z = 1.0 - (np.sin(x) ** 10 + np.cos(10 + y * x) * np.cos(x))
+
+    m = ax.pcolormesh(xx, yy, z)
+
+    new_xx = np.linspace(2.0, 6.0, N + 1)
+    m.set_xdata(new_xx)
+
+    bbox = m.get_bbox()
+    assert bbox["left"] == new_xx[0]
+    assert bbox["right"] == new_xx[-1]
+
+
+def test_pcolormesh_1d_set_ydata():
+    _, ax = plt.subplots()
+
+    M, N = 80, 100
+    xx = np.linspace(1.0, 5.0, N + 1)
+    yy = np.linspace(1.0, 5.0, M + 1)
+    x, y = np.meshgrid(xx[:-1], yy[:-1])
+    z = 1.0 - (np.sin(x) ** 10 + np.cos(10 + y * x) * np.cos(x))
+
+    m = ax.pcolormesh(xx, yy, z)
+
+    new_yy = np.linspace(3.0, 7.0, M + 1)
+    m.set_ydata(new_yy)
+
+    bbox = m.get_bbox()
+    assert bbox["bottom"] == new_yy[0]
+    assert bbox["top"] == new_yy[-1]
+
+
+def test_pcolormesh_2d_set_xdata():
+    _, ax = plt.subplots()
+
+    M, N = 80, 100
+    xx = np.linspace(1.0, 5.0, N + 1)
+    yy = np.linspace(1.0, 5.0, M + 1)
+    x, y = np.meshgrid(xx, yy)
+    z = 1.0 - (
+        np.sin(x[:-1, :-1]) ** 10
+        + np.cos(10 + y[:-1, :-1] * x[:-1, :-1]) * np.cos(x[:-1, :-1])
+    )
+
+    m = ax.pcolormesh(x, y, z)
+
+    new_x = x + 3.0
+    m.set_xdata(new_x)
+
+    bbox = m.get_bbox()
+    assert bbox["left"] == new_x[0, 0]
+    assert bbox["right"] == new_x[0, -1]
+
+
+def test_pcolormesh_2d_set_ydata():
+    _, ax = plt.subplots()
+
+    M, N = 80, 100
+    xx = np.linspace(1.0, 5.0, N + 1)
+    yy = np.linspace(1.0, 5.0, M + 1)
+    x, y = np.meshgrid(xx, yy)
+    z = 1.0 - (
+        np.sin(x[:-1, :-1]) ** 10
+        + np.cos(10 + y[:-1, :-1] * x[:-1, :-1]) * np.cos(x[:-1, :-1])
+    )
+
+    m = ax.pcolormesh(x, y, z)
+
+    new_y = y + 4.0
+    m.set_ydata(new_y)
+
+    bbox = m.get_bbox()
+    assert bbox["bottom"] == new_y[0, 0]
+    assert bbox["top"] == new_y[-1, 0]