diff --git a/CLAUDE.md b/CLAUDE.md index 8271b14..1ee7e3c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -146,9 +146,9 @@ The package is organized around a pipeline: define sections → map to grid → ### Core Modules -- **`section.py`** — Section definition and grid path algorithms. `Section` holds named waypoint coordinates; `GriddedSection` extends it with grid index information. `grid_section()` is the main entry point that maps geographic waypoints to grid vorticity-point indices `(i_c, j_c)` by walking the grid between waypoints along the `curve` requested. There are three: `"great circle"` (the default — every segment follows the geodesic), `"latitude circle"` (every segment follows a parallel, and a segment whose endpoints do not share a latitude — meridional ones included — raises), and `"latitude and great circle"` (decided per segment: constant-latitude segments follow the parallel, all others the geodesic). Under all three, every segment follows the **shortest** path between its two waypoints — raw longitudes are never read as a request to go the long way round, so a `"latitude circle"` segment written `0 -> 270` runs 90° *west* — which is why encircling the globe takes intermediate waypoints (e.g. `0 -> 120 -> 240 -> 360`). `_check_segment_span` resolves each segment's curve (via `_is_constant_latitude`, the one classification the metric choice and the legality check share) and raises for ill-posed ones: endpoints written exactly half a circle apart (neither way round is shorter — one error shared by all curves, quoting degrees of longitude along the parallel or degrees of arc, as applicable), or, under `"latitude circle"`, endpoints that do not share a latitude. The walk is deterministic and direction-independent: it admits neighbors strictly closer to the endpoint plus any seam twin of the current cell (so periodic/fold-seam crossings do not depend on floating-point rounding), and breaks ties by index. Grid topology is inferred entirely from `xgcm.Grid` metadata — each axis' `boundary` (periodic wrap, fill/extend wall, or a single-tile bipolar north fold `{"Y": {"fold": ...}}`) and `face_connections` for multi-tile grids — so there is no `topology` keyword. The pathfinder consumes topology-aware neighbor maps built by `gridutils.build_neighbor_maps`. +- **`section.py`** — Section definition and grid path algorithms. `Section` holds named waypoint coordinates; `GriddedSection` extends it with grid index information. `grid_section()` is the main entry point that maps geographic waypoints to grid vorticity-point indices `(i_c, j_c)` by walking the grid between waypoints along the `curve` requested. There are three: `"great circle"` (the default — every segment follows the geodesic), `"latitude circle"` (every segment follows a parallel, and a segment whose endpoints do not share a latitude — meridional ones included — raises), and `"latitude and great circle"` (decided per segment: constant-latitude segments follow the parallel, all others the geodesic). Under all three, every segment follows the **shortest** path between its two waypoints — raw longitudes are never read as a request to go the long way round, so a `"latitude circle"` segment written `0 -> 270` runs 90° *west* — which is why encircling the globe takes intermediate waypoints (e.g. `0 -> 120 -> 240 -> 360`). `_check_segment_span` resolves each segment's curve (via `_is_constant_latitude`, the one classification the metric choice and the legality check share) and raises for ill-posed ones: endpoints written exactly half a circle apart (neither way round is shorter — one error shared by all curves, quoting degrees of longitude along the parallel or degrees of arc, as applicable), or, under `"latitude circle"`, endpoints that do not share a latitude. The walk is deterministic and direction-independent: it admits neighbors strictly closer to the endpoint plus any seam twin of the current cell (so periodic/fold-seam crossings do not depend on floating-point rounding), and breaks ties by index. That seam-twin step visits the same physical corner twice, so it spans no grid cell and is not a velocity face; **`grid_section()` removes the redundant index right after the walk, via `drop_repeated_corners()`** (`create_section_composite()`, the lower-level entry point, does the same when passed `grid=`). Every section `grid_section()` returns therefore satisfies the invariant that **no two consecutive corners are the same physical point** — every consecutive pair is a real velocity face — and, on a multi-tile grid, that **every corner is the canonical native index of its corner node** (`_OuterTopology.node_native`). That is the only place the normalization happens: `transports.uvindices_from_qindices` *checks* the invariant and raises rather than re-establishing it, so hand-built index arrays are the caller's to fix (by calling `drop_repeated_corners` themselves). Which index of a repeated corner survives is not free: it must name the same two velocity faces, with the same transport signs, as the run it replaces — `_collapse_single_tile_runs` checks that with the very arithmetic `uvindices_from_qindices` will apply, preferring the last index (the frame the path continues in, which is what a seam hand-off needs). Where a whole grid column degenerates to one point — a bipolar cap's pole, stored as every corner of a column — a section can arrive along one index and leave along another; then **no single index names both flanking faces and it raises**, naming the corner, rather than emit a face the section does not have. (On a correctly declared fold grid the walk crosses the seam instead and this does not arise; it shows up on a tripolar grid whose fold is left undeclared.) `_check_supported_topology()` additionally rejects, at the front door, a multi-tile grid that registers no cell-center dimension on both horizontal axes: the corner-node topology is rebuilt by fingerprinting the tracer *cells* around each corner, and velocities are staggered onto the center dimensions (U at (X-corner, Y-center), V at (X-center, Y-corner)), so such a grid can carry no velocities and each of its seam corners cannot be resolved to one index. Only the center **dimensions** are needed — `_OuterTopology` numbers cells with `np.arange`, so no tracer longitude/latitude values are required. Single-tile grids are unaffected: their seam twins are detected geometrically from the corner coordinates alone, so a corner-only single-tile grid still traces *and* still enumerates its velocity faces (it just carries no velocity data to put on them). Grid topology is inferred entirely from `xgcm.Grid` metadata — each axis' `boundary` (periodic wrap, fill/extend wall, or a single-tile bipolar north fold `{"Y": {"fold": ...}}`) and `face_connections` for multi-tile grids — so there is no `topology` keyword. The pathfinder consumes topology-aware neighbor maps built by `gridutils.build_neighbor_maps`. -- **`transports.py`** — Transport computation along sections. `uvindices_from_qindices()` converts vorticity-point indices to U/V velocity-point indices using a per-position corner offset (`gridutils.corner_offset`) covering all three C-grid staggerings: 'outer', 'right', and 'left' (see "Corner staggering" below). `convergent_transport()` is the main function: it lazily computes signed normal transports with configurable orientation (positive inward to the polygon defined by the section). +- **`transports.py`** — Transport computation along sections. `uvindices_from_qindices()` converts vorticity-point indices to U/V velocity-point indices using a per-position corner offset (`gridutils.corner_offset`) covering all three C-grid staggerings: 'outer', 'right', and 'left' (see "Corner staggering" below). It emits exactly one face per consecutive pair of corners: it relies on the section-finding invariant above rather than filtering degenerate pairs itself, so an edge whose normal velocity is stored on neither the source nor the destination face raises instead of silently returning a zero-flux placeholder. Single-tile faces come from `_single_tile_face()`, the one implementation of that index arithmetic (shared with `section._collapse_single_tile_runs`); it reads a step of more than one index along an axis as a seam crossing **only on an axis that is actually periodic**, since on a walled axis the same step is a jump between two indices of one physical corner. `convergent_transport()` is the main function: it lazily computes signed normal transports with configurable orientation (positive inward to the polygon defined by the section). - **`tracers.py`** — `extract_tracer()` interpolates tracer data to U/V points along a section path for cross-section plotting. diff --git a/sectionate/gridutils.py b/sectionate/gridutils.py index 9accc51..92a8fff 100644 --- a/sectionate/gridutils.py +++ b/sectionate/gridutils.py @@ -142,9 +142,32 @@ def coord_dict(grid): dict Dictionary containing names of "X" and "Y" dimension variables, at both cell 'center' position and the corner position ('outer', 'right', or 'left'; see `corner_position`). + + Raises + ------ + ValueError + If either horizontal axis registers no 'center' position. Velocities are staggered + onto the center dimensions -- U lives at (X-corner, Y-center) and V at (X-center, + Y-corner) -- so a grid without them cannot locate any velocity, and anything that + derives faces, transports or tracers from a section needs this mapping. """ corner_pos = corner_position(grid) + missing = [ax for ax in ("X", "Y") if "center" not in grid.axes[ax].coords] + if missing: + one = len(missing) == 1 + raise ValueError( + f"The {' and '.join(missing)} {'axis' if one else 'axes'} of this grid " + f"{'registers' if one else 'register'} no 'center' position, so sectionate " + "cannot locate its velocities: U is stored at (X-corner, Y-center) and V at " + "(X-center, Y-corner). Declare the cell-center dimensions, e.g. " + "coords={'X': {'center': 'xh', 'outer': 'xq'}, " + "'Y': {'center': 'yh', 'outer': 'yq'}}. (A single-tile grid can still be " + "traced by `grid_section` from its corner coordinates alone, and its velocity " + "faces enumerated; a multi-tile grid cannot, because its corner topology is " + "rebuilt from the tracer cells around each corner.)" + ) + return { "X": { "center": grid.axes["X"].coords["center"], @@ -231,7 +254,10 @@ def build_neighbor_maps(grid, geocorners): (shared-corner) corner topology, resolved to native indices -- see `_OuterTopology`. Each physical corner appears under a single canonical native index (a shared 'outer' seam corner is not stepped through twice), - and the returned maps have the same format and native index frame. + and the returned maps have the same format and native index frame. That + construction reads the tracer cells around each corner, so a multi-tile grid + must register a cell-center dimension on both horizontal axes; `coord_dict` + raises by name if it does not. Parameters ---------- @@ -253,20 +279,14 @@ def build_neighbor_maps(grid, geocorners): multitile = facedim is not None if multitile: - has_centers = all("center" in grid.axes[ax].coords for ax in ("X", "Y")) - if has_centers: - return outer_topology(grid).maps - # Without tracer-center coordinates the cell-identity construction is - # unavailable (and neither are velocities, so only walking is needed): - # fall back to reading xgcm's corner halos directly. Only shared-corner - # ('outer') tilings are safe here -- staggered corner arrays can pad one - # corner off across rotated/reversed seams. - if corner_position(grid) != "outer": - raise ValueError( - "Multi-tile grids with staggered ('left'/'right') corners require " - "tracer-center coordinates to derive their corner topology." - ) - return _multitile_padded_maps(grid, geocorners) + # The cell-fingerprint construction identifies each corner by the tracer cells + # around it, so it needs the cell-center *dimensions* (their names and lengths; + # no tracer longitude/latitude values are read). There is no substitute: xgcm's + # halo padding of corner arrays produces one copy of a shared seam corner per + # face -- and, on staggered ('left'/'right') lattices, can land a cross-seam + # neighbor one corner off -- so maps read off it would not step through a single + # canonical index per physical corner. + return outer_topology(grid).maps da = geocorners["X"] Ydim, Xdim = da.dims[-2], da.dims[-1] ny, nx = da.sizes[Ydim], da.sizes[Xdim] @@ -310,89 +330,6 @@ def pad(a): return maps -def _multitile_padded_maps(grid, geocorners): - """ - Fallback multi-tile neighbor maps for shared-corner ('outer') tilings that - carry no tracer-center coordinates: pad index-valued corner arrays with - xgcm's `face_connections` halos and read the neighbors off directly. A - shared seam corner is stored on every face that touches it, so these maps - step through each face's own copy (seam-twin semantics); `_validate_reciprocity` - refuses topologies xgcm's padding cannot represent consistently. - """ - facedim = grid._facedim - da = geocorners["X"] - Ydim, Xdim = da.dims[-2], da.dims[-1] - ny, nx = da.sizes[Ydim], da.sizes[Xdim] - nf = da.sizes[facedim] - dims = (facedim, Ydim, Xdim) - shape = (nf, ny, nx) - iarr = xr.DataArray(np.broadcast_to(np.arange(nx), shape).astype(float), dims=dims) - jarr = xr.DataArray(np.broadcast_to(np.arange(ny)[:, None], shape).astype(float), dims=dims) - farr = xr.DataArray(np.broadcast_to(np.arange(nf)[:, None, None], shape).astype(float), dims=dims) - own_f = np.broadcast_to(np.arange(nf)[:, None, None], shape) - own_j = np.broadcast_to(np.arange(ny)[:, None], shape) - own_i = np.broadcast_to(np.arange(nx), shape) - - 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) - - pj, pi, pf = pad(jarr), pad(iarr), pad(farr) - - interior = slice(1, -1) - slices = { - "right": (interior, slice(2, None)), - "left": (interior, slice(0, -2)), - "up": (slice(2, None), interior), - "down": (slice(0, -2), interior), - } - - maps = {} - for d, (ysl, xsl) in slices.items(): - sel = {Ydim: ysl, Xdim: xsl} - jmap = pj.isel(sel).values - imap = pi.isel(sel).values - fmap = pf.isel(sel).values - wall = np.isnan(fmap) | np.isnan(jmap) | np.isnan(imap) - fmap = np.where(wall, own_f, fmap).astype(np.int64) - jmap = np.where(wall, own_j, jmap).astype(np.int64) - imap = np.where(wall, own_i, imap).astype(np.int64) - maps[d] = (fmap, jmap, imap) - - _validate_reciprocity(maps, own_f, own_j, own_i) - return maps - - -def _validate_reciprocity(maps, own_f, own_j, own_i): - """ - Verify that the neighbor maps describe a consistent topology: if B is a - (non-wall) neighbor of A, then A must be one of B's four neighbors. An - inconsistent map would yield silently-wrong sections, so we detect it and - refuse rather than return garbage neighbors. - """ - for d, (fmap, jmap, imap) in maps.items(): - same = (fmap == own_f) & (jmap == own_j) & (imap == own_i) - not_wall = ~same - # Gather each neighbor's own four neighbors and look for the point back. - reciprocated = np.zeros(jmap.shape, dtype=bool) - for (f2, j2, i2) in maps.values(): - back = ( - (f2[fmap, jmap, imap] == own_f) - & (j2[fmap, jmap, imap] == own_j) - & (i2[fmap, jmap, imap] == own_i) - ) - reciprocated |= back - if np.any(not_wall & ~reciprocated): - raise NotImplementedError( - "Could not derive a consistent neighbor topology from this grid's " - "`face_connections` metadata (the multi-tile neighbor maps are not " - "reciprocal). Sections on grids with simpler topology are supported." - ) - - def outer_topology(grid): """ The complete corner-point topology of a multi-tile grid, built on its @@ -1108,7 +1045,7 @@ def coincident(a, b): node_reps[n].append((f, J, I)) self.node_reps = node_reps - # No `_validate_reciprocity` here: the maps project an *undirected* node + # The maps need no reciprocity check: they project an *undirected* node # graph, so node-level reciprocity holds by construction. Index-level # reciprocity deliberately does not: where one physical point is stored # natively more than once (e.g. the LLC south-boundary fold, whose rows diff --git a/sectionate/section.py b/sectionate/section.py index dcbffab..c9d99c2 100644 --- a/sectionate/section.py +++ b/sectionate/section.py @@ -5,6 +5,7 @@ get_geo_corners, get_facedim, build_neighbor_maps, + outer_topology, ) # Two corner indices map to the same physical point on seams that fold or wrap (e.g. @@ -317,10 +318,10 @@ def grid_section(grid, lons, lats, curve="great circle"): # All topologies -- periodic wrap, fill/extend walls, multi-tile face # connections, and the bipolar north fold -- are derived uniformly from the # grid's xgcm metadata by padding index arrays (see `build_neighbor_maps`). - # Where a single physical corner carries two indices (a periodic seam, a shared + # Where one physical corner is stored under two indices (a periodic seam, a shared # multi-tile boundary corner, or the fold seam), the walk simply steps through - # both; the resulting zero-length section edge carries no flux and is dropped when - # faces are derived (see `transports.uvindices_from_qindices`). + # both; that hand-off step spans no grid cell, so the redundant index is dropped + # right here, by `drop_repeated_corners`, before the section is returned. neighbor_maps = build_neighbor_maps(grid, geocorners) return create_section_composite( @@ -330,9 +331,198 @@ def grid_section(grid, lons, lats, curve="great circle"): lats, neighbor_maps=neighbor_maps, curve=curve, + grid=grid, ) +def drop_repeated_corners(grid, i_c, j_c, f_c=None): + """ + Normalise a traced grid path so that consecutive corners are always distinct + physical points. + + A single physical vorticity point can be stored under more than one index: the + duplicated seam column (or row) of a symmetric ('outer') periodic grid, a corner + shared by several tiles of a multi-tile grid, the two coincident lips of a grid cut, + or the whole degenerate column where a bipolar cap converges on one of its poles. + `infer_grid_path_from_geo` deliberately admits the "seam twin" of the corner it is + standing on, so that a seam crossing is decided by the grid's topology rather than by + floating-point rounding. The resulting hand-off step spans no grid cell, so that pair + of corners does not define a velocity face. + + `grid_section` calls this immediately after the path is traced, so that every section + it returns satisfies: + + * no two consecutive corners are the same physical point -- every consecutive pair + is a real velocity face; and + * on a multi-tile grid, every corner is the canonical native index of the physical + corner point it denotes (`_OuterTopology.node_native`), which is what lets the + velocity-face attribution look a corner up in either the source or the + destination face's frame. + + Of a repeated corner's several indices, the one kept is the one that names the same + two velocity faces as the run it replaces: the faces flanking the run keep exactly + the indices they would have had if the walk had never changed frame. On a multi-tile + grid that is automatic, since every index of a corner resolves to the same canonical + native one. On a single tile the LAST index of the run -- the frame the path continues + in -- is preferred and is what a seam hand-off needs; the earlier ones are tried in + turn if it does not serve, and a run that no single index can stand in for raises + (see `_collapse_single_tile_runs`). + + `uvindices_from_qindices` requires this invariant and checks it, so hand-built index + arrays that were not produced by `grid_section` can be passed through here first. + + PARAMETERS: + ----------- + grid: xgcm.Grid + Grid the path was traced on. + i_c, j_c: array-like of int + Vorticity-point indices along the "X" and "Y" dimensions. + f_c: array-like of int or None + Face indices for multi-tile grids (`face_connections`); None for single-tile. + + RETURNS: + -------- + i_c, j_c, f_c, lons_c, lats_c: np.ndarray + The normalised path and the coordinates of its corners (`f_c` is None for + single-tile grids). + """ + i_c = np.asarray(i_c, dtype=np.int64) + j_c = np.asarray(j_c, dtype=np.int64) + f_c = None if f_c is None else np.asarray(f_c, dtype=np.int64) + + geocorners = get_geo_corners(grid) + glon = np.asarray(geocorners["X"].values) + glat = np.asarray(geocorners["Y"].values) + + if f_c is not None: + # Multi-tile: which physical point an index denotes is exact and purely + # topological -- two indices denote the same point iff they resolve to the same + # node of the corner graph. + ot = outer_topology(grid) + nodes = ot.node_id[f_c, j_c + ot.t, i_c + ot.t] + if (nodes < 0).any(): + raise ValueError( + "Section contains indices that are not grid corners." + ) + native = ot.node_native[nodes] + if (native[:, 0] < 0).any(): + raise ValueError( + "Section passes through a corner point that is stored on no face of " + "the grid (e.g. a cube vertex or an unstored cap/cut corner), so it " + "has no native index." + ) + f_c, j_c, i_c = native[:, 0].copy(), native[:, 1].copy(), native[:, 2].copy() + # Every index of a corner has just been rewritten to the node's one canonical + # native index, so the members of a run are now identical: which one survives + # cannot matter, and the neighbour maps link canonical indices only. + repeated = nodes[1:] == nodes[:-1] + keep = np.ones(i_c.size, dtype=bool) + keep[:-1] = ~repeated + else: + # Single tile: the seam twins of a periodic/fold boundary are separate indices + # with nothing to compare them by, so coincidence is measured physically. The + # tolerance sits far below any real grid spacing (see COINCIDENT_TOLERANCE_M). + lo, la = glon[j_c, i_c], glat[j_c, i_c] + repeated = distance_on_unit_sphere( + lo[:-1], la[:-1], lo[1:], la[1:] + ) < COINCIDENT_TOLERANCE_M + keep = _collapse_single_tile_runs(grid, i_c, j_c, repeated) + + i_c, j_c = i_c[keep], j_c[keep] + f_c = None if f_c is None else f_c[keep] + + if f_c is not None: + lons_c, lats_c = glon[f_c, j_c, i_c], glat[f_c, j_c, i_c] + else: + lons_c, lats_c = glon[j_c, i_c], glat[j_c, i_c] + return i_c, j_c, f_c, np.asarray(lons_c), np.asarray(lats_c) + + +def _collapse_single_tile_runs(grid, i_c, j_c, repeated): + """ + Boolean mask keeping exactly one index out of each run of repeated corners in a + single-tile path (`repeated[k]` is True where corners k and k+1 are the same + physical point). + + Collapsing a run replaces the index the path arrived at with the index it is left + under, so it is only safe when the surviving index still names the two velocity + faces that flank the run. That is checked, not assumed: the flanking faces are + derived from the raw path and from each candidate index, with + `transports._single_tile_face` -- the same arithmetic + `transports.uvindices_from_qindices` will apply -- and only an index reproducing + both is kept. Compared are the velocity (var, i, j) and the direction flag that + carries that face's transport sign ("Xinc" for a V face, "Yinc" for a U face; see + `convergent_transport`), so collapsing changes neither a face of the section nor + the sign of its transport. + + Candidates are tried last-index-first, because for a seam hand-off (the common case, + where a run is just a corner's two seam-twin indices) the last one is the frame the + path continues in. + + Where a whole grid column degenerates to one point -- the pole of a bipolar cap, + stored as every corner of a column -- a section can enter the point along one of its + indices and leave along another, far apart in the index lattice. Then no single index + names both flanking faces, and this raises rather than emit a face the section does + not have. + """ + # Imported here, not at module scope: `transports` imports this module for the + # coincidence test, and the face arithmetic must have exactly one implementation. + from .transports import _single_tile_face, _single_tile_face_context + + keep = np.ones(i_c.size, dtype=bool) + if not np.any(repeated): + return keep + + context = _single_tile_face_context(grid) + + def face(a, b): + """The velocity face between path positions a and b, with its transport sign.""" + var, vi, vj, Xinc, Yinc = _single_tile_face( + int(i_c[a]), int(j_c[a]), int(i_c[b]), int(j_c[b]), *context + ) + return var, vi, vj, (Xinc if var == "V" else Yinc) + + n = i_c.size + start = 0 + while start < n: + end = start + while end < n - 1 and repeated[end]: + end += 1 + if end > start: + before = start - 1 if start > 0 else None + after = end + 1 if end + 1 < n else None + face_in = None if before is None else face(before, start) + face_out = None if after is None else face(end, after) + chosen = None + for r in range(end, start - 1, -1): + if before is not None and face(before, r) != face_in: + continue + if after is not None and face(r, after) != face_out: + continue + chosen = r + break + if chosen is None: + raise ValueError( + f"Section corners {start}..{end} -- (i,j) " + f"{(int(i_c[start]), int(j_c[start]))} through " + f"{(int(i_c[end]), int(j_c[end]))} -- are all the same physical grid " + "corner, a column that degenerates to a single point (a bipolar " + "cap's pole). The section enters it from " + f"{(int(i_c[before]), int(j_c[before]))} and leaves it towards " + f"{(int(i_c[after]), int(j_c[after]))}, and no single index of that " + "corner is adjacent to both, so no section with distinct consecutive " + "corners can carry both of the velocity faces on either side. If this " + "grid has a bipolar/tripolar fold, declare it -- " + "padding={'Y': {'fold': 'corner'}} -- so the walk crosses the cap " + "through the fold seam rather than through the pole column; that is " + "what this arises from. Otherwise route the section around the pole." + ) + keep[start:end + 1] = False + keep[chosen] = True + start = end + 1 + return keep + + def _check_supported_topology(grid): """ Raise if the multi-tile `grid` requires topology features sectionate does not yet support. @@ -341,6 +531,16 @@ def _check_supported_topology(grid): connects to itself along the "Y" axis) is not supported. Use the single-tile bipolar-fold padding instead -- ``padding={"X": "periodic", "Y": {"fold": "corner"}}`` (xgcm >= the bipolar-fold release) -- which sectionate handles natively via xgcm's fold padding. + + A multi-tile grid must also register a cell-center dimension on both horizontal axes. + That is what the corner-node topology is reconstructed from (`gridutils.outer_topology` + fingerprints each corner by the tracer cells around it, because cells -- unlike corner + arrays -- pad reliably across any seam), and it is also where half of each velocity's + dimensions live: U sits at (X-corner, Y-center) and V at (X-center, Y-corner). Only the + *dimensions* are needed, not tracer longitudes/latitudes. Without them a multi-tile grid + can carry no velocities and each of its seam corners cannot be resolved to one index, so + a section across a seam could not be traced correctly; refuse it at the front door rather + than return a path whose seam steps are wrong. """ facedim = grid._facedim connections = (getattr(grid, "_face_connections", None) or {}).get(facedim, {}) @@ -358,6 +558,20 @@ def _check_supported_topology(grid): "padding instead: padding={'X':'periodic','Y':{'fold':'corner'}}." ) + missing = [ax for ax in ("X", "Y") if "center" not in grid.axes[ax].coords] + if missing: + raise ValueError( + "Multi-tile grids (`face_connections`) must register a cell-center dimension on " + f"both horizontal axes; {' and '.join(missing)} declare none. Sectionate rebuilds " + "the shared-corner topology from the tracer cells around each corner, and the " + "velocities themselves are staggered onto the center dimensions (U at (X-corner, " + "Y-center), V at (X-center, Y-corner)), so a grid without them carries no " + "velocities and its seam corners cannot be identified. Declare them, e.g. " + "coords={'X': {'center': 'xh', 'outer': 'xq'}, " + "'Y': {'center': 'yh', 'outer': 'yq'}} -- the center entries need only name " + "dimensions of the dataset, they need not carry coordinate values." + ) + def create_section_composite( gridlon, gridlat, @@ -365,6 +579,7 @@ def create_section_composite( lats, neighbor_maps, curve="great circle", + grid=None, ): """ Compute composite section along velocity faces, as defined by coordinates of vorticity points (gridlon, gridlat), @@ -391,6 +606,12 @@ def create_section_composite( "latitude circle", or "latitude and great circle". Each segment is resolved and checked independently; `sectionate.grid_section` documents what the three options mean and which segments they reject. + grid: xgcm.Grid or None + The grid `neighbor_maps` and (gridlon, gridlat) came from. Pass it to get a section + that satisfies the invariant the rest of the API expects -- no two consecutive + corners at the same physical point -- by normalizing the raw walk with + `drop_repeated_corners`. Without it the raw walk is returned unchanged, which + `transports.uvindices_from_qindices` refuses; `grid_section` always passes it. RETURNS: ------- @@ -399,6 +620,10 @@ def create_section_composite( (i_c, j_c[, f_c]) correspond to indices of vorticity points that define velocity faces; the face index f_c is only returned for multi-tile grids. (lons_c, lats_c) are the corresponding longitude and latitudes. + + Without `grid` this is the raw walk: where a section crosses a seam whose corner is + stored under two indices, both are still present, and the step between them is not + a velocity face. """ # A face dimension (3-D corner arrays) marks a multi-tile grid, whose sections @@ -444,9 +669,15 @@ def create_section_composite( lats_c = np.concatenate([lats_c, [lats_c_seg[-1]]], axis=0) if multitile: f_c = np.concatenate([f_c, [f_c_seg[-1]]], axis=0) - return i_c.astype(np.int64), j_c.astype(np.int64), f_c.astype(np.int64), lons_c, lats_c - return i_c.astype(np.int64), j_c.astype(np.int64), lons_c, lats_c + i_c, j_c = i_c.astype(np.int64), j_c.astype(np.int64) + f_c = f_c.astype(np.int64) if multitile else None + if grid is not None: + i_c, j_c, f_c, lons_c, lats_c = drop_repeated_corners(grid, i_c, j_c, f_c) + if multitile: + return i_c, j_c, f_c, lons_c, lats_c + + return i_c, j_c, lons_c, lats_c def create_section(gridlon, gridlat, lonstart, latstart, lonend, latend, neighbor_maps, curve="great circle"): """ diff --git a/sectionate/tests/test_ecco_llc90.py b/sectionate/tests/test_ecco_llc90.py index 27a764c..9970fb4 100644 --- a/sectionate/tests/test_ecco_llc90.py +++ b/sectionate/tests/test_ecco_llc90.py @@ -83,6 +83,37 @@ def test_llc90_polar_cut_lips_share_nodes(): assert shared, "no node shares native storage across tiles 0 and 3" +def test_llc90_sections_contain_no_repeated_corner(): + """Sections traced on the real LLC90 grid satisfy the public invariant: consecutive + corners are always distinct physical points -- so every consecutive pair is a real + velocity face -- and every corner is the canonical native index of its corner node, + which is what lets the face attribution read a corner in either face's frame.""" + from sectionate.gridutils import outer_topology + from sectionate.section import grid_section + from sectionate.transports import uvindices_from_qindices + + grid = _load_grid() + ot = outer_topology(grid) + + for lat in (-70., 10., 70.): + i_c, j_c, f_c, _, _ = grid_section( + grid, [0., 90., 180., 270., 360.], [lat] * 5, curve="latitude circle" + ) + nodes = ot.node_id[f_c, j_c + ot.t, i_c + ot.t] + assert (nodes >= 0).all() + assert np.all(nodes[1:] != nodes[:-1]), f"repeated corner at lat {lat}" + + native = ot.node_native[nodes] + assert np.array_equal(native[:, 0], f_c) + assert np.array_equal(native[:, 1], j_c) + assert np.array_equal(native[:, 2], i_c) + + # every consecutive pair yields a face, and all of them are real velocities + uv = uvindices_from_qindices(grid, i_c, j_c, f_c=f_c) + assert uv["var"].size == i_c.size - 1 + assert set(np.unique(uv["var"]).tolist()) <= {"U", "V"} + + def test_llc90_face_corners_resolve_topologically(): """Every native corner appears in the neighbour maps with the correct arity: building the maps requires resolving the four-tile-junction face corners via diff --git a/sectionate/tests/test_section.py b/sectionate/tests/test_section.py index cf82d2c..b94f861 100644 --- a/sectionate/tests/test_section.py +++ b/sectionate/tests/test_section.py @@ -7,6 +7,7 @@ get_geo_corners, build_neighbor_maps, ) +from sectionate.section import drop_repeated_corners # define simple lat-lon grid @@ -167,17 +168,13 @@ def test_infer_grid_path_requires_neighbor_maps(): infer_grid_path_from_geo(0, 0, 1, 1, lon, lat, None) -def test_zero_length_seam_face_dropped(): - """On a symmetric periodic grid the seam vertex (360 == 0) carries two indices, so a - seam-crossing section has a doubled corner in i_c. The zero-length edge between the two - identities is not emitted as a velocity face, while the real faces on either side are.""" - from sectionate.section import grid_section, distance_on_unit_sphere - from sectionate.transports import uvindices_from_qindices - N = 8 +def _symmetric_periodic_grid(N=8): + """A symmetric ('outer') X-periodic grid: the seam meridian is stored twice, once as + corner column 0 (lon 0) and once as corner column N (lon 360).""" xq = np.linspace(0., 360., N + 1); xh = 0.5 * (xq[:-1] + xq[1:]) yq = np.array([-30., 0., 30.]); yh = np.array([-15., 15.]) lon_c, lat_c = np.meshgrid(xq, yq); lon2, lat2 = np.meshgrid(xh, yh) - g = xgcm.Grid( + return xgcm.Grid( xr.Dataset(coords={ "xq": np.arange(N + 1), "yq": np.arange(3), "xh": np.arange(N), "yh": np.arange(2), "geolon_c": (("yq", "xq"), lon_c), "geolat_c": (("yq", "xq"), lat_c), @@ -186,18 +183,222 @@ def test_zero_length_seam_face_dropped(): coords={"X": {"center": "xh", "outer": "xq"}, "Y": {"center": "yh", "outer": "yq"}}, padding={"X": "periodic", "Y": "extend"}, autoparse_metadata=False, ) - # short way from 315 deg to 45 deg runs east across the periodic seam - i_c, j_c, lons_c, lats_c = grid_section(g, [315., 45.], [0., 0.]) - edge_len = distance_on_unit_sphere(lons_c[:-1], lats_c[:-1], lons_c[1:], lats_c[1:]) - n_degenerate = int(np.sum(edge_len < 1.e-3)) - assert n_degenerate == 1 # one zero-length edge at the seam vertex - assert i_c.size >= 4 # the doubled corner is retained in i_c + +def test_zero_length_seam_corner_dropped_by_section_finding(): + """On a symmetric periodic grid the seam vertex (360 == 0) carries two indices, and the + walk steps through both. That pair spans no grid cell, so it is not a velocity face: + `grid_section` drops the redundant index, so the section it returns never contains a + consecutive pair of corners at the same physical point, and every consecutive pair + becomes a face. The faces are the same ones the walk implied, in both directions.""" + from sectionate.section import grid_section, distance_on_unit_sphere + from sectionate.transports import uvindices_from_qindices + g = _symmetric_periodic_grid() + # short way from 315 deg to 45 deg runs east across the periodic seam; the reverse + # section runs west across it. + for lons, expected_i, expected_faces in [ + ([315., 45.], [7, 0, 1], [7, 0]), + ([45., 315.], [1, 8, 7], [0, 7]), + ]: + i_c, j_c, lons_c, lats_c = grid_section(g, lons, [0., 0.]) + + # The seam vertex survives exactly once: no consecutive pair is the same point. + edge_len = distance_on_unit_sphere(lons_c[:-1], lats_c[:-1], lons_c[1:], lats_c[1:]) + assert np.all(edge_len > 1.e-3) + assert i_c.tolist() == expected_i + + uv = uvindices_from_qindices(g, i_c, j_c) + # every consecutive pair of corners is now a face, and they are the two real + # faces flanking the seam vertex (index arithmetic wraps the seam crossing). + assert uv["var"].size == i_c.size - 1 + assert np.all(uv["var"] == "V") + assert uv["i"].tolist() == expected_faces + + # normalising an already-normalised path is a no-op + again = drop_repeated_corners(g, i_c, j_c) + assert again[0].tolist() == i_c.tolist() + assert again[1].tolist() == j_c.tolist() + + +def test_repeated_corner_raises_rather_than_inventing_a_face(): + """`uvindices_from_qindices` requires the section-finding invariant instead of + re-establishing it, so it must CHECK rather than assume: a hand-built path that steps + through both of the seam vertex's indices has to raise. Silently deriving a face from + that pair would invent a duplicate of a neighbouring face and double-count its flux.""" + from sectionate.section import grid_section, drop_repeated_corners + from sectionate.transports import uvindices_from_qindices + + g = _symmetric_periodic_grid() + i_c, j_c, _, _ = grid_section(g, [315., 45.], [0., 0.]) + assert i_c.tolist() == [7, 0, 1] + + # the raw walk, stepping through both indices of the seam vertex (360 == 0) + bad_i, bad_j = [7, 8, 0, 1], [1, 1, 1, 1] + with pytest.raises(ValueError, match="same physical point"): + uvindices_from_qindices(g, bad_i, bad_j) + + # `drop_repeated_corners` is the documented remedy, and recovers the traced section + fixed_i, fixed_j, fixed_f, _, _ = drop_repeated_corners(g, bad_i, bad_j) + assert fixed_f is None + assert fixed_i.tolist() == i_c.tolist() + assert fixed_j.tolist() == j_c.tolist() + + +def test_create_section_composite_normalises_when_given_a_grid(): + """`create_section_composite` is public and documented as the lower-level entry point, + so what it returns must be usable by the rest of the public API. Given the grid it + normalises exactly as `grid_section` does; without one the caller still gets the raw + walk, which `uvindices_from_qindices` refuses.""" + from sectionate.section import create_section_composite, grid_section + from sectionate.gridutils import get_geo_corners, build_neighbor_maps + from sectionate.transports import uvindices_from_qindices + + g = _symmetric_periodic_grid() + geocorners = get_geo_corners(g) + maps = build_neighbor_maps(g, geocorners) + args = (geocorners["X"], geocorners["Y"], [315., 45.], [0., 0.]) + + raw_i, raw_j, _, _ = create_section_composite(*args, neighbor_maps=maps) + assert raw_i.tolist() == [7, 8, 0, 1] # the seam vertex under both indices + with pytest.raises(ValueError, match="same physical point"): + uvindices_from_qindices(g, raw_i, raw_j) + + i_c, j_c, _, _ = create_section_composite(*args, neighbor_maps=maps, grid=g) + assert i_c.tolist() == grid_section(g, [315., 45.], [0., 0.])[0].tolist() + assert uvindices_from_qindices(g, i_c, j_c)["var"].size == i_c.size - 1 + + +def _idealized_grid(xpad, ypad, N=8): + """A symmetric ('outer') grid whose two axes can each be declared periodic or walled, + for exercising the seam-crossing index arithmetic on either topology.""" + xq = np.linspace(0., 360., N + 1); xh = 0.5 * (xq[:-1] + xq[1:]) + yq = np.linspace(-40., 40., N + 1); yh = 0.5 * (yq[:-1] + yq[1:]) + lon_c, lat_c = np.meshgrid(xq, yq); lon2, lat2 = np.meshgrid(xh, yh) + return xgcm.Grid( + xr.Dataset(coords={ + "xq": np.arange(N + 1), "yq": np.arange(N + 1), + "xh": np.arange(N), "yh": np.arange(N), + "geolon_c": (("yq", "xq"), lon_c), "geolat_c": (("yq", "xq"), lat_c), + "geolon": (("yh", "xh"), lon2), "geolat": (("yh", "xh"), lat2), + }), + coords={"X": {"center": "xh", "outer": "xq"}, "Y": {"center": "yh", "outer": "yq"}}, + padding={"X": xpad, "Y": ypad}, autoparse_metadata=False, + ) + + +def test_wrap_rule_applies_only_to_periodic_axes(): + """A step of more than one index along an axis means the section crossed that axis' + seam and went the short way round -- so its direction is the opposite of what the index + difference says, and the face is the one on the near side of the seam. That reading is + only valid on an axis that WRAPS. On a walled axis the same step is a jump between two + indices of one physical corner (the degenerate column of a bipolar cap), and reading it + as a wrap puts the face on the wrong side and its cell index out of range. Each rule is + therefore gated on its own axis' periodicity.""" + from sectionate.transports import uvindices_from_qindices + + periodic = _idealized_grid("periodic", "periodic") + walled = _idealized_grid("extend", "extend") + + # A zonal step of -7: eastward across the "X" seam if it wraps, else a jump west. + uv = uvindices_from_qindices(periodic, [7, 0], [1, 1]) + assert (uv["var"][0], uv["i"][0], bool(uv["Xinc"][0])) == ("V", 7, True) + uv = uvindices_from_qindices(walled, [7, 0], [1, 1]) + assert (uv["var"][0], uv["i"][0], bool(uv["Xinc"][0])) == ("V", 0, False) + + # ...and the same, one axis over: a meridional step of -7 across the "Y" seam. + uv = uvindices_from_qindices(periodic, [1, 1], [7, 0]) + assert (uv["var"][0], uv["j"][0], bool(uv["Yinc"][0])) == ("U", 7, True) + uv = uvindices_from_qindices(walled, [1, 1], [7, 0]) + assert (uv["var"][0], uv["j"][0], bool(uv["Yinc"][0])) == ("U", 0, False) + + +def _pole_column_grid(N=6, M=4, jp=2): + """A miniature bipolar cap: an X-periodic symmetric grid whose corner column `i=0` + (and its seam twin `i=N`) collapses to a single physical point from row `jp` up, the + way a bipolar cap's grid column converges on its pole. Rows `jp..M` of that column are + one corner stored 2*(M-jp+1) times.""" + xq = np.linspace(0., 360., N + 1); xh = 0.5 * (xq[:-1] + xq[1:]) + yq = np.linspace(-30., 30., M + 1); yh = 0.5 * (yq[:-1] + yq[1:]) + lon_c, lat_c = np.meshgrid(xq, yq); lon2, lat2 = np.meshgrid(xh, yh) + for i in (0, N): + lon_c[jp:, i] = lon_c[jp, 0] + lat_c[jp:, i] = lat_c[jp, 0] + return xgcm.Grid( + xr.Dataset(coords={ + "xq": np.arange(N + 1), "yq": np.arange(M + 1), + "xh": np.arange(N), "yh": np.arange(M), + "geolon_c": (("yq", "xq"), lon_c), "geolat_c": (("yq", "xq"), lat_c), + "geolon": (("yh", "xh"), lon2), "geolat": (("yh", "xh"), lat2), + }), + coords={"X": {"center": "xh", "outer": "xq"}, "Y": {"center": "yh", "outer": "yq"}}, + padding={"X": "periodic", "Y": "extend"}, autoparse_metadata=False, + ) + + +def test_degenerate_column_collapses_to_the_index_that_names_both_faces(): + """A path that walks up a degenerate column and turns off it enters and leaves the pole + on the SAME index, so one index stands in for the whole run: the collapsed path must + yield exactly the two faces the raw path did. Which index survives is not free -- it is + the one that reproduces them.""" + from sectionate.section import drop_repeated_corners + from sectionate.transports import uvindices_from_qindices + + g = _pole_column_grid() + raw_i, raw_j = [6, 6, 6, 6, 5], [1, 2, 3, 4, 4] + raw_faces = [("U", 6, 1), ("V", 5, 4)] # the only two real faces of that path + + i_c, j_c, f_c, lons_c, lats_c = drop_repeated_corners(g, raw_i, raw_j) + assert f_c is None + assert i_c.tolist() == [6, 6, 5] and j_c.tolist() == [1, 4, 4] uv = uvindices_from_qindices(g, i_c, j_c) - # the degenerate edge is not a face; the flanking real faces are - assert uv["var"].size == (i_c.size - 1) - n_degenerate - assert np.all(uv["var"] != "0") + assert uv["var"].size == i_c.size - 1 + assert list(zip(uv["var"], uv["i"].tolist(), uv["j"].tolist())) == raw_faces + + +def test_degenerate_column_entered_and_left_on_different_indices_raises(): + """The pole of a bipolar cap is one physical corner stored under a whole column of + indices. A section can reach it along one of them and leave along another, far apart in + the index lattice; then no single index is adjacent to both neighbours, so no path with + distinct consecutive corners carries both flanking faces. That must raise and name the + corner, not silently emit a face the section does not have (collapsing to the last + index here would turn the leading V face into a U face).""" + from sectionate.section import drop_repeated_corners + + g = _pole_column_grid() + # In along row 4 (a V face), out along row 2 (another V face), via the pole. + raw_i, raw_j = [1, 0, 0, 0, 6, 5], [4, 4, 3, 2, 2, 2] + with pytest.raises(ValueError, match="degenerate"): + drop_repeated_corners(g, raw_i, raw_j) + with pytest.raises(ValueError, match=r"enters it from \(1, 4\)"): + drop_repeated_corners(g, raw_i, raw_j) + + +def test_gridded_section_roundtrips_through_save_load(tmp_path): + """A section written by `save_gridded_section` already satisfies the invariant, so + loading it neither needs to nor does change it, and it yields identical faces.""" + import json + from sectionate.section import Section, GriddedSection + from sectionate.utils import save_gridded_section, load_gridded_section + from sectionate.transports import uvindices_from_qindices + + g = _symmetric_periodic_grid() + gs = GriddedSection(Section("seam", ([315., 45.], [0., 0.])), g) + + path = str(tmp_path / "gs.json") + save_gridded_section(path, gs) + with open(path) as fh: + assert json.load(fh)["i_c"] == [7, 0, 1] # the file holds the normalised path + + gs2 = load_gridded_section(path, g) + assert np.asarray(gs2.i_c).tolist() == np.asarray(gs.i_c).tolist() + assert np.asarray(gs2.j_c).tolist() == np.asarray(gs.j_c).tolist() + assert np.allclose(gs2.lons_c, gs.lons_c) + + uv1 = uvindices_from_qindices(g, gs.i_c, gs.j_c) + uv2 = uvindices_from_qindices(g, gs2.i_c, gs2.j_c) + for key in uv1: + assert np.array_equal(uv1[key], uv2[key]) def test_coincident_twin_termination(): diff --git a/sectionate/tests/test_section_cornercases.py b/sectionate/tests/test_section_cornercases.py index e3e658c..bf219a6 100644 --- a/sectionate/tests/test_section_cornercases.py +++ b/sectionate/tests/test_section_cornercases.py @@ -56,17 +56,35 @@ def test_periodic_grid_section(): lonseg = np.array([300, 60]) latseg = np.array([0, 0]) i, j, lons, lats = grid_section(grid, lonseg, latseg) - # The seam vertex (360 == 0) is symmetric+periodic, so it carries two indices (6 and 0): - # the path steps through both, leaving a doubled corner. The zero-length edge between - # them carries no flux and is dropped when faces are derived (see uvindices_from_qindices). + # The seam vertex (360 == 0) is symmetric+periodic, so it is stored under two indices + # (6 and 0). The walk steps through both, but that hand-off spans no grid cell, so only + # one of them survives into the section (`drop_repeated_corners` keeps the later one, + # index 0 -- the frame the path continues in). The seam vertex therefore appears exactly + # once, and every consecutive pair of corners is a real velocity face. assert np.all([ - modequal(i, np.array([5, 6, 0, 1])), - modequal(j, np.array([2, 2, 2, 2])), - modequal(lons, np.array([300., 360., 0., 60.])), - modequal(lats, np.array([0., 0., 0., 0.])) + modequal(i, np.array([5, 0, 1])), + modequal(j, np.array([2, 2, 2])), + modequal(lons, np.array([300., 0., 60.])), + modequal(lats, np.array([0., 0., 0.])) ]) +def test_corner_only_grid_still_yields_velocity_faces(): + """This module's grid declares corner coordinates only -- no cell-center dimensions at + all. It can carry no velocity data, but a section traced on it must still enumerate the + velocity faces its corners define: the corner->face derivation needs the center + dimensions only to wrap a periodic seam crossing back into the centers' range, so it + must look them up where it uses them rather than demand them up front.""" + from sectionate.section import grid_section + from sectionate.transports import uvindices_from_qindices + i, j, lons, lats = grid_section(grid, np.array([300., 60.]), np.array([0., 0.])) + uv = uvindices_from_qindices(grid, i, j) + assert uv["var"].size == i.size - 1 + assert uv["var"].tolist() == ["V", "V"] + assert uv["i"].tolist() == [5, 0] # the two faces flanking the seam vertex + assert uv["j"].tolist() == [2, 2] + + def test_latitude_circle_zero_length_closure(): """A latitude-circle section closing back to its start (a 360 -> 0 segment) must trace as a zero-length step, not be rejected as a 360-degree arc. This is the @@ -97,15 +115,18 @@ def test_latitude_circle_takes_shortest_path_west(): # Unit level: a raw 270-degree change is a 90-degree westward one; not an error. _check_segment_span(0., 0., 270., 0., "latitude circle") - # End to end: 0 -> 270 leaves lon=0 westward across the periodic seam (the seam - # vertex carries both index 6 (lon 360) and index 0 (lon 0)); the endpoint snaps to - # the nearest corner at lon=240. The path never visits the eastern half. + # End to end: 0 -> 270 leaves lon=0 westward across the periodic seam; the endpoint + # snaps to the nearest corner at lon=240. The path never visits the eastern half. + # The seam vertex is one physical corner stored under two indices -- 0 (lon 0) and + # 6 (lon 360) -- and the walk starts on the first, then hands off to the second to + # head west. `drop_repeated_corners` removes the redundant index right after the + # walk, so the returned path starts at index 6 rather than carrying both. i, j, lons, lats = grid_section(grid, [0., 270.], [0., 0.], curve="latitude circle") assert np.all([ - modequal(i, np.array([0, 6, 5, 4])), - modequal(j, np.array([2, 2, 2, 2])), - modequal(lons, np.array([0., 360., 300., 240.])), - modequal(lats, np.array([0., 0., 0., 0.])), + modequal(i, np.array([6, 5, 4])), + modequal(j, np.array([2, 2, 2])), + modequal(lons, np.array([360., 300., 240.])), + modequal(lats, np.array([0., 0., 0.])), ]) assert not np.any(np.isin(np.mod(lons, 360.), [60., 120., 180.])) @@ -254,7 +275,8 @@ def test_parallel_is_held_where_the_geodesic_bows(): """The test that actually separates the two metrics. Between (0, 40N) and (120, 40N) the geodesic bows a long way poleward -- it is the shorter path -- while the parallel does not. On a 2-degree grid the constant-latitude walk holds lat 40 for all 61 of its - points; the great-circle walk climbs to 60N and takes 83.""" + points; the great-circle walk climbs to 60N and takes 81 (it crosses the periodic + seam diagonally, and `drop_repeated_corners` removes the redundant indices).""" from sectionate.section import grid_section fine = _fine_global_grid() @@ -266,7 +288,7 @@ def test_parallel_is_held_where_the_geodesic_bows(): i_gc, j_gc, lons_gc, lats_gc = grid_section(fine, [0., 120.], [40., 40.]) assert lats_gc.max() == 60. # bows ~20 degrees poleward - assert len(i_gc) == 83 + assert len(i_gc) == 81 # The combined option classifies this segment as constant-latitude, so it must # reproduce the parallel exactly, not the geodesic. diff --git a/sectionate/tests/test_section_fold.py b/sectionate/tests/test_section_fold.py index a6f681d..46d66e7 100644 --- a/sectionate/tests/test_section_fold.py +++ b/sectionate/tests/test_section_fold.py @@ -35,7 +35,8 @@ def _fold_grid(nx=8, ny=5, pivot="corner"): def _xgcm_supports_fold(): - """The bipolar-fold boundary is only available in newer xgcm (xgcm/xgcm#711).""" + """Whether the installed xgcm accepts `padding={"Y": {"fold": ...}}`, the bipolar + north-fold boundary these tests are about. Older releases reject it.""" try: _fold_grid(4, 3) return True @@ -46,7 +47,7 @@ def _xgcm_supports_fold(): # Skip the whole module on xgcm releases without the bipolar-fold boundary. pytestmark = pytest.mark.skipif( not _xgcm_supports_fold(), - reason="installed xgcm lacks the bipolar north-fold boundary (xgcm/xgcm#711)", + reason="installed xgcm does not accept padding={'Y': {'fold': ...}}", ) @@ -111,3 +112,104 @@ def test_fold_section_crosses_arctic_on_real_grid(): assert np.all(np.abs(np.asarray(lats_c) - 80.0) < 1.0) # It reaches the northern fold row. assert j_c.max() == ds.yq.size - 1 + + +@pytest.mark.skipif( + not os.path.exists(_mom6_example_path()), + reason="MOM6 example dataset not present (download via examples/load_example_model_grid.py)", +) +def test_fold_grid_sections_contain_no_repeated_corner(): + """On the real MOM6 tripolar grid, sections that cross the periodic seam used to carry + the seam meridian twice (it is stored as both corner column 0 and corner column nx). + `grid_section` now drops the redundant index, so no consecutive pair of corners is + the same physical point and every consecutive pair is a real velocity face.""" + from sectionate.section import distance_on_unit_sphere, COINCIDENT_TOLERANCE_M + from sectionate.transports import uvindices_from_qindices + + ds = xr.open_dataset(_mom6_example_path()) + coords = {"X": {"center": "xh", "outer": "xq"}, "Y": {"center": "yh", "outer": "yq"}} + fold = xgcm.Grid(ds, coords=coords, padding={"X": "periodic", "Y": {"fold": "corner"}}, + autoparse_metadata=False) + + cases = [ + # a closed Arctic polygon that crosses the seam twice + (dict(lons=[40., 60., 80., 40.], lats=[60., 60., 75., 60.]), {}), + # a full southern latitude circle, crossing the seam once + (dict(lons=[0., 120., 240., 360.], lats=[-60.] * 4), dict(curve="latitude circle")), + # a high-latitude circle that reaches the fold row + (dict(lons=np.arange(0., 365., 5.).tolist(), lats=[80.] * 73), {}), + ] + for section, kwargs in cases: + i_c, j_c, lons_c, lats_c = grid_section( + fold, section["lons"], section["lats"], **kwargs + ) + edge = distance_on_unit_sphere(lons_c[:-1], lats_c[:-1], lons_c[1:], lats_c[1:]) + assert np.all(edge > COINCIDENT_TOLERANCE_M), "section repeats a physical corner" + + uv = uvindices_from_qindices(fold, i_c, j_c) + assert uv["var"].size == i_c.size - 1 + assert set(np.unique(uv["var"]).tolist()) <= {"U", "V"} + # cell-center indices stay within the grid + assert np.all(uv["i"][uv["var"] == "V"] < ds.xh.size) + assert np.all(uv["j"][uv["var"] == "U"] < ds.yh.size) + + +@pytest.mark.skipif( + not os.path.exists(_mom6_example_path()), + reason="MOM6 example dataset not present (download via examples/load_example_model_grid.py)", +) +def test_meridional_section_through_pole_column_stays_in_range(): + """A meridional section along 120W runs straight into the corner column where the + bipolar cap converges on one of its poles: 41 corner rows that are all one physical + point. Collapsing that run turns the step into the pole into a jump of 41 rows along + "Y" -- which must NOT be read as a periodic wrap, because a bipolar-fold grid's "Y" + axis does not wrap. Read as one, the U face lands on `j = ny`, one past the last cell + center, and reading its transport goes out of bounds. Assert the invariant rather than + the indices: every velocity index sits inside its own centre axis.""" + from sectionate.transports import uvindices_from_qindices, convergent_transport + + ds = xr.open_dataset(_mom6_example_path()) + coords = {"X": {"center": "xh", "outer": "xq"}, "Y": {"center": "yh", "outer": "yq"}} + fold = xgcm.Grid(ds, coords=coords, padding={"X": "periodic", "Y": {"fold": "corner"}}, + autoparse_metadata=False) + + # 120W is one of this grid's two bipoles; its neighbourhood must behave the same way. + for lon in (-121., -120.5, -120., -119.5, -119.): + i_c, j_c, lons_c, lats_c = grid_section(fold, [lon, lon], [50., 85.]) + uv = uvindices_from_qindices(fold, i_c, j_c) + # U lives at (X-corner, Y-center) and V at (X-center, Y-corner): each index must + # be inside the dimension it is staggered onto. + assert uv["j"][uv["var"] == "U"].max() < ds.sizes["yh"] + assert uv["i"][uv["var"] == "U"].max() < ds.sizes["xq"] + assert uv["i"][uv["var"] == "V"].max() < ds.sizes["xh"] + assert uv["j"][uv["var"] == "V"].max() < ds.sizes["yq"] + # ...so the transport can actually be read. + conv = convergent_transport(fold, i_c, j_c, utr="umo", vtr="vmo") + assert np.isfinite(np.nansum(conv["conv_mass_transport"].values)) + + +@pytest.mark.skipif( + not os.path.exists(_mom6_example_path()), + reason="MOM6 example dataset not present (download via examples/load_example_model_grid.py)", +) +def test_section_through_a_pole_column_raises_rather_than_naming_a_wrong_face(): + """A section can reach a bipole along one index of its degenerate column and leave along + another, far apart in the index lattice -- here in along the 60E meridian at corner row + 180 and out along corner row 140. No single index of that corner is adjacent to both + neighbours, so no section with distinct consecutive corners carries both flanking faces + (keeping either one turns the face on the other side from a V into a U). Refuse it and + name the corner instead of emitting a face the section does not have.""" + ds = xr.open_dataset(_mom6_example_path()) + coords = {"X": {"center": "xh", "outer": "xq"}, "Y": {"center": "yh", "outer": "yq"}} + # The fold left undeclared: the walk cannot cross the seam, so it is pushed through the + # pole column instead of around it. + extend = xgcm.Grid(ds, coords=coords, padding={"X": "periodic", "Y": "extend"}, + autoparse_metadata=False) + with pytest.raises(ValueError, match="degenerate"): + grid_section(extend, [188.437, 32.017], [84.876, 13.851]) + + # Declaring the fold routes the same section around the pole, and it traces. + fold = xgcm.Grid(ds, coords=coords, padding={"X": "periodic", "Y": {"fold": "corner"}}, + autoparse_metadata=False) + i_c, j_c, lons_c, lats_c = grid_section(fold, [188.437, 32.017], [84.876, 13.851]) + assert i_c.size > 1 diff --git a/sectionate/tests/test_section_multitile.py b/sectionate/tests/test_section_multitile.py index e42f766..57aed99 100644 --- a/sectionate/tests/test_section_multitile.py +++ b/sectionate/tests/test_section_multitile.py @@ -11,26 +11,39 @@ from sectionate.transports import convergent_transport -def _make_grid(lon, lat, face_connections): - """Build a symmetric ('outer' corner) multi-tile grid from (face, yg, xg) coords.""" +def _make_grid(lon, lat, face_connections, centers=False): + """Build a symmetric ('outer' corner) multi-tile grid from (face, yg, xg) coords. + + `centers=True` also registers the cell-center dimensions, which a multi-tile grid needs + for anything at all: the shared-corner topology is rebuilt from the tracer cells around + each corner, so both `build_neighbor_maps` and `grid_section` (see + `_check_supported_topology`) require them. The default leaves them out, to check that + such a grid is refused by name rather than mis-traced. + """ nf, ng, _ = lon.shape - ds = xr.Dataset({}, coords={ + coords = { "xg": (("xg",), np.arange(ng)), "yg": (("yg",), np.arange(ng)), "face": (("face",), np.arange(nf)), "geolon_c": (("face", "yg", "xg"), lon), "geolat_c": (("face", "yg", "xg"), lat), - }) + } + axes = {"X": {"outer": "xg"}, "Y": {"outer": "yg"}} + if centers: + coords["xc"] = (("xc",), np.arange(ng - 1)) + coords["yc"] = (("yc",), np.arange(ng - 1)) + axes = {"X": {"outer": "xg", "center": "xc"}, + "Y": {"outer": "yg", "center": "yc"}} return xgcm.Grid( - ds, - coords={"X": {"outer": "xg"}, "Y": {"outer": "yg"}}, + xr.Dataset({}, coords=coords), + coords=axes, padding="fill", fill_value=np.nan, face_connections=face_connections, autoparse_metadata=False, ) -def two_face_x_to_x(Nc=6): +def two_face_x_to_x(Nc=6, centers=False): """Two faces side-by-side in longitude: face0 [0,90], face1 [90,180].""" ng = Nc + 1 lat1d = np.linspace(-45, 45, ng) @@ -41,7 +54,7 @@ def two_face_x_to_x(Nc=6): lat[1] = lat[0] fc = {"face": {0: {"X": (None, (1, "X", False))}, 1: {"X": ((0, "X", False), None)}}} - return _make_grid(lon, lat, fc) + return _make_grid(lon, lat, fc, centers=centers) def left_two_tile_x_to_y(Nc=5): @@ -67,35 +80,6 @@ def left_two_tile_x_to_y(Nc=5): padding="fill", fill_value=np.nan, face_connections=fc, autoparse_metadata=False) -def two_face_x_to_y(Nc=4): - """face0 right-X connects to face1 Y (a 90-degree rotation).""" - ng = Nc + 1 - lon = np.zeros((2, ng, ng)); lat = np.zeros((2, ng, ng)) - # Encode (face, i) so we can read topology back out; geometry is not used here. - for f in range(2): - for j in range(ng): - for i in range(ng): - lon[f, j, i] = f * 1000 + i - lat[f, j, i] = j - fc = {"face": {0: {"X": (None, (1, "Y", False))}, - 1: {"Y": ((0, "X", False), None)}}} - return _make_grid(lon, lat, fc) - - -def cubed_sphere(Nc=4): - ng = Nc + 1 - lon = np.zeros((6, ng, ng)); lat = np.zeros((6, ng, ng)) - fc = {"face": { - 0: {"X": ((3, "X", False), (1, "X", False)), "Y": ((4, "Y", False), (5, "Y", False))}, - 1: {"X": ((0, "X", False), (2, "X", False)), "Y": ((4, "X", False), (5, "X", True))}, - 2: {"X": ((1, "X", False), (3, "X", False)), "Y": ((4, "Y", True), (5, "Y", True))}, - 3: {"X": ((2, "X", False), (0, "X", False)), "Y": ((4, "X", True), (5, "X", False))}, - 4: {"X": ((3, "Y", True), (1, "Y", False)), "Y": ((2, "Y", True), (0, "Y", False))}, - 5: {"X": ((3, "Y", False), (1, "Y", True)), "Y": ((0, "Y", False), (2, "Y", True))}, - }} - return _make_grid(lon, lat, fc) - - def is_neighbor(prev, cur, maps): f, j, i = prev for d in NEIGHBOR_DIRECTIONS: @@ -119,33 +103,6 @@ def assert_path_invariant(i_c, j_c, f_c, maps): # Neighbor-map unit tests (topology only) # --------------------------------------------------------------------------- -def test_neighbor_maps_x_to_x_seam(): - grid = two_face_x_to_x() - maps = build_neighbor_maps(grid, get_geo_corners(grid)) - nx = maps["right"][0].shape[-1] - # Right edge of face 0 -> face 1, landing on its left edge (i=0), same row. - fmap, jmap, imap = maps["right"] - assert np.all(fmap[0][:, -1] == 1) - assert np.all(imap[0][:, -1] == 0) - assert np.all(jmap[0][:, -1] == np.arange(grid._ds.sizes["yg"])) - # Interior steps right by one on the same face. - assert np.all(fmap[0][:, :-1] == 0) - assert np.all(imap[0][:, :-1] == np.arange(1, nx)) - - -def test_neighbor_maps_x_to_y_rotation_and_reversal(): - grid = two_face_x_to_y() - maps = build_neighbor_maps(grid, get_geo_corners(grid)) - fmap, jmap, imap = maps["right"] - n = grid._ds.sizes["yg"] - # Right edge of face0 lands on face1... - assert np.all(fmap[0][:, -1] == 1) - # ...on its bottom row (Y=0, the X->Y rotation)... - assert np.all(jmap[0][:, -1] == 0) - # ...with the tangential index reversed (j -> n-1-j). - assert np.all(imap[0][:, -1] == (n - 1) - np.arange(n)) - - def test_neighbor_maps_left_grid_rotated_seam_are_edge_adjacent(): """On a native 'left'-staggered grid, the vorticity lattice sits half a cell from the centers a face seam is defined on, so xgcm's corner-array padding lands @@ -183,31 +140,12 @@ def cells(f, j, i): # the (up to four) cells touching corner (f, j, i) assert len(a & b) >= 2, f"{d}: ({f},{j},{i}) -> {nb} only diagonally adjacent" -def test_cubed_sphere_reciprocal_or_raises(): - # xgcm's padding is unreliable (hash-seed dependent) for the cubed sphere's - # complex reversed/rotated connections. build_neighbor_maps must therefore - # EITHER produce a fully reciprocal (correct) topology OR refuse with a clear - # error -- it must never silently return an inconsistent topology. This holds - # regardless of which (buggy) halo xgcm happens to produce. - grid = cubed_sphere() - try: - maps = build_neighbor_maps(grid, get_geo_corners(grid)) - except NotImplementedError: - return - # If it did not raise, the maps must be self-consistent: when face 0's right - # neighbor is face 1, the connection definition (X right -> 1) must hold. - assert np.all(maps["right"][0][0][:, -1] == 1) - assert np.all(maps["left"][0][0][:, 0] == 3) - assert np.all(maps["up"][0][0][-1, :] == 5) - assert np.all(maps["down"][0][0][0, :] == 4) - - # --------------------------------------------------------------------------- # End-to-end section tests # --------------------------------------------------------------------------- def test_cross_face_section_returns_face_and_crosses_seam(): - grid = two_face_x_to_x() + grid = two_face_x_to_x(centers=True) out = grid_section(grid, [15., 165.], [0., 0.]) assert len(out) == 5 # multi-tile grids return the face index too i_c, j_c, f_c, lons_c, lats_c = out @@ -217,7 +155,7 @@ def test_cross_face_section_returns_face_and_crosses_seam(): def test_cross_face_section_path_invariant(): - grid = two_face_x_to_x() + grid = two_face_x_to_x(centers=True) maps = build_neighbor_maps(grid, get_geo_corners(grid)) i_c, j_c, f_c, lons_c, lats_c = grid_section(grid, [15., 165.], [0., 0.]) assert_path_invariant(i_c, j_c, f_c, maps) @@ -547,6 +485,103 @@ def test_save_load_roundtrip_preserves_face_indices(tmp_path): assert np.isclose(t1, t2) +def test_section_never_repeats_a_shared_seam_corner(): + """The corner where two tiles meet is stored on both of them. A section crossing the + seam must contain that physical corner exactly ONCE -- consecutive corners are always + distinct points, so every consecutive pair is a real velocity face and no zero-flux + placeholder is ever emitted.""" + from sectionate.gridutils import outer_topology + from sectionate.transports import uvindices_from_qindices + + grid = _two_face_transport_grid() + i_c, j_c, f_c, _, _ = grid_section(grid, [15., 165.], [0., 0.]) + assert set(np.unique(f_c).tolist()) == {0, 1} # it really crosses the seam + + ot = outer_topology(grid) + nodes = ot.node_id[f_c, j_c + ot.t, i_c + ot.t] + assert (nodes >= 0).all() + assert np.all(nodes[1:] != nodes[:-1]) # no repeated physical corner + # every corner is the canonical native index of its node + native = ot.node_native[nodes] + assert np.array_equal(native[:, 0], f_c) + assert np.array_equal(native[:, 1], j_c) + assert np.array_equal(native[:, 2], i_c) + + uv = uvindices_from_qindices(grid, i_c, j_c, f_c=f_c) + assert uv["var"].size == i_c.size - 1 # every pair is a face + assert set(np.unique(uv["var"]).tolist()) <= {"U", "V"} + + +def test_repeated_shared_corner_raises_rather_than_inventing_a_face(): + """A multi-tile path that steps through BOTH tiles' copies of a shared boundary corner + must raise rather than yield a face for that zero-length step. `drop_repeated_corners` + is the documented remedy and recovers the traced section exactly, including the + transport, so the check costs nothing that cannot be recovered.""" + from sectionate.gridutils import outer_topology + from sectionate.section import drop_repeated_corners + from sectionate.transports import uvindices_from_qindices + + grid = _two_face_transport_grid() + i_c, j_c, f_c, _, _ = grid_section(grid, [15., 165.], [0., 0.]) + + # Re-insert, after each corner that is stored on two tiles, the other tile's copy of + # it -- the zero-length hand-off step the walk takes and section-finding removes. + ot = outer_topology(grid) + reps = {} + for f in range(ot.nf): + for j in range(ot.nyq): + for i in range(ot.nxq): + n = int(ot.node_id[f, j + ot.t, i + ot.t]) + if n >= 0: + reps.setdefault(n, []).append((f, j, i)) + + raw_i, raw_j, raw_f = [], [], [] + for k in range(i_c.size): + cur = (int(f_c[k]), int(j_c[k]), int(i_c[k])) + raw_f.append(cur[0]); raw_j.append(cur[1]); raw_i.append(cur[2]) + twins = [r for r in reps[int(ot.node_id[cur[0], cur[1] + ot.t, cur[2] + ot.t])] + if r != cur] + if twins: + raw_f.append(twins[0][0]); raw_j.append(twins[0][1]); raw_i.append(twins[0][2]) + assert len(raw_i) > i_c.size # the raw path really is longer + + with pytest.raises(ValueError, match="same physical point"): + uvindices_from_qindices(grid, raw_i, raw_j, f_c=raw_f) + + fixed_i, fixed_j, fixed_f, _, _ = drop_repeated_corners(grid, raw_i, raw_j, raw_f) + assert fixed_i.tolist() == i_c.tolist() + assert fixed_j.tolist() == j_c.tolist() + assert fixed_f.tolist() == f_c.tolist() + + kw = dict(utr="u", vtr="v", geometry="cartesian") + t_fresh = convergent_transport(grid, i_c, j_c, f_c, **kw)["conv_mass_transport"].values + t_fixed = convergent_transport(grid, fixed_i, fixed_j, fixed_f, **kw)["conv_mass_transport"].values + np.testing.assert_array_equal(t_fresh, t_fixed) + + +def test_multitile_grid_without_center_dims_raises_at_the_front_door(): + """A multi-tile grid that registers no cell-center dimension carries no velocities (U + is at (X-corner, Y-center), V at (X-center, Y-corner)) and its shared-seam corners + cannot be resolved to one identity, because the corner topology is rebuilt from the + cells around each corner. `grid_section` must refuse it immediately, naming what is + missing, rather than return a path whose seam steps are wrong.""" + grid = two_face_x_to_x() # corner coordinates only + with pytest.raises(ValueError, match="cell-center dimension"): + grid_section(grid, [15., 165.], [0., 0.]) + + # ...and the same grid with the center dimensions declared traces fine. Only the + # dimensions are needed: they carry no coordinate values. + grid_ok = two_face_x_to_x(centers=True) + assert "xc" not in grid_ok._ds.coords or grid_ok._ds["xc"] is not None + i_c, j_c, f_c, _, _ = grid_section(grid_ok, [15., 165.], [0., 0.]) + assert set(np.unique(f_c).tolist()) == {0, 1} + + # Deriving velocities from a center-less grid names the same missing piece. + from sectionate.transports import uvindices_from_qindices + with pytest.raises(ValueError, match="no 'center' position"): + uvindices_from_qindices(grid, i_c, j_c, f_c=f_c) + + # --------------------------------------------------------------------------- # A registered vertical axis must not affect horizontal topology # --------------------------------------------------------------------------- @@ -577,12 +612,12 @@ def assert_maps_equal(a, b): @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 + pytest.param(lambda: two_face_x_to_x(centers=True), id="outer_corners"), + pytest.param(left_two_tile_x_to_y, id="left_corners"), ]) 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 + """The multi-tile neighbor maps 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)) diff --git a/sectionate/transports.py b/sectionate/transports.py index 5e24a91..f965b41 100644 --- a/sectionate/transports.py +++ b/sectionate/transports.py @@ -5,7 +5,7 @@ from .gridutils import ( corner_offset, coord_dict, get_geo_corners, get_facedim, build_neighbor_maps, - outer_topology, NEIGHBOR_DIRECTIONS, + NEIGHBOR_DIRECTIONS, ) from .section import distance_on_unit_sphere, COINCIDENT_TOLERANCE_M @@ -82,16 +82,20 @@ def _left_sign(var, fv, jc, ic, A, B, glon, glat): def _uv_for_edge(A, B, neighbor_maps, offset, ranges, glon, glat): """ Velocity face for the directed section edge from corner A=(fA,jA,iA) to B=(fB,jB,iB). - Returns (var, i, j, face, Lsign), where Lsign is +1 if the stored velocity's positive - direction points left of travel (a geographic sign -- see `_left_sign`); var is "0" for a - degenerate edge that carries no flux. + Returns (var, i, j, face, Lsign), where var is "U" or "V" and Lsign is +1 if the stored + velocity's positive direction points left of travel (a geographic sign -- see `_left_sign`). The velocity index is read in a single face's frame, so no velocity is rotated across the seam: the SOURCE face when the edge's normal velocity lives there (the usual case, and where a rotated connection needs no rotation since the edge is a normal X/Y edge on that face); otherwise the DESTINATION face (e.g. the trailing edge of a crossing whose normal velocity - lives there); otherwise the edge is degenerate (a crossing through a shared boundary corner - of an 'outer' tiling). The sign is always geographic, so rotated seams orient correctly. + lives there). The sign is always geographic, so rotated seams orient correctly. + + A and B must be distinct physical points: `section.drop_repeated_corners` removes the + zero-length hand-off step between a seam corner's two indices when the section is traced, + so such a pair never reaches here. If neither face stores the edge's normal velocity, the + edge is not a velocity face of this grid and a ValueError is raised rather than a + zero-flux placeholder returned, which would silently drop real transport. """ fA, jA, iA = A fB, jB, iB = B @@ -111,8 +115,73 @@ def _uv_for_edge(A, B, neighbor_maps, offset, ranges, glon, glat): Lsign = _left_sign(var_d, fB, jB, iB, A, B, glon, glat) return var_d, int(vi_d), int(vj_d), int(fB), Lsign - # 3. degenerate crossing through a shared boundary corner -- carries no flux. - return "0", 0, 0, int(fB), 0 + # 3. Neither face stores this edge's normal velocity. + raise ValueError( + f"The section edge from corner (face={fA}, j={jA}, i={iA}) to corner " + f"(face={fB}, j={jB}, i={iB}) is not a velocity face of this grid: its normal " + f"velocity is stored neither on the source face {fA} nor on the destination " + f"face {fB}, so no transport can be read across it. This happens where a face " + "edge is not backed by any stored velocity (e.g. an unstored cap or grid-cut " + "lip). Move the waypoints so the section is routed around that corner." + ) + + +def _center_size(grid, axis): + """ + Length of `axis`' cell-center dimension, or None if the grid registers no 'center' + position on it. A corner-only grid stores no velocities, so it has no cell-center + index range to wrap a periodic seam crossing back into -- but it can still be traced, + and its corner-to-corner faces still enumerated. + """ + name = grid.axes[axis].coords.get("center") + return None if name is None else grid._ds[name].size + + +def _single_tile_face_context(grid): + """Grid-derived inputs of `_single_tile_face`: (offset, periodic, centers).""" + return ( + corner_offset(grid), + {ax: grid.axes[ax].padding == "periodic" for ax in ("X", "Y")}, + {ax: _center_size(grid, ax) for ax in ("X", "Y")}, + ) + + +def _single_tile_face(i0, j0, i1, j1, offset, periodic, centers): + """ + Velocity face for the single-tile section edge from corner (i0, j0) to corner (i1, j1), + as (var, i, j, Xinc, Yinc). `offset` is the corner->velocity index shift from + `gridutils.corner_offset`; `periodic` and `centers` come from + `_single_tile_face_context`. + + A step along one axis of more than one index means the section crossed that axis' + periodic seam and the path went the short way round, so its direction is the opposite + of what the index difference suggests. That reading is only valid on an axis that + actually wraps: on a non-periodic axis such a step is a jump between two indices of one + physical corner (e.g. the degenerate column where a bipolar cap converges on its pole), + and there the index difference already gives the direction. Hence both rules are gated + on their own axis' periodicity. + """ + zonal = j1 == j0 + Xinc = i1 > i0 + Yinc = j1 > j0 + if periodic["X"]: + if (i1 - i0) > 1: Xinc = False + elif (i1 - i0) < -1: Xinc = True + if periodic["Y"]: + if (j1 - j0) > 1: Yinc = False + elif (j1 - j0) < -1: Yinc = True + + var = "V" if zonal else "U" + vi = (i1 if (not(Xinc) and zonal) else i0) + (offset if zonal else 0) + vj = (j1 if (not(Yinc) and not(zonal)) else j0) + (offset if not(zonal) else 0) + # The velocity's along-face index is a cell-center index ("i" for a V face, "j" for a + # U face); on a periodic axis a seam crossing can land on the duplicated corner index, + # one past the last center, so wrap it back into the centers' range. + if zonal and periodic["X"] and centers["X"] is not None: + vi = vi % centers["X"] + elif not(zonal) and periodic["Y"] and centers["Y"] is not None: + vj = vj % centers["Y"] + return var, int(vi), int(vj), bool(Xinc), bool(Yinc) def uvindices_from_qindices(grid, i_c, j_c, f_c=None): @@ -121,6 +190,14 @@ def uvindices_from_qindices(grid, i_c, j_c, f_c=None): N vorticity points, automatically reading the vorticity corner position from `grid` metadata ('outer', 'right', or 'left'; see `gridutils.corner_offset`). + Every consecutive pair of corners must be a real velocity face, i.e. two *distinct* + physical points. Sections produced by `grid_section` satisfy this: the zero-length + hand-off step between a seam corner's two indices is dropped when the section is traced + (see `section.drop_repeated_corners`). This is checked, not assumed -- an index array + that repeats a physical corner raises, because deriving a face from that pair would + invent one that does not exist. Pass hand-built indices through + `section.drop_repeated_corners` first. + PARAMETERS: ----------- grid: xgcm.Grid @@ -155,22 +232,29 @@ def uvindices_from_qindices(grid, i_c, j_c, f_c=None): glon = np.asarray(geocorners["X"].values) glat = np.asarray(geocorners["Y"].values) - if f_c is not None: - # Resolve each corner to its canonical native representation on the - # grid's corner topology, dropping consecutive corners that are the - # same physical point. Sections saved before multi-tile paths became - # twin-free step through both native copies of a shared seam corner: - # that zero-length edge is not a velocity face, and the neighbor maps - # used below link only canonical representations. - ot = outer_topology(grid) - f_c = np.asarray(f_c) - nodes = ot.node_id[f_c, j_c + ot.t, i_c + ot.t] - if (nodes < 0).any(): - raise ValueError("Section contains indices that are not grid corners.") - keep = np.ones(nodes.size, dtype=bool) - keep[1:] = nodes[1:] != nodes[:-1] - nat = ot.node_native[nodes[keep]] - f_c, j_c, i_c = nat[:, 0].copy(), nat[:, 1].copy(), nat[:, 2].copy() + # Enforce the section-finding invariant. A pair of corners at the same physical point + # spans no grid cell, so it is not a velocity face; the index arithmetic below would + # nonetheless produce one (a spurious duplicate of a neighbouring face), so refuse it. + if i_c.size > 1: + if f_c is not None: + fq = np.asarray(f_c) + clon, clat = glon[fq, j_c, i_c], glat[fq, j_c, i_c] + else: + clon, clat = glon[j_c, i_c], glat[j_c, i_c] + # Strictly less than, matching `drop_repeated_corners` and the walker exactly: a + # pair this check rejects must be a pair normalisation would have removed, or a + # section straight from `grid_section` could fail the invariant it establishes. + repeated = distance_on_unit_sphere( + clon[:-1], clat[:-1], clon[1:], clat[1:] + ) < COINCIDENT_TOLERANCE_M + if repeated.any(): + k = int(np.flatnonzero(repeated)[0]) + raise ValueError( + f"Section corners {k} and {k+1} are the same physical point, so they do " + "not define a velocity face. Sections from `grid_section` never contain " + "such a pair; pass hand-built indices through " + "`sectionate.drop_repeated_corners(grid, i_c, j_c, f_c)` first." + ) nsec = i_c.size uvindices = { @@ -208,38 +292,23 @@ def uvindices_from_qindices(grid, i_c, j_c, f_c=None): uvindices["face"][k] = face uvindices["Lsign"][k] = Lsign else: + # A section that crosses a periodic seam takes one step of a whole domain width + # (or height): the seam corner's two indices are collapsed to one when the path is + # traced (`section.drop_repeated_corners`), so the face on the near side of the + # seam is expressed as a wrapped step -- which `_single_tile_face` reads as such, + # on periodic axes only. + face_context = _single_tile_face_context(grid) for k in range(0, nsec-1): - zonal = not(j_c[k+1] != j_c[k]) - Xinc = i_c[k+1] > i_c[k] - Yinc = j_c[k+1] > j_c[k] - # Handle corner cases for wrapping boundaries - if (i_c[k+1] - i_c[k])>1: Xinc = False - elif (i_c[k+1] - i_c[k])<-1: Xinc = True - uvindex = { - "var": "V" if zonal else "U", - "i": i_c[k+(1 if not(Xinc) and zonal else 0)], - "j": j_c[k+(1 if not(Yinc) and not(zonal) else 0)], - "Yinc": Yinc, - "Xinc": Xinc, - } - uvindex["i"] += (offset if zonal else 0) - uvindex["j"] += (offset if not(zonal) else 0) - for (key, v) in uvindices.items(): - v[k] = uvindex[key] - - # Drop zero-length faces. A single physical corner can carry two indices -- at a - # periodic seam, a shared multi-tile boundary corner, or the bipolar fold seam -- so - # the section path can contain a consecutive pair whose endpoints are the same point. - # That edge spans no grid cell and carries no flux, so it is not a velocity face. Both - # corners are kept (they anchor the real faces on either side), but the degenerate edge - # between them emits no face. Done here, in the deterministic corner->face derivation, - # so it applies identically when a saved section is reloaded from its (i_c, j_c[, f_c]). - if f_c is not None: - clon, clat = glon[f_c, j_c, i_c], glat[f_c, j_c, i_c] - else: - clon, clat = glon[j_c, i_c], glat[j_c, i_c] - keep = distance_on_unit_sphere(clon[:-1], clat[:-1], clon[1:], clat[1:]) > COINCIDENT_TOLERANCE_M - return {key: np.asarray(val)[keep] for key, val in uvindices.items()} + var, vi, vj, Xinc, Yinc = _single_tile_face( + int(i_c[k]), int(j_c[k]), int(i_c[k+1]), int(j_c[k+1]), *face_context + ) + uvindices["var"][k] = var + uvindices["i"][k] = vi + uvindices["j"][k] = vj + uvindices["Xinc"][k] = Xinc + uvindices["Yinc"][k] = Yinc + + return uvindices def uvcoords_from_uvindices(grid, uvindices): """ @@ -296,9 +365,10 @@ def uvcoords_from_uvindices(grid, uvindices): for p in range(len(uvindices["var"])): var, i, j = uvindices["var"][p], uvindices["i"][p], uvindices["j"][p] if var not in ("U", "V"): - # Degenerate edge (e.g. a seam crossing through a shared corner): no point. - lons[p], lats[p] = np.nan, np.nan - continue + raise ValueError( + f"`uvindices['var'][{p}]` is {var!r}; every section face is a 'U' or 'V' " + "velocity point." + ) # On multi-tile grids, also select the velocity point's face. fsel = {facedim: int(faces[p])} if (facedim is not None and faces is not None) else {} if var == "U":