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, 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)

• 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 centre grid is assumed to span exactly one period.

Type

float | 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 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

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 centre grid of size nlon in longitude, U lives on the nlon - 1 east faces between adjacent cells:

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

(1)

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)$$

(2)

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.

Centre (tracer) fields inherit lon_period. U/V faces always get None: their coordinate arrays no longer span a full period, so periodic wrapping would be ill-defined at first order.

Parameters

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

Return type

float | None