Fix contour coordinates to use DataArray coordinate space

contours() returned raw array indices, which produced misaligned
contour lines when plotted on xarray imshow axes that use the
DataArray's coordinate values.  Transform output coordinates from
array indices to the DataArray's y/x coordinate space via np.interp.
Update tests, contour_explorer, and user guide notebook accordingly.

Author: brendancol

examples/contour_explorer.py

@@ -156,8 +156,8 @@ def draw_contours():
             fill_cmap = "gist_earth"
         # Use matplotlib's contourf for fills (our contour lines overlay on top)
         cf = ax.contourf(
-            elev_vals, levels=levels_list, cmap=fill_cmap,
-            origin="lower", alpha=0.35, extend="both",
+            xs, ys, elev_vals, levels=levels_list, cmap=fill_cmap,
+            alpha=0.35, extend="both",
         )
         fill_artists.extend(cf.collections)
 
@@ -252,16 +252,24 @@ def update_status():
 for spine in ax.spines.values():
     spine.set_color("white")
 
+# Map the image extent to the DataArray's coordinate space so that
+# contour coordinates (which are in coordinate space) align correctly.
+xs = elevation.coords["x"].values
+ys = elevation.coords["y"].values
+img_extent = [xs[0], xs[-1], ys[0], ys[-1]]
+
 # Hillshade layer
 hillshade_img = ax.imshow(
     hillshade_vals, cmap="gray", origin="lower",
     aspect="equal", interpolation="bilinear", alpha=1.0,
+    extent=img_extent,
 )
 
 # Terrain colour layer
 terrain_img = ax.imshow(
     elev_vals, cmap="gist_earth", origin="lower",
     aspect="equal", interpolation="bilinear", alpha=0.4,
+    extent=img_extent,
 )
 
 # Elevation readout under cursor
@@ -287,12 +295,18 @@ def update_status():
 
 # -- Event handlers ------------------------------------------------------------
 
+def _coord_to_pixel(xdata, ydata):
+    """Convert coordinate-space position to nearest pixel indices."""
+    col = int(round(np.interp(xdata, xs, np.arange(len(xs)))))
+    row = int(round(np.interp(ydata, ys, np.arange(len(ys)))))
+    return row, col
+
+
 def on_click(event):
     """Left-click: add contour at that elevation. Right-click: remove nearest."""
     if event.inaxes != ax:
         return
-    col = int(round(event.xdata))
-    row = int(round(event.ydata))
+    row, col = _coord_to_pixel(event.xdata, event.ydata)
     if not (0 <= row < GRID_H and 0 <= col < GRID_W):
         return
 
@@ -304,7 +318,7 @@ def on_click(event):
         snapped = round(elev / 10) * 10
         if snapped not in custom_levels:
             custom_levels.append(snapped)
-            marker_positions.append((col, row))
+            marker_positions.append((event.xdata, event.ydata))
             print(f"  Added contour at elevation {snapped:.0f} "
                   f"(clicked {elev:.1f} at pixel {row}, {col})")
             _update_markers()
@@ -353,8 +367,7 @@ def on_motion(event):
         elev_text.set_text("")
         fig.canvas.draw_idle()
         return
-    col = int(round(event.xdata))
-    row = int(round(event.ydata))
+    row, col = _coord_to_pixel(event.xdata, event.ydata)
     if 0 <= row < GRID_H and 0 <= col < GRID_W:
         e = elev_vals[row, col]
         elev_text.set_text(f"elev: {e:.1f} m  ({row}, {col})")

examples/user_guide/26_Contour_Lines.ipynb

