"""Retrieve and store daily ECCO2 cube92 ocean state estimate data.
Unlike the GLORYS12/Mercator THREDDS archive (one file per day, per
variable, discovered through a rolling filename), the APDRC GDS server
exposes each ECCO2 cube92 variable as a single OPeNDAP dataset spanning
the full 1992-01-04 -- 2025-12-31 daily record. A requested date is
selected from that record by nearest-neighbour lookup on the remote
``time`` coordinate and cached locally as a single-day Zarr store so
repeated access does not re-hit the network.
ECCO2 cube92 is also a plain regular lat/lon grid (0.25 deg, global),
not a curvilinear NEMO grid, so there are no separate mesh_hgr/mesh_zgr
files to derive geometry from. :func:`grid_definition` instead computes
cell areas analytically from the lat/lon spacing and derives the bottom
level index ``kmt`` from the land/below-seafloor pattern of a reference
``uvel`` snapshot (see :func:`calc_kmt_dzt`).
The server's native ``lon`` runs 0.125..359.875 (Prime Meridian at the
array seam); :func:`_center_on_prime_meridian` rolls every cached field to
-179.875..179.875 so 0 deg sits at the array centre, matching the
convention used elsewhere in the codebase.
"""
import pathlib
import time
import warnings
from concurrent.futures import ThreadPoolExecutor
import dask
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
EARTH_RADIUS = 6_371_000.0 # meters
# (connect, read) timeout, in seconds, for every OPeNDAP request made against
# the APDRC GDS 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)
# These three warnings are permanent, unfixable quirks of the APDRC GDS
# server (and of the local Zarr caching this module builds around it) —
# never a sign of an actual problem — so they are silenced globally rather
# than scoped per call site. A `with warnings.catch_warnings():` block is
# not reliable here: the Zarr writes below run on dask's threaded scheduler,
# which can execute the lazy `time` CF-decode in a worker thread after a
# `with` block in the calling thread has already restored the filters —
# `catch_warnings` is documented as not thread-safe.
warnings.filterwarnings("ignore", category=xr.SerializationWarning)
warnings.filterwarnings(
"ignore", message="Consolidated metadata is currently not part"
)
warnings.filterwarnings(
"ignore", message="PyDAP was unable to determine the DAP protocol"
)
settings = dict(
host="http://apdrc.soest.hawaii.edu:80/dods",
product="public_data/ECCO/ECCO2/cube92_daily",
timestep=pd.Timedelta(1, "D"),
disk="raid",
)
vardict = dict(
uvel=["U", "uvel"],
vvel=["V", "vvel"],
wvel=["W", "wvel"],
temp=["T", "theta"],
salt=["S", "salt"],
)
def basedir():
path = pathlib.Path(f"/data/{settings['disk']}/ecco2/cube92_daily")
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 date 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 ``CUBE92_<date>_<data_var>.zarr``.
"""
dtm = pd.to_datetime(dtm)
return f"CUBE92_{dtm.date()}_{data_var}.zarr"
def _center_on_prime_meridian(ds):
"""Roll the lon axis so the grid is centred on 0 deg (Prime Meridian).
ECCO2 cube92 serves ``lon`` as 0.125..359.875: the array seam falls on
the Prime Meridian and the antimeridian sits at the array centre.
Rolling to -179.875..179.875 matches the convention used elsewhere in
the codebase (e.g. ``gcm.glorys``) and keeps Atlantic/European regions
contiguous for ``i1:i2`` subsetting instead of wrapping around the
array edge.
"""
n_roll = int((ds.lon.values > 180).argmax())
ds = ds.roll(lon=-n_roll, roll_coords=True)
ds = ds.assign_coords(lon=(((ds.lon + 180) % 360) - 180))
return ds
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
chunks = {"lev": 10, "lat": 360, "lon": 720}
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)
url = f"{settings['host']}/{settings['product']}/{vardict[data_var][1]}"
# PydapDataStore's own suggested fix for the DAP-protocol warning
# (prefixing the URL with `dap2://`) makes pydap rewrite the scheme to
# `https` internally, which would bypass the `session.mount("http://",
# ...)` retry adapter above — so the warning is silenced globally
# instead (see the module-level filterwarnings calls).
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).normalize()
iloc = np.argmin(np.abs(dtm_index - dtm))
da = ds[[vardict[data_var][1]]].isel(time=iloc)
da = _center_on_prime_meridian(da)
return da.to_zarr(str(ddir / fname), mode="w", compute=False)
def _compute_with_retry(*jobs, max_retries=3, backoff=10, **kwargs):
"""Compute dask jobs, retrying on transient OPeNDAP connection drops.
ECCO2 cube92 daily fetches are large (order 100-200 MB per variable)
over a slow remote link, and occasionally fail mid-transfer with a
``requests.exceptions.ChunkedEncodingError``/``IncompleteRead`` from
the APDRC server. ``session.mount("http://", adapter)`` in
:func:`_prepare_write_job` only retries failures during connection
setup (via urllib3's ``Retry``); a dropped connection while streaming
an already-established response — which is what happens here, deep
inside pydap's ``r.content`` — is not covered by that and otherwise
kills a long loop over many dates partway through.
Parameters
----------
*jobs
Dask-computable objects, as passed to ``dask.compute``.
max_retries : int, optional
Number of retries after the initial attempt. Default 3.
backoff : float, optional
Seconds to wait before retry *n*, scaled by *n*. Default 10.
**kwargs
Forwarded to ``dask.compute`` (e.g. ``scheduler``, ``num_workers``).
"""
for attempt in range(max_retries + 1):
try:
return dask.compute(*jobs, **kwargs)
except requests.exceptions.RequestException as exc:
if attempt == max_retries:
raise
wait = backoff * (attempt + 1)
print(
f"OPeNDAP connection dropped mid-transfer ({exc!r}); "
f"retrying in {wait}s ({attempt + 1}/{max_retries}) …",
flush=True,
)
time.sleep(wait)
[docs]
def retrieve(dtm="2020-01-01", data_var="uvel", force=False):
"""Download a single daily field from the APDRC GDS server and save it as Zarr.
Opens the remote ECCO2 cube92 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 ProgressBar():
_compute_with_retry(write_job, scheduler="threads", num_workers=8)
def _cell_edges(centers):
"""Return cell-boundary coordinates bracketing a 1-D array of centres."""
centers = np.asarray(centers)
edges = np.empty(centers.size + 1)
edges[1:-1] = 0.5 * (centers[:-1] + centers[1:])
edges[0] = centers[0] - (edges[1] - centers[0])
edges[-1] = centers[-1] + (centers[-1] - edges[-2])
return edges
[docs]
def calc_grid_spacing(lat, lon):
"""Compute T-cell face widths on a regular lat/lon grid.
Parameters
----------
lat : np.ndarray
1-D array of cell-centre latitudes in degrees, shape ``(jmt,)``.
lon : np.ndarray
1-D array of cell-centre longitudes in degrees, shape ``(imt,)``.
Returns
-------
dxv : np.ndarray
Zonal cell width in meters (face crossed by meridional flux),
shape ``(jmt, imt)``.
dyu : np.ndarray
Meridional cell width in meters (face crossed by zonal flux),
shape ``(jmt, imt)``.
"""
lat_edges = _cell_edges(lat)
lon_edges = _cell_edges(lon)
dlat = np.deg2rad(np.diff(lat_edges)) # (jmt,)
dlon = np.deg2rad(np.diff(lon_edges)) # (imt,)
jmt, imt = lat.size, lon.size
dyu = np.broadcast_to(EARTH_RADIUS * dlat[:, None], (jmt, imt))
dxv = EARTH_RADIUS * np.cos(np.deg2rad(lat))[:, None] * dlon[None, :]
return dxv, np.ascontiguousarray(dyu)
[docs]
def calc_kmt_dzt(uvel0, lev):
"""Derive bottom-level index and layer thickness from a reference snapshot.
ECCO2 cube92 has no separate bathymetry file; land and below-seafloor
cells carry an exact-zero velocity (see Notes), so the ocean depth at
each column is recovered by counting non-zero levels.
Parameters
----------
uvel0 : np.ndarray
Zonal velocity at a single reference time, shape ``(km, jmt, imt)``.
lev : np.ndarray
1-D array of depth-level centres in meters, shape ``(km,)``.
Returns
-------
kmt : np.ndarray
1-based bottom level index per column, shape ``(jmt, imt)``.
dzt : np.ndarray
Layer thickness, zeroed below the seafloor, shape ``(km, jmt, imt)``.
Notes
-----
The APDRC GDS gateway declares ``missing_value``/``_FillValue`` as
``-1e23`` but actually serves exact ``0.0`` for land and below-seafloor
cells (verified against known land and shallow-shelf columns), so the
mask is built from exact equality rather than NaN/CF decoding.
"""
kmt = (uvel0 != 0.0).sum(axis=0).astype("int32") # (jmt, imt)
edges = _cell_edges(lev)
edges[0] = 0.0 # surface
dz = np.diff(edges) # (km,)
km, jmt, imt = uvel0.shape
dzt = np.broadcast_to(dz[:, None, None], (km, jmt, imt)).copy()
k_index = np.arange(1, km + 1)[:, None, None]
dzt[k_index > kmt[None, :, :]] = 0.0
return kmt, dzt
[docs]
def grid_definition(dtm="2020-01-01", dask=True, force=False):
"""Build (and cache) the ECCO2 cube92 grid geometry.
Parameters
----------
dtm : str or datetime-like, optional
Reference date used to derive ``kmt`` from the land/underwater NaN
pattern of a single ``uvel`` snapshot. Default ``"2020-01-01"``.
dask : bool, optional
If True (default), open the cached grid file with Dask chunking.
force : bool, optional
If True, rebuild the grid file even if it already exists.
Returns
-------
xarray.Dataset
Dataset with ``dxdy``, ``dxv``, ``dyu``, ``dzt``, ``kmt``, ``lon``,
``lat``.
"""
gridfile = basedir() / "trm_cube92_grid_def.zarr"
if not gridfile.is_dir() or force:
ref = open_dataset(dtm=dtm, data_var="uvel", dask=False)
lat = ref["y"].values
lon = ref["x"].values
lev = ref["z"].values
kmt, dzt = calc_kmt_dzt(ref["uvel"].values, lev)
dxv, dyu = calc_grid_spacing(lat, lon)
ds = xr.Dataset(
{
"dxdy": (("y", "x"), dxv * dyu),
"dxv": (("y", "x"), dxv),
"dyu": (("y", "x"), dyu),
"dzt": (("z", "y", "x"), dzt),
"kmt": (("y", "x"), kmt),
},
coords={"lon": ("x", lon), "lat": ("y", lat), "z": ("z", lev)},
)
ds.to_zarr(gridfile)
zarr.consolidate_metadata(gridfile)
chunks = {} if dask else None
return xr.open_dataset(gridfile, chunks=chunks)
[docs]
def landmask(k: int = 1) -> xr.DataArray:
"""ECCO2 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 ECCO2 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):
"""Open an ECCO2 cube92 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'). If ``None``,
both velocity components are fetched and merged.
force : bool
If True, re-download the file even if it already exists locally.
dask : bool
If True (default), open the cached Zarr with Dask chunking.
Returns
-------
xarray.Dataset
Dataset containing the requested variable(s) for the given date,
with dims renamed to ``(z, y, x)``.
"""
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")
)
if n_missing:
print(
f"Connecting to OPeNDAP for {n_missing} file(s) ({dtm_ts.date()}) …",
flush=True,
)
# Open both OPeNDAP connections concurrently (network-bound).
with ThreadPoolExecutor(max_workers=2) as pool:
futs = {
var: pool.submit(_prepare_write_job, dtm_ts, var, force)
for var in ("uvel", "vvel")
}
write_jobs = [futs[var].result() for var in ("uvel", "vvel")]
pending = [j for j in write_jobs if j is not None]
if pending:
print("Writing …", flush=True)
with ProgressBar():
_compute_with_retry(*pending, scheduler="threads", num_workers=8)
ds = open_dataset(dtm=dtm, data_var="uvel", dask=dask)
ds["vvel"] = open_dataarray(dtm=dtm, data_var="vvel", dask=dask)
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)
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, "lev": "z", "lat": "y", "lon": "x"})
ds.attrs["timestep"] = settings["timestep"]
ds = ds.assign_coords({"time": dtm})
return ds
[docs]
def open_dataarray(dtm="2020-01-01", data_var="uvel", force=False, dask=True):
"""Open an ECCO2 cube92 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.
dask : bool
If True (default), open the cached Zarr with Dask chunking.
Returns
-------
xarray.DataArray
Data array for the requested variable and date.
"""
ds = open_dataset(dtm=dtm, data_var=data_var, force=force, dask=dask)
return ds[data_var]
