"""Retrieve and store daily Scripps Gulf of Guinea (GoG) MITgcm output.
The Scripps GoG regional MITgcm configuration (``GG_i001``) is served as
plain NetCDF files over HTTPS from https://sose.ucsd.edu/GoG/DATA/, one
file per day and variable group, e.g.
``daily/GG_i001_2024-01-02_1day_UVEL_VVEL_top200m.nc``. Unlike ECCO2
there is no OPeNDAP server, so :func:`retrieve` downloads whole files and
caches them unchanged under :func:`datadir`; :func:`open_dataset` opens
the cached copy.
The model uses an MITgcm Arakawa C-grid on a regular 1/64 deg lat/lon
mesh (1376 x 385 T-cells, 11W--10.5E, 1--7N). The velocity files cover
only the 320 T-cell rows 2--7N, keep MITgcm's staggered dimensions
(``UVEL`` on ``(Z, YC, XG)``: the **west** face of each T-cell; ``VVEL``
on ``(Z, YG, XC)``: the **south** face, plus one extra row for the
northernmost north face), and carry their own face metrics (``dyG``,
``dxG``) and open fractions (``hFacW``, ``hFacS``). Moving those faces
onto TRACMASS's east/north convention is done in
:mod:`tracmass.gcm.scrippsgog`.
There is no ``hFacC`` in the velocity files. :func:`grid_definition`
rebuilds it from the binary ``Depth`` and ``delR`` grid files with
MITgcm's partial-cell rule (``hFacMin``, see :func:`calc_hfac`) and
checks the result against the files' ``hFacW``/``hFacS``.
"""
import pathlib
import re
import time
import numpy as np
import pandas as pd
import requests
import xarray as xr
import zarr
from tracmass.seagrid import coast
settings = dict(
host="https://sose.ucsd.edu/GoG/DATA",
product="PRELIM",
run="i001",
timestep=pd.Timedelta(1, "D"),
disk="raid",
# MITgcm partial-cell minimum fraction of the run; reproduces the
# velocity files' hFacW/hFacS from Depth (checked in grid_definition).
hfac_min=0.2,
)
# key -> (file variable group, variable name inside the file)
vardict = dict(
uvel=("UVEL_VVEL_top200m", "UVEL"),
vvel=("UVEL_VVEL_top200m", "VVEL"),
temp=("THETA", "THETA"),
salt=("SALT", "SALT"),
)
# Binary MITgcm grid files (big-endian float32) under ``<product>/grid/``.
_GRID_FILES = ("XC", "YC", "Depth", "RC")
_DELR_FILE = "delR_v2.bin"
# (connect, read) timeout in seconds for every HTTP request.
_HTTP_TIMEOUT = (10, 60)
_CHUNK_BYTES = 1 << 20
[docs]
def basedir():
"""Return the local root directory for the configured product.
Returns
-------
pathlib.Path
``/data/<disk>/scrippsgog/<product>``, created if needed.
"""
path = pathlib.Path(f"/data/{settings['disk']}/scrippsgog/{settings['product']}")
path.mkdir(parents=True, exist_ok=True)
return path
[docs]
def datadir(dtm="2024-01-02"):
"""Return the local data directory for a given date.
Parameters
----------
dtm : str or datetime-like
Date string or datetime-like object.
Returns
-------
pathlib.Path
Path to the year-specific daily data directory.
"""
dtm = pd.to_datetime(dtm)
return basedir() / "daily" / str(dtm.year)
[docs]
def griddir():
"""Return the local directory holding the binary grid files."""
return basedir() / "grid"
[docs]
def filename(dtm="2024-01-02", data_var="uvel"):
"""Build the remote (and cached) 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', 'temp', 'salt'). ``uvel`` and
``vvel`` live in the same file.
Returns
-------
str
e.g. ``GG_i001_2024-01-02_1day_UVEL_VVEL_top200m.nc``.
"""
dtm = pd.to_datetime(dtm)
group = vardict[data_var][0]
return f"GG_{settings['run']}_{dtm.date()}_1day_{group}.nc"
[docs]
def url(dtm="2024-01-02", data_var="uvel"):
"""Return the download URL of a daily file.
Parameters
----------
dtm : str or datetime-like
Date string or datetime-like object.
data_var : str
Variable key (one of 'uvel', 'vvel', 'temp', 'salt').
Returns
-------
str
Full HTTPS URL.
"""
base = f"{settings['host']}/{settings['product']}/daily"
return f"{base}/{filename(dtm=dtm, data_var=data_var)}"
[docs]
def available_dates(data_var="uvel"):
"""List the dates the server has files for.
Parameters
----------
data_var : str
Variable key (one of 'uvel', 'vvel', 'temp', 'salt').
Returns
-------
pandas.DatetimeIndex
Sorted dates parsed from the server's directory listing.
"""
listing = f"{settings['host']}/{settings['product']}/daily/"
resp = requests.get(listing, timeout=_HTTP_TIMEOUT)
resp.raise_for_status()
group = re.escape(vardict[data_var][0])
pattern = rf'href="GG_{settings["run"]}_(\d{{4}}-\d{{2}}-\d{{2}})_1day_{group}\.nc"'
return pd.DatetimeIndex(sorted(set(re.findall(pattern, resp.text))))
def _download(src, dest, max_retries=3, backoff=10):
"""Stream *src* to *dest* via a temporary file, retrying dropped transfers.
The file only appears at *dest* once it is complete, so an interrupted
download never leaves a truncated file that looks cached.
"""
dest = pathlib.Path(dest)
dest.parent.mkdir(parents=True, exist_ok=True)
tmp = dest.with_name(dest.name + ".part")
for attempt in range(max_retries + 1):
try:
with requests.get(src, stream=True, timeout=_HTTP_TIMEOUT) as resp:
if resp.status_code == 404:
raise FileNotFoundError(f"Not on the server: {src}")
resp.raise_for_status()
with open(tmp, "wb") as fh:
for chunk in resp.iter_content(chunk_size=_CHUNK_BYTES):
fh.write(chunk)
tmp.replace(dest)
return dest
except requests.exceptions.RequestException as exc:
tmp.unlink(missing_ok=True)
if attempt == max_retries:
raise
wait = backoff * (attempt + 1)
print(
f"Download of {src} failed ({exc!r}); "
f"retrying in {wait}s ({attempt + 1}/{max_retries}) …",
flush=True,
)
time.sleep(wait)
[docs]
def retrieve(dtm="2024-01-02", data_var="uvel", force=False):
"""Download a single daily file and cache it locally.
Parameters
----------
dtm : str or datetime-like
Date of the daily file.
data_var : str
Variable key (one of 'uvel', 'vvel', 'temp', 'salt'). ``uvel`` and
``vvel`` share one file.
force : bool
If True, re-download even if the file is already cached.
Returns
-------
pathlib.Path
Path to the cached file.
Raises
------
FileNotFoundError
If the server has no file for *dtm* (see :func:`available_dates`).
Examples
--------
>>> from tracmass.seagrid import scrippsgog
>>> for dtm in pd.date_range("2024-01-02", "2024-01-31"):
... scrippsgog.retrieve(dtm)
"""
dtm = pd.to_datetime(dtm)
dest = datadir(dtm=dtm) / filename(dtm=dtm, data_var=data_var)
if dest.exists() and not force:
return dest
print(f"Downloading {dest.name} …", flush=True)
return _download(url(dtm=dtm, data_var=data_var), dest)
[docs]
def retrieve_grid(force=False):
"""Download the binary MITgcm grid files (``XC``, ``YC``, ``Depth``, ...).
Parameters
----------
force : bool
If True, re-download even if the files are already cached.
Returns
-------
pathlib.Path
The local grid directory (:func:`griddir`).
"""
base = f"{settings['host']}/{settings['product']}/grid"
names = [f"{n}.{ext}" for n in _GRID_FILES for ext in ("data", "meta")]
for name in names + [_DELR_FILE]:
dest = griddir() / name
if force or not dest.exists():
_download(f"{base}/{name}", dest)
return griddir()
def _read_meta(path):
"""Return the array shape ``(..., ny, nx)`` declared in an MITgcm .meta file."""
text = pathlib.Path(path).read_text()
dims = [int(v) for v in re.search(r"dimList = \[([^\]]*)\]", text)[1].split(",")]
sizes = dims[0::3] # each dim is listed as (size, first, last), fastest first
return tuple(reversed(sizes))
[docs]
def read_grid_field(name):
"""Read one binary MITgcm grid field (``XC``, ``YC``, ``Depth``, ``RC``).
Parameters
----------
name : str
Field name; ``<name>.data`` / ``<name>.meta`` must be in
:func:`griddir` (see :func:`retrieve_grid`).
Returns
-------
np.ndarray
float32 array, shape ``(385, 1376)`` for 2-D fields (y first).
"""
shape = _read_meta(griddir() / f"{name}.meta")
data = np.fromfile(griddir() / f"{name}.data", dtype=">f4")
return data.reshape([s for s in shape if s > 1] or [1]).astype("float32")
[docs]
def calc_hfac(depth, drf, hfac_min=0.2, hfac_min_dr=0.0):
"""Rebuild MITgcm's T-cell open fraction ``hFacC`` from bathymetry.
Parameters
----------
depth : np.ndarray
Ocean depth (positive, m), shape ``(jmt, imt)``; 0 on land.
drf : np.ndarray
Layer thicknesses (m), shape ``(km,)``, surface first.
hfac_min, hfac_min_dr : float, optional
MITgcm ``hFacMin`` / ``hFacMinDr``.
Returns
-------
np.ndarray
``hFacC``, shape ``(km, jmt, imt)``, in ``[0, 1]``.
Notes
-----
Follows MITgcm ``ini_masks_etc``: the raw fraction of each layer above
the sea floor is rounded to 0 below half the minimum size, and up to
the minimum size below it.
"""
drf = np.asarray(drf, dtype="float64")
z_top = np.concatenate([[0.0], np.cumsum(drf)[:-1]])
hfac = np.clip((depth[None] - z_top[:, None, None]) / drf[:, None, None], 0, 1)
hfac_mn = np.maximum(hfac_min, np.minimum(hfac_min_dr / drf, 1.0))[:, None, None]
hfac = np.where(hfac < hfac_mn, np.where(hfac < 0.5 * hfac_mn, 0.0, hfac_mn), hfac)
return hfac.astype("float32")
[docs]
def calc_kmt_dzt(hfac, drf):
"""Derive bottom-level index and layer thickness from ``hFacC``.
Parameters
----------
hfac : np.ndarray
Open fraction, shape ``(km, jmt, imt)``.
drf : np.ndarray
Layer thicknesses (m), shape ``(km,)``.
Returns
-------
kmt : np.ndarray
Number of open layers per column, shape ``(jmt, imt)``, int32.
Land is 0.
dzt : np.ndarray
``drf * hfac`` (partial bottom cells), shape ``(km, jmt, imt)``.
"""
kmt = (hfac > 0).sum(axis=0).astype("int32")
dzt = (np.asarray(drf)[:, None, None] * hfac).astype("float32")
return kmt, dzt
[docs]
def grid_definition(dtm=None, dask=True, force=False):
"""Build (and cache) the GoG grid geometry on the velocity files' T-cells.
Parameters
----------
dtm : str or datetime-like or None, optional
Date of the reference velocity file whose rows, layers and face
metrics define the grid. Default: the first file on the server's
start date (``2024-01-02``).
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 on ``(z, y, x)`` T-cells with ``kmt``, ``dzt``, ``hfac``,
``dxdy``, ``dxv``, ``dyu``, the 2-D ``lons``/``lats`` and 1-D
``lon``/``lat``/``z`` coordinates.
Raises
------
ValueError
If the rebuilt ``hFacC`` does not reproduce the velocity file's
``hFacW``/``hFacS`` (wrong ``settings["hfac_min"]`` or misaligned
rows).
"""
gridfile = basedir() / "trm_gog_grid_def.zarr"
if not gridfile.is_dir() or force:
retrieve_grid()
ref = open_dataset(dtm=dtm or "2024-01-02", dask=False)
xc, yc = read_grid_field("XC"), read_grid_field("YC")
depth = read_grid_field("Depth")
# Velocity files hold a row subset of the binary grid; locate it.
j0 = int(np.argmin(np.abs(yc[:, 0] - ref["YC"].values[0])))
rows = slice(j0, j0 + ref.sizes["YC"])
if not np.allclose(yc[rows, 0], ref["YC"].values, atol=1e-4):
raise ValueError("Velocity file rows do not match the grid's YC")
drf = ref["drF"].values.astype("float64")
# hFacS needs the T-cells on both sides of every YG face: one row
# south of the velocity rows and one north of them.
if j0 == 0 or rows.stop >= yc.shape[0]:
raise ValueError("Velocity rows touch the edge of the binary grid")
ext = slice(j0 - 1, rows.stop + 1)
hfac_ext = calc_hfac(depth[ext], drf, settings["hfac_min"])
_check_hfac(hfac_ext, ref)
hfac = hfac_ext[:, 1:-1]
kmt, dzt = calc_kmt_dzt(hfac, drf)
dyg = ref["dyG"].values # (y, x) west face, constant along x
dxg = ref["dxG"].values # (y+1, x) south faces + northernmost north face
dxv = dxg[1:] # north face of each T-cell
dyu = np.concatenate([dyg[:, 1:], dyg[:, -1:]], axis=1) # east face
dxdy = 0.5 * (dxg[:-1] + dxg[1:]) * dyg
ds = xr.Dataset(
{
"kmt": (("y", "x"), kmt),
"dzt": (("z", "y", "x"), dzt),
"hfac": (("z", "y", "x"), hfac),
"dxdy": (("y", "x"), dxdy.astype("float32")),
"dxv": (("y", "x"), dxv.astype("float32")),
"dyu": (("y", "x"), dyu.astype("float32")),
"lons": (("y", "x"), xc[rows]),
"lats": (("y", "x"), yc[rows]),
},
coords={
"lon": ("x", ref["XC"].values),
"lat": ("y", ref["YC"].values),
"z": ("z", -ref["z"].values),
},
)
ds.to_zarr(gridfile, mode="w")
zarr.consolidate_metadata(gridfile)
chunks = {} if dask else None
return xr.open_dataset(gridfile, chunks=chunks)
[docs]
def landmask(k: int = 1) -> xr.DataArray:
"""GoG 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 GoG 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`.
Notes
-----
The GoG domain is regional: land outside it is unknown, so distances
near the open boundaries can be too large.
"""
return coast.distance_to_land(landmask(k))
def _check_hfac(hfac_ext, ref, atol=1e-4):
"""Raise if rebuilt hFacC disagrees with the velocity file's hFacW / hFacS.
*hfac_ext* covers the velocity rows plus one extra row on each side.
MITgcm sets a face's open fraction to the smaller of its two T-cells'.
"""
hfac = hfac_ext[:, 1:-1]
hfac_w = np.minimum(hfac[:, :, :-1], hfac[:, :, 1:])
err_w = np.abs(hfac_w - ref["hFacW"].values[:, :, 1:]).max()
hfac_s = np.minimum(hfac_ext[:, :-1], hfac_ext[:, 1:])
err_s = np.abs(hfac_s - ref["hFacS"].values).max()
if max(err_w, err_s) > atol:
raise ValueError(
f"Rebuilt hFacC does not match the velocity file (|dhFacW| = "
f"{err_w:.3g}, |dhFacS| = {err_s:.3g}); check settings['hfac_min']"
)
[docs]
def open_dataset(dtm="2024-01-02", data_var=None, force=False, dask=True):
"""Open a cached daily GoG file, downloading it first if necessary.
Parameters
----------
dtm : str or datetime-like
Date of the daily file.
data_var : str or None
Variable key (one of 'uvel', 'vvel', 'temp', 'salt'). ``None``
(default) opens the velocity file with both ``uvel`` and ``vvel``.
force : bool
If True, re-download the file even if it already exists locally.
dask : bool
If True (default), open with Dask chunking (the file's own chunks).
Returns
-------
xarray.Dataset
The file's single time step, with the data variable(s) renamed to
their keys and ``Z`` renamed to ``z``. MITgcm's staggered
horizontal dims (``XC``/``XG``, ``YC``/``YG``) and face metrics are
kept as is.
"""
data_var = data_var or "uvel"
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)
ds = xr.open_dataset(fname, chunks={} if dask else None)
group = vardict[data_var][0]
renames = {var: key for key, (grp, var) in vardict.items() if grp == group}
ds = ds.rename({**renames, "Z": "z"}).isel(time=0)
ds.attrs["timestep"] = settings["timestep"]
return ds
[docs]
def open_dataarray(dtm="2024-01-02", data_var="uvel", force=False, dask=True):
"""Open one variable of a daily GoG file, downloading it if necessary.
Parameters
----------
dtm : str or datetime-like
Date of the daily file.
data_var : str
Variable key (one of 'uvel', 'vvel', 'temp', 'salt').
force : bool
If True, re-download the file even if it already exists locally.
dask : bool
If True (default), open 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]
