Support seam-crossing regional longitude grids (open periodic axis)#38
Closed
vadmbertr wants to merge 1 commit into
Closed
Support seam-crossing regional longitude grids (open periodic axis)#38vadmbertr wants to merge 1 commit into
vadmbertr wants to merge 1 commit into
Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
PR #36 made
lon_periodrequire 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 running170°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-1interior faces, clip indexing, extrapolation past its two ends). The only periodic thing it needs is folding a query given in either convention (-175or185) into the grid's window. This PR lets you declare circularity once vialon_periodat 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):lon_period is Nonelon_periodset, spans one period(i+1) % n(unchanged, byte-identical)lon_periodset, spans lessKey point: open = bounded grid + a query pre-fold, so there are no interpolation-kernel changes. The open path is realized at the
Field/Datasetdispatch sites by folding the query and reusing the existinglon_period=Nonepath.Changes
grid.py: new staticlon_closedfield andclosed_for(stagger)accessor;u_face_coordsgivesnlon-1interior midpoint faces on an open grid (vs the seam-inclusivenlonfaces 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.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_closedproperty; the open-grid query fold applied atinterp,neighborhood, thevelocity_interppartial-slip path, andstandalone.Behaviour changes
lon_periodset is now a valid open grid rather than an error (this is the point). Period overshoot (duplicate wrap endpoint) and descending axes still raise.Field.lon_coordsreads back e.g.190where-170was passed. Queries may still be given in any convention; they are folded automatically.Tests
TestPeriodicLonSpantests that rejected sub-period grids now assert they load as open grids (lon_closed is False); overshoot/descending still raise.TestOpenSeamCrossingLon: seam continuity, representation invariance (-175==185), open-vs-closed extrapolation,jit/vmap/gradsafety (grad finite and≈1across the seam),neighborhoodclamp, and thelon_closed=override underjit.test_grid.pycoverage for open C-grid interior-midpoint faces andclosed_for.Full suite: 338 passed, 2 skipped (the closed/global periodic path is unchanged and byte-identical).