Skip to content

Support seam-crossing regional longitude grids (open periodic axis)#38

Closed
vadmbertr wants to merge 1 commit into
mainfrom
claude/lon-periodicity-regional-subsets-l57bhz
Closed

Support seam-crossing regional longitude grids (open periodic axis)#38
vadmbertr wants to merge 1 commit into
mainfrom
claude/lon-periodicity-regional-subsets-l57bhz

Conversation

@vadmbertr

@vadmbertr vadmbertr commented Jul 8, 2026

Copy link
Copy Markdown
Owner

Motivation

PR #36 made lon_period require the grid to span exactly one period (nlon * dlon == lon_period). That is correct for a global grid whose wrap identifies the cell past the last centre with the first, but it rejects a regional subset that crosses the antimeridian — e.g. longitudes running 170°E … 180° … -170° (= 190°E), spanning ~20°, not 360°.

Such a slab is topologically an open interval with the seam interior to it, so it behaves exactly like a bounded grid (nlon-1 interior faces, clip indexing, extrapolation past its two ends). The only periodic thing it needs is folding a query given in either convention (-175 or 185) into the grid's window. This PR lets you declare circularity once via lon_period at grid/dataset creation and get seam continuity automatically — no manual unwrapping or query folding.

Approach

A three-way classification of the longitude axis, resolved statically at trace time so interpolation stays jit/vmap/grad-safe (never a data-dependent branch):

axis condition behaviour
bounded lon_period is None non-periodic (unchanged)
closed / global lon_period set, spans one period wraps (i+1) % n (unchanged, byte-identical)
open / seam-crossing lon_period set, spans less centred query fold + bounded path (new)

Key point: open = bounded grid + a query pre-fold, so there are no interpolation-kernel changes. The open path is realized at the Field/Dataset dispatch sites by folding the query and reusing the existing lon_period=None path.

Changes

  • grid.py: new static lon_closed field and closed_for(stagger) accessor; u_face_coords gives nlon-1 interior midpoint faces on an open grid (vs the seam-inclusive nlon faces on a closed one).
  • forcing.py:
    • _check_periodic_lon_classify_and_unwrap_lon: classifies closed vs open, unwraps a seam-crossing axis to a monotonic ascending representation (so every downstream consumer stays seam-agnostic), still rejects descending axes and period overshoot, and runs in float64 on the host before the storage cast.
    • Optional lon_closed= override on all loaders (from_arrays, from_arrays_cgrid, from_xarray, from_xarray_cgrid) — authoritative when passed, so an open grid can be built under tracing where the axis can't be classified.
    • Field.lon_closed property; the open-grid query fold applied at interp, neighborhood, the velocity_interp partial-slip path, and standalone.
    • The fold centres its modulo branch cut in the uncovered arc opposite the grid, so both near-edge extrapolation zones stay well-behaved and differentiable (a cut at the western edge would fold a query just west of the grid nearly a full period east).

Behaviour changes

  • A sub-period grid with lon_period set is now a valid open grid rather than an error (this is the point). Period overshoot (duplicate wrap endpoint) and descending axes still raise.
  • Periodic grids are stored unwrapped (strictly ascending), so Field.lon_coords reads back e.g. 190 where -170 was passed. Queries may still be given in any convention; they are folded automatically.
  • Queries outside an open grid's covered arc extrapolate (like a bounded grid) rather than wrapping to the far side.

Tests

  • The two PR Validate periodic longitude grids span exactly one period #36 TestPeriodicLonSpan tests that rejected sub-period grids now assert they load as open grids (lon_closed is False); overshoot/descending still raise.
  • New TestOpenSeamCrossingLon: seam continuity, representation invariance (-175 == 185), open-vs-closed extrapolation, jit/vmap/grad safety (grad finite and ≈1 across the seam), neighborhood clamp, and the lon_closed= override under jit.
  • New test_grid.py coverage for open C-grid interior-midpoint faces and closed_for.

Full suite: 338 passed, 2 skipped (the closed/global periodic path is unchanged and byte-identical).

PR #36 required a periodic longitude grid to span exactly one period
(nlon * dlon == lon_period), which rejected a regional subset crossing the
antimeridian (e.g. 170°…-170°, spanning ~20°). Such a slab is topologically
an open interval with the seam interior to it: it behaves like a bounded
grid (nlon-1 interior faces, clip indexing, extrapolation past its ends).
The only periodic operation it needs is folding a query given in either
convention (-175 or 185) into the grid's window.

Introduce a three-way classification, resolved statically at trace time so
interpolation stays jit/vmap/grad-safe:

  lon_period is None                 -> bounded (unchanged)
  lon_period set, spans one period   -> closed / global periodic (unchanged)
  lon_period set, spans less         -> open / seam-crossing regional (new)

- Grid gains a static `lon_closed` bool and a `closed_for` accessor; the
  closed/global path is byte-identical to before.
- `_check_periodic_lon` becomes `_classify_and_unwrap_lon`: it classifies the
  axis (closed vs open), unwraps a seam-crossing axis to a monotonic ascending
  representation so every downstream consumer stays seam-agnostic, still
  rejects descending axes and period overshoot, and runs in float64 on the
  host before the storage cast. A `lon_closed=` override on the loaders lets an
  open grid be built under tracing, where the axis can't be classified.
- The open path is realized at the Field/Dataset dispatch sites (interp,
  neighborhood, velocity_interp partial-slip, standalone) by folding the query
  and reusing the existing bounded path — no interpolation kernel changes. The
  fold centres its branch cut in the uncovered arc opposite the grid so both
  near-edge extrapolation zones stay well-behaved (and differentiable).
- C-grid U faces on an open grid are the nlon-1 interior midpoints, not the
  seam-inclusive nlon faces.

The two PR #36 tests that rejected sub-period grids now assert they load as
valid open grids; overshoot and descending axes still raise. Adds coverage
for seam continuity, representation invariance, open-vs-closed extrapolation,
jit/vmap/grad safety, the neighborhood clamp, and open C-grid faces.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DTaB6gfoum7XU5WbMMTndn
@vadmbertr vadmbertr closed this Jul 9, 2026
@vadmbertr
vadmbertr deleted the claude/lon-periodicity-regional-subsets-l57bhz branch July 9, 2026 02:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants