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

roms_tools/setup/utils.py

@@ -4,12 +4,11 @@ import time
 import typing
 from collections.abc import Sequence
 from dataclasses import asdict, fields, is_dataclass
-from datetime import datetime, timedelta
+from datetime import datetime
 from enum import StrEnum
 from pathlib import Path
 from typing import Any, Literal, TypeAlias
 
-import cftime
 import numba as nb
 import numpy as np
 import pandas as pd
@@ -128,32 +127,6 @@ def substitute_nans_by_fillvalue(field, fill_value=0.0) -> xr.DataArray:
     return field.fillna(fill_value)
 
 
-def one_dim_fill(da: xr.DataArray, dim: str, direction="forward") -> xr.DataArray:
-    """Fill NaN values in a DataArray along a specified dimension.
-
-    Parameters
-    ----------
-    da : xr.DataArray
-        The input DataArray with NaN values to be filled, which must include the specified dimension.
-    dim : str
-        The name of the dimension along which to fill NaN values (e.g., 'depth' or 'time').
-    direction : str, optional
-        The filling direction; either "forward" to propagate non-NaN values downward or "backward" to propagate them upward.
-        Defaults to "forward".
-
-    Returns
-    -------
-    xr.DataArray
-        A new DataArray with NaN values filled in the specified direction, leaving the original data unchanged.
-    """
-    if dim in da.dims:
-        if direction == "forward":
-            return da.ffill(dim=dim)
-        elif direction == "backward":
-            return da.bfill(dim=dim)
-    return da
-
-
 def assign_dates_to_climatology(ds: xr.Dataset, time_dim: str) -> xr.Dataset:
     """Assigns climatology dates to the dataset's time dimension.
 
@@ -182,240 +155,6 @@ def assign_dates_to_climatology(ds: xr.Dataset, time_dim: str) -> xr.Dataset:
     return ds
 
 
-def interpolate_cyclic_time(
-    data_array: xr.DataArray,
-    time_dim_name: str,
-    day_of_year: int | float | np.ndarray | xr.DataArray | Sequence[int | float],
-) -> xr.DataArray:
-    """Interpolates a DataArray cyclically across the start and end of the year.
-
-    This function extends the data cyclically by appending the last time step
-    (shifted back by one year) at the beginning and the first time step
-    (shifted forward by one year) at the end. It then performs linear interpolation
-    to match the specified `day_of_year` values.
-
-    Parameters
-    ----------
-    data_array : xr.DataArray
-        The input data array containing a time-like dimension.
-    time_dim_name : str
-        The name of the time dimension in the dataset.
-    day_of_year : Union[int, float, np.ndarray, xr.DataArray, Sequence[Union[int, float]]]
-        The target day(s) of the year for interpolation. This can be:
-        - A single integer or float representing the day of the year.
-        - A NumPy array or xarray DataArray containing multiple days.
-        - A list or tuple of integers or floats for multiple target days.
-
-    Returns
-    -------
-    xr.DataArray
-        The interpolated DataArray, ensuring cyclic continuity across year boundaries.
-
-    Notes
-    -----
-    - This function is useful for interpolating climatological data, where the time axis
-      represents a repeating annual cycle.
-    - The `day_of_year` values should be within the range [1, 365] or [1, 366] for leap years.
-    """
-    # Concatenate across the beginning and end of the year
-    time_concat = xr.concat(
-        [
-            data_array[time_dim_name][-1] - 365.25,  # Shift last time backward
-            data_array[time_dim_name],
-            data_array[time_dim_name][0] + 365.25,  # Shift first time forward
-        ],
-        dim=time_dim_name,
-    )
-
-    data_array_concat = xr.concat(
-        [
-            data_array.isel(
-                **{time_dim_name: -1}
-            ),  # Append last value at the beginning
-            data_array,
-            data_array.isel(**{time_dim_name: 0}),  # Append first value at the end
-        ],
-        dim=time_dim_name,
-    )
-    data_array_concat[time_dim_name] = time_concat
-
-    # Interpolate to specified times
-    data_array_interpolated = data_array_concat.interp(
-        **{time_dim_name: day_of_year}, method="linear"
-    )
-
-    return data_array_interpolated
-
-
-def interpolate_from_climatology(
-    field: xr.DataArray | xr.Dataset,
-    time_dim_name: str,
-    time: xr.DataArray | pd.DatetimeIndex,
-) -> xr.DataArray | xr.Dataset:
-    """Interpolates a climatological field to specified time points.
-
-    This function interpolates the input `field` based on `day_of_year` values
-    extracted from the provided `time` points. If `field` is an `xarray.Dataset`,
-    interpolation is applied to all its data variables individually.
-
-    Parameters
-    ----------
-    field : xarray.DataArray or xarray.Dataset
-        The input field to be interpolated.
-        - If `field` is an `xarray.DataArray`, it must have a time dimension identified by `time_dim_name`.
-        - If `field` is an `xarray.Dataset`, all variables within the dataset are interpolated along `time_dim_name`.
-        The time dimension is assumed to represent `day_of_year` for climatological purposes.
-    time_dim_name : str
-        The name of the time dimension in `field`. This dimension is used for interpolation.
-    time : xarray.DataArray or pandas.DatetimeIndex
-        The target time points for interpolation. These are internally converted to `day_of_year`
-        before performing interpolation.
-
-    Returns
-    -------
-    xarray.DataArray or xarray.Dataset
-        The interpolated field, maintaining the same type (`xarray.DataArray` or `xarray.Dataset`)
-        but aligned to the specified `time` values.
-
-    Notes
-    -----
-    - This function assumes that `field` represents a climatological dataset, where time is expressed as `day_of_year` (1-365).
-    - The `time` input is automatically converted to `day_of_year`, so manual conversion is not required before calling this function.
-    """
-
-    def interpolate_single_field(data_array: xr.DataArray) -> xr.DataArray:
-        if isinstance(time, xr.DataArray):
-            # Extract day of year from xarray.DataArray
-            day_of_year = time.dt.dayofyear
-        else:
-            if np.size(time) == 1:
-                # Convert single datetime64 object to pandas.Timestamp
-                date = pd.Timestamp(time)
-                day_of_year = (
-                    date.dayofyear
-                    + (date.hour / 24)
-                    + (date.minute / 1440)
-                    + (date.second / 86400)
-                )
-            else:
-                # Convert each datetime64 object in the array to pandas.Timestamp and compute fractional day of year
-                day_of_year = np.array(
-                    [
-                        pd.Timestamp(t).dayofyear
-                        + (pd.Timestamp(t).hour / 24)
-                        + (pd.Timestamp(t).minute / 1440)
-                        + (pd.Timestamp(t).second / 86400)
-                        for t in time
-                    ]
-                )
-
-        data_array_interpolated = interpolate_cyclic_time(
-            data_array, time_dim_name, day_of_year
-        )
-
-        if np.size(time) == 1:
-            data_array_interpolated = data_array_interpolated.expand_dims(
-                {time_dim_name: 1}
-            )
-        return data_array_interpolated
-
-    if isinstance(field, xr.DataArray):
-        return interpolate_single_field(field)
-    elif isinstance(field, xr.Dataset):
-        interpolated_data_vars = {
-            var: interpolate_single_field(data_array)
-            for var, data_array in field.data_vars.items()
-        }
-        return xr.Dataset(interpolated_data_vars, attrs=field.attrs)
-
-    else:
-        raise TypeError("Input 'field' must be an xarray.DataArray or xarray.Dataset.")
-
-
-def get_time_type(data_array: xr.DataArray) -> str:
-    """Determines the type of time values in the xarray DataArray.
-
-    Parameters
-    ----------
-    data_array : xr.DataArray
-        The xarray DataArray to be checked for time data types.
-
-    Returns
-    -------
-    str
-        A string indicating the type of the time data: 'cftime', 'datetime', or 'int'.
-
-    Raises
-    ------
-    TypeError
-        If the values in the DataArray are not of type numpy.ndarray or list.
-    """
-    # List of cftime datetime types
-    cftime_types = (
-        cftime.DatetimeNoLeap,
-        cftime.DatetimeJulian,
-        cftime.DatetimeGregorian,
-        cftime.Datetime360Day,
-        cftime.DatetimeProlepticGregorian,
-    )
-
-    # Check if any of the coordinate values are of cftime, datetime, or integer type
-    if isinstance(data_array.values, np.ndarray | list):
-        # Check if the data type is numpy datetime64, indicating standard datetime objects
-        if data_array.values.dtype == "datetime64[ns]":
-            return "datetime"
-
-        # Check if any values in the array are instances of cftime types
-        if any(isinstance(value, cftime_types) for value in data_array.values):
-            return "cftime"
-
-        # Check if all values are of integer type (e.g., for indices or time steps)
-        if np.issubdtype(data_array.values.dtype, np.integer):
-            return "int"
-
-        # If none of the above conditions are met, raise a ValueError
-        raise ValueError("Unsupported data type for time values in input dataset.")
-
-    # Handle unexpected types
-    raise TypeError("DataArray values must be of type numpy.ndarray or list.")
-
-
-def convert_cftime_to_datetime(data_array: np.ndarray) -> np.ndarray:
-    """Converts cftime datetime objects to numpy datetime64 objects in a numpy ndarray.
-
-    Parameters
-    ----------
-    data_array : np.ndarray
-        The numpy ndarray containing cftime datetime objects to be converted.
-
-    Returns
-    -------
-    np.ndarray
-        The ndarray with cftime datetimes converted to numpy datetime64 objects.
-
-    Notes
-    -----
-    This function is intended to be used with numpy ndarrays. If you need to convert
-    cftime datetime objects in an xarray.DataArray, please use the appropriate function
-    to handle xarray.DataArray conversions.
-    """
-    # List of cftime datetime types
-    cftime_types = (
-        cftime.DatetimeNoLeap,
-        cftime.DatetimeJulian,
-        cftime.DatetimeGregorian,
-    )
-
-    # Define a conversion function for cftime to numpy datetime64
-    def convert_datetime(dt):
-        if isinstance(dt, cftime_types):
-            # Convert to ISO format and then to nanosecond precision
-            return np.datetime64(dt.isoformat(), "ns")
-        return np.datetime64(dt, "ns")
-
-    return np.vectorize(convert_datetime)(data_array)
-
-
 def get_variable_metadata():
     """Retrieves metadata for commonly used variables in the dataset.
 
