"""xhycom — xarray interface for HYCOM a.b binary output files.
Public API
----------
open_dataset(path, ...) Open any HYCOM .ab file pair (archv, grid, bathy).
open_mfdataset(paths, ...) Open a time series of archive snapshots.
"""
from __future__ import annotations
import warnings
from typing import Iterable, Union
import xarray as xr
from ._abfile import ABFile
from ._discovery import find_archv_files
from ._postprocess import postprocess
from ._reader import (
_build_mf_lazy,
_read_archv_meta,
detect_filetype,
read_archv,
read_ave,
read_bathy,
read_grid,
)
from ._regrid import (
regrid,
regrid_horizontal,
regrid_to_hycom,
regrid_vertical,
velocities_east_north,
)
# Private alias so the public `postprocess` name can also be a keyword argument
# on open_dataset / open_mfdataset without shadowing the function.
_postprocess_ds = postprocess
__version__ = "0.1.0"
__all__ = [
"open_dataset",
"open_mfdataset",
"postprocess",
"read_ave",
"regrid",
"regrid_horizontal",
"regrid_to_hycom",
"regrid_vertical",
"velocities_east_north",
]
# A grid argument is either a path to ``regional.grid`` or a pre-loaded Dataset.
GridArg = Union[str, xr.Dataset, None]
# ``chunks`` is forwarded to ``Dataset.chunk`` (int, mapping, "auto", or None).
Chunks = Union[int, dict, str, None]
# Layer-velocity names and their barotropic partners. ``postprocess`` turns the
# baroclinic ``archv`` velocities into the total current by adding the barotropic
# part, so when velocities are requested via ``variables=`` we must also read it.
_VEL_TOTAL = ("u-vel.", "v-vel.")
_VEL_BTROP = ("u_btrop", "v_btrop")
def _augment_velocity_vars(variables: list[str] | None, postprocess: bool):
"""Pull in the barotropic velocity when velocities are requested + postprocess.
Returns ``(variables_to_read, auto_added)``; *auto_added* are barotropic
names added only to build the total current, which the caller drops after
postprocess so an explicit ``variables=`` list is honoured.
"""
if not postprocess or variables is None:
return variables, []
if not set(variables) & set(_VEL_TOTAL):
return variables, []
auto = [b for b in _VEL_BTROP if b not in variables]
return list(variables) + auto, auto
def _drop_auto(ds: xr.Dataset, auto: list[str]) -> xr.Dataset:
"""Drop the barotropic variables that were auto-added just to form the total."""
drop = [v for v in auto if v in ds.variables]
return ds.drop_vars(drop) if drop else ds
def _load_grid(grid: GridArg, endian: str) -> xr.Dataset | None:
"""Accept a path string or pre-loaded Dataset; return a Dataset."""
if grid is None:
return None
if isinstance(grid, xr.Dataset):
return grid
return open_dataset(grid, endian=endian)
[docs]
def open_dataset(
path: str,
grid: GridArg = None,
endian: str = "big",
chunks: Chunks = None,
variables: list[str] | None = None,
postprocess: bool = False,
) -> xr.Dataset:
"""Open a HYCOM ``.ab`` file pair as an ``xr.Dataset``.
Automatically detects the file type (archive, grid, or bathymetry) from
the ``.b`` header, so the same function works for all HYCOM output files.
If *path* is a glob pattern or directory, it is forwarded to
:func:`open_mfdataset` automatically.
Parameters
----------
path : str
Path to the file. The ``.a`` / ``.b`` extension is optional.
Glob patterns (``*``, ``?``, ``[``) and directory paths are forwarded
to :func:`open_mfdataset`.
grid : str or xr.Dataset, optional
Path to ``regional.grid`` (without extension), or a Dataset already
returned by a previous ``open_dataset`` call on a grid file.
* For **archive** files: attaches ``lon`` / ``lat`` as non-dimension
coordinates on every variable.
* For **bathymetry** files: required (grid dimensions and coordinates
are not stored in the bathymetry file itself).
* For **grid** files: ignored.
endian : str
Byte order: ``"big"`` (default), ``"little"``, or ``"native"``.
chunks : int, dict, or "auto", optional
If provided, the returned Dataset is chunked with Dask. Passed
directly to ``ds.chunk()``. Example: ``chunks={"k": 1}`` to chunk
one layer at a time.
postprocess : bool
If ``True``, convert native units to physical ones and add derived
fields via :func:`xhycom.postprocess` — e.g. sea-surface height and
layer thicknesses in metres, plus ``area`` / ``landmask``. Default
``False`` (data are returned exactly as stored on disk).
Returns
-------
xr.Dataset
Contents depend on file type:
**Archive** (``archv.YYYY_DDD_HH``)
* ``time`` dimension of size 1.
* 2-D fields on ``(time, y, x)``.
* Layered fields on ``(time, k, y, x)`` with ``k`` (layer index,
1-based) and ``dens`` (target sigma-2 density) coordinates.
* Global attributes ``iversn``, ``iexpt``, ``yrflag``.
**Grid** (``regional.grid``)
* All 19 grid variables on ``(y, x)``: ``plon``, ``plat``,
``ulon``, ``ulat``, ``vlon``, ``vlat``, ``qlon``, ``qlat``,
``pang``, ``scpx``, ``scpy``, ``scqx``, ``scqy``, ``scux``,
``scuy``, ``scvx``, ``scvy``, ``cori``, ``pasp``.
**Bathymetry** (``depth_*``)
* Single ``depth`` variable (metres) on ``(y, x)``.
``lon`` / ``lat`` non-dimension coordinates are attached to every
variable when *grid* is supplied (archive and bathymetry files).
Raises
------
ValueError
If the file type cannot be detected, or if *grid* is not provided for
a bathymetry file.
Examples
--------
Open the grid:
>>> grid = xhycom.open_dataset("topo/regional.grid")
Open the bathymetry (grid required for dimensions and coordinates):
>>> bathy = xhycom.open_dataset("topo/depth_TP2a0.10_04",
... grid="topo/regional.grid")
Open a single archive snapshot with grid coordinates:
>>> ds = xhycom.open_dataset("data/archv.2020_001_00",
... grid="topo/regional.grid")
Re-use a pre-loaded grid Dataset to avoid reading the file twice:
>>> grid = xhycom.open_dataset("topo/regional.grid")
>>> bathy = xhycom.open_dataset("topo/depth_TP2a0.10_04", grid=grid)
>>> ds = xhycom.open_dataset("data/archv.2020_001_00", grid=grid)
"""
path = str(path)
# Forward globs and directories to open_mfdataset.
import os as _os
if any(c in path for c in "*?[") or _os.path.isdir(path):
return open_mfdataset(
path,
grid=grid,
endian=endian,
chunks=chunks,
variables=variables,
postprocess=postprocess,
)
basename = ABFile.strip_ab_ending(path)
filetype = detect_filetype(basename)
grid_ds = _load_grid(grid, endian)
if filetype == "grid":
ds = read_grid(basename, endian=endian)
elif filetype == "archv":
# chunks is handled inside read_archv: data is never loaded eagerly
# when chunks is set — Dask tasks are created instead.
aug, auto = _augment_velocity_vars(variables, postprocess)
ds = read_archv(
basename, grid_ds=grid_ds, endian=endian, chunks=chunks, variables=aug
)
if postprocess:
ds = _drop_auto(_postprocess_ds(ds), auto)
return ds
elif filetype == "ave":
ds = read_ave(
basename, grid_ds=grid_ds, endian=endian, chunks=chunks, variables=variables
)
return ds
elif filetype == "bathy":
if grid_ds is None:
raise ValueError(
"grid= is required to open a bathymetry file — it provides "
"the grid dimensions (idm, jdm) and lon/lat coordinates.\n"
"Example: open_dataset('depth_...', grid='regional.grid')"
)
ds = read_bathy(basename, grid_ds=grid_ds, endian=endian)
else:
raise ValueError(f"Unsupported file type {filetype!r} for open_dataset.")
if postprocess:
ds = _postprocess_ds(ds)
return ds.chunk(chunks) if chunks is not None else ds
[docs]
def open_mfdataset(
paths: str | Iterable[str],
grid: GridArg = None,
endian: str = "big",
skip_errors: bool = False,
chunks: Chunks = None,
variables: list[str] | None = None,
postprocess: bool = False,
) -> xr.Dataset:
"""Open multiple HYCOM archive ``.ab`` file pairs as a single ``xr.Dataset``.
Snapshots are concatenated along a ``time`` dimension in chronological
order.
Parameters
----------
paths : str or list of str
One of:
* A directory path — all ``archv.`` / ``archm.YYYY_DDD_HH.[ab]``
pairs found inside are used.
* A glob pattern such as ``"data/archm.1993_*.a"``.
* An explicit list of archive basenames or filenames.
grid : str or xr.Dataset, optional
Grid file path or pre-loaded Dataset. Loaded once and shared across
all files.
endian : str
Byte order.
skip_errors : bool
If ``True``, files that fail to open are skipped with a warning
rather than raising an exception. Default ``False``.
chunks : int, dict, or "auto", optional
If provided, the returned Dataset is chunked with Dask. Passed
directly to ``ds.chunk()``. Example: ``chunks={"time": 1}``.
variables : list of str, optional
If provided, only these variables are included in the returned Dataset.
Reduces the Dask graph size proportionally — useful when working with
large archives that contain many variables (e.g. BGC runs).
Variables not present in the archive are skipped with a warning.
postprocess : bool
If ``True``, apply :func:`xhycom.postprocess` to the combined Dataset
(native-unit conversions + derived fields). Default ``False``.
Returns
-------
xr.Dataset
Combined Dataset with a ``time`` dimension spanning all snapshots.
Examples
--------
Open all snapshots in a directory:
>>> ds = xhycom.open_mfdataset("data/", grid="topo/regional.grid")
Open a subset using a glob:
>>> ds = xhycom.open_mfdataset("data/archv.2020_*.a",
... grid="topo/regional.grid")
Compute time-mean surface salinity:
>>> ds["saln"].isel(k=0).mean("time").plot(x="lon", y="lat")
"""
if isinstance(paths, str):
basenames = find_archv_files(paths)
else:
basenames = [ABFile.strip_ab_ending(str(p)) for p in paths]
grid_ds = _load_grid(grid, endian)
# When velocities are requested with postprocess, also read the barotropic
# part so the total current can be formed, then drop it afterwards.
aug, auto = _augment_velocity_vars(variables, postprocess)
if chunks is not None:
# Lazy path: parse all .b headers in parallel (no .a I/O), then build
# a combined Dask Dataset in one pass — avoids xr.concat overhead.
try:
import dask # noqa: F401
except ImportError:
raise ImportError(
"Dask is required for lazy/chunked loading. "
"Install it with: pip install dask"
)
from concurrent.futures import ThreadPoolExecutor, as_completed
meta_map: dict = {}
with ThreadPoolExecutor() as executor:
future_to_base = {
executor.submit(_read_archv_meta, bn, endian): bn for bn in basenames
}
for future in as_completed(future_to_base):
bn = future_to_base[future]
try:
meta_map[bn] = future.result()
except Exception as exc:
if skip_errors:
warnings.warn(f"Skipping {bn!r}: {exc}", stacklevel=2)
else:
raise
# Restore chronological order, dropping any skipped files.
valid_basenames = [bn for bn in basenames if bn in meta_map]
metas = [meta_map[bn] for bn in valid_basenames]
if not metas:
raise RuntimeError("No files were successfully opened.")
# Extract the integer time chunk size so the graph is built with the
# right granularity — avoids creating 1-file tasks and then rechunking.
time_chunk = 1
if isinstance(chunks, dict) and isinstance(chunks.get("time"), int):
time_chunk = chunks["time"]
elif isinstance(chunks, int):
time_chunk = chunks
ds = _build_mf_lazy(
valid_basenames,
metas,
grid_ds,
endian,
variables=aug,
time_chunk=time_chunk,
)
if postprocess:
ds = _drop_auto(_postprocess_ds(ds), auto)
return ds.chunk(chunks)
else:
# Eager path: read each file and concatenate.
datasets = []
for basename in basenames:
try:
datasets.append(
read_archv(basename, grid_ds=grid_ds, endian=endian, variables=aug)
)
except Exception as exc:
if skip_errors:
warnings.warn(f"Skipping {basename!r}: {exc}", stacklevel=2)
else:
raise
if not datasets:
raise RuntimeError("No files were successfully opened.")
ds = xr.concat(
datasets,
dim="time",
data_vars="minimal",
coords="minimal",
compat="override",
)
if postprocess:
return _drop_auto(_postprocess_ds(ds), auto)
return ds