"""Retrieve and store daily ocean reanalysis data from Mercator Ocean."""
import pathlib
import warnings
from concurrent.futures import ThreadPoolExecutor
import numpy as np
import pandas as pd
import requests
import xarray as xr
import zarr
from dask.diagnostics import ProgressBar
from xarray.backends import PydapDataStore
from tracmass.seagrid import coast
from tracmass.seagrid.debug import trace_lines
settings = dict(
host="dap2://tds.mercator-ocean.fr",
product="glorys12v1",
timestep=pd.Timedelta(1, "D"),
disk="raid",
)
vardict = dict(
uvel=["U", "vozocrtx"],
vvel=["V", "vomecrty"],
wvel=["W", "vovecrtz"],
temp=["T", "votemper"],
salt=["S", "vosaline"],
)
# (connect, read) timeout, in seconds, for every OPeNDAP request made against
# the Mercator THREDDS server. A short connect timeout fails fast against an
# unreachable host; the read timeout bounds how long a single stalled chunk
# request (one dask chunk's worth of array data) may block once connected.
# Paired with a low retry count below, a genuinely dead/unreachable server
# now fails within roughly a minute instead of retrying for 10+ minutes per
# chunk (each retry re-running the full connect+read timeout).
_OPENDAP_TIMEOUT = (10, 60)
# NetCDF default fill value (9.96921e36). The Mercator OPeNDAP velocities carry it at
# masked (land) points, and xarray does not decode it there, so it lands in
# the cached Zarrs as data. Anything at or above _FILL_THRESHOLD is fill.
_FILL_THRESHOLD = 1e30
def _mask_fill(da, other):
"""Replace fill values in *da* with *other* (lazy; keeps Dask chunks)."""
return da.where(abs(da) < _FILL_THRESHOLD, other)
[docs]
def mask_fill(da):
"""Replace the fill values (land) in a GLORYS variable (lazy).
Masked points hold the fill value: no flow through land (velocities
``uvel`` / ``vvel`` / ``wvel`` → 0); tracers have no land value (→ NaN).
Parameters
----------
da : xr.DataArray
A variable from :func:`open_dataset`, named by its key (``"uvel"``,
``"temp"``, ...).
Returns
-------
xr.DataArray
Same shape and Dask chunks.
"""
return _mask_fill(da, 0 if da.name in ("uvel", "vvel", "wvel") else np.nan)
def basedir():
path = pathlib.Path(f"/data/{settings['disk']}/mercator/{settings['product']}")
path.mkdir(parents=True, exist_ok=True)
return path
[docs]
def datadir(dtm="2020-01-01"):
"""Return the local data directory for a given date, creating it if needed.
Parameters
----------
dtm : str or datetime-like
Date string or datetime-like object.
Returns
-------
pathlib.Path
Path to the year-specific data directory.
"""
dtm = pd.to_datetime(dtm)
path = basedir() / str(dtm.year)
path = pathlib.Path(path)
return path
[docs]
def filename(dtm="2020-01-01", data_var="uvel"):
"""Build the Zarr filename for a given dat and variable.
Parameters
----------
dtm : str or datetime-like
Date string or datetime-like object.
data_var : str
Variable key (one of 'uvel', 'vvel', 'wvel', 'temp', 'salt').
Returns
-------
str
Filename string in the form ``GLORYS12V1_<date>_grid<X>.zarr``.
"""
def dtmstr(dtm):
return f"{dtm.year}{dtm.month:02}{dtm.day:02}"
dtm = pd.to_datetime(dtm)
chr = vardict[data_var][0]
if settings["product"] == "psy4v3r1":
# ext-PSY4V3R1_1dAV_20121123_20121124_gridW_R20121128.nc
fn = (
"ext-PSY4V3R1_1dAV_"
+ f"{dtmstr(dtm)}_{dtmstr(dtm + pd.DateOffset(1, 'D'))}_grid{chr}_"
+ f"R{dtmstr(pd.Timestamp.now())}.zarr"
)
else:
fn = f"{settings['product'].upper()}_{dtm.date()}_grid{chr}.zarr"
return fn
def _prepare_write_job(dtm, data_var, force=False):
"""Open the remote OPeNDAP dataset and return a pending dask write job.
Returns ``None`` if the local Zarr file already exists and ``force`` is
False. Intended to be called from a thread pool so that multiple
OPeNDAP connections can be set up concurrently.
Parameters
----------
dtm : pd.Timestamp
Date for the desired time step (already converted).
data_var : str
Variable key (one of 'uvel', 'vvel', 'wvel', 'temp', 'salt').
force : bool
If True, re-download even if the file already exists.
Returns
-------
dask.delayed or None
Pending ``to_zarr`` compute job, or ``None`` if the file is cached.
"""
fname = filename(dtm=dtm, data_var=data_var)
ddir = datadir(dtm=dtm)
if (ddir / fname).exists() and not force:
return None
# wvel is on depthw: without its own key it got one 50-level chunk.
chunks = {"time_counter": 1, "deptht": 10, "depthw": 10, "y": 1000, "x": 1000}
session = requests.Session()
adapter = requests.adapters.HTTPAdapter(
max_retries=requests.adapters.Retry(total=2, backoff_factor=1),
pool_connections=8,
pool_maxsize=16,
)
session.mount("http://", adapter)
path = f"/thredds/dodsC/{settings['product']}-daily-grid"
url = settings["host"] + path + vardict[data_var][0]
store = PydapDataStore.open(url, session=session, timeout=_OPENDAP_TIMEOUT)
ds = xr.open_dataset(store, chunks=chunks).astype("float32")
dtm_index = pd.to_datetime(ds.time_counter).normalize()
iloc = np.argmin(np.abs(dtm_index - dtm))
da = ds[vardict[data_var][1]].isel(time_counter=iloc).to_dataset()
return da.to_zarr(str(ddir / fname), mode="w", compute=False)
[docs]
def retrieve(dtm="2020-01-01", data_var="uvel", force=False):
"""Download a single daily field from the Mercator THREDDS server and save it as Zarr.
Opens the remote dataset via OPeNDAP, selects the time step closest to
``dtm``, and writes the result to a local Zarr store.
Parameters
----------
dtm : str or datetime-like
Date string or datetime-like object for the desired time step.
data_var : str
Variable key (one of 'uvel', 'vvel', 'wvel', 'temp', 'salt').
force : bool
If True, re-download the file even if it already exists locally.
"""
dtm = pd.to_datetime(dtm)
write_job = _prepare_write_job(dtm, data_var, force)
if write_job is None:
return
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
message="Consolidated metadata is currently not part",
)
with ProgressBar():
write_job.compute(scheduler="threads", num_workers=8)
def calc_e3t(ds):
e3t_0 = ds.e3t_0.isel(t=0).values # (z,)
e3t_ps = ds.e3t_ps.isel(t=0).values # (y, x)
mbathy = ds.mbathy.isel(t=0).values # (y, x), 1-based level index of bottom cell
nz = len(ds.nav_lev)
ny, nx = e3t_ps.shape
# k_index = np.arange(1, nz + 1)[:, None, None] * np.ones((nz, ny, nx))
k_index = np.arange(1, nz + 1)[:, None, None] # shape (nz, 1, 1)
# e3t_3d = np.broadcast_to(e3t_0[:, None, None], (nz, ny, nx)).copy()
e3t_3d = np.empty((nz, ny, nx))
e3t_3d[:] = e3t_0[:, None, None]
bottom_mask = k_index == mbathy[None, :, :]
e3t_3d[bottom_mask] = np.broadcast_to(e3t_ps[None, :, :], (nz, ny, nx))[bottom_mask]
e3t_3d[k_index > mbathy[None, :, :]] = 0.0
e3t = xr.DataArray(
e3t_3d,
dims=["z", "y", "x"],
coords={"nav_lev": ds.nav_lev},
attrs={"units": "m", "long_name": "T-cell thickness"},
)
return e3t
def grid_definition(dask=True, force=False):
gridfile = basedir() / f"trm_{settings['product']}_grid_def.zarr"
if not gridfile.is_dir() or force:
ds = xr.Dataset()
chunks = {} if dask else None
zgr = xr.open_dataset(basedir() / "ext-PSY4V3R1_mesh_zgr.nc", chunks=chunks)
hgr = xr.open_dataset(basedir() / "ext-PSY4V3R1_mesh_hgr.nc", chunks=chunks)
ds["dxv"] = hgr["e1v"]
ds["dyu"] = hgr["e2u"]
ds["dxdy"] = ds["dxv"] * ds["dyu"]
ds["lon"] = hgr["nav_lon"]
ds["lat"] = hgr["nav_lat"]
ds["kmt"] = zgr["mbathy"]
ds = ds.set_coords(["lon", "lat"])
ds["dzt"] = calc_e3t(zgr)
ds.to_zarr(gridfile)
zarr.consolidate_metadata(gridfile)
# chunks={} honours the zarr store's native layout (200×200 after rechunking).
# chunks="auto" would collapse each variable into a single global slab, forcing
# expensive Dask rechunks when multiplying grid fields against the velocity arrays.
ds = xr.open_dataset(gridfile, chunks={})
# Grid files built before 2026-10-03 had every zero (land kmt, land dzt)
# overwritten with the fill value; land must be kmt == 0 / dzt == 0.
return ds.map(_mask_fill, keep_attrs=True, other=0)
[docs]
def landmask(k: int = 1) -> xr.DataArray:
"""GLORYS12 land mask on the T-cells of :func:`grid_definition`.
Parameters
----------
k : int, optional
1-based model level. Default ``1`` (surface): land is ``kmt == 0``.
Deeper, the sea floor counts as land too (``kmt < k``).
Returns
-------
xr.DataArray
bool ``(y, x)``, ``True`` on land, with ``lon`` / ``lat``
coordinates. See :func:`tracmass.seagrid.coast.landmask`.
"""
return coast.landmask(grid_definition(), k)
[docs]
def distance_to_land(k: int = 1) -> xr.DataArray:
"""Distance (km) from each GLORYS12 T-cell centre to the nearest land cell.
Parameters
----------
k : int, optional
1-based model level whose land mask is used (see :func:`landmask`).
Default ``1`` (surface).
Returns
-------
xr.DataArray
float32 ``(y, x)`` great-circle distance in km, 0 on land, with
``lon`` / ``lat`` coordinates. See
:func:`tracmass.seagrid.coast.distance_to_land`.
"""
return coast.distance_to_land(landmask(k))
# @trace_lines
[docs]
def open_dataset(dtm="2020-01-01", data_var=None, force=False, dask=True, mask=True):
"""Open a mercator dataset, downloading it first if necessary.
Parameters
----------
dtm : str or datetime-like
Date string or datetime-like object for the desired time step.
data_var : str
Variable key (one of 'uvel', 'vvel', 'wvel', 'temp', 'salt').
force : bool
If True, re-download the file even if it already exists locally.
mask : bool, optional
Replace fill values (land) with 0 for velocities and NaN for
tracers. Default True. ``False`` leaves them in, for callers that
first cut out a region and then mask just that with
:func:`mask_fill` (masking a whole Zarr chunk to keep a small part
of it is several times slower).
Returns
-------
xarray.Dataset
Dataset containing the requested variable for the given date.
"""
if data_var is None:
dtm_ts = pd.to_datetime(dtm)
n_missing = sum(
not (datadir(dtm=dtm_ts) / filename(dtm=dtm_ts, data_var=var)).exists()
for var in ("uvel", "vvel", "wvel")
)
if n_missing:
print(
f"Connecting to OPeNDAP for {n_missing} file(s) ({dtm_ts.date()}) …",
flush=True,
)
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
message="Consolidated metadata is currently not part",
)
# Open all OPeNDAP connections concurrently (network-bound).
# The warning filter must be active before threads start because
# zarr fires ZarrUserWarning during to_zarr(..., compute=False).
with ThreadPoolExecutor(max_workers=3) as pool:
futs = {
var: pool.submit(_prepare_write_job, dtm_ts, var, force)
for var in ("uvel", "vvel", "wvel")
}
write_jobs = [futs[var].result() for var in ("uvel", "vvel", "wvel")]
# Filter out cached files, then compute all pending writes together
# under a single ProgressBar so output is never interleaved.
pending = [j for j in write_jobs if j is not None]
if pending:
import dask
print("Writing …", flush=True)
with ProgressBar():
dask.compute(*pending, scheduler="threads", num_workers=8)
ds = open_dataset(dtm=dtm, data_var="uvel", mask=mask)
ds["vvel"] = open_dataarray(dtm=dtm, data_var="vvel", mask=mask)
ds["wvel"] = open_dataarray(dtm=dtm, data_var="wvel", mask=mask)
ds = ds.rename_dims({"deptht": "z", "depthw": "z"})
return ds
dtm = pd.to_datetime(dtm)
fname = datadir(dtm=dtm) / filename(dtm=dtm, data_var=data_var)
if not fname.exists() or force:
retrieve(dtm=dtm, data_var=data_var, force=force)
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
message="Consolidated metadata is currently not part",
)
if dask:
ds = xr.open_dataset(fname, chunks={}, consolidated=True)
else:
ds = xr.open_dataset(fname, consolidated=True)
ds = ds.rename(
{
vardict[data_var][1]: data_var,
"nav_lat": "lat_" + data_var[0],
"nav_lon": "lon_" + data_var[0],
}
)
ds.attrs["timestep"] = settings["timestep"]
ds = ds.assign_coords({"time": dtm}).drop_vars("time_counter")
ds = ds.assign_coords({"time": dtm})
if mask:
ds[data_var] = mask_fill(ds[data_var])
return ds
[docs]
def rechunk_wvel(dtm1, dtm2=None, depth_chunk: int = 10, path=None) -> list:
"""Rewrite cached ``wvel`` Zarr stores with *depth_chunk*-level chunks.
``wvel`` stores downloaded before 2026-10-08 hold each 1000 x 1000 tile
as one 50-level chunk (the download chunked ``deptht`` but ``wvel`` is
on ``depthw``), so reading a few levels decompresses all 50 in one
thread. The x/y chunks are kept. Each store is written next to the
original, checked to hold identical values, and only then swapped in;
stores that already have *depth_chunk* levels per chunk are skipped.
Do not run it on days a tracking run is reading at the same time. A
run opens each day's store (and its chunk layout) a window before it
reads the data; if the store is swapped in between, the read fails
with ``cannot reshape array of size 10000000 into shape
(50,1000,1000)``. A run is usually faster than the rechunk, so a run
that starts behind it catches up. Rechunk first, or rechunk days the
run has already passed.
Parameters
----------
dtm1 : str or datetime-like
First day.
dtm2 : str or datetime-like or None, optional
Last day (inclusive). Default: only *dtm1*.
depth_chunk : int, optional
Levels per chunk. Default 10, as ``uvel`` / ``vvel``.
path : str or pathlib.Path or None, optional
Rechunk this one store instead of the cached ``wvel`` of
*dtm1* .. *dtm2* (for testing on a copy).
Returns
-------
list of pathlib.Path
The stores that were rewritten.
Raises
------
RuntimeError
If a rewritten store does not match its original; the original is
then left untouched.
Examples
--------
>>> rechunk_wvel("2010-06-01", "2011-06-01")
"""
import shutil
if path is not None:
paths = [pathlib.Path(path)]
else:
days = pd.date_range(dtm1, dtm2 if dtm2 is not None else dtm1, freq="D")
paths = [datadir(dtm=d) / filename(dtm=d, data_var="wvel") for d in days]
var = vardict["wvel"][1]
done = []
for i, store in enumerate(paths, 1):
if not store.exists():
print(f"[{i}/{len(paths)}] {store.name}: missing, skipped")
continue
src = zarr.open_group(str(store), mode="r")
if src[var].chunks[0] == depth_chunk:
print(f"[{i}/{len(paths)}] {store.name}: already rechunked")
continue
print(f"[{i}/{len(paths)}] {store.name} ...", flush=True)
tmp = store.with_name(store.name + ".tmp")
if tmp.exists():
shutil.rmtree(tmp)
try:
with warnings.catch_warnings():
warnings.filterwarnings("ignore")
ds = xr.open_dataset(store, chunks={}, consolidated=True)
da = ds[var]
chunks = (depth_chunk,) + tuple(src[var].chunks[1:])
ds = ds.chunk(dict(zip(da.dims, chunks)))
for v in ds.variables:
ds[v].encoding.pop("chunks", None)
ds[v].encoding.pop("preferred_chunks", None)
ds.to_zarr(
str(tmp), mode="w", consolidated=True,
encoding={var: {"chunks": chunks}},
zarr_format=src.metadata.zarr_format, # keep v2 stores v2
)
new = xr.open_dataset(
tmp, chunks={}, consolidated=True, engine="zarr"
)[var]
same = bool(
((new == da) | (new.isnull() & da.isnull())).all().compute()
)
except BaseException:
shutil.rmtree(tmp, ignore_errors=True)
raise
if not same:
shutil.rmtree(tmp)
raise RuntimeError(
f"rechunk_wvel: {tmp} differs from {store}; kept original"
)
# Swap by renames, so the store is missing only between two renames
# (deleting the original first left it missing for the whole rmtree).
old = store.with_name(store.name + ".old")
if old.exists():
shutil.rmtree(old)
store.rename(old)
tmp.rename(store)
shutil.rmtree(old)
done.append(store)
return done
[docs]
def open_dataarray(dtm="2020-01-01", data_var="uvel", force=False, mask=True):
"""Open a mercator dataarray, downloading it first if necessary.
Parameters
----------
dtm : str or datetime-like
Date string or datetime-like object for the desired time step.
data_var : str
Variable key (one of 'uvel', 'vvel', 'wvel', 'temp', 'salt').
force : bool
If True, re-download the file even if it already exists locally.
mask : bool, optional
Mask fill values, see :func:`open_dataset`. Default True.
Returns
-------
xarray.Dataset
Dataset containing the requested variable for the given date.
"""
ds = open_dataset(dtm=dtm, data_var=data_var, force=force, mask=mask)
return ds[data_var]
