xhycom 0.1.0__py3-none-any.whl

xhycom/__init__.py

@@ -0,0 +1,400 @@
+"""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)
+
+
+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
+
+
+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