pastax.interpolation

Bilinear interpolation on equally-spaced rectilinear A-grids.

linear_interp_1d¶

pastax.interpolation.linear_interp_1d(values, coords, x)

Linearly interpolate a 1-D field on an equally-spaced grid.

Parameters

• values (Float[jaxlib._jax.Array, 'n']) – Field values at each grid node, shape (n,).

• coords (Float[jaxlib._jax.Array, 'n']) – Equally-spaced 1-D grid coordinates, shape (n,).

• x (Float[jaxlib._jax.Array, '']) – Query coordinate.

Returns

Interpolated scalar value at x. For x outside the grid the result is the linear extrapolation from the nearest cell.

Return type

Float[jaxlib._jax.Array, ‘’]

bilinear_interp_2d¶

pastax.interpolation.bilinear_interp_2d(values, lat_coords, lon_coords, lat, lon, lon_period=None, mask=None)

Bilinearly interpolate a 2-D field on an equally-spaced rectilinear grid.

Parameters

• values (Float[jaxlib._jax.Array, 'lat lon']) – Field values, shape (n_lat, n_lon).

• lat_coords (Float[jaxlib._jax.Array, 'lat']) – Equally-spaced latitude coordinates in degrees, shape (n_lat,).

• lon_coords (Float[jaxlib._jax.Array, 'lon']) – Equally-spaced longitude coordinates in degrees, shape (n_lon,).

• lat (Float[jaxlib._jax.Array, '']) – Query latitude in degrees.

• lon (Float[jaxlib._jax.Array, '']) – Query longitude in degrees.

• lon_period (float | None) – If given (e.g. 360.0), the longitude axis is treated as periodic with that period; the grid is assumed to span exactly one period (n_lon * dlon == lon_period) and the cell at lon_coords[-1] + dlon is identified with lon_coords[0]. None (default) reproduces the non-wrapping behaviour with linear extrapolation past the boundary.

• mask (Bool[jaxlib._jax.Array, 'lat lon'] | None) –

  Optional 2-D boolean land mask, same shape as values. True marks a land cell. Behaviour by mixed-corner count:

  • All four corners ocean → standard bilinear (bit-exact identical to the mask=None path).

  • Mixed corners → inverse-distance partial-cell weighting on the ocean corners in normalised cell coordinates ($d^2 = \alpha^2 + \beta^2 + \varepsilon$ where $\alpha$, $\beta$ are the fractional distances along the cell axes); land corners are dropped.

  • All four corners land → returns 0 (zero velocity for fully-grounded cells).

  The $\varepsilon$ floor and safe_divide() keep both forward and backward passes finite for queries on or near a corner.

Returns

Interpolated scalar value at (lat, lon).

Return type

Float[jaxlib._jax.Array, ‘’]

spatiotemporal_interp¶

pastax.interpolation.spatiotemporal_interp(values, t_coords, lat_coords, lon_coords, t, lat, lon, lon_period=None, mask=None)

Trilinearly interpolate a field in time and space on an A-grid.

Performs bilinear_interp_2d() at the two bounding time slices, then linearly blends the two results in time.

Parameters

• values (Float[jaxlib._jax.Array, 'time lat lon']) – Field values, shape (n_time, n_lat, n_lon).

• t_coords (Float[jaxlib._jax.Array, 'time']) – Equally-spaced time coordinates in seconds, shape (n_time,).

• lat_coords (Float[jaxlib._jax.Array, 'lat']) – Equally-spaced latitude coordinates in degrees.

• lon_coords (Float[jaxlib._jax.Array, 'lon']) – Equally-spaced longitude coordinates in degrees.

• t (Float[jaxlib._jax.Array, '']) – Query time in seconds.

• lat (Float[jaxlib._jax.Array, '']) – Query latitude in degrees.

• lon (Float[jaxlib._jax.Array, '']) – Query longitude in degrees.

