Pad only over the axes the array actually spans

[hdrake]

Fixes MOM6-community#52.

The neighbor-map builders requested a one-point halo on every axis the grid registers:

padding_width = {ax: (1, 1) for ax in grid.axes}

but the arrays they pad are the horizontal corner/center index arrays constructed a few lines above each call. xgcm.padding._pad_basic iterates the whole padding_width mapping and only continues on a zero width, so an axis the array has no dimension for reaches axis._get_position_name(da) and raises. Any grid registering a vertical axis — every real model grid — was therefore untraceable:

KeyError: "None of the DataArray's dims ('yq', 'xq') were found in axis coords."

Changes

• New private helper gridutils._pad_axes(grid, dims) returning the grid axes that dims spans.
• Applied at all three sites, one per neighbor-map path: build_neighbor_maps (single-tile), _multitile_padded_maps (multi-tile, 'outer' corners without tracer centers) and _OuterTopology.__init__ (multi-tile with tracer centers — LLC90, cubed-sphere). The issue as first reported named only the first two; the third fails identically, so the bug was not confined to grids without tracer centers.
• The two single-axis pads just below the third site (GX/GY) already wrote {ax: (1, 1) if ax == "X" else (0, 0) for ax in grid.axes} and so were safe — but only via the zero-width continue. They now draw on the same axis list, so nothing in the module depends on that accident.

Why derive from the array's dims rather than hardcode ("X", "Y"). Both are correct here, since the arrays are built in-function and their dims are known. I went with the dim-derived form because it states the actual precondition — pad only over axes this array is defined on — rather than restating it as a naming coincidence, and because it is uniform across all three sites (the third pads a center-position array, the others corner-position, one of them with a face dim). The face dimension is not an xgcm axis, so it drops out on its own. Hardcoding would also silently rot if the axes were ever renamed.

Tests

• test_section.py::test_vertical_axis_does_not_affect_neighbor_maps — neighbor maps for the lat-lon grid are identical with and without a registered Z axis (not merely non-raising).
• test_section.py::test_grid_section_on_grid_with_vertical_axis — end-to-end: the same section traced on the X/Y/Z grid and on its horizontal-only view.
• test_section_multitile.py::test_vertical_axis_does_not_affect_multitile_neighbor_maps — parametrized over two_face_x_to_x (no centers → _multitile_padded_maps) and left_two_tile_x_to_y (with centers → _OuterTopology), covering the other two sites.

_latlon_neighbor_maps gained a vertical_axis=False keyword and its grid construction moved into a new _latlon_grid helper; existing callers are unchanged.

All 4 new cases fail on the current branch and pass with the fix.

Verification

• Full suite in an env with xgcm 0.10.1, data/ symlinked from an existing checkout: 65 passed, 0 skipped, 0 failed (61 before this PR, plus the 4 new cases).

