"""Scripps Gulf of Guinea (GoG) MITgcm interface for TRACMASS JAX.
This module provides a ``fluxes`` callable that loads the daily Scripps
GoG regional MITgcm velocities (top 200 m) and converts them to
volume-flux arrays for the TRACMASS particle-tracking framework.
The files keep MITgcm's C-grid staggering: ``UVEL[i]`` sits on the
**west** face of T-cell ``i`` and ``VVEL[j]`` on the **south** face of
T-cell row ``j`` (with one extra row for the northernmost north face).
TRACMASS reads ``uflux[i]`` / ``vflux[:, j]`` as the **east** / **north**
faces, so :func:`fluxes` shifts both by one face. The east face of the
last column is not in the files and is set to zero (it is land).
Transports use the files' own face metrics and open fractions,
``uflux = UVEL * dyG * drF * hFacW`` and ``vflux = VVEL * dxG * drF *
hFacS``. The western (11 W) and southern (2 N) open boundaries carry
inflow in the files, but TRACMASS treats a grid edge as closed, so those
faces are zero in the tracking (particles cannot leave or enter there).
The files hold no vertical velocity, so tracking is horizontal-only.
"""
import types
from typing import NewType
import dask.array as dsa
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 scrippsgog as gog_source
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,
),
)
def _clean_velocity(da: xr.DataArray) -> xr.DataArray:
"""Zero out NaN (the files' ``_FillValue``) in a velocity field."""
return da.fillna(0.0)
def _bare(da: xr.DataArray) -> xr.DataArray:
"""*da* without coordinates, so it aligns with the velocity fields by position."""
return xr.DataArray(da.data, dims=da.dims)
[docs]
def face_transports(vels: xr.Dataset) -> tuple[xr.DataArray, xr.DataArray]:
"""Volume transports through the east and north face of every T-cell.
Parameters
----------
vels : xarray.Dataset
A velocity file from :func:`tracmass.seagrid.scrippsgog.open_dataset`
(staggered ``(z, YC, XG)`` / ``(z, YG, XC)`` dims).
Returns
-------
uflux, vflux : xarray.DataArray
m³/s on T-cell dims ``(z, y, x)``. ``uflux[..., i]`` is the east
face of cell ``i`` (zero for the last column); ``vflux[:, j]`` is
the north face of row ``j``.
"""
dz = vels["drF"].data[:, None, None]
u = _clean_velocity(vels["uvel"]) * vels["hFacW"] * vels["dyG"]
v = _clean_velocity(vels["vvel"]) * vels["hFacS"] * vels["dxG"]
# West face of cell i+1 is the east face of cell i; nothing east of the
# last column, so pad it with a closed (zero) face.
u_east = dsa.pad(dsa.asarray(u.data)[:, :, 1:] * dz, ((0, 0), (0, 0), (0, 1)))
v_north = dsa.asarray(v.data)[:, 1:, :] * dz
dims = ("z", "y", "x")
return xr.DataArray(u_east, dims=dims), xr.DataArray(v_north, dims=dims)
[docs]
def fluxes(dtm: pd.Timestamp | str | None = None) -> FluxDataset:
"""Load Scripps GoG flux fields for a single time window.
Reads the daily velocity file for the requested time (downloading it
if needed), computes zonal and meridional volume fluxes on the
TRACMASS east/north face convention, 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 day. Default ``None`` — no ``dtm`` is passed to
:func:`tracmass.seagrid.scrippsgog.open_dataset`, so its own
default applies.
Returns
-------
FluxDataset
xarray Dataset with ``uflux``, ``vflux``, ``kmt``, ``dzt``,
``dxdy``, ``timestep``, ``dtm``, ``reverse``, ``imt``, ``jmt``,
``km``, ``iperio``, ``jperio`` and ``zstot``. There is no
``wflux`` (the files have no vertical velocity).
"""
if dtm is None:
vels = gog_source.open_dataset()
else:
vels = gog_source.open_dataset(dtm=pd.to_datetime(dtm))
dtm = pd.Timestamp(vels.time.item())
subset = settings["subset"]
# k1/k2 are 1-based TRACMASS layers; i1/i2/j1/j2 raw 0-based slice ends.
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"])
try:
chunks = vels.uvel.chunks
_zarr_cx, _zarr_cy = int(chunks[-1][0]), int(chunks[-2][0])
except Exception:
_zarr_cx, _zarr_cy = 1, 1
uflux, vflux = face_transports(vels)
grd = gog_source.grid_definition()
dsdict = {
"uflux": uflux,
"vflux": vflux,
"kmt": _bare(grd["kmt"]),
"dzt": _bare(grd["dzt"]).expand_dims("t", axis=-1),
"dxdy": _bare(grd["dxdy"]),
}
ds = xr.Dataset(dsdict).isel(x=x_slice, y=y_slice, z=k_slice)
ds["timestep"] = vels.timestep
ds["dtm"] = dtm
ds["reverse"] = False
ds["km"], ds["jmt"], ds["imt"] = ds["uflux"].shape
# A regional, closed-box domain: never periodic.
ds["iperio"], ds["jperio"] = 0, 0
ds["zstot"] = ds["dxdy"] * 0 + 1
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 _gog_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 = gog_source.grid_definition(dask=False)
# Stored (y, x); transpose to (x, y) = (imt, jmt) to match tracmass
return types.SimpleNamespace(
lon=grd["lons"].values.T,
lat=grd["lats"].values.T,
kmt=np.asarray(grd["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 GoG 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 = _gog_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 in a land cell
(``kmt == 0``) is placed at the centre of the nearest ocean cell.
Parameters
----------
lon, lat : array-like
Longitude and latitude in degrees.
Returns
-------
x, y : np.ndarray
Fractional 0-based full-grid (absolute) TRACMASS positions, same
shape as *lon*/*lat*. Subtract ``grid.i_offset`` /
``grid.j_offset`` for positions relative to a subset grid.
Examples
--------
>>> from tracmass.gcm import scrippsgog as gcm
>>> x, y = gcm.lonlat_to_xy(np.array([0.0, 5.0]), np.array([3.0, 4.0]))
"""
return _lonlat_to_xy(lon, lat, _gog_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.
Parameters
----------
x, y : array-like
Fractional 0-based full-grid positions as stored in the trajectory
Parquet file (columns ``"x"`` and ``"y"``).
Returns
-------
lon, lat : np.ndarray
Longitude and latitude in degrees, same shape as *x* and *y*.
Examples
--------
>>> df = pd.read_parquet("trm_gog.parquet", columns=["x", "y"])
>>> lon, lat = gcm.scrippsgog.xy_to_lonlat(df["x"], df["y"])
"""
return _xy_to_lonlat(x, y, _gog_grid())
