roms-tools 3.2.0__py3-none-any.whl → 3.3.0__py3-none-any.whl

roms_tools/setup/utils.py

@@ -4,10 +4,10 @@ import time
 import typing
 from collections.abc import Sequence
 from dataclasses import asdict, fields, is_dataclass
-from datetime import datetime
+from datetime import datetime, timedelta
 from enum import StrEnum
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, TypeAlias
 
 import cftime
 import numba as nb
@@ -31,6 +31,8 @@ yaml.SafeDumper.add_multi_representer(
 HEADER_SZ = 96
 HEADER_CHAR = "="
 
+RawDataSource: TypeAlias = dict[str, str | Path | list[str | Path] | bool]
+
 
 def log_the_separator() -> None:
     """Log a separator line using HEADER_CHAR repeated HEADER_SZ times."""
@@ -1963,3 +1965,262 @@ def validate_names(
         names = names[:max_to_plot]
 
     return names
+
+
+def check_dataset(
+    ds: xr.Dataset,
+    dim_names: dict[str, str],
+    var_names: dict[str, str],
+    opt_var_names: dict[str, str] | None = None,
+) -> None:
+    """Check if the dataset contains the specified variables and dimensions.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The xarray Dataset to check.
+    dim_names: Dict[str, str], optional
+        Dictionary specifying the names of dimensions in the dataset.
+    var_names: Dict[str, str]
+        Dictionary of variable names that are required in the dataset.
+    opt_var_names : Optional[Dict[str, str]], optional
+        Dictionary of optional variable names.
+        These variables are not strictly required, and the function will not raise an error if they are missing.
+        Default is None, meaning no optional variables are considered.
+
+
+    Raises
+    ------
+    ValueError
+        If the dataset does not contain the specified variables or dimensions.
+    """
+    missing_dims = [dim for dim in dim_names.values() if dim not in ds.dims]
+    if missing_dims:
+        raise ValueError(
+            f"Dataset does not contain all required dimensions. The following dimensions are missing: {missing_dims}"
+        )
+
+    missing_vars = [var for var in var_names.values() if var not in ds.data_vars]
+    if missing_vars:
+        raise ValueError(
+            f"Dataset does not contain all required variables. The following variables are missing: {missing_vars}"
+        )
+
+    if opt_var_names:
+        missing_optional_vars = [
+            var for var in opt_var_names.values() if var not in ds.data_vars
+        ]
+        if missing_optional_vars:
+            logging.warning(
+                f"Optional variables missing (but not critical): {missing_optional_vars}"
+            )
+
+
+def select_relevant_times(
+    ds: xr.Dataset,
+    time_dim: str,
+    start_time: datetime,
+    end_time: datetime | None = None,
+    climatology: bool = False,
+    allow_flex_time: bool = False,
+) -> xr.Dataset:
+    """
+    Select a subset of the dataset based on time constraints.
+
+    This function supports two main use cases:
+
+    1. **Time range selection (start_time + end_time provided):**
+       - Returns all records strictly between `start_time` and `end_time`.
+       - Ensures at least one record at or before `start_time` and one record at or
+         after `end_time` are included, even if they fall outside the strict range.
+
+    2. **Initial condition selection (start_time provided, end_time=None):**
+       - Delegates to `_select_initial_time`, which reduces the dataset to exactly one
+         time entry.
+       - If `allow_flex_time=True`, a +24-hour buffer around `start_time` is allowed,
+         and the closest timestamp is chosen.
+       - If `allow_flex_time=False`, requires an exact timestamp match.
+
+    Additional behavior:
+    - If `climatology=True`, the dataset must contain exactly 12 time steps. If valid,
+      the climatology dataset is returned without further filtering.
+    - If the dataset uses `cftime` datetime objects, these are converted to
+      `np.datetime64` before filtering.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset to filter. Must contain a valid time dimension.
+    time_dim : str
+        Name of the time dimension in `ds`.
+    start_time : datetime
+        Start time for filtering.
+    end_time : datetime or None
+        End time for filtering. If `None`, the function assumes an initial condition
+        use case and selects exactly one timestamp.
+    climatology : bool, optional
+        If True, requires exactly 12 time steps and bypasses normal filtering.
+        Defaults to False.
+    allow_flex_time : bool, optional
+        Whether to allow a +24h search window after `start_time` when `end_time`
+        is None. If False (default), requires an exact match.
+
+    Returns
+    -------
+    xr.Dataset
+        A filtered dataset containing only the selected time entries.
+
+    Raises
+    ------
+    ValueError
+        - If `climatology=True` but the dataset does not contain exactly 12 time steps.
+        - If `climatology=False` and the dataset contains integer time values.
+        - If no valid records are found within the requested range or window.
+
+    Warns
+    -----
+    UserWarning
+        - If no records exist at or before `start_time` or at or after `end_time`.
+        - If the specified time dimension does not exist in the dataset.
+
+    Notes
+    -----
+    - For initial conditions (end_time=None), see `_select_initial_time` for details
+      on strict vs. flexible selection behavior.
+    - Logs warnings instead of failing hard when boundary records are missing, and
+      defaults to using the earliest or latest available time in such cases.
+    """
+    if time_dim not in ds.variables:
+        logging.warning(
+            f"Dataset does not contain time dimension '{time_dim}'. "
+            "Please check variable naming or dataset structure."
+        )
+        return ds
+
+    time_type = get_time_type(ds[time_dim])
+
+    if climatology:
+        if len(ds[time_dim]) != 12:
+            raise ValueError(
+                f"The dataset contains {len(ds[time_dim])} time steps, but the climatology flag is set to True, which requires exactly 12 time steps."
+            )
+    else:
+        if time_type == "int":
+            raise ValueError(
+                "The dataset contains integer time values, which are only supported when the climatology flag is set to True. However, your climatology flag is set to False."
+            )
+    if time_type == "cftime":
+        ds = ds.assign_coords({time_dim: convert_cftime_to_datetime(ds[time_dim])})
+
+    if not end_time:
+        # Assume we are looking for exactly one time record for initial conditions
+        return _select_initial_time(
+            ds, time_dim, start_time, climatology, allow_flex_time
+        )
+
+    if climatology:
+        return ds
+
+    # Identify records before or at start_time
+    before_start = ds[time_dim] <= np.datetime64(start_time)
+    if before_start.any():
+        closest_before_start = ds[time_dim].where(before_start, drop=True)[-1]
+    else:
+        logging.warning(f"No records found at or before the start_time: {start_time}.")
+        closest_before_start = ds[time_dim][0]
+
+    # Identify records after or at end_time
+    after_end = ds[time_dim] >= np.datetime64(end_time)
+    if after_end.any():
+        closest_after_end = ds[time_dim].where(after_end, drop=True).min()
+    else:
+        logging.warning(f"No records found at or after the end_time: {end_time}.")
+        closest_after_end = ds[time_dim].max()
+
+    # Select records within the time range and add the closest before/after
+    within_range = (ds[time_dim] > np.datetime64(start_time)) & (
+        ds[time_dim] < np.datetime64(end_time)
+    )
+    selected_times = ds[time_dim].where(
+        within_range
+        | (ds[time_dim] == closest_before_start)
+        | (ds[time_dim] == closest_after_end),
+        drop=True,
+    )
+    ds = ds.sel({time_dim: selected_times})
+
+    return ds
+
+
+def _select_initial_time(
+    ds: xr.Dataset,
+    time_dim: str,
+    ini_time: datetime,
+    climatology: bool,
+    allow_flex_time: bool = False,
+) -> xr.Dataset:
+    """Select exactly one initial time from dataset.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The input dataset with a time dimension.
+    time_dim : str
+        Name of the time dimension.
+    ini_time : datetime
+        The desired initial time.
+    allow_flex_time : bool
+        - If True: allow a +24h window and pick the closest available timestamp.
+        - If False (default): require an exact match, otherwise raise ValueError.
+
+    Returns
+    -------
+    xr.Dataset
+        Dataset reduced to exactly one timestamp.
+
+    Raises
+    ------
+    ValueError
+        If no matching time is found (when `allow_flex_time=False`), or no entries are
+        available within the +24h window (when `allow_flex_time=True`).
+    """
+    if climatology:
+        # Convert from timedelta64[ns] to fractional days
+        ds["time"] = ds["time"] / np.timedelta64(1, "D")
+        # Interpolate from climatology for initial conditions
+        return interpolate_from_climatology(ds, time_dim, ini_time)
+
+    if allow_flex_time:
+        # Look in time range [ini_time, ini_time + 24h)
+        end_time = ini_time + timedelta(days=1)
+        times = (np.datetime64(ini_time) <= ds[time_dim]) & (
+            ds[time_dim] < np.datetime64(end_time)
+        )
+
+        if np.all(~times):
+            raise ValueError(
+                f"No time entries found between {ini_time} and {end_time}."
+            )
+
+        ds = ds.where(times, drop=True)
+        if ds.sizes[time_dim] > 1:
+            # Pick the time closest to start_time
+            ds = ds.isel({time_dim: 0})
+
+        logging.warning(
+            f"Selected time entry closest to the specified start_time in +24 hour range: {ds[time_dim].values}"
+        )
+
+    else:
+        # Strict match required
+        if not (ds[time_dim].values == np.datetime64(ini_time)).any():
+            raise ValueError(
+                f"No exact match found for initial time {ini_time}. Consider setting allow_flex_time to True."
+            )
+
+        ds = ds.sel({time_dim: np.datetime64(ini_time)})
+
+    if time_dim not in ds.dims:
+        ds = ds.expand_dims(time_dim)
+
+    return ds