oceanarray 0.1.0__py3-none-any.whl

oceanarray/utilities.py

@@ -0,0 +1,538 @@
+"""Shared utility helpers used across the oceanarray processing pipeline."""
+
+import math
+from datetime import datetime
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional
+
+import numpy as np
+import xarray as xr
+
+from oceanarray import logger
+
+
+def extract_inline_instruments(
+    inline_list: List[Dict[str, Any]],
+) -> List[Dict[str, Any]]:
+    """Extract processable instrument entries from the mooring YAML ``inline`` list.
+
+    Most ``inline`` entries describe passive hardware (ropes, shackles, floats).
+    This function filters for entries that have an ``instrument`` field and either
+    a ``filename`` (to be processed) or ``skip: true`` (to be reported as skipped).
+
+    Normalisation applied to each matching entry:
+
+    - ``hab`` is set from ``hab_bottom`` if ``hab`` is not already present.
+      For a downward-looking instrument the transducer is at the bottom of the
+      housing (``hab_bottom``); for an upward-looking one use ``hab_top``.
+    - ``serial`` is split on the first comma and the first token is used as the
+      primary serial number.  Any remaining tokens are joined and stored under
+      ``beacon_id`` in the returned dict.
+    - ``source: "inline"`` is added to distinguish these entries from ``clamp``
+      entries in downstream logging.
+
+    **Serial parsing is fragile**: the convention ``serial: 16430, R01-024``
+    assumes the *first* comma-separated token is the instrument serial and the
+    rest is a transponder/beacon ID.  If a YAML author places the beacon serial
+    first, the wrong value will be used as the output filename stem.  The
+    validator emits a WARNING for any inline entry whose serial contains a comma
+    so operators can confirm the ordering is correct.
+
+    Args:
+        inline_list: The raw list parsed from the ``inline`` YAML key.
+
+    Returns:
+        List of normalised instrument-config dicts ready for stage processing.
+
+    """
+    result: List[Dict[str, Any]] = []
+    for entry in inline_list:
+        if not isinstance(entry, dict):
+            continue
+        if "instrument" not in entry:
+            continue
+        if not entry.get("filename") and not entry.get("skip"):
+            continue
+
+        normalized: Dict[str, Any] = dict(entry)
+
+        if "hab" not in normalized and "hab_bottom" in normalized:
+            normalized["hab"] = normalized["hab_bottom"]
+
+        serial_raw = str(entry.get("serial", ""))
+        if "," in serial_raw:
+            parts = [p.strip() for p in serial_raw.split(",")]
+            normalized["serial"] = parts[0]
+            normalized["beacon_id"] = ", ".join(parts[1:])
+
+        normalized["source"] = "inline"
+        result.append(normalized)
+    return result
+
+
+def concat_with_scalar_vars(
+    datasets: List[xr.Dataset],
+    dim: str,
+    scalar_vars: Optional[List[str]] = None,
+) -> xr.Dataset:
+    """Concatenate datasets along a dimension, preserving scalar variables.
+
+    Scalar (0-D) variables would normally be broadcast to the concatenation
+    dimension by ``xr.concat``; this helper strips them beforehand and
+    re-attaches the first occurrence after the concat so they remain 0-D.
+
+    Parameters
+    ----------
+    datasets : list of xarray.Dataset
+        Datasets to concatenate.
+    dim : str
+        Dimension along which to concatenate.
+    scalar_vars : list of str, optional
+        List of variable names to treat as scalars. If None, auto-detect
+        scalar variables (those with ndim == 0 in any dataset).
+
+    Returns
+    -------
+    xarray.Dataset
+        Concatenated dataset with scalar variables re-attached as 0-D DataArrays.
+
+    """
+    scalar_storage = {}
+
+    # Auto-detect scalar variables if not provided
+    if scalar_vars is None:
+        scalar_vars = set()
+        for ds in datasets:
+            scalar_vars.update([v for v in ds.data_vars if ds[v].ndim == 0])
+        scalar_vars = list(scalar_vars)
+
+    # Strip scalar variables and store them
+    cleaned = []
+    for _i, ds in enumerate(datasets):
+        ds_copy = ds.copy()
+        for var in scalar_vars:
+            if var in ds_copy and ds_copy[var].ndim == 0:
+                scalar_storage[(_i, var)] = ds_copy[var]
+                del ds_copy[var]
+        cleaned.append(ds_copy)
+
+    # Concatenate non-scalar parts
+    combined = xr.concat(cleaned, dim=dim)
+
+    # Re-attach scalar variables as 0-D DataArrays
+    for (_i, var), da in scalar_storage.items():
+        combined[var] = da  # safe: 0-D, no coords
+
+    return combined
+
+
+def check_necessary_variables(ds: xr.Dataset, vars: list) -> None:
+    """Check that all required variables are present in a dataset.
+
+    Parameters
+    ----------
+    ds: xarray.Dataset
+        Dataset that should be checked
+    vars: list
+        List of variables
+
+    Raises
+    ------
+    KeyError:
+        Raises an error if all vars not present in ds
+
+    Notes
+    -----
+    Original Author: Callum Rollo
+
+    """
+    missing_vars = set(vars).difference(set(ds.variables))
+    if missing_vars:
+        msg = f"Required variables {list(missing_vars)} do not exist in the supplied dataset."
+        raise KeyError(msg)
+
+
+def get_time_key(ds: xr.Dataset) -> str:
+    """Return the name of the time coordinate or variable in an xarray.Dataset.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        The dataset to inspect.
+
+    Returns
+    -------
+    str
+        The name of the time coordinate or variable.
+
+    Raises
+    ------
+    ValueError
+        If no time dimension or coordinate is found.
+
+    """
+    # Common time names to check
+    time_candidates = ["TIME", "time", "Time", "DATETIME", "datetime", "date"]
+    for name in time_candidates:
+        if name in ds.coords or name in ds.variables:
+            return name
+
+    # Fallback: look for any variable or coordinate with datetime dtype
+    for name in ds.coords:
+        if np.issubdtype(ds.coords[name].dtype, np.datetime64):
+            logger.log_warning(
+                "No standard time coordinate found. Using first datetime coordinate: %s",
+                name,
+            )
+            return name
+
+    for name in ds.variables:
+        if np.issubdtype(ds.variables[name].dtype, np.datetime64):
+            logger.log_warning(
+                "No standard time coordinate found. Using first datetime coordinate: %s",
+                name,
+            )
+            return name
+
+    raise ValueError("No valid time coordinate found in dataset.")  # noqa: TRY003
+
+
+def get_dims(
+    ds_gridded: xr.Dataset,
+) -> tuple:
+    """Extract pressure key, time key, and their dimensions from a dataset.
+
+    Parameters
+    ----------
+    ds_gridded : xarray.Dataset
+        Dataset containing the variables and dimensions.
+
+    Returns
+    -------
+    pres_key : str or None
+        Key for the pressure (or depth) variable; None if not found.
+    time_key : str
+        Key for the time variable.
+    pres_dim : str or None
+        Dimension associated with the pressure variable; None if not found.
+    time_dim : str or None
+        Dimension associated with the time variable.
+
+    """
+    if isinstance(ds_gridded, xr.DataArray):
+        ds_gridded = ds_gridded.to_dataset()
+
+    # Determine pressure key
+    pressure_candidates = [
+        "PRES_ADJUSTED",
+        "PRESSURE",
+        "pressure",
+        "PRES",
+        "PRES_INTERPOLATED",
+    ]
+    pres_key = None
+    for name in pressure_candidates:
+        if name in ds_gridded.coords or name in ds_gridded.variables:
+            pres_key = name
+            break
+
+    # Fall back to depth if no pressure variable found
+    dpth_key = None
+    if pres_key is None:
+        depth_candidates = [
+            "DEPTH_ADJUSTED",
+            "DEPTH",
+            "depth",
+            "Depth",
+            "DEPTH_INTERPOLATED",
+        ]
+        for name in depth_candidates:
+            if name in ds_gridded.coords or name in ds_gridded.variables:
+                dpth_key = name
+                break
+
+    # Get time key using helper
+    time_key = get_time_key(ds_gridded)
+
+    # Determine dimensions
+    time_dim = ds_gridded[time_key].dims[0] if ds_gridded[time_key].dims else None
+    pres_dim = None
+    if pres_key is None and dpth_key is not None:
+        pres_dims = ds_gridded[dpth_key].dims
+        if len(pres_dims) > 1:
+            pres_dim = next(dim for dim in pres_dims if dim != time_dim)
+        else:
+            pres_dim = pres_dims[0]
+
+    elif pres_key is not None:
+        pres_dims = ds_gridded[pres_key].dims
+        if len(pres_dims) > 1:
+            pres_dim = next(dim for dim in pres_dims if dim != time_dim)
+        else:
+            pres_dim = pres_dims[0]
+
+    return pres_key or dpth_key, time_key, pres_dim, time_dim
+
+
+def iso8601_duration_from_seconds(seconds: float) -> str:
+    """Convert a duration in seconds to an ISO 8601 duration string.
+
+    Parameters
+    ----------
+    seconds : float
+        Duration in seconds.
+
+    Returns
+    -------
+    str
+        ISO 8601 duration string, e.g., 'PT1H', 'PT30M', 'PT15S'.
+
+    """
+    seconds = int(round(seconds))
+    if seconds >= 86400:
+        return f"P{seconds // 86400}D"
+    elif seconds >= 3600:
+        return f"PT{seconds // 3600}H"
+    elif seconds >= 60:
+        return f"PT{seconds // 60}M"
+    else:
+        return f"PT{seconds}S"
+
+
+def is_iso8601_utc(timestr: str) -> bool:
+    """Validate whether a string is in ISO8601 UTC format: YYYY-MM-DDTHH:MM:SSZ.
+
+    Parameters
+    ----------
+    timestr : str
+        Input time string.
+
+    Returns
+    -------
+    bool
+        True if valid ISO8601 UTC format, False otherwise.
+
+    """
+    try:
+        datetime.strptime(timestr, "%Y/%m/%dT%H:%M:%SZ")  # RODB-style
+    except ValueError:
+        try:
+            datetime.strptime(timestr, "%Y-%m-%dT%H:%M:%SZ")  # ISO8601 style
+        except ValueError:
+            return False
+    return True
+
+
+def apply_defaults(default_source: str, default_files: List[str]) -> Callable:
+    """Decorate a function to apply default values for source and file_list parameters.
+
+    Parameters
+    ----------
+    default_source : str
+        Default source URL or path.
+    default_files : list of str
+        Default list of filenames.
+
+    Returns
+    -------
+    Callable
+        A wrapped function with defaults applied.
+
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        def wrapper(
+            source: Optional[str] = None,
+            file_list: Optional[List[str]] = None,
+            *args: Any,
+            **kwargs: Any,
+        ) -> Any:
+            if source is None:
+                source = default_source
+            if file_list is None:
+                file_list = default_files
+            return func(source=source, file_list=file_list, *args, **kwargs)  # noqa: B026
+
+        return wrapper
+
+    return decorator
+
+
+def _status(kind: str, msg: str, *, char: str = "=") -> None:
+    """Print a formatted status line to stdout.
+
+    kind     msg                        char (optional)
+    -------  -------------------------  ---------------
+    section  "Stage 1: dsG3_1_2026"    border char, default "="  → "=== Stage 1: dsG3_1_2026 ==="
+    instr    "<type> <serial>"          ignored  → "--- microcat 7507 ---"
+    file     relative path             ignored  → "Creating output file: proc/..."
+    skip     relative path             ignored  → "  OUTFILE EXISTS: proc/...  (use --force to overwrite)"
+    error    message                   ignored  → "ERROR: ..."
+
+    Pass ``char`` to ``section`` calls to distinguish mooring runs in a multi-mooring
+    script, e.g. ``char="*"`` gives ``"*** Stage 1: dsG3_2_2026 ***"``.
+    """
+    if kind == "section":
+        border = char * 3
+        print(f"\n{border} {msg} {border}")
+    elif kind == "instr":
+        print(f"\n--- {msg} ---")
+    elif kind == "file":
+        print(f"Creating output file: {msg}")
+    elif kind == "skip":
+        print(f"  OUTFILE EXISTS: {msg}  (use --force to overwrite)")
+    elif kind == "error":
+        print(f"ERROR: {msg}")
+
+
+def _nice_colorbar_bounds(vmin: float, vmax: float, n: int = 20) -> np.ndarray:
+    """Return a boundary array for a discrete colorbar with approximately *n* levels.
+
+    The step is rounded to 1 significant figure so tick labels land on clean values.
+    The range is centered on the midpoint of [vmin, vmax].
+
+    Examples
+    --------
+    Temperature  vmin=0.87, vmax=7.23 → step=0.3,  bounds 1.2–7.2  (20 levels)
+    Salinity     vmin=34.782, vmax=35.139 → step=0.02, bounds 34.76–35.16 (20 levels)
+
+    """
+    span = vmax - vmin
+    if span <= 0:
+        return np.linspace(vmin - 1, vmin + 1, n + 1)
+    raw_step = span / n
+    mag = 10.0 ** math.floor(math.log10(raw_step))
+    rounded = round(raw_step / mag)
+    if rounded == 0:
+        rounded = 1
+    elif rounded >= 10:
+        mag *= 10
+        rounded = 1
+    nice_step = rounded * mag
+    mid = (vmin + vmax) / 2
+    mid_aligned = round(mid / nice_step) * nice_step
+    lo = mid_aligned - (n / 2) * nice_step
+    return np.array([lo + i * nice_step for i in range(n + 1)])
+
+
+# ---------------------------------------------------------------------------
+# Output dtype helpers
+# ---------------------------------------------------------------------------
+
+
+def find_best_dtype(var_name: str, da: xr.DataArray) -> type:
+    """Determine the optimal storage dtype for a variable.
+
+    Parameters
+    ----------
+    var_name : str
+        Variable name.
+    da : xr.DataArray
+        Data array to inspect.
+
+    Returns
+    -------
+    type
+        Recommended numpy dtype.
+
+    Notes
+    -----
+    Rules applied in order:
+
+    - String / datetime / object variables: unchanged.
+    - ``time`` in name: unchanged (preserve datetime64 / float encoding).
+    - ``*_qc`` suffix or ``flag`` in name: ``int8``.
+    - ``serial_number`` or ``serial``: ``int32``.
+    - ``latitude`` / ``longitude`` in name: ``float64``.
+    - Integer input: downsize to ``int32`` if stored as ``int64``, else unchanged.
+    - ``float64`` input: ``float32``.
+    - Anything else: unchanged.
+
+    """
+    input_dtype = da.dtype.type
+    if da.dtype.kind in ("U", "S", "O", "M"):
+        return input_dtype
+    if "time" in var_name.lower():
+        return input_dtype
+    if var_name.endswith("_qc") or "flag" in var_name:
+        return np.int8
+    if var_name in ("serial_number", "serial"):
+        return np.int32
+    if "latitude" in var_name.lower() or "longitude" in var_name.lower():
+        return np.float64
+    if da.dtype.kind in ("i", "u") and da.dtype.itemsize == 8:
+        return np.int32
+    if input_dtype == np.float64:
+        return np.float32
+    return input_dtype
+
+
+def drop_all_zero_vars(ds: xr.Dataset, prefixes: list[str]) -> xr.Dataset:
+    """Drop variables whose finite values are all zero.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Dataset to filter.
+    prefixes : list of str
+        Variable name prefixes to check (e.g. ``["amplitude_beam", "analog_input_"]``).
+        Any variable whose name starts with one of these prefixes and whose
+        finite values are all zero (or has no finite values) is dropped.
+
+    Returns
+    -------
+    xr.Dataset
+        Dataset with all-zero prefix-matched variables removed.
+
+    """
+    to_drop = []
+    for vname in ds.data_vars:
+        if not any(vname.startswith(p) for p in prefixes):
+            continue
+        arr = ds[vname].values
+        finite = arr[np.isfinite(arr)] if np.issubdtype(arr.dtype, np.floating) else arr
+        if finite.size == 0 or np.all(finite == 0):
+            to_drop.append(vname)
+    return ds.drop_vars(to_drop) if to_drop else ds
+
+
+def cast_output_dtypes(ds: xr.Dataset) -> xr.Dataset:
+    """Cast every variable in *ds* to its optimal storage dtype.
+
+    Calls :func:`find_best_dtype` per variable and rebuilds only those that
+    change.  Attributes are preserved.  The input dataset is not modified.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Dataset to cast.
+
+    Returns
+    -------
+    xr.Dataset
+        New dataset with optimised dtypes, ready for NetCDF output.
+
+    """
+    updates: dict[str, xr.Variable] = {}
+    for vname in ds.data_vars:
+        var = ds[vname]
+        target = find_best_dtype(vname, var)
+        if np.dtype(target) != var.dtype:
+            if np.issubdtype(np.dtype(target), np.integer) and np.issubdtype(
+                var.dtype, np.floating
+            ):
+                # NaN cannot be represented as an integer; replace before cast.
+                # QC variables use 9 (CF "missing value"); other integer vars use 0.
+                fill_val = 9 if (vname.endswith("_qc") or "flag" in vname) else 0
+                safe_vals = np.where(np.isfinite(var.values), var.values, fill_val)
+                updates[vname] = xr.Variable(
+                    var.dims, safe_vals.astype(target), attrs=var.attrs
+                )
+            else:
+                updates[vname] = xr.Variable(
+                    var.dims, var.values.astype(target), attrs=var.attrs
+                )
+    if not updates:
+        return ds
+    return ds.assign(updates)