"""Land mask and distance to land on a seagrid model's T-cells.
Shared by the per-model ``landmask`` / ``distance_to_land`` functions in
:mod:`tracmass.seagrid.glorys`, :mod:`tracmass.seagrid.ecco2` and
:mod:`tracmass.seagrid.scrippsgog`. Each model passes its
``grid_definition()`` dataset; everything else is grid-independent.
"""
import warnings
import numpy as np
import xarray as xr
from scipy import ndimage
from scipy.spatial import cKDTree
EARTH_RADIUS_KM = 6371.0
def _kmt_2d(grid: xr.Dataset) -> xr.DataArray:
"""``kmt`` on ``(y, x)``, without the size-1 ``t`` of some grid files."""
kmt = grid["kmt"]
if "t" in kmt.dims:
kmt = kmt.isel(t=0, drop=True)
return kmt.transpose("y", "x")
[docs]
def landmask(grid: xr.Dataset, k: int = 1) -> xr.DataArray:
"""Land mask of a ``grid_definition()`` dataset.
Parameters
----------
grid : xr.Dataset
A seagrid ``grid_definition()`` with ``kmt`` (the number of ocean
levels per column; 0 on land) and ``lon`` / ``lat`` coordinates.
k : int, optional
1-based model level. Default ``1`` (the surface): land is
``kmt == 0``. At deeper levels the sea floor counts as land too:
``kmt < k``.
Returns
-------
xr.DataArray
bool, ``(y, x)``, ``True`` on land, with the grid's ``lon`` /
``lat`` coordinates. NumPy-backed.
Raises
------
ValueError
If *k* is below 1.
"""
if k < 1:
raise ValueError(f"k is a 1-based level, got {k}")
kmt = _kmt_2d(grid)
mask = (kmt.astype("int32") < k).compute()
mask = mask.assign_coords(
{c: grid[c].compute() for c in ("lon", "lat") if c in grid.coords}
)
mask.name = "landmask"
mask.attrs = dict(long_name="land mask", level=int(k), description=(
"True where the column has fewer than `level` ocean levels (kmt < level)"
))
return mask
def _unit_xyz(lon: np.ndarray, lat: np.ndarray) -> np.ndarray:
"""Points on the unit sphere, shape ``(n, 3)``."""
lon, lat = np.deg2rad(lon), np.deg2rad(lat)
coslat = np.cos(lat)
return np.column_stack([coslat * np.cos(lon), coslat * np.sin(lon), np.sin(lat)])
def _lonlat_2d(mask: xr.DataArray) -> tuple[np.ndarray, np.ndarray]:
"""Broadcast the mask's 1-D or 2-D ``lon`` / ``lat`` to ``(y, x)``."""
lon, lat = xr.broadcast(mask["lon"], mask["lat"])
lon, lat = lon.transpose("y", "x"), lat.transpose("y", "x")
return np.asarray(lon, dtype="float64"), np.asarray(lat, dtype="float64")
[docs]
def distance_to_land(mask: xr.DataArray) -> xr.DataArray:
"""Great-circle distance from every cell centre to the nearest land cell.
Parameters
----------
mask : xr.DataArray
A land mask from :func:`landmask` (bool, ``(y, x)``, ``True`` on
land, with ``lon`` / ``lat`` coordinates in degrees).
Returns
-------
xr.DataArray
float32, ``(y, x)``, in km, with the mask's coordinates. 0 on
land; NaN everywhere when the grid has no land cell.
Notes
-----
Distances run between cell centres: an ocean cell next to the coast is
about one cell width from land, not half. The search uses a KD-tree on
unit-sphere coordinates, so it is exact on curvilinear grids and across
the date line. Only land cells with an ocean neighbour (8-neighbour)
go into the tree: any other land cell is farther from the ocean. On a
regional grid, land outside the domain is unknown, so distances near
the open boundaries can be too large.
"""
land = np.asarray(mask, dtype=bool)
lon, lat = _lonlat_2d(mask)
ocean = ~land
coast = land & ndimage.binary_dilation(ocean, structure=np.ones((3, 3), bool))
dist = np.zeros(land.shape, dtype="float32")
if not land.any():
warnings.warn("grid has no land cells; distance_to_land is NaN",
stacklevel=3)
dist[:] = np.nan
elif coast.any(): # else all land: 0 everywhere
tree = cKDTree(_unit_xyz(lon[coast], lat[coast]))
chord, _ = tree.query(_unit_xyz(lon[ocean], lat[ocean]), workers=-1)
dist[ocean] = 2 * EARTH_RADIUS_KM * np.arcsin(np.minimum(chord / 2, 1.0))
out = xr.DataArray(dist, dims=("y", "x"), coords=mask.coords,
name="distance_to_land")
out.attrs = dict(
units="km", long_name="distance to nearest land cell centre",
level=mask.attrs.get("level", 1),
)
return out
