"""GLORYS12 reanalysis GCM interface for TRACMASS JAX.
This module provides a ``fluxes`` callable that loads GLORYS12 ocean
reanalysis velocity fields and converts them to volume-flux arrays
suitable for use with the TRACMASS particle-tracking framework.
"""
from typing import NewType
import numpy as np
import pandas as pd
import xarray as xr
import types
from tracmass.gcm.flux_grid import subset_offset_attrs
from tracmass.seagrid import glorys as glorys_gcm
from tracmass.trajs.coords import lonlat_to_xy as _lonlat_to_xy
from tracmass.trajs.coords import xy_to_lonlat as _xy_to_lonlat
FluxDataset = NewType("FluxDataset", xr.Dataset)
settings = dict(
subset=dict(
i1=None,
i2=None,
istep=None,
j1=None,
j2=None,
jstep=None,
k1=None,
k2=None,
kstep=None,
),
include_wflux=False,
)
# x-halo trims that put velocity column and grid column i on the same T-cell;
# see the comment in :func:`fluxes`.
_GRID_XTRIM = slice(1, -1) # mesh files, x=4320
_VEL_XTRIM = slice(3, -1) # velocity Zarrs, x=4322
[docs]
def fluxes(
dtm: pd.Timestamp | str | None = None, include_wflux: bool = True
) -> FluxDataset:
"""Load GLORYS12 flux fields for a single time window.
Reads GLORYS12 reanalysis velocities and grid geometry for the
requested time, computes zonal, meridional, and vertical volume
fluxes, and returns them as an ``xr.Dataset`` compatible with
:func:`tracmass.gcm.fluxes_to_grid`.
Backward-in-time tracking is handled at the :func:`~tracmass.gcm.new_grid`
/ :func:`~tracmass.gcm.new_flux` level by passing ``dtm_end < dtm_start``; this function always returns
forward (positive) fluxes.
Parameters
----------
dtm : pd.Timestamp or str or None, optional
Target time for the flux snapshot. Default ``None`` — no ``dtm``
is passed to :func:`tracmass.seagrid.glorys.open_dataset` at all,
so its own default applies.
Returns
-------
FluxDataset
xarray Dataset containing the variables required by
``GridFields``: ``uflux``, ``vflux``, ``wflux``, ``kmt``,
``dzt``, ``dxdy``, ``timestep``, ``dtm``, ``reverse``,
``imt``, ``jmt``, ``km``, ``iperio``, ``jperio``, and
``zstot``.
"""
if dtm is None:
vels = glorys_gcm.open_dataset(mask=False)
else:
dtm = pd.to_datetime(dtm)
vels = glorys_gcm.open_dataset(dtm=dtm, mask=False)
dtm = pd.Timestamp(vels.time.item())
# Grid geometry files (mesh_hgr/mesh_zgr) have x=4320 with one halo column
# on each side; slice(1, -1) gives 4318, so TRACMASS cell i is mesh
# T-cell i+1. The velocity Zarrs (U, V, W) have x=4322: two extra columns
# on the *west* only — their lon_u/lon_v/lon_w equal the mesh
# glamu/glamv/glamt exactly at velocity column c+2 = mesh column c. So
# velocity column i+3 belongs to cell i: slice(3, -1). With that, uflux[i]
# is the east face of cell i and vflux[:, j] its north face, as the kernels
# expect. (Two slice(1, -1) trims used to put every velocity one cell
# west of its T-cell.) The halos are trimmed *before* the subset is
# applied, so i1/i2 index TRACMASS cells and velocity and geometry stay
# the same width.
# .data is used on geometry arrays to avoid xarray coordinate-alignment
# errors when mixing index coordinates from different source files.
sub = _subset_isel()
vels = vels.isel(x=_VEL_XTRIM).isel(**sub) # 4322 → 4318
# Mask land fill values only now, after the subset: masking before it
# processes every whole Zarr chunk the subset touches (~7x slower).
vels = vels.assign(
{v: glorys_gcm.mask_fill(vels[v]) for v in ("uvel", "vvel", "wvel")}
)
grd = glorys_gcm.grid_definition().isel(x=_GRID_XTRIM).isel(**sub) # 4320 → 4318
dxdy = grd["dxdy"].isel(t=0)
# Capture zarr chunk sizes from the velocity DataArray *before* arithmetic
# (multiplying by grid arrays with different chunks creates irregular dask
# chunks that no longer reflect zarr storage boundaries). The first x/y
# chunk may be smaller due to the halo trim, so take the second chunk size
# (or fall back to the first if only one chunk exists).
try:
_ux_chunks = vels.uvel.chunks[-1] # x-dim chunks (z, y, x ordering)
_uy_chunks = vels.uvel.chunks[-2] # y-dim chunks
_zarr_cx = int(_ux_chunks[1] if len(_ux_chunks) > 1 else _ux_chunks[0])
_zarr_cy = int(_uy_chunks[1] if len(_uy_chunks) > 1 else _uy_chunks[0])
except Exception:
_zarr_cx, _zarr_cy = 1, 1
_uvel = vels.uvel
_vvel = vels.vvel
_dzt = grd.dzt.data # (z, y, x)
_dyu = grd.dyu.isel(t=0).data # (y, x)
_dxv = grd.dxv.isel(t=0).data # (y, x)
dsdict = {
"uflux": _uvel * _dzt * _dyu,
"vflux": _vvel * _dzt * _dxv,
"kmt": grd["kmt"].isel(t=0),
# dzt has no t-dim in the mesh file; add a trailing t-dim so that
# xr.concat propagates it into slot index (imt, jmt, km, 2)
"dzt": grd["dzt"].expand_dims("t", axis=-1),
"dxdy": dxdy,
"timestep": vels.timestep,
"dtm": dtm,
"reverse": False,
}
if include_wflux:
# .data avoids coordinate mismatch between velocity and mesh files
dsdict["wflux"] = vels.wvel * dxdy.data
dsdict["km"], dsdict["jmt"], dsdict["imt"] = dsdict["uflux"].shape[:3]
# Zonally periodic only on the full globe; an x-subset is a closed box.
x_subset = sub["x"].start is not None or sub["x"].stop is not None
dsdict["iperio"], dsdict["jperio"] = (0 if x_subset else 1), 0
dsdict["zstot"] = dxdy * 0 + 1
ds = xr.Dataset(dsdict)
ds = ds.transpose("x", "y", ...)
ds.attrs["_zarr_chunk_x"] = _zarr_cx
ds.attrs["_zarr_chunk_y"] = _zarr_cy
ds.attrs.update(subset_offset_attrs(settings["subset"]))
return ds
def _subset_isel() -> dict[str, slice]:
"""Return ``isel`` slices for ``settings["subset"]``.
*i1*/*i2*/*j1*/*j2* are 0-based slice endpoints in TRACMASS cell space
(i.e. after the halo trim), ``None`` meaning no limit. *k1*/*k2* use
1-based TRACMASS layer indexing (``k1=1`` is the first depth layer), so
they become ``slice(k1 - 1, k2)``.
Returns
-------
dict[str, slice]
Slices keyed by the source dimension names ``x``, ``y`` and ``z``.
"""
subset = settings["subset"]
k1 = subset["k1"]
return dict(
x=slice(subset["i1"], subset["i2"], subset["istep"]),
y=slice(subset["j1"], subset["j2"], subset["jstep"]),
z=slice((k1 - 1) if k1 is not None else None, subset["k2"], subset["kstep"]),
)
def _glorys_grid(
i_slice: slice = slice(None), j_slice: slice = slice(None), kmt: bool = True
) -> types.SimpleNamespace:
"""Return a SimpleNamespace with lon, lat (and kmt) arrays for the full grid.
The x/y subset is deliberately *not* applied: recorded trajectories hold
full-grid (absolute) positions, see ``GridFields.i_offset``. *i_slice*
/ *j_slice* (in TRACMASS cells) read just a window of it.
"""
grd_sub = glorys_gcm.grid_definition().isel(x=_GRID_XTRIM)
grd_sub = grd_sub.isel(x=i_slice, y=j_slice)
# NEMO stores (y, x); transpose to (x, y) = (imt, jmt) to match tracmass
return types.SimpleNamespace(
lon=np.asarray(grd_sub["lon"].values).T,
lat=np.asarray(grd_sub["lat"].values).T,
**({"kmt": np.asarray(grd_sub["kmt"].isel(t=0).values).T} if kmt else {}),
)
[docs]
def lonlat_grid(grid=None) -> tuple[np.ndarray, np.ndarray]:
"""Cell-centre longitude / latitude of a run's grid.
Parameters
----------
grid : GridFields or None, optional
The run grid; its ``i_offset`` / ``j_offset`` / ``imt`` / ``jmt``
select the matching window of the full GLORYS12 grid (a ``subset``).
Default: the full grid.
Returns
-------
lon, lat : np.ndarray
Degrees, shape ``(imt, jmt)`` — indexed like ``grid.kmt``.
"""
if grid is None:
i_slice = j_slice = slice(None)
else:
i_slice = slice(grid.i_offset, grid.i_offset + grid.imt)
j_slice = slice(grid.j_offset, grid.j_offset + grid.jmt)
g = _glorys_grid(i_slice, j_slice, kmt=False)
return g.lon, g.lat
# Cells of margin around the positions when xy_to_lonlat reads a window of
# the grid: bilinear interpolation only needs the neighbouring centres.
_WINDOW_MARGIN = 2
[docs]
def lonlat_to_xy(
lon: np.ndarray | float,
lat: np.ndarray | float,
) -> tuple[np.ndarray, np.ndarray]:
"""Convert geographic coordinates to fractional TRACMASS (x, y) grid positions.
The exact inverse of :func:`xy_to_lonlat` (see
:func:`tracmass.trajs.lonlat_to_xy`): a point off-centre within a cell
gets the matching fractional offset rather than being snapped to the
cell centre. A point in a land cell (``kmt == 0``) is placed at the
centre of the nearest ocean cell instead.
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.
Returns
-------
x, y : np.ndarray
Fractional 0-based TRACMASS grid positions, same shape as *lon*/*lat*.
Notes
-----
The same x-halo trim as :func:`fluxes` is applied.
Positions are full-grid (absolute) indices — the same frame as saved
trajectory files — regardless of any ``subset``. Subtract
``grid.i_offset`` / ``grid.j_offset`` to get positions relative to a
subset grid (e.g. for seeding).
Examples
--------
>>> import numpy as np
>>> from tracmass.gcm import glorys as gcm
>>> lon = np.array([-30.0, 10.0])
>>> lat = np.array([45.0, 55.0])
>>> x, y = gcm.lonlat_to_xy(lon, lat)
"""
return _lonlat_to_xy(lon, lat, _glorys_grid())
[docs]
def xy_to_lonlat(
x: np.ndarray | float,
y: np.ndarray | float,
) -> tuple[np.ndarray, np.ndarray]:
"""Convert fractional TRACMASS (x, y) positions to longitude/latitude.
Applies the same x-halo trim that :func:`fluxes` uses (but not the
``subset``: *x*/*y* are full-grid indices, as saved by the trajectory
recorder), then bilinearly interpolates the curvilinear grid's lon/lat arrays at the
fractional particle positions. Works correctly for curvilinear grids
(e.g. NEMO ORCA) because interpolation is done in Cartesian (x, y, z)
space on the unit sphere to avoid longitude discontinuities at ±180°.
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 column ``"ia" - 1`` / ``"ja" - 1``); its centre, where the
grid's lon/lat are defined, is at ``x = floor(x) + 0.5``.
Returns
-------
lon, lat : np.ndarray
Longitude and latitude in degrees, same shape as *x* and *y*.
Notes
-----
The grid-definition file caches longitude **and** latitude from the NEMO
mesh_hgr file. If the file was built before the ``nav_lat`` bug-fix was
applied, delete it so it is regenerated::
from tracmass.seagrid import glorys as _g
_g.grid_definition(force=True)
Examples
--------
>>> import pandas as pd
>>> df = pd.read_parquet("trm_omi.parquet", columns=["x", "y"])
>>> lon, lat = gcm.glorys.xy_to_lonlat(df["x"], df["y"])
"""
x = np.asarray(x, dtype=float)
y = np.asarray(y, dtype=float)
if x.size == 0:
return np.empty_like(x), np.empty_like(y)
# Read only the window of the (large) global grid the positions need;
# positions near the zonal edge use the full grid so the seam is handled.
grd = glorys_gcm.grid_definition()
imt, jmt = grd.sizes["x"] - 2, grd.sizes["y"]
i0 = int(np.floor(np.nanmin(x) - 0.5)) - _WINDOW_MARGIN
i1 = int(np.floor(np.nanmax(x) - 0.5)) + _WINDOW_MARGIN + 2
if i0 <= 0 or i1 >= imt:
return _xy_to_lonlat(x, y, _glorys_grid(kmt=False))
j0 = max(int(np.floor(np.nanmin(y) - 0.5)) - _WINDOW_MARGIN, 0)
j1 = min(int(np.floor(np.nanmax(y) - 0.5)) + _WINDOW_MARGIN + 2, jmt)
window = _glorys_grid(slice(i0, i1), slice(j0, j1), kmt=False)
return _xy_to_lonlat(x - i0, y - j0, window, periodic=False)