@@ -1861,38 +1600,6 @@ def get_boundary_coords():
     return bdry_coords
 
 
-def wrap_longitudes(grid_ds, straddle):
-    """Adjusts longitude values in a dataset to handle dateline crossing.
-
-    Parameters
-    ----------
-    grid_ds : xr.Dataset
-        The dataset containing longitude variables to adjust.
-    straddle : bool
-        If True, adjusts longitudes to the range [-180, 180] for datasets
-        that straddle the dateline. If False, adjusts longitudes to the
-        range [0, 360].
-
-    Returns
-    -------
-    xr.Dataset
-        The dataset with adjusted longitude values.
-    """
-    for lon_dim in ["lon_rho", "lon_u", "lon_v"]:
-        if straddle:
-            grid_ds[lon_dim] = xr.where(
-                grid_ds[lon_dim] > 180,
-                grid_ds[lon_dim] - 360,
-                grid_ds[lon_dim],
-            )
-        else:
-            grid_ds[lon_dim] = xr.where(
-                grid_ds[lon_dim] < 0, grid_ds[lon_dim] + 360, grid_ds[lon_dim]
-            )
-
-    return grid_ds
-
-
 def to_float(val):
     """Convert a value or list of values to float.
 
@@ -1967,260 +1674,101 @@ def validate_names(
     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.
+def check_and_set_boundaries(
+    boundaries: dict[str, bool] | None,
+    mask: xr.DataArray,
+) -> dict[str, bool]:
+    """
+    Validate and finalize the `boundaries` dictionary.
 
     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.
