Pad only over the axes the array actually spans

sectionate/gridutils.py

@@ -20,6 +20,36 @@ def get_facedim(grid):
     return getattr(grid, "_facedim", None)
 
 
+def _pad_axes(grid, dims):
+    """
+    Names of the `grid` axes that `dims` spans.
+
+    `xgcm.padding.pad` iterates the whole `padding_width` mapping it is given and
+    looks each axis' position up in the array being padded, so an axis the array
+    has no dimension for is an error rather than a no-op. The arrays padded in
+    this module are the horizontal corner/center index arrays built a few lines
+    above each call, so passing `grid.axes` wholesale makes them fail on any grid
+    that registers a vertical axis -- which every real model grid does. Deriving
+    the list from the array's own dims keeps padding independent of whatever
+    *other* axes the grid happens to carry.
+
+    Parameters
+    ----------
+    grid: xgcm.Grid
+    dims: iterable of str
+        Dimension names of the array about to be padded.
+
+    Returns
+    -------
+    list of str
+    """
+    dims = set(dims)
+    return [
+        name for name, axis in grid.axes.items()
+        if dims & set(axis.coords.values())
+    ]
+
+
 def corner_position(grid):
     """
     Return the C-grid vorticity ("corner") position shared by the X and Y axes:
@@ -248,8 +278,9 @@ def build_neighbor_maps(grid, geocorners):
     own_j = np.broadcast_to(np.arange(ny)[:, None], shape)
     own_i = np.broadcast_to(np.arange(nx), shape)
 
-    padding = {ax: grid.axes[ax].padding for ax in grid.axes}
-    padding_width = {ax: (1, 1) for ax in grid.axes}
+    axes = _pad_axes(grid, dims)
+    padding = {ax: grid.axes[ax].padding for ax in axes}
+    padding_width = {ax: (1, 1) for ax in axes}
 
     def pad(a):
         return _module_pad(a, grid, padding_width, padding=padding, fill_value=np.nan)
@@ -302,8 +333,9 @@ def _multitile_padded_maps(grid, geocorners):
     own_j = np.broadcast_to(np.arange(ny)[:, None], shape)
     own_i = np.broadcast_to(np.arange(nx), shape)
 
-    padding = {ax: grid.axes[ax].padding for ax in grid.axes}
-    padding_width = {ax: (1, 1) for ax in grid.axes}
+    axes = _pad_axes(grid, dims)
+    padding = {ax: grid.axes[ax].padding for ax in axes}
+    padding_width = {ax: (1, 1) for ax in axes}
 
     def pad(a):
         return _module_pad(a, grid, padding_width, padding=padding, fill_value=np.nan)
@@ -531,12 +563,13 @@ def __init__(self, grid):
         # topology; every non-seam boundary pads NaN so walls stay walls ---
         def _seam_or_fill(b):
             return b if b == "periodic" else "fill"
-        padding = {ax: _seam_or_fill(grid.axes[ax].padding) for ax in grid.axes}
         cid = xr.DataArray(
             np.arange(nf * Nyc * Nxc, dtype=float).reshape(nf, Nyc, Nxc),
             dims=(facedim, Yc, Xc),
         )
-        bw = {ax: (1, 1) for ax in grid.axes}
+        axes = _pad_axes(grid, cid.dims)
+        padding = {ax: _seam_or_fill(grid.axes[ax].padding) for ax in axes}
+        bw = {ax: (1, 1) for ax in axes}
         C = _module_pad(cid, grid, bw, padding=padding, fill_value=np.nan)
         C = C.transpose(facedim, ..., Yc, Xc).values  # (nf, Nyc+2, Nxc+2)
         # A diagonally-padded halo cell is a pad of a pad: across two seams it
@@ -560,11 +593,11 @@ def _seam_or_fill(b):
         # left NaN and is resolved later by the 3-cell junction match (if it has a
         # native storage) or falls through to the coordinate-free `by_junction`.
         GX = _module_pad(
-            cid, grid, {ax: (1, 1) if ax == "X" else (0, 0) for ax in grid.axes},
+            cid, grid, {ax: (1, 1) if ax == "X" else (0, 0) for ax in axes},
             padding=padding, fill_value=np.nan,
         ).transpose(facedim, ..., Yc, Xc).values           # (nf, Nyc, Nxc+2)
         GY = _module_pad(
-            cid, grid, {ax: (1, 1) if ax == "Y" else (0, 0) for ax in grid.axes},
+            cid, grid, {ax: (1, 1) if ax == "Y" else (0, 0) for ax in axes},
             padding=padding, fill_value=np.nan,
         ).transpose(facedim, ..., Yc, Xc).values           # (nf, Nyc+2, Nxc)
 

sectionate/tests/test_section.py

@@ -16,26 +16,62 @@
 ds["lat"] = xr.DataArray(lat, dims=("y", "x"))
 
 
-def _latlon_neighbor_maps():
-    """Neighbor maps for the lat-lon grid above (X-periodic, Y clipped), built the
-    same way `grid_section` does -- from an xgcm.Grid. The low-level pathfinder always
-    requires these; a grid is the only source of topology-aware connectivity."""
+def _latlon_grid(vertical_axis=False):
+    """The lat-lon grid above (X-periodic, Y clipped), optionally also registering a
+    vertical axis -- as every real model grid does, and as sections never traverse."""
     ny, nx = lat.shape
-    g = xgcm.Grid(
-        xr.Dataset(coords={
-            "xq": np.arange(nx), "yq": np.arange(ny),
-            "xh": np.arange(nx) + 0.5, "yh": np.arange(ny) + 0.5,
-            "geolon_c": (("yq", "xq"), lon.astype(float)),
-            "geolat_c": (("yq", "xq"), lat.astype(float)),
-            "geolon": (("yh", "xh"), lon.astype(float)),
-            "geolat": (("yh", "xh"), lat.astype(float)),
-        }),
-        coords={"X": {"center": "xh", "right": "xq"}, "Y": {"center": "yh", "right": "yq"}},
-        padding={"X": "periodic", "Y": "extend"},
-        autoparse_metadata=False,
-    )
+    ds = xr.Dataset(coords={
+        "xq": np.arange(nx), "yq": np.arange(ny),
+        "xh": np.arange(nx) + 0.5, "yh": np.arange(ny) + 0.5,
+        "geolon_c": (("yq", "xq"), lon.astype(float)),
+        "geolat_c": (("yq", "xq"), lat.astype(float)),
+        "geolon": (("yh", "xh"), lon.astype(float)),
+        "geolat": (("yh", "xh"), lat.astype(float)),
+    })
+    coords = {"X": {"center": "xh", "right": "xq"}, "Y": {"center": "yh", "right": "yq"}}
+    if vertical_axis:
+        ds = ds.assign_coords({"z_l": ("z_l", np.array([5., 15.])),
+                               "z_i": ("z_i", np.array([0., 10., 20.]))})
+        coords["Z"] = {"center": "z_l", "outer": "z_i"}
+    return xgcm.Grid(ds, coords=coords, padding={"X": "periodic", "Y": "extend"},
+                     autoparse_metadata=False)
+
+
+def _latlon_neighbor_maps(vertical_axis=False):
+    """Neighbor maps for the lat-lon grid above, built the same way `grid_section`
+    does -- from an xgcm.Grid. The low-level pathfinder always requires these; a grid
+    is the only source of topology-aware connectivity."""
+    g = _latlon_grid(vertical_axis=vertical_axis)
     return build_neighbor_maps(g, get_geo_corners(g))
 
+
+def test_vertical_axis_does_not_affect_neighbor_maps():
+    """Horizontal connectivity must not depend on whether the grid also registers a
+    vertical axis. `build_neighbor_maps` padded its index arrays over *every* axis of
+    the grid, and xgcm's `pad` raises on an axis the array has no dimension for, so a
+    Z axis made an otherwise ordinary grid untraceable."""
+    plain = _latlon_neighbor_maps()
+    with_z = _latlon_neighbor_maps(vertical_axis=True)
+    assert set(plain) == set(with_z)
+    for d in plain:
+        for a, b in zip(plain[d], with_z[d]):
+            if a is None or b is None:
+                assert a is None and b is None
+            else:
+                np.testing.assert_array_equal(a, b)
+
+
+def test_grid_section_on_grid_with_vertical_axis():
+    """End-to-end: a section traced on a grid registering X, Y and Z is identical to
+    the same section traced on the horizontal-only view of that grid."""
+    from sectionate.section import grid_section
+    lons, lats = [10., 40.], [-10., 20.]
+    plain = grid_section(_latlon_grid(), lons, lats)
+    with_z = grid_section(_latlon_grid(vertical_axis=True), lons, lats)
+    for a, b in zip(plain, with_z):
+        np.testing.assert_array_equal(a, b)
+
+
 def test_distance_on_unit_sphere():
     from sectionate.section import distance_on_unit_sphere
 

sectionate/tests/test_section_multitile.py

@@ -545,3 +545,46 @@ def test_save_load_roundtrip_preserves_face_indices(tmp_path):
     t1 = convergent_transport(grid, gs.i_c, gs.j_c, gs.f_c, **kw)["conv_mass_transport"].sum().values
     t2 = convergent_transport(grid, gs2.i_c, gs2.j_c, gs2.f_c, **kw)["conv_mass_transport"].sum().values
     assert np.isclose(t1, t2)
+
+
+# ---------------------------------------------------------------------------
+# A registered vertical axis must not affect horizontal topology
+# ---------------------------------------------------------------------------
+
+def _with_vertical_axis(grid):
+    """Rebuild `grid` with an extra registered Z axis and nothing else changed."""
+    coords = {name: dict(axis.coords) for name, axis in grid.axes.items()}
+    coords["Z"] = {"center": "z_l", "outer": "z_i"}
+    ds = grid._ds.assign_coords({"z_l": ("z_l", np.array([5., 15.])),
+                                 "z_i": ("z_i", np.array([0., 10., 20.]))})
+    return xgcm.Grid(
+        ds, coords=coords,
+        padding={**{name: axis.padding for name, axis in grid.axes.items()}, "Z": "extend"},
+        fill_value=np.nan,
+        face_connections=grid._face_connections,
+        autoparse_metadata=False,
+    )
+
+
+def assert_maps_equal(a, b):
+    assert set(a) == set(b)
+    for d in a:
+        for x, y in zip(a[d], b[d]):
+            if x is None or y is None:
+                assert x is None and y is None
+            else:
+                np.testing.assert_array_equal(x, y)
+
+
+@pytest.mark.parametrize("fixture", [
+    two_face_x_to_x,        # 'outer' corners, no tracer centers -> _multitile_padded_maps
+    left_two_tile_x_to_y,   # 'left' corners with centers        -> _OuterTopology
+])
+def test_vertical_axis_does_not_affect_multitile_neighbor_maps(fixture):
+    """Both multi-tile neighbor-map paths padded their index arrays over *every* axis
+    of the grid, so a Z axis -- which every real model grid registers -- made the grid
+    untraceable. The maps must be identical with and without one."""
+    grid = fixture()
+    plain = build_neighbor_maps(grid, get_geo_corners(grid))
+    with_z = _with_vertical_axis(grid)
+    assert_maps_equal(plain, build_neighbor_maps(with_z, get_geo_corners(with_z)))