@@ -546,9 +546,10 @@ def contours(
 
     return_type : str, default "numpy"
         Output format.  ``"numpy"`` returns a list of ``(level, coords)``
-        tuples where *coords* is an Nx2 array of ``(row, col)``
-        coordinates.  ``"geopandas"`` returns a GeoDataFrame with
-        ``level`` and ``geometry`` columns (requires geopandas/shapely).
+        tuples where *coords* is an Nx2 array of ``(y, x)`` coordinates
+        in the DataArray's coordinate space.  ``"geopandas"`` returns a
+        GeoDataFrame with ``level`` and ``geometry`` columns (requires
+        geopandas/shapely).
 
     Returns
     -------
@@ -608,6 +609,20 @@ def contours(
     )
     results = mapper(agg)(agg.data, levels)
 
+    # Transform from array indices to the DataArray's coordinate values.
+    y_coords = agg.coords[agg.dims[0]].values
+    x_coords = agg.coords[agg.dims[1]].values
+    y_idx = np.arange(len(y_coords), dtype=np.float64)
+    x_idx = np.arange(len(x_coords), dtype=np.float64)
+
+    transformed = []
+    for level, coords in results:
+        out = np.empty_like(coords)
+        out[:, 0] = np.interp(coords[:, 0], y_idx, y_coords)
+        out[:, 1] = np.interp(coords[:, 1], x_idx, x_coords)
+        transformed.append((level, out))
+    results = transformed
+
     if return_type == "numpy":
         return results
     elif return_type == "geopandas":

examples/user_guide/images/contour_lines_preview.png

@@ -52,10 +52,12 @@ def test_single_level_ramp(self):
         agg = create_test_raster(data, backend='numpy')
         result = contours(agg, levels=[2.5])
         assert len(result) > 0
+        # With res=0.5 x_coords = [0, 0.5, 1.0, 1.5, 2.0, 2.5].
+        # The crossing between col indices 2 and 3 maps to x = 1.25.
+        expected_x = 1.25
         for level, coords in result:
             assert level == 2.5
-            # All points should have col ~= 2.5 (crossing between col 2 and 3).
-            np.testing.assert_allclose(coords[:, 1], 2.5, atol=1e-10)
+            np.testing.assert_allclose(coords[:, 1], expected_x, atol=1e-10)
 
     def test_multiple_levels_ramp(self):
         """Multiple levels on a ramp produce one line per level."""
@@ -90,13 +92,17 @@ def test_peak_contour_at_1_5(self):
         agg = create_test_raster(data, backend='numpy')
         result = contours(agg, levels=[1.5])
         assert len(result) >= 1
+        y_min = float(agg.coords[agg.dims[0]].values.min())
+        y_max = float(agg.coords[agg.dims[0]].values.max())
+        x_min = float(agg.coords[agg.dims[1]].values.min())
+        x_max = float(agg.coords[agg.dims[1]].values.max())
         for level, coords in result:
             assert level == 1.5
-            # All points should be inside the raster bounds.
-            assert np.all(coords[:, 0] >= 0)
-            assert np.all(coords[:, 0] <= 4)
-            assert np.all(coords[:, 1] >= 0)
-            assert np.all(coords[:, 1] <= 4)
+            # All points should be inside the raster coordinate bounds.
+            assert np.all(coords[:, 0] >= y_min - 1e-10)
+            assert np.all(coords[:, 0] <= y_max + 1e-10)
+            assert np.all(coords[:, 1] >= x_min - 1e-10)
+            assert np.all(coords[:, 1] <= x_max + 1e-10)
 
 
 # ---------------------------------------------------------------------------
@@ -119,9 +125,11 @@ def test_partial_nan(self):
         agg = create_test_raster(data, backend='numpy')
         result = contours(agg, levels=[2.5])
         assert len(result) > 0
+        # y_coords = [2.0, 1.5, 1.0, 0.5, 0.0] (decreasing with res=0.5).
+        # NaN row is row 0 (y=2.0).  All contour points must stay at y <= 1.5.
+        nan_row_y = agg.coords[agg.dims[0]].values[0]
         for level, coords in result:
-            # No contour point should be in the NaN row.
-            assert np.all(coords[:, 0] >= 1.0 - 1e-10)
+            assert np.all(coords[:, 0] < nan_row_y + 1e-10)
 
 
 # ---------------------------------------------------------------------------
@@ -367,8 +375,22 @@ def test_matches_skimage(self):
         assert len(our_result) > 0
         assert len(sk_contours) > 0
 
-        # Collect all our contour points into a set for proximity checking.
-        our_points = np.vstack([coords for _, coords in our_result])
+        # Transform our coordinate-space points back to array indices for
+        # comparison with skimage (which returns array index coordinates).
+        y_coords = agg.coords[agg.dims[0]].values
+        x_coords = agg.coords[agg.dims[1]].values
+        y_idx = np.arange(len(y_coords), dtype=np.float64)
+        x_idx = np.arange(len(x_coords), dtype=np.float64)
+
+        our_points_idx = []
+        for _, coords in our_result:
+            pts = np.empty_like(coords)
+            pts[:, 0] = np.interp(coords[:, 0], y_coords[::-1], y_idx[::-1]) \
+                if y_coords[0] > y_coords[-1] \
+                else np.interp(coords[:, 0], y_coords, y_idx)
+            pts[:, 1] = np.interp(coords[:, 1], x_coords, x_idx)
+            our_points_idx.append(pts)
+        our_points = np.vstack(our_points_idx)
 
         # Every skimage contour point should be close to some point of ours.
         for sk_line in sk_contours: