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

roms_tools/setup/utils.py

@@ -1,11 +1,13 @@
 import importlib.metadata
 import logging
+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
+from typing import Any, Literal, TypeAlias
 
 import cftime
 import numba as nb
@@ -18,11 +20,55 @@ from pydantic import BaseModel
 from roms_tools.constants import R_EARTH
 from roms_tools.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
 
+if typing.TYPE_CHECKING:
+    from roms_tools.setup.grid import Grid
+
 yaml.SafeDumper.add_multi_representer(
     StrEnum,
     yaml.representer.SafeRepresenter.represent_str,
 )
 
+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."""
+    logging.info(HEADER_CHAR * HEADER_SZ)
+
+
+class Timed:
+    """Context manager to time a block and log messages."""
+
+    def __init__(self, message: str = "", verbose: bool = True) -> None:
+        """
+        Initialize the context manager.
+
+        Parameters
+        ----------
+        message : str, optional
+            A log message printed at the start of the block (default: "").
+        verbose : bool, optional
+            Whether to log timing information (default: True).
+        """
+        self.message = message
+        self.verbose = verbose
+        self.start: float | None = None
+
+    def __enter__(self) -> "Timed":
+        if self.verbose:
+            self.start = time.time()
+            if self.message:
+                logging.info(self.message)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        if self.verbose and self.start is not None:
+            logging.info(f"Total time: {time.time() - self.start:.3f} seconds")
+            log_the_separator()
+
 
 def nan_check(field, mask, error_message=None) -> None:
     """Checks for NaN values at wet points in the field.
