Match layer/interface through the grid, not string substitution

sectionate/tests/test_convergent_transport.py

@@ -183,3 +183,168 @@ def test_convergent_transport_convention():
     )['conv_mass_transport'].sum().values
 
     assert np.equal(-3., conv) and np.equal(-3, conv_rev)
+
+
+def initialize_minimal_vertical_grid(layer, interface, register_z=True):
+    """The minimal 1x1-cell 'outer' grid of `initialize_minimal_outer_grid`, plus a
+    single-layer vertical coordinate named (`layer`, `interface`). `register_z` selects
+    whether the grid declares that vertical axis or (as most grids handed to sectionate
+    do, since sections are only ever traced horizontally) only its horizontal ones."""
+    xh, yh = np.array([0.5]), np.array([0.5])
+    xq, yq = np.array([0., 1.]), np.array([0., 1.])
+
+    lon, lat = np.meshgrid(xh, yh)
+    lon_c, lat_c = np.meshgrid(xq, yq)
+    ds = xr.Dataset({}, coords={
+        "xh": xr.DataArray(xh, dims=("xh",)),
+        "yh": xr.DataArray(yh, dims=("yh",)),
+        "xq": xr.DataArray(xq, dims=("xq",)),
+        "yq": xr.DataArray(yq, dims=("yq",)),
+        layer: xr.DataArray(np.array([0.5]), dims=(layer,)),
+        interface: xr.DataArray(np.array([0., 1.]), dims=(interface,)),
+        "geolon": xr.DataArray(lon, dims=("yh", "xh")),
+        "geolat": xr.DataArray(lat, dims=("yh", "xh")),
+        "geolon_c": xr.DataArray(lon_c, dims=("yq", "xq",)),
+        "geolat_c": xr.DataArray(lat_c, dims=("yq", "xq",)),
+    })
+    ds["u"] = xr.DataArray(np.array([[[1., -np.sqrt(2.)]]]), dims=(layer, "yh", "xq"))
+    ds["v"] = xr.DataArray(np.array([[[0.], [np.pi]]]), dims=(layer, "yq", "xh"))
+    coords = {
+        'X': {'outer': 'xq', 'center': 'xh'},
+        'Y': {'outer': 'yq', 'center': 'yh'},
+    }
+    if register_z:
+        coords['Z'] = {'center': layer, 'outer': interface}
+    return xgcm.Grid(ds, coords=coords, padding={'X': "extend", 'Y': "extend"},
+                     autoparse_metadata=False)
+
+
+# The closed path of vorticity points around the single cell of the minimal grid,
+# i.e. what `grid_section` returns for the square [0,1]x[0,1] (see
+# `test_convergent_transport`); hardcoded so these tests exercise only
+# `convergent_transport`'s layer/interface handling.
+MINIMAL_LOOP_I = np.array([0, 1, 1, 0, 0])
+MINIMAL_LOOP_J = np.array([0, 0, 1, 1, 0])
+
+
+def transport_around_minimal_cell(grid, layer, interface):
+    from sectionate.transports import convergent_transport
+    return convergent_transport(
+        grid,
+        MINIMAL_LOOP_I,
+        MINIMAL_LOOP_J,
+        utr="u",
+        vtr="v",
+        layer=layer,
+        interface=interface,
+        geometry="cartesian",
+    )
+
+
+@pytest.mark.parametrize("register_z", [True, False])
+@pytest.mark.parametrize("layer,interface", [
+    ("z_l", "z_i"),           # the canonical MOM6 depth coordinate
+    ("sigma2_l", "sigma2_i"),  # a stem with no "l" in it
+    ("lam_l", "lam_i"),        # a stem containing an "l" ...
+    ("level_l", "level_i"),    # ... and one containing two
+])
+def test_layer_interface_pairs_accepted(layer, interface, register_z):
+    """A consistent (layer, interface) pair must be accepted whatever its stem spells.
+    Matching them by rewriting every "l" in `layer` into an "i" spuriously rejected any
+    stem that itself contains an "l"."""
+    grid = initialize_minimal_vertical_grid(layer, interface, register_z=register_z)
+    dsout = transport_around_minimal_cell(grid, layer, interface)
+
+    # both vertical coordinates are carried through to the output ...
+    assert layer in dsout.coords
+    assert interface in dsout.coords
+    # ... and the transport is the depth-independent result of `test_convergent_transport`
+    conv = dsout['conv_mass_transport'].sum().values
+    assert np.isclose(1. + 0. + np.sqrt(2.) - np.pi, conv, rtol=1.e-14)
+
+
+def test_layer_interface_from_grid_axis_without_naming_convention():
+    """Names the grid itself pairs on one axis are accepted even when they follow no
+    "_l"/"_i" naming convention: the grid, not the spelling, is the authority."""
+    grid = initialize_minimal_vertical_grid("MyCenters", "MyEdges", register_z=True)
+    dsout = transport_around_minimal_cell(grid, "MyCenters", "MyEdges")
+    assert "MyCenters" in dsout.coords and "MyEdges" in dsout.coords
+
+
+@pytest.mark.parametrize("layer,interface", [
+    ("z_l", "sigma2_i"),  # two different vertical coordinates
+    ("ml_l", "mi_i"),     # different stems ("ml" vs "mi"), but a whole-string
+                          # "l"->"i" substitution would have accepted them
+    ("z_l", "z_l"),       # the layer coordinate passed twice
+])
+def test_inconsistent_layer_interface_rejected(layer, interface):
+    """Genuinely mismatched pairs are still rejected, and the message names them. The
+    grid registers no vertical axis here, so the names are all there is to go on; a
+    grid that *does* register one rejects a contradicting name earlier and more
+    specifically (see `test_layer_interface_contradicting_z_axis_raises`)."""
+    grid = initialize_minimal_vertical_grid("z_l", "z_i", register_z=False)
+    with pytest.raises(ValueError, match="do not describe the same vertical axis"):
+        transport_around_minimal_cell(grid, layer, interface)
+
+
+# ---------------------------------------------------------------------------
+# layer/interface are taken from the grid's vertical axis when it has one
+# ---------------------------------------------------------------------------
+
+def test_layer_interface_derived_from_grid_z_axis():
+    """A grid that registers a "Z" axis already names its vertical coordinates, so the
+    caller should not have to repeat them."""
+    grid = initialize_minimal_vertical_grid("sigma2_l", "sigma2_i", register_z=True)
+    dsout = transport_around_minimal_cell(grid, None, None)  # nothing passed
+    assert "sigma2_l" in dsout.coords
+    assert "sigma2_i" in dsout.coords
+    conv = dsout['conv_mass_transport'].sum().values
+    assert np.isclose(1. + 0. + np.sqrt(2.) - np.pi, conv, rtol=1.e-14)
+
+
+def test_layer_interface_derived_matches_explicitly_passed():
+    """Passing the same names the "Z" axis carries -- what a caller that read them off
+    `grid.axes["Z"].coords` does -- is equivalent to passing nothing."""
+    grid = initialize_minimal_vertical_grid("lam_l", "lam_i", register_z=True)
+    zc = grid.axes["Z"].coords["center"]
+    zi = grid.axes["Z"].coords["outer"]
+    explicit = transport_around_minimal_cell(grid, zc, zi)
+    derived = transport_around_minimal_cell(grid, None, None)
+    xr.testing.assert_identical(explicit, derived)
+
+
+@pytest.mark.parametrize("layer,interface", [
+    ("z_l", None),        # layer contradicts the axis
+    (None, "z_i"),        # interface contradicts the axis
+    ("z_l", "z_i"),       # both do
+])
+def test_layer_interface_contradicting_z_axis_raises(layer, interface):
+    """An explicit name that disagrees with the grid's own vertical axis is an error,
+    not an override: the output would be labelled with a coordinate that need not
+    describe the data."""
+    grid = initialize_minimal_vertical_grid("sigma2_l", "sigma2_i", register_z=True)
+    with pytest.raises(ValueError, match="contradicts the grid"):
+        transport_around_minimal_cell(grid, layer, interface)
+
+
+def test_no_vertical_axis_and_no_names_attaches_nothing():
+    """Without a "Z" axis and without explicit names there is no vertical coordinate to
+    attach -- and, unlike the old `layer="z_l"` default, no lookup of a name the grid
+    has never heard of."""
+    grid = initialize_minimal_outer_grid()
+    grid._ds['u'] = xr.DataArray(np.array([[1., -np.sqrt(2.)]]), dims=("yh", "xq",))
+    grid._ds['v'] = xr.DataArray(np.array([[0], [np.pi]]), dims=("yq", "xh",))
+    dsout = transport_around_minimal_cell(grid, None, None)
+    assert not any(c.endswith(("_l", "_i")) for c in dsout.coords)
+    conv = dsout['conv_mass_transport'].sum().values
+    assert np.isclose(1. + 0. + np.sqrt(2.) - np.pi, conv, rtol=1.e-14)
+
+
+def test_z_axis_center_without_coordinate_values_is_not_derived():
+    """A "Z" axis position may name a bare dimension carrying no coordinate values;
+    such a name cannot be attached to the output, so it is not derived."""
+    grid = initialize_minimal_vertical_grid("z_l", "z_i", register_z=True)
+    grid._ds = grid._ds.drop_vars("z_i")  # keep the dim, drop its values
+    dsout = transport_around_minimal_cell(grid, None, None)
+    assert "z_l" in dsout.coords
+    assert "z_i" not in dsout.coords

