"""ECCO2 cube92 GCM interface for TRACMASS JAX.
This module provides a ``fluxes`` callable that loads ECCO2 cube92 daily
ocean state estimate velocity fields and converts them to volume-flux
arrays suitable for use with the TRACMASS particle-tracking framework.
Unlike GLORYS12/NEMO, ECCO2 cube92 is a plain global regular lat/lon grid
(no curvilinear coordinates, no halo columns), so :func:`fluxes` needs no
x-halo trimming and cell face widths (``dxv``, ``dyu``) come from
:func:`tracmass.seagrid.ecco2.grid_definition` rather than a mesh file.
"""
import types
from typing import NewType
import numpy as np
import pandas as pd
import xarray as xr
from tracmass.gcm.flux_grid import subset_offset_attrs
from tracmass.seagrid import ecco2 as ecco2_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)
# ECCO2 cube92 ocean currents are O(1) m/s at most; the APDRC gateway's
# declared missing_value/_FillValue (-1e23) sometimes survives undecoded
# (rather than becoming NaN), and independently of that, wvel carries its
# own invalid mask that is *not* implied by uvel/kmt (e.g. under ice
# shelves wvel is invalid at levels where uvel/vvel are valid). Any
# magnitude this large is unambiguously a fill artefact, not data.
_FILL_MAGNITUDE = 1e10
def _clean_velocity(da: xr.DataArray) -> xr.DataArray:
"""Zero out NaN and undecoded fill-sentinel values in a velocity field."""
return da.where(np.abs(da) < _FILL_MAGNITUDE, 0.0)
settings = dict(
subset=dict(
i1=None,
i2=None,
istep=None,
j1=None,
j2=None,
jstep=None,
k1=None,
k2=None,
kstep=None,
),
)
[docs]
def fluxes(
dtm: pd.Timestamp | str | None = None, include_wflux: bool | None = None
) -> FluxDataset:
"""Load ECCO2 cube92 flux fields for a single time window.
Reads ECCO2 cube92 velocities and grid geometry for the requested
time, computes zonal, meridional, and (optionally) 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.ecco2.open_dataset` at all,
so its own default applies.
include_wflux : bool or None, optional
If True, also fetch ``wvel`` and compute ``wflux``. If False, the
``wvel`` OPeNDAP dataset is not touched at all — use this with
``advection_scheme.piecewise(grid, no_wflux=True)`` for
horizontal-only advection. If ``None`` (default), taken from the
``move_vertically`` setting (default ``True``); see
:func:`tracmass.config.move_vertically`.
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 include_wflux is None:
from tracmass.config import move_vertically
include_wflux = move_vertically()
if dtm is None:
vels = ecco2_gcm.open_dataset()
else:
dtm = pd.to_datetime(dtm)
vels = ecco2_gcm.open_dataset(dtm=dtm)
dtm = pd.Timestamp(vels.time.item())
subset = settings["subset"]
# settings["subset"]["k1"]/k2 use 1-based TRACMASS layer indexing (k1=1
# means the first depth layer). Python's isel/slice uses 0-based
# exclusive-end indexing, so convert: slice(k1-1, k2).
# i1/i2/j1/j2 are left as raw 0-based slice endpoints (None = no subset).
k1 = subset["k1"]
k_start = (k1 - 1) if k1 is not None else None
k_slice = slice(k_start, subset["k2"], subset["kstep"])
x_slice = slice(subset["i1"], subset["i2"], subset["istep"])
y_slice = slice(subset["j1"], subset["j2"], subset["jstep"])
vels = vels.isel(x=x_slice, y=y_slice, z=k_slice)
grd = ecco2_gcm.grid_definition()
grd = grd.isel(x=x_slice, y=y_slice, z=k_slice)
# 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).
try:
_zarr_cx = int(vels.uvel.chunks[-1][0]) # x-dim chunk (z, y, x ordering)
_zarr_cy = int(vels.uvel.chunks[-2][0]) # y-dim chunk
except Exception:
_zarr_cx, _zarr_cy = 1, 1
# .data is used on geometry arrays to avoid xarray coordinate-alignment
# errors when mixing index coordinates from the velocity Zarr with
# those from the separately-cached grid-definition Zarr.
_dzt = grd.dzt.data # (z, y, x)
_dyu = grd.dyu.data # (y, x)
_dxv = grd.dxv.data # (y, x)
_uvel = _clean_velocity(vels.uvel)
_vvel = _clean_velocity(vels.vvel)
dsdict = {
"uflux": _uvel * _dzt * _dyu,
"vflux": _vvel * _dzt * _dxv,
"kmt": grd["kmt"],
"dzt": grd["dzt"].expand_dims("t", axis=-1),
"dxdy": grd["dxdy"],
"timestep": vels.timestep,
"dtm": dtm,
"reverse": False,
}
if include_wflux:
wvel = ecco2_gcm.open_dataarray(dtm=dtm, data_var="wvel")
wvel = wvel.isel(x=x_slice, y=y_slice, z=k_slice)
dsdict["wflux"] = _clean_velocity(wvel) * grd["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 = subset["i1"] is not None or subset["i2"] is not None
dsdict["iperio"], dsdict["jperio"] = (0 if x_subset else 1), 0
dsdict["zstot"] = grd["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(subset))
return ds
def _ecco2_grid() -> types.SimpleNamespace:
"""Return a SimpleNamespace with lon, lat, 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``.
"""
grd_sub = ecco2_gcm.grid_definition()
# ECCO2 stores (y, x); transpose to (x, y) = (imt, jmt) to match tracmass
lon2d, lat2d = np.meshgrid(grd_sub["lon"].values, grd_sub["lat"].values)
return types.SimpleNamespace(
lon=lon2d.T,
lat=lat2d.T,
kmt=np.asarray(grd_sub["kmt"].values).T,
)
[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 ECCO2 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 = _ecco2_grid()
return g.lon[i_slice, j_slice], g.lat[i_slice, j_slice]
[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
-----
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 ecco2 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, _ecco2_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.
*x*/*y* are full-grid indices (as saved by the trajectory recorder),
whatever the ``subset``. Bilinearly interpolates the grid's lon/lat
arrays at the fractional particle positions.
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*.
Examples
--------
>>> import pandas as pd
>>> df = pd.read_parquet("trm_ecco2.parquet", columns=["x", "y"])
>>> lon, lat = gcm.ecco2.xy_to_lonlat(df["x"], df["y"])
"""
return _xy_to_lonlat(x, y, _ecco2_grid())
