pastax.grid

Grid metadata for Arakawa A- and C-grid forcing layouts.

The Grid object describes the centre (tracer) coordinates of a forcing dataset together with its stagger type, and — for C-grids — the staggered U/V-face coordinates. It is the single owner of coordinates: it is held on pastax.Dataset and every pastax.Field references it, reading the coordinates appropriate to its stagger role via Grid.coords_for().

Each Field therefore carries no coordinates of its own; interpolation reads them from the shared Grid. This keeps the Field.interp path (bilinear on equally-spaced coords) valid for both A- and C-grid fields, since shifting an equally-spaced grid by half a cell preserves equal spacing.

Curvilinear (2-D) coordinate arrays are accepted structurally for forward compatibility but not yet supported by interpolation.

Grid¶

class pastax.grid.Grid(t_coords, lat_coords, lon_coords, grid_type='rectilinear', stagger_type='A', lon_period=None, lon_closed=True, u_lat_coords=None, u_lon_coords=None, v_lat_coords=None, v_lon_coords=None)

Bases: Module

Grid metadata: centre coordinates plus stagger and topology type.

Parameters

• t_coords (Float[jaxlib._jax.Array, 'time'])

• lat_coords (Float[jaxlib._jax.Array, '...'])

• lon_coords (Float[jaxlib._jax.Array, '...'])

• grid_type (Literal['rectilinear', 'curvilinear'])

• stagger_type (Literal['A', 'C'])

• lon_period (float | None)

• lon_closed (bool)

• u_lat_coords (Float[jaxlib._jax.Array, '...'] | None)

• u_lon_coords (Float[jaxlib._jax.Array, '...'] | None)

• v_lat_coords (Float[jaxlib._jax.Array, '...'] | None)

• v_lon_coords (Float[jaxlib._jax.Array, '...'] | None)

t_coords

1-D time coordinates in seconds, equally spaced.

Type

jaxtyping.Float[jaxlib._jax.Array, ‘time’]

lat_coords

Latitude of cell centres. 1-D for rectilinear grids, 2-D (lat, lon) for curvilinear grids.

Type

jaxtyping.Float[jaxlib._jax.Array, ‘…’]

lon_coords

Longitude of cell centres. 1-D for rectilinear grids, 2-D (lat, lon) for curvilinear grids.

Type

jaxtyping.Float[jaxlib._jax.Array, ‘…’]

grid_type

"rectilinear" (default) or "curvilinear". Curvilinear grids are accepted structurally but not yet supported by Field.interp.

Type

Literal[‘rectilinear’, ‘curvilinear’]

stagger_type

"A" (default — all variables at cell centres) or "C" (NEMO-convention Arakawa C-grid: U on east faces, V on north faces).

Type

Literal[‘A’, ‘C’]

lon_period

If set (e.g. 360.0), longitude is treated as periodic with that period. The grid then either spans exactly one period (a global grid — see lon_closed) or is a seam-crossing regional subset.

Type

float | None

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.

Type

bool

u_lat_coords, u_lon_coords

U-face coordinates (NEMO C-grid). Populated by the C-grid loaders (derived from the centre grid via u_face_coords() or supplied explicitly); None on A-grids.

v_lat_coords, v_lon_coords

V-face coordinates (NEMO C-grid). Populated by the C-grid loaders (derived via v_face_coords() or supplied explicitly); None on A-grids.

t_coords: Float[jaxlib._jax.Array, 'time']

lat_coords: Float[jaxlib._jax.Array, '...']

lon_coords: Float[jaxlib._jax.Array, '...']

grid_type: Literal['rectilinear', 'curvilinear'] = 'rectilinear'

stagger_type: Literal['A', 'C'] = 'A'

lon_period: float | None = None

lon_closed: bool = True

u_lat_coords: Float[jaxlib._jax.Array, '...'] | None = None

u_lon_coords: Float[jaxlib._jax.Array, '...'] | None = None

v_lat_coords: Float[jaxlib._jax.Array, '...'] | None = None

v_lon_coords: Float[jaxlib._jax.Array, '...'] | None = None

u_face_coords()

Return (lat_u, lon_u) — coordinates of U-face centres (NEMO C-grid).

For a bounded centre grid of size nlon in longitude, U lives on the nlon - 1 interior east faces between adjacent cells:

$$\mathrm{lon}_u[i] = \tfrac{1}{2}\left(\mathrm{lon}_c[i] + \mathrm{lon}_c[i+1]\right)$$

(1)

When lon_period is set and the grid is closed (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:

$$\mathrm{lon}_u[i] = \mathrm{lon}_c[i] + \tfrac{1}{2}\,\Delta\lambda$$

(2)

This nlon-face axis spans exactly one period, so U interpolation wraps across the seam with the same first-order periodic scheme used for centre fields (see 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: $\mathrm{lat}_u = \mathrm{lat}_c$.

Raises

NotImplementedError – For curvilinear grids.

Return type

tuple[Float[jaxlib._jax.Array, ‘lat’], Float[jaxlib._jax.Array, ‘lon_u’]]

v_face_coords()

Return (lat_v, lon_v) — coordinates of V-face centres (NEMO C-grid).

For a centre grid of size nlat in latitude, V lives on the nlat - 1 north faces between adjacent cells:

$$\mathrm{lat}_v[j] = \tfrac{1}{2}\left(\mathrm{lat}_c[j] + \mathrm{lat}_c[j+1]\right)$$

(3)

Longitude is unchanged: $\mathrm{lon}_v = \mathrm{lon}_c$.

Raises

NotImplementedError – For curvilinear grids.

Return type

tuple[Float[jaxlib._jax.Array, ‘lat_v’], Float[jaxlib._jax.Array, ‘lon’]]

coords_for(stagger)

Return the (lat, lon) coordinates a field at stagger reads.

"center" returns the centre (tracer) coordinates; "u_face" and "v_face" return the stored staggered coordinates (populated by the C-grid loaders). This is the single source of coordinates consulted by pastax.Field.interp().

Raises

ValueError – For "u_face" / "v_face" on a grid that carries no staggered coordinates, or for an unknown stagger.

Parameters

stagger (Literal['center', 'u_face', 'v_face'])

Return type

tuple[Float[jaxlib._jax.Array, ‘…’], Float[jaxlib._jax.Array, ‘…’]]

period_for(stagger)

Return the longitude period a field at stagger should use.

Every stagger role inherits lon_period: centre longitudes span one period by construction, V-face longitudes equal the centre longitudes, and periodic U faces are built seam-inclusive (nlon faces spanning one period; see u_face_coords()). When lon_period is None all roles get None (bounded grid, no wrapping). Latitude is never periodic, so this concerns only the longitude axis.

Parameters

stagger (Literal['center', 'u_face', 'v_face'])

Return type

float | None

closed_for(stagger)

Return whether the longitude axis a field at stagger reads is closed.

Every stagger role inherits the centre grid’s 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 period_for(); on a bounded grid the value is ignored.

Parameters

stagger (Literal['center', 'u_face', 'v_face'])

Return type

bool