sectionate/transports.py

@@ -382,15 +382,121 @@ def uvcoords_from_qindices(grid, i_c, j_c, f_c=None):
         uvindices_from_qindices(grid, i_c, j_c, f_c=f_c),
     )
 
+# xgcm positions at which a vertical *interface* (cell-edge) coordinate can sit,
+# as opposed to the "center" position where the paired layer coordinate sits.
+_INTERFACE_POSITIONS = ("outer", "inner", "left", "right")
+
+
+def _layer_interface_axis(grid, layer, interface):
+    """Name of the `grid` axis that registers `layer` at its "center" position and
+    `interface` at one of its interface positions, or None if no single axis does."""
+    for name, axis in grid.axes.items():
+        coords = axis.coords
+        if coords.get("center") != layer:
+            continue
+        if any(coords.get(pos) == interface for pos in _INTERFACE_POSITIONS):
+            return name
+    return None
+
+
+def _suffix_paired(layer, interface):
+    """Whether `layer`/`interface` follow the `<stem>l`/`<stem>i` naming convention
+    (e.g. "z_l"/"z_i", "sigma2_l"/"sigma2_i"), anchored at the *end* of the name."""
+    return (
+        layer.endswith("l")
+        and interface.endswith("i")
+        and layer[:-1] == interface[:-1]
+    )
+
+
+def _validate_layer_interface(grid, layer, interface):
+    """Raise unless `layer` and `interface` describe the same vertical axis.
+
+    The grid already knows the answer, so ask it: accept the pair if any axis of
+    `grid` registers `layer` at its "center" position and `interface` at an
+    interface position. Grids handed to sectionate frequently declare only their
+    horizontal axes, though (the vertical one is never traced over), so a pair the
+    grid says nothing about falls back to the `<stem>l`/`<stem>i` naming convention.
+
+    Note this is a suffix-anchored fallback, not a whole-string substitution: the
+    latter rewrites every "l" in the name and so rejects perfectly consistent pairs
+    whose stem happens to contain one (e.g. "lam_l"/"lam_i").
+    """
+    if _layer_interface_axis(grid, layer, interface) is not None:
+        return
+    if _suffix_paired(layer, interface):
+        return
+    offered = ", ".join(
+        f"{name}={dict(axis.coords)}" for name, axis in grid.axes.items()
+    ) or "none"
+    raise ValueError(
+        f"Inconsistent layer and interface grid variables: layer={layer!r} and "
+        f"interface={interface!r} do not describe the same vertical axis. No axis "
+        f"of the grid registers {layer!r} at its 'center' position together with "
+        f"{interface!r} at an interface position "
+        f"({'/'.join(_INTERFACE_POSITIONS)}), and the two names do not follow the "
+        f"'<stem>l'/'<stem>i' convention (e.g. 'z_l'/'z_i') either. "
+        f"Grid axes: {offered}."
+    )
+
+
+def _vertical_axis_pair(grid):
+    """The (center, interface) coordinate names of the `grid`'s vertical ("Z") axis,
+    or (None, None) if it registers none.
+
+    Only names that are actually present in `grid._ds` are returned: xgcm requires an
+    axis' positions to name *dimensions*, but a dimension need not carry coordinate
+    values, and these names are attached to the output by lookup.
+    """
+    axis = grid.axes.get("Z")
+    if axis is None:
+        return None, None
+    coords = axis.coords
+    center = coords.get("center")
+    interface = next((coords[pos] for pos in _INTERFACE_POSITIONS if pos in coords), None)
+    return (center if center in grid._ds else None,
+            interface if interface in grid._ds else None)
+
+
+def _resolve_layer_interface(grid, layer, interface):
+    """Decide which vertical coordinate names label the output.
+
+    The grid is the authority whenever it has one to give: if it registers a "Z" axis,
+    that axis' center and interface coordinates are used, and the caller need not name
+    them at all. An explicitly passed name that *contradicts* the grid is an error
+    rather than an override -- a caller who disagrees with the grid about which
+    vertical coordinate its transports live on is confused, and quietly preferring
+    either one would label the output with a coordinate that may not describe the data.
+    Passing names that agree with the grid stays valid, since that is how callers that
+    read them off `grid.axes["Z"].coords` in the first place invoke this.
+
+    Without a "Z" axis the passed names are used as given, checked against each other
+    by `_validate_layer_interface`.
+    """
+    zc, zi = _vertical_axis_pair(grid)
+    for passed, derived, what in ((layer, zc, "layer"), (interface, zi, "interface")):
+        if passed is not None and derived is not None and passed != derived:
+            raise ValueError(
+                f"{what}={passed!r} contradicts the grid: its 'Z' axis registers "
+                f"{derived!r} at that position ({dict(grid.axes['Z'].coords)}). Pass "
+                f"{what}={derived!r}, or omit it and let the axis supply it."
+            )
+    layer = zc if zc is not None else layer
+    interface = zi if zi is not None else interface
+    if (layer is not None) and (interface is not None):
+        _validate_layer_interface(grid, layer, interface)
+    return layer, interface
+
+
 def convergent_transport(
     grid,
     i_c,
     j_c,
     f_c=None,
     utr="umo",
     vtr="vmo",
-    layer="z_l",
-    interface="z_i",
+    layer=None,
+    interface=None,
     outname="conv_mass_transport",
     sect_coord="sect",
     geometry="spherical",
@@ -420,10 +526,23 @@ def convergent_transport(
     vtr: str
         Name of "Y"-direction tracer transport
     layer : str or None
-        Name of the vertical layer (cell-center) coordinate, or None for grids without one.
+        Name of the vertical layer (cell-center) coordinate. Default: None, meaning
+        "take it from the grid" — if `grid` registers a "Z" axis, that axis' "center"
+        coordinate is used and nothing need be passed. Only grids that do not declare
+        their vertical axis need to name it here; if there is no "Z" axis and nothing
+        is passed, the output simply carries no layer coordinate.
     interface : str or None
-        Name of the vertical interface coordinate, or None. If both are given, they must be
-        consistent (`layer` is `interface` with "l" in place of "i").
+        Name of the vertical interface coordinate, resolved the same way from the "Z"
+        axis' interface position ("outer"/"inner"/"left"/"right"). Default: None.
+
+        A name that *contradicts* the grid's "Z" axis raises, rather than overriding it.
+        Passing names that agree with the axis is fine, which is how callers that read
+        them off `grid.axes["Z"].coords` invoke this.
+
+        When the grid declares no vertical axis and both names are given, they must
+        still describe the same one: either some axis of the grid registers `layer` at
+        "center" and `interface` at an interface position, or the names follow the
+        `<stem>l`/`<stem>i` convention (e.g. "z_l"/"z_i", "sigma2_l"/"sigma2_i").
     outname : str
         Name of output xr.DataArray variable. Default: "conv_mass_transport".
     sect_coord: str
@@ -453,11 +572,9 @@ def convergent_transport(
         as well as some useful metadata, such as whether each point corresponds to a "U" or "V" velocity and whether
         the sign of the transport had to be flipped to make it point inwards.
     """
-
-    if (layer is not None) and (interface is not None):
-        if layer.replace("l", "i") != interface:
-            raise ValueError("Inconsistent layer and interface grid variables!")
-
+
+    layer, interface = _resolve_layer_interface(grid, layer, interface)
+
     # On a multi-tile grid the contributing velocity face varies along the section
     # (`uvindices["face"]`); it is selected pointwise in every `.isel` below.
     facedim = get_facedim(grid) if f_c is not None else None