@@ -437,157 +483,193 @@ def get_variable_metadata():
             "long_name": "dissolved inorganic phosphate",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "NO3": {
             "long_name": "dissolved inorganic nitrate",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "SiO3": {
             "long_name": "dissolved inorganic silicate",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "NH4": {
             "long_name": "dissolved ammonia",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "Fe": {
             "long_name": "dissolved inorganic iron",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "Lig": {
             "long_name": "iron binding ligand",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "O2": {
             "long_name": "dissolved oxygen",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "DIC": {
             "long_name": "dissolved inorganic carbon",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "DIC_ALT_CO2": {
             "long_name": "dissolved inorganic carbon, alternative CO2",
             "units": "mmol/m^3",
             "flux_units": "meq/s",
+            "integrated_units": "meq",
+        },
+        "ALK": {
+            "long_name": "alkalinity",
+            "units": "meq/m^3",
+            "flux_units": "meq/s",
+            "integrated_units": "meq",
         },
-        "ALK": {"long_name": "alkalinity", "units": "meq/m^3", "flux_units": "meq/s"},
         "ALK_ALT_CO2": {
             "long_name": "alkalinity, alternative CO2",
             "units": "meq/m^3",
             "flux_units": "meq/s",
+            "integrated_units": "meq",
         },
         "DOC": {
             "long_name": "dissolved organic carbon",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "DON": {
             "long_name": "dissolved organic nitrogen",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "DOP": {
             "long_name": "dissolved organic phosphorus",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "DOCr": {
             "long_name": "refractory dissolved organic carbon",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "DONr": {
             "long_name": "refractory dissolved organic nitrogen",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "DOPr": {
             "long_name": "refractory dissolved organic phosphorus",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "zooC": {
             "long_name": "zooplankton carbon",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "spChl": {
             "long_name": "small phytoplankton chlorophyll",
             "units": "mg/m^3",
             "flux_units": "mg/s",
+            "integrated_units": "mg",
         },
         "spC": {
             "long_name": "small phytoplankton carbon",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "spP": {
             "long_name": "small phytoplankton phosphorous",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "spFe": {
             "long_name": "small phytoplankton iron",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "spCaCO3": {
             "long_name": "small phytoplankton CaCO3",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "diatChl": {
             "long_name": "diatom chloropyll",
             "units": "mg/m^3",
             "flux_units": "mg/s",
+            "integrated_units": "mg",
         },
         "diatC": {
             "long_name": "diatom carbon",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "diatP": {
             "long_name": "diatom phosphorus",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "diatFe": {
             "long_name": "diatom iron",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "diatSi": {
             "long_name": "diatom silicate",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "diazChl": {
             "long_name": "diazotroph chloropyll",
             "units": "mg/m^3",
             "flux_units": "mg/s",
+            "integrated_units": "mg",
         },
         "diazC": {
             "long_name": "diazotroph carbon",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "diazP": {
             "long_name": "diazotroph phosphorus",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "diazFe": {
             "long_name": "diazotroph iron",
             "units": "mmol/m^3",
             "flux_units": "mmol/s",
+            "integrated_units": "mmol",
         },
         "pco2_air": {"long_name": "atmospheric pCO2", "units": "ppmv"},
         "pco2_air_alt": {
@@ -720,7 +802,10 @@ def compute_missing_surface_bgc_variables(bgc_data):
     return bgc_data
 
 
-def get_tracer_metadata_dict(include_bgc=True, with_flux_units=False):
+def get_tracer_metadata_dict(
+    include_bgc: bool = True,
+    unit_type: Literal["concentration", "flux", "integrated"] = "concentration",
+):
     """Generate a dictionary containing metadata for model tracers.
 
     The returned dictionary maps tracer names to their associated units and long names.
@@ -732,9 +817,8 @@ def get_tracer_metadata_dict(include_bgc=True, with_flux_units=False):
         If True (default), includes biogeochemical tracers in the output.
         If False, returns only physical tracers (e.g., temperature, salinity).
 
-    with_flux_units : bool, optional
-        If True, uses units appropriate for tracer fluxes (e.g., mmol/s).
-        If False (default), uses units appropriate for tracer concentrations (e.g., mmol/m³).
+    unit_type : str
+        One of "concentration" (default), "flux", or "integrated".
 
     Returns
     -------
@@ -784,15 +868,19 @@ def get_tracer_metadata_dict(include_bgc=True, with_flux_units=False):
 
     metadata = get_variable_metadata()
 
-    tracer_dict = {
-        tracer: {
-            "units": metadata[tracer]["flux_units"]
-            if with_flux_units
-            else metadata[tracer]["units"],
+    tracer_dict = {}
+    for tracer in tracer_names:
+        if unit_type == "flux":
+            unit = metadata[tracer]["flux_units"]
+        elif unit_type == "integrated":
+            unit = metadata[tracer].get("integrated_units", None)
+        else:  # default to concentration units
+            unit = metadata[tracer]["units"]
+
+        tracer_dict[tracer] = {
+            "units": unit,
             "long_name": metadata[tracer]["long_name"],
         }
-        for tracer in tracer_names
-    }
 
     return tracer_dict
 
@@ -819,7 +907,8 @@ def add_tracer_metadata_to_ds(ds, include_bgc=True, with_flux_units=False):
     xarray.Dataset
         The dataset with added tracer metadata.
     """
-    tracer_dict = get_tracer_metadata_dict(include_bgc, with_flux_units)
+    unit_type = "flux" if with_flux_units else "concentration"
+    tracer_dict = get_tracer_metadata_dict(include_bgc, unit_type=unit_type)
 
     tracer_names = list(tracer_dict.keys())
     tracer_units = [tracer_dict[tracer]["units"] for tracer in tracer_names]
@@ -1045,22 +1134,47 @@ def group_by_year(ds, filepath):
     return dataset_list, output_filenames
 
 
-def get_target_coords(grid, use_coarse_grid=False):
-    """Retrieves longitude and latitude coordinates from the grid, adjusting them based
-    on longitude range.
+def get_target_coords(
+    grid: "Grid", use_coarse_grid: bool = False
+) -> dict[str, xr.DataArray | bool | None]:
+    """
+    Retrieve longitude, latitude, and auxiliary grid coordinates, adjusting for
+    longitude ranges and coarse grid usage.
 
     Parameters
     ----------
     grid : Grid
-        Object representing the grid information used for the model.
+        Grid object.
     use_coarse_grid : bool, optional
-        Use coarse grid data if True. Defaults to False.
+        If True, use the coarse grid variables (`lat_coarse`, `lon_coarse`, etc.)
+        instead of the native grid. Defaults to False.
 
     Returns
     -------
-    dict
-        Dictionary containing the longitude, latitude, angle and mask arrays, along with a boolean indicating
-        if the grid straddles the meridian.
+    dict[str, xr.DataArray | bool | None]
+        Dictionary containing the following keys:
+
+        - `"lat"` : xr.DataArray
+            Latitude at rho points.
+        - `"lon"` : xr.DataArray
+            Longitude at rho points, adjusted to -180 to 180 or 0 to 360 range.
+        - `"lat_psi"` : xr.DataArray | None
+            Latitude at psi points, if available.
+        - `"lon_psi"` : xr.DataArray | None
+            Longitude at psi points, if available.
+        - `"angle"` : xr.DataArray
+            Grid rotation angle.
+        - `"mask"` : xr.DataArray | None
+            Land/sea mask at rho points.
+        - `"straddle"` : bool
+            True if the grid crosses the Greenwich meridian, False otherwise.
+
+    Notes
+    -----
+    - If `grid.straddle` is False and the ROMS domain lies more than 5° from
+      the Greenwich meridian, longitudes are adjusted to 0-360 range.
+    - Renaming of coarse grid dimensions is applied to match the rho-point
+      naming convention (`eta_rho`, `xi_rho`) for compatibility with ROMS-Tools.
     """
     # Select grid variables based on whether the coarse grid is used
     if use_coarse_grid:
@@ -1851,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