+    boundaries : dict[str, bool] or None
+        User-supplied dictionary controlling which boundaries are active.
+        Keys may include any subset of {"south", "east", "north", "west"}.
+        Missing keys will be filled from mask-based defaults.
+        If None, all boundaries are inferred from the land mask.
 
+    mask : xr.DataArray
+        2D land/sea mask on rho-points. Used to determine which boundaries
+        contain at least one ocean point.
 
-    Raises
-    ------
-    ValueError
-        If the dataset does not contain the specified variables or dimensions.
+    Returns
+    -------
+    dict[str, bool]
+        Completed and validated boundary configuration.
     """
-    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}"
+    valid_keys = {"south", "east", "north", "west"}
+
+    # --------------------------------------------
+    # Case 1: boundaries not provided → infer them
+    # --------------------------------------------
+    if boundaries is None:
+        inferred = _infer_valid_boundaries_from_mask(mask)
+        logging.info(f"No `boundaries` provided. Using mask-based defaults: {inferred}")
+        return inferred
+
+    # --------------------------------------------
+    # Case 2: boundaries provided → validate
+    # --------------------------------------------
+    if not isinstance(boundaries, dict):
+        raise TypeError(
+            "`boundaries` must be a dict mapping boundary names to booleans."
         )
 
-    missing_vars = [var for var in var_names.values() if var not in ds.data_vars]
-    if missing_vars:
+    # Unknown keys?
+    unknown_keys = set(boundaries) - valid_keys
+    if unknown_keys:
         raise ValueError(
-            f"Dataset does not contain all required variables. The following variables are missing: {missing_vars}"
+            f"`boundaries` contains invalid keys: {unknown_keys}. "
+            "Allowed keys are: 'south', 'east', 'north', 'west'."
         )
 
-    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}"
+    # Type-check provided values
+    for key, val in boundaries.items():
+        if not isinstance(val, bool):
+            raise TypeError(f"Boundary '{key}' must be a boolean.")
+
+    # Fill missing boundaries using defaults
+    inferred_defaults = _infer_valid_boundaries_from_mask(mask)
+    completed = boundaries.copy()
+
+    for key in valid_keys:
+        if key not in completed:
+            completed[key] = inferred_defaults[key]
+            logging.info(
+                f"`boundaries[{key!r}]` not provided — defaulting to "
+                f"{inferred_defaults[key]}"
             )
 
+    logging.info(f"Using boundary configuration: {completed}")
+    return completed
 
-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.
+def _infer_valid_boundaries_from_mask(mask: xr.DataArray) -> dict[str, bool]:
     """