• Real-grid check on the CM4p25 example file with the vertical axis registered ('Z': {'center': 'sigma2_l', 'outer': 'sigma2_i'}), which is what a budget grid actually looks like:

  	stock v0.4.0rc1	this branch
  grid_section	KeyError	ok, 36 corner points
  regionate.GriddedRegion (0.6.0rc1)	KeyError	ok, 55 boundary points
  convergent_transport	ok	ok, dims ('time', 'sigma2_l', 'sect')

  Correction: an earlier revision of this table listed convergent_transport as
  failing on stock. It does not, and I had measured it wrongly — my script called
  grid_section and convergent_transport inside one try, so the first raised and
  the second never ran. Re-measured with indices obtained from a horizontal-only grid,
  convergent_transport succeeds on stock with a Z axis registered, both with the
  vertical padding set and with it unset. That is structural rather than lucky: it
  takes explicit i_c/j_c and so never builds neighbor maps. The affected entry
  points are grid_section and anything that calls it (which includes
  regionate.GriddedRegion) -- not convergent_transport, and not
  regionate.MaskRegions, which goes through regionate.boundaries (already fixed
  there, see MaskRegions(2d_mask, 3d_grid) fails: xgcm.pad iterates all grid axes regionate#24).

I did not re-execute the example notebooks. They construct horizontal-only grids (examples/load_example_model_grid.py registers X and Y only), which is the path the existing suite already covers and which this change leaves untouched — _pad_axes returns exactly ["X", "Y"] for them. Worth a run before merge regardless.

Relationship to #7

Independent, not stacked. Both branch from topology-driven-neighbors; this one touches gridutils.py / test_section.py / test_section_multitile.py, #7 touches transports.py / test_convergent_transport.py. No overlap, so they merge in either order.

They are worth reading together, though: #7 is why convergent_transport rejects some valid layer/interface pairs, and this one is why the grid those names come from could not be handed to sectionate in the first place. Together they remove the need for the horizontal_grid() workaround I added in hdrake/xwmb#43.

_{Drafted with AI assistance (Claude Code). I have read the diff and ran the tests, the reproducer and the real-grid check myself.}

---

Example notebooks

Re-executed all five in a dedicated env built from docs/environment.yml with this branch pip install -e'd (xgcm 0.10.1, sectionate 0.4.0rc2.dev1), data symlinked from an existing checkout:

notebook	result
1_creating_an_OSNAP_section.ipynb	executed cleanly
2_OSNAP_transports_CM4p25.ipynb	executed cleanly
3_Labrador_convergence_CM4p25.ipynb	executed cleanly
4_sections_on_global_tripolar_grid.ipynb	executed cleanly
5_MOC_transports_ECCOv4r4.ipynb	executed cleanly

As expected, this change is a no-op for them: every notebook grid registers only X and Y, so _pad_axes returns exactly ["X", "Y"] — the same mapping the old code built. Comparing every textual output against the committed ones, the only differences are the version banner (0.3.3 → 0.4.0rc2.dev1), the absent "Downloading … from Zenodo" lines, and warning paths. Notebook 5's overturning streamfunction reproduces bit-for-bit (psi range (Sv): -58.09992975038242 to 53.35870875408872). Textual outputs are identical to those produced on #7's branch except for warning paths and line numbers.

The re-executed notebooks are committed (last commit on this branch). They were re-run fresh against this branch's current HEAD, so the committed outputs are the ones this branch's committed code actually produces. Two caveats are recorded in that commit message: the version banner reads 0.4.0rc2.dev1, a hatch-vcs dev string from an editable install of an untagged branch, and warning text cites build-worktree paths rather than a normal checkout. Both are artifacts of how the notebooks were run, not of this change, and both resolve on a refresh from a tagged release. See the merge-order note at the end.

Coverage: why no test caught this

Counted across the topology-driven-neighbors tree and regionate@topology-overhaul — 33 xgcm.Grid constructions in sectionate (30 tests, 3 examples), of which zero register a Z axis; 15 in regionate, of which exactly one does. Details and the regionate#24 history are in MOM6-community#52.

_{This section drafted with AI assistance (Claude Code); I ran the notebooks and the comparisons myself.}

---

⚠️ Merge order: #7 and #8 now conflict on the notebooks

The re-executed notebooks are committed here, and the same is true on the other PR, so #7 and #8 no longer merge in either order without a conflict. I verified this rather than assuming it — a test merge of the two pushed refs conflicts on exactly the five notebooks:

examples/1_creating_an_OSNAP_section.ipynb
examples/2_OSNAP_transports_CM4p25.ipynb
examples/3_Labrador_convergence_CM4p25.ipynb
examples/4_sections_on_global_tripolar_grid.ipynb
examples/5_MOC_transports_ECCOv4r4.ipynb

Everything else merges cleanly: sectionate/transports.py, sectionate/gridutils.py and both test files are disjoint between the two branches. The conflict is entirely re-executed output, not code.

Suggested order: #8 first, then #7

• Pad only over the axes the array actually spans #8 is the smaller, self-contained fix (gridutils.py only) and its notebook refresh is a no-op by construction — every notebook grid registers only X and Y, so _pad_axes returns exactly ["X", "Y"], the same mapping the old code built. Its outputs differ from the current ones only in the version banner and warning paths.
• Match layer/interface through the grid, not string substitution #7 changes convergent_transport's API (the layer/interface defaults), so its notebook outputs are the ones with any semantic content. Merging it second means the surviving copy is the one produced by the later code.

Resolving the conflict

Don't hand-merge the notebook JSON. Take one side wholesale and, ideally, re-execute once on the merge result:

git checkout --theirs examples/   # after merging #8, take #7's notebooks
git add examples/
# then, best-effort:
cd examples && jupyter nbconvert --to notebook --execute --inplace \
  --ExecutePreprocessor.timeout=1800 --ExecutePreprocessor.kernel_name=python3 *.ipynb

Since neither branch changes any notebook result (verified: notebook 5's streamfunction is bit-identical on both, psi range (Sv): -58.09992975038242 to 53.35870875408872), taking either side is numerically safe; re-executing afterwards just makes the banner and warning paths consistent with the merged tree.

If you would rather not carry the churn at all, dropping the notebook commit from both branches (git revert of the last commit on each) restores clean either-order merging — which is what I had originally recommended. Either is fine; this note is just so nothing is a surprise at merge time.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB