From 46ae6bad32f28a41f2b0f79eacd0ebd4e04a8979 Mon Sep 17 00:00:00 2001 From: Vadim Bertrand <36510417+vadmbertr@users.noreply.github.com> Date: Thu, 9 Jul 2026 02:28:29 +0000 Subject: [PATCH] Support seam-crossing regional longitude grids (open periodic axis) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Requiring a periodic longitude grid to span exactly one period (nlon * dlon == lon_period) 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, and still rejects descending axes and period overshoot. A `lon_closed=` override on the loaders lets an open grid be built under tracing, where the axis can't be classified; passing `lon_closed=False` there also unwraps a seam so raw seam coordinates load under jit/vmap without pre-unwrapping. - 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 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, open C-grid faces, and the in-jit auto-unwrap. --- src/pastax/forcing.py | 294 +++++++++++++++++++++++++++++++----------- src/pastax/grid.py | 46 ++++++- tests/test_forcing.py | 219 ++++++++++++++++++++++++++++--- tests/test_grid.py | 35 +++++ 4 files changed, 497 insertions(+), 97 deletions(-) diff --git a/src/pastax/forcing.py b/src/pastax/forcing.py index 7493a0b..1b92041 100644 --- a/src/pastax/forcing.py +++ b/src/pastax/forcing.py @@ -169,6 +169,17 @@ def lon_period(self) -> float | None: """Longitude period for this field's stagger role (``None`` on faces).""" return self.grid.period_for(self.stagger) + @property + def lon_closed(self) -> bool: + """Whether this field's longitude axis spans one full period (wraps). + + ``True`` for a global periodic grid (interpolation wraps across the + seam), ``False`` for an open seam-crossing regional slab (continuous + across the antimeridian but not wrapping). Only meaningful when + :attr:`lon_period` is set. + """ + return self.grid.closed_for(self.stagger) + @classmethod def standalone( cls, @@ -187,27 +198,31 @@ def standalone( :class:`Grid` in the slot matching ``stagger``, and the returned field reads them back through the usual grid-backed properties. - With ``lon_period`` set, the passed ``lon_coords`` are taken to span - exactly one period and the field wraps in longitude — for a face - stagger this means the coordinates must be the seam-inclusive face - axis (``nlon`` points; see :meth:`pastax.Grid.u_face_coords`), not the - ``nlon - 1`` interior faces. + With ``lon_period`` set, the passed ``lon_coords`` are classified like + the loaders: an axis spanning exactly one period is a global (closed) + grid that wraps — for a face stagger the coordinates must then be the + seam-inclusive face axis (``nlon`` points; see + :meth:`pastax.Grid.u_face_coords`), not the ``nlon - 1`` interior faces — + while a shorter seam-crossing axis is an open grid, stored unwrapped and + continuous across the antimeridian without wrapping. """ + lon_out, lon_closed = _classify_and_unwrap_lon(lon_coords, lon_period) + lon_coords = jnp.asarray(lon_out) if stagger == "center": grid = Grid( t_coords=t_coords, lat_coords=lat_coords, lon_coords=lon_coords, - lon_period=lon_period, + lon_period=lon_period, lon_closed=lon_closed, ) elif stagger == "u_face": grid = Grid( t_coords=t_coords, lat_coords=lat_coords, lon_coords=lon_coords, - stagger_type="C", lon_period=lon_period, + stagger_type="C", lon_period=lon_period, lon_closed=lon_closed, u_lat_coords=lat_coords, u_lon_coords=lon_coords, ) elif stagger == "v_face": grid = Grid( t_coords=t_coords, lat_coords=lat_coords, lon_coords=lon_coords, - stagger_type="C", lon_period=lon_period, + stagger_type="C", lon_period=lon_period, lon_closed=lon_closed, v_lat_coords=lat_coords, v_lon_coords=lon_coords, ) else: @@ -233,8 +248,13 @@ def interp( Returns: Interpolated scalar value at the query point. Outside the grid the interpolation extrapolates linearly (clamping to grid boundaries - beyond one cell). When ``lon_period`` is set, longitude wraps - instead of extrapolating. When ``self.mask`` is set, coastal + beyond one cell). When ``lon_period`` is set on a global (closed) + grid, longitude wraps instead of extrapolating. On an open + seam-crossing regional grid the query is folded into the grid's + window (so it may be given in any longitude convention) and + interpolated continuously across the antimeridian; queries outside + the covered arc extrapolate like a bounded grid rather than wrapping + to the far side. When ``self.mask`` is set, coastal cells use inverse-distance partial-cell weighting and fully land-bound cells return ``0`` (see :func:`pastax.interpolation.bilinear_interp_2d`) — the right @@ -242,9 +262,15 @@ def interp( tracers (an SST of ``0`` over land is a value, not a gap). """ lat_coords, lon_coords = self.grid.coords_for(self.stagger) + lon_period = self.grid.period_for(self.stagger) + if lon_period is not None and not self.grid.closed_for(self.stagger): + # Open seam-crossing grid: fold the query into the stored (unwrapped, + # ascending) window and interpolate as a bounded grid — no wrap. + lon = _fold_into_window(lon_coords, lon, lon_period) + lon_period = None return spatiotemporal_interp( self.values, self.grid.t_coords, lat_coords, lon_coords, - t, lat, lon, lon_period=self.grid.period_for(self.stagger), + t, lat, lon, lon_period=lon_period, mask=self.mask, ) @@ -270,8 +296,10 @@ def neighborhood( Returns: Array of shape (2*t_window+1, 2*lat_window+1, 2*lon_window+1). Time and latitude windows are clamped to the grid boundary near the - edges. The longitude window wraps modulo ``lon_period`` when that - attribute is set, otherwise it is clamped like the others. + edges. The longitude window wraps modulo ``lon_period`` on a global + (closed) periodic grid; on an open seam-crossing grid the query is + folded into the stored window and the window is clamped (no wrap), + and without ``lon_period`` it is clamped like the others. """ nt = self.t_coords.shape[0] nlat = self.lat_coords.shape[0] @@ -287,23 +315,29 @@ def neighborhood( it_start = jnp.clip(it - t_window, 0, nt - wt) ilat_start = jnp.clip(ilat - lat_window, 0, nlat - wlat) - if self.lon_period is None: - ilon = _nearest_idx(self.lon_coords, lon, nlon) - ilon_start = jnp.clip(ilon - lon_window, 0, nlon - wlon) - return jax.lax.dynamic_slice( + if self.lon_period is not None and self.lon_closed: + # Closed global periodic grid — wrap the window across the seam. + ilon = _nearest_idx_periodic(self.lon_coords, lon, nlon, self.lon_period) + block = jax.lax.dynamic_slice( self.values, - (it_start, ilat_start, ilon_start), - (wt, wlat, wlon), + (it_start, ilat_start, jnp.astype(0, jnp.int32)), + (wt, wlat, nlon), ) - - ilon = _nearest_idx_periodic(self.lon_coords, lon, nlon, self.lon_period) - block = jax.lax.dynamic_slice( + lon_idx = (ilon - lon_window + jnp.arange(wlon)) % nlon + return block[:, :, lon_idx] + + # Bounded OR open seam-crossing grid — clamped window, no wrap. An open + # grid folds the query into its stored (unwrapped) window first. + lon_q = lon + if self.lon_period is not None: + lon_q = _fold_into_window(self.lon_coords, lon, self.lon_period) + ilon = _nearest_idx(self.lon_coords, lon_q, nlon) + ilon_start = jnp.clip(ilon - lon_window, 0, nlon - wlon) + return jax.lax.dynamic_slice( self.values, - (it_start, ilat_start, jnp.astype(0, jnp.int32)), - (wt, wlat, nlon), + (it_start, ilat_start, ilon_start), + (wt, wlat, wlon), ) - lon_idx = (ilon - lon_window + jnp.arange(wlon)) % nlon - return block[:, :, lon_idx] class Dataset(eqx.Module): @@ -460,12 +494,18 @@ def velocity_interp( "(NaN-inferred or explicit)." ) joint_mask = u_field.mask & v_field.mask + lon_q = lon + lon_period = u_field.lon_period + if lon_period is not None and not u_field.lon_closed: + # Open seam-crossing grid: fold into the window, no wrap. + lon_q = _fold_into_window(u_field.lon_coords, lon, lon_period) + lon_period = None u, v = spatiotemporal_velocity_partialslip( u_field.values, v_field.values, u_field.t_coords, u_field.lat_coords, u_field.lon_coords, - t, lat, lon, joint_mask, + t, lat, lon_q, joint_mask, slip_a=slip_a, slip_b=slip_b, - lon_period=u_field.lon_period, + lon_period=lon_period, ) return jnp.stack([u, v]) @@ -482,6 +522,7 @@ def from_arrays( lon: Array, dtype: DTypeLike = jnp.float32, lon_period: float | None = None, + lon_closed: bool | None = None, masks: dict[str, Array] | None = None, ) -> Dataset: """Build a Dataset from numpy or JAX arrays. @@ -496,11 +537,24 @@ def from_arrays( Double precision should be used to avoid truncation errors. lat: 1-D latitude coordinate array (degrees), equally spaced. lon: 1-D longitude coordinate array (degrees), equally spaced. - dtype: JAX dtype for all arrays, except for the time coordinate + dtype: JAX dtype for all arrays, except for the time coordinate (default float32). - lon_period: If set (e.g. ``360.0``), all fields are constructed - with periodic longitude wrapping. The grid must span exactly - one period. + lon_period: If set (e.g. ``360.0``), longitude is treated as circular + with this period. A grid spanning exactly one period is a global + (closed) grid that wraps; a shorter grid that crosses the + antimeridian (e.g. ``170°…-170°``) is an open regional slab, + stored unwrapped (ascending) and continuous across the seam + without wrapping. Seam-crossing input may be given in any + longitude convention; queries are folded automatically. + lon_closed: Optional override for the closed/open classification. + Normally derived from the coordinates (``nlon*dlon == lon_period`` + → closed). Pass explicitly to build a grid under ``jit``/``vmap`` + tracing, where concrete coordinates are unavailable and the + classification would otherwise default to closed. Passing + ``lon_closed=False`` under tracing additionally unwraps a + seam-crossing axis for you, so raw seam coordinates can be loaded + inside ``jit``/``vmap`` without pre-unwrapping. Ignored when + ``lon_period`` is ``None``. masks: Optional ``{field_name: 2-D bool array of shape (lat, lon)}`` land masks. ``True`` marks a land cell. When a field appears in ``masks``, that mask is used. Otherwise a mask is inferred @@ -516,8 +570,8 @@ def from_arrays( t = _coerce_time_to_seconds(t) t_arr = _asarray_time_int64(t) lat_arr = jnp.asarray(lat, dtype=dtype) - lon_arr = jnp.asarray(lon, dtype=dtype) - _check_periodic_lon(lon_arr, lon_period) + lon_out, lon_closed = _classify_and_unwrap_lon(lon, lon_period, lon_closed) + lon_arr = jnp.asarray(lon_out, dtype=dtype) nt = int(t_arr.shape[0]) nlat = int(lat_arr.shape[0]) nlon = int(lon_arr.shape[0]) @@ -528,6 +582,7 @@ def from_arrays( grid_type="rectilinear", stagger_type="A", lon_period=lon_period, + lon_closed=lon_closed, ) masks = masks or {} unknown = set(masks) - set(fields) @@ -566,6 +621,7 @@ def from_xarray( coordinates: dict[str, str], dtype: DTypeLike = jnp.float32, lon_period: float | None = None, + lon_closed: bool | None = None, masks: dict[str, Array] | None = None, ) -> Dataset: """Load a Dataset from an xarray Dataset (zarr or netCDF backed). @@ -575,9 +631,11 @@ def from_xarray( fields: Mapping {internal_name: xarray_variable_name}. coordinates: Mapping with keys "time", "lat", "lon" → xarray coord names. dtype: JAX dtype for all arrays (default float32). - lon_period: If set (e.g. ``360.0``), all fields are constructed - with periodic longitude wrapping. The grid must span exactly - one period. + lon_period: If set (e.g. ``360.0``), longitude is treated as circular + with this period (global wrap or seam-crossing regional slab); + see :meth:`from_arrays`. + lon_closed: Optional override for the closed/open classification; see + :meth:`from_arrays`. masks: Optional land masks keyed by internal field name; see :meth:`from_arrays` for semantics. If omitted, masks are inferred from NaN — which matches the CMEMS / CF @@ -596,6 +654,7 @@ def from_xarray( lon=ds[coordinates["lon"]].values, dtype=dtype, lon_period=lon_period, + lon_closed=lon_closed, masks=masks, ) @@ -613,6 +672,7 @@ def from_arrays_cgrid( v_lon: Array | None = None, dtype: DTypeLike = jnp.float32, lon_period: float | None = None, + lon_closed: bool | None = None, masks: dict[str, Array] | None = None, ) -> Dataset: """Build a Dataset on a NEMO-convention Arakawa C-grid. @@ -694,8 +754,8 @@ def from_arrays_cgrid( t = _coerce_time_to_seconds(t) t_arr = _asarray_time_int64(t) lat_arr = jnp.asarray(center_lat, dtype=dtype) - lon_arr = jnp.asarray(center_lon, dtype=dtype) - _check_periodic_lon(lon_arr, lon_period) + lon_out, lon_closed = _classify_and_unwrap_lon(center_lon, lon_period, lon_closed) + lon_arr = jnp.asarray(lon_out, dtype=dtype) nt = int(t_arr.shape[0]) nlat = int(lat_arr.shape[0]) @@ -708,14 +768,16 @@ def from_arrays_cgrid( grid_type="rectilinear", stagger_type="C", lon_period=lon_period, + lon_closed=lon_closed, ) derived_u_lat, derived_u_lon = centre_grid.u_face_coords() derived_v_lat, derived_v_lon = centre_grid.v_face_coords() - # A periodic centre grid has one U face per centre cell (the seam face - # included), so U is nlon-wide; a bounded grid has nlon - 1 interior - # faces. V is unaffected (latitude is never periodic). - nlon_u = nlon if lon_period is not None else nlon - 1 + # A closed (global periodic) centre grid has one U face per centre cell + # (the seam face included), so U is nlon-wide; a bounded OR open + # (seam-crossing regional) grid has nlon - 1 interior faces. V is + # unaffected (latitude is never periodic). + nlon_u = nlon if (lon_period is not None and lon_closed) else nlon - 1 u_lat_arr = jnp.asarray(u_lat, dtype=dtype) if u_lat is not None else derived_u_lat u_lon_arr = jnp.asarray(u_lon, dtype=dtype) if u_lon is not None else derived_u_lon @@ -735,6 +797,7 @@ def from_arrays_cgrid( grid_type="rectilinear", stagger_type="C", lon_period=lon_period, + lon_closed=lon_closed, u_lat_coords=u_lat_arr, u_lon_coords=u_lon_arr, v_lat_coords=v_lat_arr, @@ -824,6 +887,7 @@ def from_xarray_cgrid( staggered_coordinates: dict[str, str] | None = None, dtype: DTypeLike = jnp.float32, lon_period: float | None = None, + lon_closed: bool | None = None, masks: dict[str, Array] | None = None, ) -> Dataset: """Load a C-grid Dataset from an xarray Dataset. @@ -858,6 +922,7 @@ def from_xarray_cgrid( vector. dtype: JAX dtype for all arrays (default float32). lon_period: Forwarded to :meth:`from_arrays_cgrid`. + lon_closed: Forwarded to :meth:`from_arrays_cgrid`. masks: Forwarded to :meth:`from_arrays_cgrid`. Keys must match the field names declared in ``vectors`` and ``tracers``. @@ -893,6 +958,7 @@ def from_xarray_cgrid( v_lon=ds[stag["v_lon"]].values if "v_lon" in stag else None, dtype=dtype, lon_period=lon_period, + lon_closed=lon_closed, masks=masks, ) @@ -907,49 +973,127 @@ def _is_traced(x: Array) -> bool: return isinstance(x, jax.core.Tracer) -def _check_periodic_lon( - lon_arr: Float[Array, "lon"], lon_period: float | None -) -> None: - """Validate longitude coordinates when periodic wrapping is requested. +def _fold_into_window( + lon_coords: Float[Array, "lon"], lon: Float[Array, ""], period: float +) -> Float[Array, ""]: + """Fold a query longitude into an open grid's window on a circle of ``period``. - Two preconditions of the periodic index arithmetic are checked: + An open (seam-crossing) grid stores an unwrapped, strictly-ascending axis + ``[x0, ..., x1]`` spanning less than one period. A query may be given in any + convention (e.g. ``-175`` or ``185`` for the same point), so it must be + folded to the grid's representation before the ordinary bounded index + arithmetic runs. - 1. **Ascending axis.** The query is folded with ``% lon_period`` above - ``lon_coords[0]``, which assumes an ascending axis; a descending axis - silently produces wrong indices and weights. (Without ``lon_period``, - descending coordinates work — the spacing sign cancels.) - 2. **Spans exactly one period.** The wrap identifies the cell at - ``lon_coords[-1] + dlon`` with ``lon_coords[0]``, so the grid must span - exactly one period: ``nlon * dlon == lon_period``. A grid that does not - (e.g. a regional slab mislabelled periodic, or one that includes the - duplicate wrap endpoint) yields silently wrong wrapped indices. - - The check reads concrete values, so it is a no-op when ``lon_arr`` is a - tracer (the loader was called from inside ``jit`` / ``vmap``): correctness - of the coordinate axis is then the caller's responsibility. + The modulo branch cut is centred in the *uncovered* arc opposite the grid + (around ``x0 + period/2``) rather than at the western edge ``x0``: with the + cut at ``x0`` a query one cell west of the grid would fold almost a full + period east and extrapolate catastrophically (and non-differentiably) right + at the boundary. Centring it makes both near-grid extrapolation zones behave + like a plain bounded grid. Pure arithmetic — safe under ``jit``/``vmap``/``grad``. + """ + x0 = lon_coords[0] + x1 = lon_coords[-1] + mid = 0.5 * (x0 + x1) + return mid + jnp.mod(lon - mid + 0.5 * period, period) - 0.5 * period + + +def _classify_and_unwrap_lon( + lon: Array, lon_period: float | None, lon_closed: bool | None = None +) -> tuple[Array, bool]: + """Validate/unwrap a periodic longitude axis and classify it closed vs open. + + Returns ``(lon_out, lon_closed)`` where ``lon_out`` is the coordinate array + to store and ``lon_closed`` says whether the grid spans exactly one period. + + Two topologies are supported when ``lon_period`` is set: + + * **closed** — an ascending axis spanning exactly one period + (``nlon * dlon == lon_period``): the global periodic grid, wrap + ``(i+1) % n``. Ascending input is returned unchanged (the closed path is + byte-identical to before this function existed). + * **open** — a seam-crossing regional slab (e.g. ``170°…-170°``) spanning + *less* than one period. The single antimeridian wrap is unwrapped into a + strictly-ascending axis (``lon[0] + cumsum(diff % period)``) so every + downstream consumer stays seam-agnostic; the grid does not wrap. + + The *classification* reads concrete values, so under tracing + (``jit``/``vmap``) it cannot run: the closed/open split then defaults to + ``closed`` (today's behaviour) unless an explicit ``lon_closed`` override is + passed. The *unwrap*, however, is pure arithmetic and works under tracing — + so when the caller declares the grid open (``lon_closed=False``) a seam is + unwrapped even under a tracer, letting raw seam coords be loaded inside + ``jit``/``vmap`` without pre-unwrapping. + + Computation uses ``jnp`` and so honours the caller's precision config: under + the default ``jax_enable_x64=False`` it runs in float32, which the ``1e-3`` + classification tolerance comfortably absorbs. One edge: the seam + equal-spacing check compares float32 ``diff``s, so an ultra-fine grid at the + dateline (spacing ``dlon`` at or below ~0.02°, where float32 cancellation + near 180° approaches the tolerance) may be spuriously rejected — pass + ``lon_closed=False`` (which skips host-side validation) or pre-unwrap. """ if lon_period is None: - return - if _is_traced(lon_arr): - return - if not bool(jnp.all(jnp.diff(lon_arr) > 0)): + return lon, True + if _is_traced(lon): + # No concrete values → cannot classify or validate (caller owns + # correctness). The unwrap is traceable, so honour an explicit "open" + # declaration and unwrap a possible seam; closed/None grids need no + # monotone axis (the periodic index uses only x0 and dx) and are left + # untouched, preserving the closed byte-identity. + if lon_closed is False: + step = jnp.mod(jnp.diff(lon), lon_period) + lon = lon[0] + jnp.cumsum( + jnp.concatenate([jnp.zeros((1,), lon.dtype), step]) + ) + return lon, False + return lon, True if lon_closed is None else bool(lon_closed) + + lon_j = jnp.asarray(lon) + nlon = int(lon_j.shape[0]) + if nlon < 2: + raise ValueError( + "lon_period requires at least 2 longitude points to define the " + f"grid spacing; got nlon={nlon}." + ) + diffs = jnp.diff(lon_j) + if bool(jnp.all(diffs > 0)): + lon_out = lon # already ascending — pass through unchanged (byte-identical) + dlon = float(lon_j[1] - lon_j[0]) + elif bool(jnp.all(diffs < 0)): raise ValueError( "lon_period requires strictly ascending longitude coordinates; " "the periodic wrap arithmetic is undefined on a descending axis. " "Flip the longitude axis (and the field values) before loading." ) - nlon = int(lon_arr.shape[0]) - dlon = float(lon_arr[1] - lon_arr[0]) - implied_period = nlon * dlon - if abs(implied_period - lon_period) > 1e-3 * abs(lon_period): + else: + # Seam-crossing: unwrap the single antimeridian wrap into a monotonic + # ascending axis. Each step, taken modulo the period, is the forward arc. + step = jnp.mod(diffs, lon_period) + if not bool(jnp.allclose(step, step[0], rtol=1e-3)): + raise ValueError( + "lon_period longitudes must be equally spaced and ascending, or " + "a single antimeridian-crossing wrap of an equally-spaced axis; " + "got an axis that is neither (non-uniform spacing after " + "unwrapping the seam)." + ) + lon_out = lon_j[0] + jnp.cumsum( + jnp.concatenate([jnp.zeros((1,), lon_j.dtype), step]) + ) + dlon = float(step[0]) + + span = nlon * dlon + if span > lon_period * (1.0 + 1e-3): raise ValueError( - f"lon_period={lon_period} requires the grid to span exactly one " - f"period (nlon * dlon == lon_period), but nlon={nlon} points with " - f"spacing dlon={dlon:g} span {implied_period:g}. The wrap identifies " - "the cell past the last centre with the first, so a periodic grid " - "must not include the duplicate wrap endpoint, and its extent must " - "equal lon_period. Check the period, or drop/append a column." + f"lon_period={lon_period} but the grid spans {span:g} (nlon={nlon} " + f"points with spacing dlon={dlon:g}), overshooting one period. A " + "periodic grid must not include the duplicate wrap endpoint; drop " + "the repeated column or fix the period." ) + if lon_closed is None: + lon_closed = abs(span - lon_period) <= 1e-3 * abs(lon_period) + else: + lon_closed = bool(lon_closed) + return lon_out, lon_closed def _check_cgrid_shape(name: str, got: tuple[int, ...], expected: tuple[int, ...]) -> None: diff --git a/src/pastax/grid.py b/src/pastax/grid.py index a46e17f..ad28387 100644 --- a/src/pastax/grid.py +++ b/src/pastax/grid.py @@ -44,8 +44,19 @@ class Grid(eqx.Module): ``"C"`` (NEMO-convention Arakawa C-grid: U on east faces, V on north faces). lon_period: If set (e.g. ``360.0``), longitude is treated as periodic - with that period. The centre grid is assumed to span exactly one - period. + with that period. The grid then either spans exactly one period + (a global grid — see :attr:`lon_closed`) or is a seam-crossing + regional subset. + lon_closed: Whether a periodic (``lon_period`` set) grid spans exactly + one full period. ``True`` (default) is a **closed** global grid: + the cell past the last centre is identified with the first and + interpolation wraps (``(i+1) % n``). ``False`` is an **open** + seam-crossing regional slab (e.g. ``170°…-170°``): longitude is + still circular for query folding, but the grid does not wrap — it + behaves like a bounded grid with ``nlon - 1`` interior faces and + extrapolates past its two ends. Derived at load from + ``nlon * dlon`` vs ``lon_period``; irrelevant when ``lon_period`` + is ``None``. u_lat_coords, u_lon_coords: U-face coordinates (NEMO C-grid). Populated by the C-grid loaders (derived from the centre grid via :meth:`u_face_coords` or supplied explicitly); ``None`` on A-grids. @@ -62,6 +73,7 @@ class Grid(eqx.Module): ) stagger_type: Literal["A", "C"] = eqx.field(static=True, default="A") lon_period: float | None = eqx.field(static=True, default=None) + lon_closed: bool = eqx.field(static=True, default=True) u_lat_coords: Float[Array, "..."] | None = None u_lon_coords: Float[Array, "..."] | None = None v_lat_coords: Float[Array, "..."] | None = None @@ -77,10 +89,11 @@ def u_face_coords(self) -> tuple[Float[Array, "lat"], Float[Array, "lon_u"]]: \mathrm{lon}_u[i] = \tfrac{1}{2}\left(\mathrm{lon}_c[i] + \mathrm{lon}_c[i+1]\right) - When :attr:`lon_period` is set the grid is periodic, so every centre - cell has an east face — including the seam face between the last centre - and the first-plus-a-period. There are then ``nlon`` U faces, each a - half-cell east of its centre: + When :attr:`lon_period` is set **and** the grid is closed + (:attr:`lon_closed`), the grid is a global periodic grid, so every + centre cell has an east face — including the seam face between the last + centre and the first-plus-a-period. There are then ``nlon`` U faces, + each a half-cell east of its centre: .. math:: @@ -90,6 +103,12 @@ def u_face_coords(self) -> tuple[Float[Array, "lat"], Float[Array, "lon_u"]]: wraps across the seam with the same first-order periodic scheme used for centre fields (see :meth:`period_for`). + An **open** seam-crossing regional grid (``lon_period`` set but + ``lon_closed`` false) has no wrap and therefore no seam face: like a + bounded grid it has ``nlon - 1`` interior east faces at the centre + midpoints. The midpoints are correct because the centre longitudes are + stored unwrapped (strictly ascending) by the loaders. + Latitude is unchanged: :math:`\mathrm{lat}_u = \mathrm{lat}_c`. Raises: @@ -97,7 +116,7 @@ def u_face_coords(self) -> tuple[Float[Array, "lat"], Float[Array, "lon_u"]]: """ self._check_rectilinear("u_face_coords") lon_c = self.lon_coords - if self.lon_period is None: + if self.lon_period is None or not self.lon_closed: lon_u = 0.5 * (lon_c[:-1] + lon_c[1:]) else: dlon = lon_c[1] - lon_c[0] @@ -174,6 +193,19 @@ def period_for( """ return self.lon_period + def closed_for( + self, stagger: Literal["center", "u_face", "v_face"] + ) -> bool: + """Return whether the longitude axis a field at ``stagger`` reads is closed. + + Every stagger role inherits the centre grid's :attr:`lon_closed`: a + U/V face of an open (seam-crossing) centre grid is itself an open + interior axis, and the faces of a closed global grid are closed. Only + meaningful together with a non-``None`` :meth:`period_for`; on a bounded + grid the value is ignored. + """ + return self.lon_closed + def _check_rectilinear(self, method: str) -> None: if self.grid_type != "rectilinear": raise NotImplementedError( diff --git a/tests/test_forcing.py b/tests/test_forcing.py index a9c3475..040c4aa 100644 --- a/tests/test_forcing.py +++ b/tests/test_forcing.py @@ -684,15 +684,22 @@ def test_cgrid_descending_lon_with_period_raises(self): class TestPeriodicLonSpan: - """lon_period demands the grid span exactly one period (nlon*dlon == period).""" + """lon_period classifies the grid: span == period is a closed (global) grid + that wraps; span < period is an open seam-crossing regional grid; span > + period (a duplicate wrap endpoint) is rejected.""" - def test_regional_grid_mislabelled_periodic_raises(self): + def test_regional_subperiod_grid_accepted_as_open(self): + """A sub-period grid is a valid *open* grid (span < period), not an error.""" t = np.linspace(0.0, 3600.0, 2) lat = np.array([0.0, 1.0]) - lon = np.array([0.0, 1.0, 2.0, 3.0]) # spans 4 deg, not 360 + lon = np.array([0.0, 1.0, 2.0, 3.0]) # spans 4 deg < 360 -> open u = np.ones((2, 2, 4), dtype=np.float32) - with pytest.raises(ValueError, match="span exactly one|span exactly"): - Dataset.from_arrays({"u": u}, t=t, lat=lat, lon=lon, lon_period=360.0) + ds = Dataset.from_arrays({"u": u}, t=t, lat=lat, lon=lon, lon_period=360.0) + assert ds.grid.lon_closed is False + assert ds["u"].lon_period == 360.0 + # interpolation is continuous within the window (constant field -> 1.0) + v = ds["u"].interp(jnp.array(0.0), jnp.array(1.5), jnp.array(0.5)) + assert float(v) == pytest.approx(1.0) def test_duplicate_wrap_endpoint_raises(self): """Including the wrap endpoint (nlon+1 points, last == first + period) @@ -701,29 +708,211 @@ def test_duplicate_wrap_endpoint_raises(self): lat = np.array([0.0, 1.0]) lon = np.array([0.0, 90.0, 180.0, 270.0, 360.0]) # 5 pts incl. wrap u = np.ones((2, 2, 5), dtype=np.float32) - with pytest.raises(ValueError, match="span exactly one|duplicate wrap"): + with pytest.raises(ValueError, match="overshoot|duplicate wrap|spans"): Dataset.from_arrays({"u": u}, t=t, lat=lat, lon=lon, lon_period=360.0) - def test_correct_span_accepted(self): + def test_correct_span_accepted_as_closed(self): t = np.linspace(0.0, 3600.0, 2) lat = np.array([0.0, 1.0]) lon = np.array([0.0, 90.0, 180.0, 270.0]) # 4 * 90 == 360 u = np.ones((2, 2, 4), dtype=np.float32) ds = Dataset.from_arrays({"u": u}, t=t, lat=lat, lon=lon, lon_period=360.0) assert ds["u"].lon_period == 360.0 + assert ds.grid.lon_closed is True + + def test_seam_crossing_grid_unwrapped_and_open(self): + """A grid crossing the antimeridian (170...-170) is stored unwrapped + (ascending) and classified open.""" + t = np.linspace(0.0, 3600.0, 2) + lat = np.array([0.0, 1.0]) + lon = np.array([170.0, 175.0, 180.0, -175.0, -170.0]) # crosses seam + u = np.ones((2, 2, 5), dtype=np.float32) + ds = Dataset.from_arrays({"u": u}, t=t, lat=lat, lon=lon, lon_period=360.0) + assert ds.grid.lon_closed is False + stored = np.asarray(ds["u"].lon_coords) + assert np.all(np.diff(stored) > 0) # strictly ascending + np.testing.assert_allclose(stored, [170.0, 175.0, 180.0, 185.0, 190.0]) - def test_cgrid_mislabelled_periodic_raises(self): + def test_cgrid_subperiod_grid_accepted_as_open(self): + """A sub-period C-grid is open: U has nlon-1 interior faces, not nlon.""" t = np.linspace(0.0, 3600.0, 2) lat = np.linspace(0.0, 4.0, 5) - lon = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) # spans 6 deg, not 360 - u = np.ones((2, 5, 5), dtype=np.float32) + lon = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) # spans 5 deg < 360 -> open + u = np.ones((2, 5, 5), dtype=np.float32) # nlon - 1 = 5 interior faces v = np.ones((2, 4, 6), dtype=np.float32) - with pytest.raises(ValueError, match="span exactly one|span exactly"): - Dataset.from_arrays_cgrid( - t, lat, lon, - vectors={"current": {"u": ("uo", u), "v": ("vo", v)}}, - lon_period=360.0, + ds = Dataset.from_arrays_cgrid( + t, lat, lon, + vectors={"current": {"u": ("uo", u), "v": ("vo", v)}}, + lon_period=360.0, + ) + assert ds.grid.lon_closed is False + assert ds.grid.u_lon_coords.shape == (5,) + assert ds.grid.closed_for("u_face") is False + + +class TestOpenSeamCrossingLon: + """Open (seam-crossing) regional grids: continuity across the antimeridian + without wrapping, in any longitude convention, and safe under tracing.""" + + @staticmethod + def _ramp_field(): + # Field value equals the (unwrapped) longitude of each column, so a + # correct interpolation returns the query longitude itself. + t = np.array([0.0, 3600.0]) + lat = np.array([0.0, 1.0]) + lon = np.array([170.0, 175.0, 180.0, -175.0, -170.0]) # -> 170..190 + uw = np.array([170.0, 175.0, 180.0, 185.0, 190.0]) + u = np.tile(uw, (2, 2, 1)).astype(np.float32) + ds = Dataset.from_arrays({"u": u}, t=t, lat=lat, lon=lon, lon_period=360.0) + return ds["u"] + + def test_interp_continuous_across_seam(self): + f = self._ramp_field() + for q in (178.0, 182.0, 185.0): + v = f.interp(jnp.array(0.0), jnp.array(q), jnp.array(0.5)) + assert float(v) == pytest.approx(q, abs=1e-2) + + def test_representation_invariance(self): + """A query given as -175 equals the same query given as 185.""" + f = self._ramp_field() + a = f.interp(jnp.array(0.0), jnp.array(-175.0), jnp.array(0.5)) + b = f.interp(jnp.array(0.0), jnp.array(185.0), jnp.array(0.5)) + assert float(a) == pytest.approx(float(b)) + assert float(a) == pytest.approx(185.0, abs=1e-2) + + def test_open_extrapolates_where_closed_would_wrap(self): + """West of the window, an open grid extrapolates (does not wrap east).""" + f = self._ramp_field() # window [170, 190] + # 169 is 1 deg west of the western edge: gentle extrapolation ~169, + # NOT a wrap to the eastern end (~190). + v = f.interp(jnp.array(0.0), jnp.array(169.0), jnp.array(0.5)) + assert float(v) == pytest.approx(169.0, abs=1e-1) + + def test_nonseam_open_matches_bounded(self): + """A non-seam open subset (lon_period=360, lon_closed=False) interpolates a + varying field identically to the same grid loaded bounded (lon_period=None) + for in-window queries — the centred fold is the identity there — and folds + alternate-convention queries into the window.""" + t = np.array([0.0, 3600.0]) + lat = np.array([0.0, 1.0]) + lon = np.array([10.0, 20.0, 30.0, 40.0]) # non-seam, spans 30 deg < 360 + ramp = np.tile(lon.astype(np.float32), (2, 2, 1)) # value == longitude + open_ds = Dataset.from_arrays( + {"u": ramp}, t=t, lat=lat, lon=lon, lon_period=360.0, lon_closed=False + ) + bounded_ds = Dataset.from_arrays({"u": ramp}, t=t, lat=lat, lon=lon) + assert open_ds.grid.lon_closed is False + for q in (12.5, 25.0, 33.7): + o = float(open_ds["u"].interp(jnp.array(0.0), jnp.array(q), jnp.array(0.5))) + b = float(bounded_ds["u"].interp(jnp.array(0.0), jnp.array(q), jnp.array(0.5))) + assert o == pytest.approx(q, abs=1e-3) # ramp: interpolated value == longitude + # open == bounded for in-window queries (up to float32 fold rounding, + # ~1e-5; a wrong cell would differ by a full grid step ~10 deg) + assert o == pytest.approx(b, abs=1e-4) + # a query in another convention (25 - 360) folds back into the window; + # the bounded grid would instead extrapolate far outside it. + alt = float(open_ds["u"].interp(jnp.array(0.0), jnp.array(25.0 - 360.0), jnp.array(0.5))) + assert alt == pytest.approx(25.0, abs=1e-3) + + def test_jit_vmap_grad_safe(self): + f = self._ramp_field() + + def fn(q): + return f.interp(jnp.array(0.0), q, jnp.array(0.5)) + # jit + assert float(jax.jit(fn)(jnp.array(182.0))) == pytest.approx(182.0, abs=1e-2) + # vmap over queries straddling the seam (incl. the -175 convention) + out = jax.vmap(fn)(jnp.array([178.0, 182.0, -175.0])) + np.testing.assert_allclose(np.asarray(out), [178.0, 182.0, 185.0], atol=1e-2) + # grad is finite and ~1 (ramp slope), including across the seam at 180 + for q in (182.0, 180.0): + g = jax.grad(fn)(jnp.array(q)) + assert np.isfinite(float(g)) + assert float(g) == pytest.approx(1.0, abs=1e-2) + + def test_neighborhood_open_clamps_no_wrap(self): + f = self._ramp_field() + # centred on -175 (=185): the three columns are 180, 185, 190 + nb = f.neighborhood( + jnp.array(0.0), jnp.array(-175.0), jnp.array(0.5), + t_window=0, lat_window=0, lon_window=1, + ) + np.testing.assert_allclose(np.asarray(nb[0, 0]), [180.0, 185.0, 190.0]) + # at the western edge it clamps (no wrap to the eastern end) + nb_edge = f.neighborhood( + jnp.array(0.0), jnp.array(170.0), jnp.array(0.5), + t_window=0, lat_window=0, lon_window=1, + ) + np.testing.assert_allclose(np.asarray(nb_edge[0, 0]), [170.0, 175.0, 180.0]) + + def test_lon_closed_override_under_jit(self): + """Under tracing the axis can't be classified; an explicit lon_closed + lets an open grid be built inside jit.""" + t = np.array([0.0, 3600.0]) + lat = np.array([0.0, 1.0]) + lon = jnp.array([170.0, 175.0, 180.0, 185.0, 190.0]) # already unwrapped + uw = np.array([170.0, 175.0, 180.0, 185.0, 190.0]) + u = jnp.asarray(np.tile(uw, (2, 2, 1)).astype(np.float32)) + + def build_and_interp(u, lon, q): + ds = Dataset.from_arrays( + {"u": u}, t=t, lat=lat, lon=lon, + lon_period=360.0, lon_closed=False, + ) + return ds["u"].interp(jnp.array(0.0), q, jnp.array(0.5)) + + out = jax.jit(build_and_interp)(u, lon, jnp.array(182.0)) + assert float(out) == pytest.approx(182.0, abs=1e-2) + + def test_raw_seam_autounwrap_under_jit(self): + """With lon_closed=False the loader unwraps a seam even under tracing, so + RAW seam coords can be loaded inside jit without pre-unwrapping.""" + t = np.array([0.0, 3600.0]) + lat = np.array([0.0, 1.0]) + lon_raw = jnp.array([170.0, 175.0, 180.0, -175.0, -170.0]) # crosses seam + uw = np.array([170.0, 175.0, 180.0, 185.0, 190.0]) + u = jnp.asarray(np.tile(uw, (2, 2, 1)).astype(np.float32)) + + def build(u, lon, q): + ds = Dataset.from_arrays( + {"u": u}, t=t, lat=lat, lon=lon, + lon_period=360.0, lon_closed=False, ) + return ds["u"].interp(jnp.array(0.0), q, jnp.array(0.5)), ds["u"].lon_coords + + val, coords = jax.jit(build)(u, lon_raw, jnp.array(-175.0)) + assert float(val) == pytest.approx(185.0, abs=1e-2) # -175 folds to 185 + coords = np.asarray(coords) + assert np.all(np.diff(coords) > 0) # stored axis unwrapped to ascending + np.testing.assert_allclose(coords, [170.0, 175.0, 180.0, 185.0, 190.0], atol=1e-3) + + # same under vmap over queries in both conventions + out = jax.vmap(lambda q: jax.jit(build)(u, lon_raw, q)[0])( + jnp.array([178.0, 182.0, -175.0]) + ) + np.testing.assert_allclose(np.asarray(out), [178.0, 182.0, 185.0], atol=1e-2) + + +class TestClassifyAndUnwrapLon: + """Eager classification/unwrap of the longitude axis (host-side).""" + + def test_ascending_coords_stored_byte_identical(self): + """A plain ascending grid is passed through unchanged (no cumsum drift): + the stored coords equal the input exactly.""" + from pastax.forcing import _classify_and_unwrap_lon + + lon = np.array([0.3, 0.4, 0.5, 0.6], dtype=np.float32) # non-exact spacing + lon_out, closed = _classify_and_unwrap_lon(lon, 360.0) + assert closed is False # spans 0.4 deg < 360 -> open + assert lon_out is lon # same object, not a reconstructed axis + + def test_seam_input_unwrapped_and_open(self): + from pastax.forcing import _classify_and_unwrap_lon + + lon = np.array([170.0, 175.0, 180.0, -175.0, -170.0]) + lon_out, closed = _classify_and_unwrap_lon(lon, 360.0) + assert closed is False + np.testing.assert_allclose(np.asarray(lon_out), [170, 175, 180, 185, 190]) class TestLoadersAreTracerSafe: diff --git a/tests/test_grid.py b/tests/test_grid.py index e45cd47..5a3e6eb 100644 --- a/tests/test_grid.py +++ b/tests/test_grid.py @@ -61,6 +61,41 @@ def test_u_face_coords_periodic_includes_seam_face(): assert jnp.allclose(jnp.diff(lon_u), 90.0) +def test_u_face_coords_open_uses_interior_midpoints(): + """An open (seam-crossing) grid has no wrap and no seam face: U lives on the + nlon-1 interior midpoints, like a bounded grid, computed on the unwrapped + (ascending) centres.""" + t = jnp.asarray([0.0, 3600.0]) + lat = jnp.asarray([0.0, 1.0]) + lon = jnp.asarray([170.0, 175.0, 180.0, 185.0, 190.0]) # unwrapped seam grid + g = Grid( + t_coords=t, lat_coords=lat, lon_coords=lon, + stagger_type="C", lon_period=360.0, lon_closed=False, + ) + lat_u, lon_u = g.u_face_coords() + assert jnp.allclose(lat_u, lat) + assert lon_u.shape == (lon.shape[0] - 1,) # nlon - 1, no seam face + assert jnp.allclose(lon_u, jnp.asarray([172.5, 177.5, 182.5, 187.5])) + + +def test_closed_for_matches_lon_closed_for_all_staggers(): + t = jnp.asarray([0.0, 3600.0]) + lat = jnp.asarray([0.0, 1.0]) + lon = jnp.asarray([170.0, 175.0, 180.0, 185.0, 190.0]) + g_open = Grid( + t_coords=t, lat_coords=lat, lon_coords=lon, + stagger_type="C", lon_period=360.0, lon_closed=False, + ) + for stagger in ("center", "u_face", "v_face"): + assert g_open.closed_for(stagger) is False + g_closed = Grid( + t_coords=t, lat_coords=lat, lon_coords=jnp.asarray([0.0, 90.0, 180.0, 270.0]), + stagger_type="C", lon_period=360.0, + ) + for stagger in ("center", "u_face", "v_face"): + assert g_closed.closed_for(stagger) is True + + def test_period_for_returns_period_for_all_staggers_when_periodic(): t = jnp.asarray([0.0, 3600.0]) lat = jnp.asarray([0.0, 1.0])