-    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
+    Determine which grid boundaries contain at least one ocean point.
 
-
-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.
+    Any boundary consisting entirely of land is considered inactive.
 
     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.
+    mask : xr.DataArray
+        2D mask array on rho-points where 1 = ocean, 0 = land.
 
     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`).
+    dict[str, bool]
+        Boolean availability for {south, east, north, west}.
     """
-    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}."
-            )
+    bdry_coords = get_boundary_coords()
+    boundaries = {}
 
-        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})
+    for direction in ["south", "east", "north", "west"]:
+        coords = bdry_coords["rho"][direction]
+        bdry_mask = mask.isel(**coords)
 
-        logging.warning(
-            f"Selected time entry closest to the specified start_time in +24 hour range: {ds[time_dim].values}"
-        )
+        # Boundary is valid if ANY ocean point exists
+        boundaries[direction] = bool(bdry_mask.values.any())
 
-    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
+    return boundaries

roms_tools/tests/test_analysis/test_cdr_ensemble.py

@@ -86,18 +86,16 @@ def test_extract_efficiency(create_member_ds: xr.Dataset) -> None:
     assert eff_rel.time.attrs.get("long_name") == "time since release start"
 
 
-def test_extract_efficiency_missing_abs_time() -> None:
-    """Test that _extract_efficiency raises an error if 'abs_time' is missing."""
+def test_extract_efficiency_missing_time() -> None:
+    """Test that _extract_efficiency raises an error if 'time' is missing."""
     times = np.array(["2000-01-01", "2000-01-02"], dtype="datetime64[ns]")
     ds = xr.Dataset(
         {"cdr_efficiency": ("time", [0.1, 0.2])},
-        coords={"time": times},  # Note: no 'abs_time' coordinate
+        coords={"abs_time": times},  # Note: no 'time' coordinate
     )
 
     ens = Ensemble.__new__(Ensemble)
-    with pytest.raises(
-        ValueError, match="Dataset must contain an 'abs_time' coordinate."
-    ):
+    with pytest.raises(ValueError, match="Dataset must contain a 'time' coordinate."):
         ens._extract_efficiency(ds)