• lon_period (float | None) – If given, treat the longitude axis as periodic with this period (see bilinear_interp_2d()).

• mask (Bool[jaxlib._jax.Array, 'lat lon'] | None) – Optional 2-D land mask shared across time (see bilinear_interp_2d()).

Returns

Interpolated scalar value at (t, lat, lon).

Return type

Float[jaxlib._jax.Array, ‘’]

bilinear_velocity_partialslip_2d¶

pastax.interpolation.bilinear_velocity_partialslip_2d(u_values, v_values, lat_coords, lon_coords, lat, lon, mask, slip_a=0.5, slip_b=0.5, lon_period=None)

Bilinear A-grid velocity interpolation with partial-slip wall correction.

Replaces the standard bilinear weights for U (tangential to a latitudinal coast) and V (tangential to a longitudinal coast) with a wall-slip-aware formula whenever an entire cell edge is land.

For a south-edge-fully-land cell with the north edge in the ocean, the naive bilinear for U is $w_l\,U_{\mathrm{along}\,N}$ where

$$U_{\mathrm{along}\,N} = (1 - w_j)\,U_{NW} + w_j\,U_{NE}$$

(1)

is the linear interpolation of U along the north (ocean) edge. The partial-slip correction replaces this with $(a + b\,w_l)\,U_{\mathrm{along}\,N}$: at the coast ($w_l = 0$) the tangential velocity is $a\,U_{\mathrm{along}\,N}$ rather than 0 (which would trap the particle). The default $a = b = 0.5$ gives a half-slip wall; $a = 1,\ b = 0$ recovers full free-slip; $a = 0,\ b = 1$ recovers the no-slip naive bilinear.

Symmetrically for north-edge-land cells (using $U_{\mathrm{along}\,S}$ and $a + b\,(1 - w_l)$) and for V across east/west land edges. Cells without a fully-land edge fall back to standard bilinear.

Parameters

• u_values (Float[jaxlib._jax.Array, 'lat lon']) – U-component values on the A-grid, shape (n_lat, n_lon).

• v_values (Float[jaxlib._jax.Array, 'lat lon']) – V-component values, same shape as u_values.

• lat_coords (Float[jaxlib._jax.Array, 'lat']) – Equally-spaced latitudes, shape (n_lat,).

• lon_coords (Float[jaxlib._jax.Array, 'lon']) – Equally-spaced longitudes, shape (n_lon,).

• lat (Float[jaxlib._jax.Array, '']) – Query latitude.

• lon (Float[jaxlib._jax.Array, '']) – Query longitude.

• mask (Bool[jaxlib._jax.Array, 'lat lon']) – Joint U/V land mask, True = land. Typically built as u_mask & v_mask so a corner is considered land only when both components are masked there.

• slip_a (float) – Slip coefficient at the wall. Default 0.5.

• slip_b (float) – Slip coefficient gradient. Default 0.5.

• lon_period (float | None) – Periodic longitude period in degrees, or None.

Returns

Tuple (u, v) of interpolated A-grid velocity components.

Return type

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

spatiotemporal_velocity_partialslip¶

pastax.interpolation.spatiotemporal_velocity_partialslip(u_values, v_values, t_coords, lat_coords, lon_coords, t, lat, lon, mask, slip_a=0.5, slip_b=0.5, lon_period=None)

Trilinear A-grid velocity interpolation with partial-slip wall correction.

Applies bilinear_velocity_partialslip_2d() at the two bounding time slabs and blends linearly in time.

Returns

Tuple (u, v) of trilinearly-interpolated A-grid velocities at (t, lat, lon).

Parameters

• u_values (Float[jaxlib._jax.Array, 'time lat lon'])

• v_values (Float[jaxlib._jax.Array, 'time lat lon'])

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

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

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

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

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

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

• mask (Bool[jaxlib._jax.Array, 'lat lon'])

• slip_a (float)

• slip_b (float)

• lon_period (float | None)

Return type

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