"""Coordinate conversion utilities for TRACMASS-JAX.
Provides :func:`xy_to_lonlat` and :func:`lonlat_to_xy` for converting between
fractional TRACMASS grid positions and geographic lon/lat coordinates. Both
functions are grid-agnostic: any object that exposes ``.lon`` and ``.lat``
arrays (and optionally ``.kmt``) can be passed as the *grid* argument.
Both use the TRACMASS position convention: cell ``i`` (0-based, ``ia - 1``)
spans ``i <= x < i + 1`` and its centre -- where the grid's ``lon``/``lat``
are defined -- sits at ``x = i + 0.5``; likewise for *y*. The
coordinate arrays may be 1-D (``lon`` ``(imt,)``, ``lat`` ``(jmt,)``) or 2-D
in either axis order — :func:`_orient_lonlat` normalises them to TRACMASS
``(imt, jmt)`` and transposes a ``(jmt, imt)`` layout with a warning.
"""
from __future__ import annotations
import warnings
import numpy as np
from scipy.spatial import cKDTree
def _orient_lonlat(lon_arr, lat_arr, kmt_arr):
"""Return ``lon``/``lat``/``kmt`` as TRACMASS ``(imt, jmt)`` arrays.
Accepts 1-D or 2-D ``lon``/``lat`` in either axis order and normalises
them to the layout the interpolators expect: first axis zonal (*x*),
second axis meridional (*y*).
- **1-D ``lon`` (imt,) and 1-D ``lat`` (jmt,)** — a regular lon/lat
grid; broadcast to ``(imt, jmt)`` with ``indexing="ij"``.
- **2-D ``lon``/``lat``** — kept as-is when ``lon`` already varies
along axis 0 and ``lat`` along axis 1; transposed (with a warning)
when they look like ``(jmt, imt)``, e.g. the output of
``np.meshgrid(lon_1d, lat_1d)`` (default ``indexing="xy"``).
``kmt`` is transposed to match whenever its shape lines up with the
transposed ``lon``/``lat``.
"""
lon_arr = np.asarray(lon_arr, dtype=float)
lat_arr = np.asarray(lat_arr, dtype=float)
if kmt_arr is not None:
kmt_arr = np.asarray(kmt_arr)
if lon_arr.ndim == 1 and lat_arr.ndim == 1:
imt, jmt = lon_arr.size, lat_arr.size
lon_arr, lat_arr = np.meshgrid(lon_arr, lat_arr, indexing="ij")
if (
kmt_arr is not None
and kmt_arr.ndim == 2
and kmt_arr.shape == (jmt, imt)
and imt != jmt
):
kmt_arr = kmt_arr.T
return lon_arr, lat_arr, kmt_arr
if lon_arr.ndim != 2 or lat_arr.ndim != 2 or lon_arr.shape != lat_arr.shape:
raise ValueError(
"lon/lat must be both 1-D (lon (imt,), lat (jmt,)) or both 2-D "
f"with matching shape; got {lon_arr.shape} and {lat_arr.shape}"
)
# Proper layout: lon varies down axis 0 (zonal), lat across axis 1
# (meridional). Median |diff| is ~grid spacing along the varying axis
# and ~0 along the constant one, and shrugs off the ±180 seam jump.
def _spread(a, axis):
return np.nanmedian(np.abs(np.diff(a, axis=axis))) if a.shape[axis] > 1 else 0.0
proper = _spread(lon_arr, 0) + _spread(lat_arr, 1)
swapped = _spread(lon_arr, 1) + _spread(lat_arr, 0)
if swapped > proper:
warnings.warn(
"lon/lat arrays look like (jmt, imt); transposing to TRACMASS "
"(imt, jmt). Pass arrays with the zonal axis first (e.g. "
"np.meshgrid(lon, lat, indexing='ij')) to silence this.",
stacklevel=3,
)
lon_arr, lat_arr = lon_arr.T, lat_arr.T
if kmt_arr is not None and kmt_arr.ndim == 2:
kmt_arr = kmt_arr.T
elif (
kmt_arr is not None
and kmt_arr.ndim == 2
and kmt_arr.shape == lon_arr.shape[::-1]
and lon_arr.shape[0] != lon_arr.shape[1]
):
kmt_arr = kmt_arr.T
return lon_arr, lat_arr, kmt_arr
def _unpack_grid(grid):
"""Return ``(lon_arr, lat_arr, kmt_arr_or_None)`` from a *grid* argument.
Accepted forms
--------------
- Object with ``.lon`` and ``.lat`` attributes (1-D or 2-D). ``.kmt``
is used when present.
- 2-tuple ``(lon_arr, lat_arr)`` — kmt is ``None`` (all cells treated
as ocean).
- 3-tuple ``(lon_arr, lat_arr, kmt_arr)``.
``lon``/``lat`` may be 1-D (``lon`` ``(imt,)``, ``lat`` ``(jmt,)``) or
2-D in either axis order; :func:`_orient_lonlat` normalises them to
TRACMASS ``(imt, jmt)``.
"""
if isinstance(grid, tuple):
if len(grid) == 2:
lon_arr, lat_arr = grid
kmt_arr = None
elif len(grid) == 3:
lon_arr, lat_arr, kmt_arr = grid
else:
raise ValueError("grid tuple must have 2 or 3 elements")
return _orient_lonlat(lon_arr, lat_arr, kmt_arr)
if not (hasattr(grid, "lon") and hasattr(grid, "lat")):
raise TypeError(
"grid must be a tuple (lon, lat) or (lon, lat, kmt), or an object "
"with .lon and .lat attributes"
)
kmt_arr = grid.kmt if hasattr(grid, "kmt") else None
return _orient_lonlat(grid.lon, grid.lat, kmt_arr)
def _to_xyz(lon_deg, lat_deg):
"""Unit-sphere Cartesian ``(..., 3)`` array from lon/lat in degrees."""
lo = np.radians(lon_deg)
la = np.radians(lat_deg)
return np.stack(
[np.cos(la) * np.cos(lo), np.cos(la) * np.sin(lo), np.sin(la)], axis=-1
)
def _is_zonally_periodic(xyz):
"""True when the last grid column wraps around onto the first.
Compares the median chord between column ``imt - 1`` and column ``0``
with the median chord between neighbouring columns: a global grid
closes its seam with an ordinary-sized step (or less, for grids with
duplicated halo columns), a regional one does not.
"""
if xyz.shape[0] < 3:
return False
step = np.nanmedian(np.linalg.norm(np.diff(xyz, axis=0), axis=-1))
gap = np.nanmedian(np.linalg.norm(xyz[0] - xyz[-1], axis=-1))
return bool(gap < 1.5 * step)
def _xy_to_xyz(x, y, xyz, periodic):
"""Unit vectors at fractional positions *x*, *y*; see :func:`xy_to_lonlat`.
Bilinear interpolation between the four surrounding cell centres, done
in Cartesian space on the unit sphere. Centre ``i`` sits at
``x = i + 0.5``. With *periodic*, *x* wraps across the zonal seam;
otherwise, and always in *y*, positions in the outer half cell (and
beyond) are linearly extrapolated from the two edge centres.
"""
imt, jmt = xyz.shape[:2]
u = np.asarray(x, dtype=float) - 0.5
v = np.asarray(y, dtype=float) - 0.5
if periodic:
i0 = np.floor(u).astype(int)
fx = u - i0
i0 = np.mod(i0, imt)
i1 = np.mod(i0 + 1, imt)
else:
i0 = np.clip(np.floor(u).astype(int), 0, max(imt - 2, 0))
fx = u - i0
i1 = np.minimum(i0 + 1, imt - 1)
j0 = np.clip(np.floor(v).astype(int), 0, max(jmt - 2, 0))
fy = v - j0
j1 = np.minimum(j0 + 1, jmt - 1)
fx = fx[..., None]
fy = fy[..., None]
c = (
(1.0 - fx) * (1.0 - fy) * xyz[i0, j0]
+ fx * (1.0 - fy) * xyz[i1, j0]
+ (1.0 - fx) * fy * xyz[i0, j1]
+ fx * fy * xyz[i1, j1]
)
r = np.linalg.norm(c, axis=-1, keepdims=True)
return c / np.where(r == 0, 1.0, r)
[docs]
def xy_to_lonlat(
x,
y,
grid,
periodic: bool | None = None,
) -> tuple[np.ndarray, np.ndarray]:
"""Convert fractional TRACMASS (x, y) positions to longitude/latitude.
Bilinearly interpolates the grid's cell-centre lon/lat at the
fractional particle positions. Cell ``i`` spans ``i <= x < i + 1``
with its centre at ``x = i + 0.5`` (likewise *y*), so a particle at a
cell centre maps exactly to that cell's ``lon``/``lat`` and one on a
cell wall maps half-way between the two neighbouring centres.
Interpolation is performed in Cartesian (x, y, z) space on the unit
sphere to handle the ±180° longitude seam correctly on curvilinear
grids.
Parameters
----------
x, y : array-like
Fractional 0-based grid positions as stored in the trajectory Parquet
file (columns ``"x"`` and ``"y"``). ``floor(x)`` is the cell
(matching ``"ia" - 1`` / ``"ja" - 1``).
grid : object or tuple
Provides the ``(imt, jmt)`` arrays of grid-cell centres in degrees:
- **Object with ``.lon`` and ``.lat`` attributes** — e.g. a
``GridFields`` extended with coordinate arrays, or a
:func:`types.SimpleNamespace`. A ``.kmt`` attribute is not required
for this function and is ignored if present.
- **2-tuple** ``(lon_arr, lat_arr)`` — lon/lat arrays directly.
- **3-tuple** ``(lon_arr, lat_arr, kmt_arr)`` — kmt is ignored here.
``lon_arr``/``lat_arr`` may be 1-D (``lon`` ``(imt,)``, ``lat``
``(jmt,)`` — a regular lon/lat grid) or 2-D. 2-D arrays are used
as ``(imt, jmt)`` (first axis zonal *x*, second meridional *y*);
a ``(jmt, imt)`` layout — e.g. ``np.meshgrid(lon, lat)`` with the
default ``indexing="xy"`` — is detected and transposed, with a
:class:`warnings.warn`.
periodic : bool or None, optional
Whether *x* wraps around (a global grid): positions between the
last and first cell centres then interpolate across the seam.
``None`` (default) detects it from the grid. Otherwise positions
in the outermost half cell are linearly extrapolated.
Returns
-------
lon, lat : np.ndarray
Longitude and latitude in degrees, same shape as *x* and *y*.
Examples
--------
>>> import pandas as pd, types
>>> from tracmass.trajs import xy_to_lonlat
>>> grid = types.SimpleNamespace(lon=lon_arr, lat=lat_arr)
>>> df = pd.read_parquet("traj.parquet", columns=["x", "y"])
>>> lon, lat = xy_to_lonlat(df["x"], df["y"], grid)
"""
lon_arr, lat_arr, _ = _unpack_grid(grid)
xyz = _to_xyz(lon_arr, lat_arr)
if periodic is None:
periodic = _is_zonally_periodic(xyz)
c = _xy_to_xyz(x, y, xyz, periodic)
lon_out = np.degrees(np.arctan2(c[..., 1], c[..., 0]))
lat_out = np.degrees(np.arcsin(np.clip(c[..., 2], -1.0, 1.0)))
return lon_out, lat_out
[docs]
def lonlat_to_xy(
lon,
lat,
grid,
periodic: bool | None = None,
) -> tuple[np.ndarray, np.ndarray]:
"""Convert geographic coordinates to fractional TRACMASS (x, y) grid positions.
The exact inverse of :func:`xy_to_lonlat`: the returned position maps
back to *lon*/*lat*, so a point off-centre within a cell gets the
matching fractional offset rather than being snapped to the centre.
Finds the nearest grid-cell centre by great-circle distance (a
:class:`scipy.spatial.cKDTree` on unit-sphere coordinates) as a
starting guess, then refines it with Newton iterations on the
bilinear map in the local tangent plane.
When a ``.kmt`` array is available on *grid*, the search is restricted
to ocean cells (``kmt > 0``), and a point that falls in a land cell
(or outside the grid) is placed at the centre of the nearest ocean
cell instead, so seeded particles always start in the ocean.
The resulting ``(x, y)`` values can be passed directly to
:func:`tracmass.seed.expand_sub` or used to initialise a
:class:`~tracmass.ParticleState`, with
``ia = floor(x) + 1``, ``ja = floor(y) + 1`` (1-based cell indices).
Parameters
----------
lon, lat : array-like
Longitude and latitude in degrees.
grid : object or tuple
Provides the ``(imt, jmt)`` arrays of grid-cell centres in degrees,
and optionally a land mask:
- **Object with ``.lon`` and ``.lat`` attributes** — e.g. a
``GridFields`` extended with coordinate arrays, or a
:func:`types.SimpleNamespace`. If the object also has a ``.kmt``
attribute, only ocean cells (``kmt > 0``) are included in the
spatial index.
- **2-tuple** ``(lon_arr, lat_arr)`` — all cells are included.
- **3-tuple** ``(lon_arr, lat_arr, kmt_arr)`` — ocean-only index.
As for :func:`xy_to_lonlat`, ``lon``/``lat`` may be 1-D or 2-D in
either axis order; a ``(jmt, imt)`` 2-D layout is transposed to
``(imt, jmt)`` (along with ``kmt``) with a warning.
periodic : bool or None, optional
Whether *x* wraps around; see :func:`xy_to_lonlat`. With
*periodic*, the returned *x* is wrapped into ``[0, imt)``.
Returns
-------
x, y : np.ndarray
Fractional 0-based TRACMASS grid positions, same shape as *lon*/*lat*.
Examples
--------
>>> import numpy as np, types
>>> from tracmass.trajs import lonlat_to_xy
>>> grid = types.SimpleNamespace(lon=lon_arr, lat=lat_arr, kmt=kmt_arr)
>>> x, y = lonlat_to_xy(np.array([-30.0, 10.0]), np.array([45.0, 55.0]), grid)
"""
lon_arr, lat_arr, kmt_arr = _unpack_grid(grid)
lon = np.asarray(lon, dtype=float)
lat = np.asarray(lat, dtype=float)
scalar = lon.ndim == 0
lon = np.atleast_1d(lon)
lat = np.atleast_1d(lat)
original_shape = lon.shape
lon = lon.ravel()
lat = lat.ravel()
imt, jmt = lon_arr.shape
xyz = _to_xyz(lon_arr, lat_arr)
if periodic is None:
periodic = _is_zonally_periodic(xyz)
if kmt_arr is not None:
ii, jj = np.where(kmt_arr > 0)
else:
ii, jj = np.mgrid[0:imt, 0:jmt]
ii = ii.ravel()
jj = jj.ravel()
# Chord length on the unit sphere is monotonic in great-circle
# distance, so a Euclidean KD-tree query finds the nearest centre.
target = _to_xyz(lon, lat)
_, nearest = cKDTree(xyz[ii, jj]).query(target)
xc = ii[nearest] + 0.5
yc = jj[nearest] + 0.5
# Newton iterations on the bilinear map, with the residual measured in
# the target's local east/north tangent plane (unit-sphere radians).
lo, la = np.radians(lon), np.radians(lat)
east = np.stack([-np.sin(lo), np.cos(lo), np.zeros_like(lo)], axis=-1)
north = np.stack(
[-np.sin(la) * np.cos(lo), -np.sin(la) * np.sin(lo), np.cos(la)], axis=-1
)
def _resid(xx, yy):
d = _xy_to_xyz(xx, yy, xyz, periodic) - target
return np.stack([(d * east).sum(-1), (d * north).sum(-1)], axis=-1)
x, y = xc.copy(), yc.copy()
h = 1e-4
for _ in range(8):
r = _resid(x, y)
jac = np.stack(
[(_resid(x + h, y) - r) / h, (_resid(x, y + h) - r) / h], axis=-1
)
det = jac[:, 0, 0] * jac[:, 1, 1] - jac[:, 0, 1] * jac[:, 1, 0]
ok = np.abs(det) > 1e-30
det = np.where(ok, det, 1.0)
dx = np.where(ok, (jac[:, 1, 1] * r[:, 0] - jac[:, 0, 1] * r[:, 1]) / det, 0)
dy = np.where(ok, (jac[:, 0, 0] * r[:, 1] - jac[:, 1, 0] * r[:, 0]) / det, 0)
# Cap the step so a poor starting guess cannot overshoot wildly.
x -= np.clip(dx, -1.0, 1.0)
y -= np.clip(dy, -1.0, 1.0)
if np.all(np.abs(dx) < 1e-9) and np.all(np.abs(dy) < 1e-9):
break
if periodic:
x = np.mod(x, imt)
# Keep the refined position only where it converged onto the target
# inside the grid (and in an ocean cell, when kmt is known); otherwise
# fall back to the nearest (ocean) cell centre.
ci = np.floor(x).astype(int)
cj = np.floor(y).astype(int)
good = (
np.isfinite(x) & np.isfinite(y)
& (ci >= 0) & (ci < imt) & (cj >= 0) & (cj < jmt)
& (np.abs(_resid(x, y)).max(axis=-1) < 1e-7)
)
if kmt_arr is not None:
good[good] = kmt_arr[ci[good], cj[good]] > 0
x = np.where(good, x, xc).reshape(original_shape)
y = np.where(good, y, yc).reshape(original_shape)
if scalar:
return float(x.item()), float(y.item())
return x, y
