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

roms_tools/setup/river_forcing.py

@@ -3,6 +3,7 @@ from collections import defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
+from typing import TypeAlias
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -18,6 +19,7 @@ from roms_tools.plot import (
 )
 from roms_tools.setup.datasets import (
     DaiRiverDataset,
+    RawDataSource,
     get_indices_of_nearest_grid_cell_for_rivers,
 )
 from roms_tools.setup.utils import (
@@ -38,6 +40,8 @@ from roms_tools.utils import save_datasets
 INCLUDE_ALL_RIVER_NAMES = "all"
 MAX_RIVERS_TO_PLOT = 20  # must be <= MAX_DISTINCT_COLORS
 
+TRiverIndex: TypeAlias = dict[tuple[int, int], list[str]]
+
 
 @dataclass(kw_only=True)
 class RiverForcing:
@@ -51,7 +55,7 @@ class RiverForcing:
         Start time of the desired river forcing data.
     end_time : datetime
         End time of the desired river forcing data.
-    source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool], optional
+    source : RawDataSource, optional
         Dictionary specifying the source of the river forcing data. Keys include:
 
           - "name" (str): Name of the data source (e.g., "DAI").
@@ -75,7 +79,7 @@ class RiverForcing:
         Whether to include BGC tracers. Defaults to `False`.
     model_reference_date : datetime, optional
         Reference date for the ROMS simulation. Default is January 1, 2000.
-    indices : dict[str, list[tuple]], optional
+    indices : dict[str, list[tuple[int, int]]], optional
         A dictionary specifying the river indices for each river to be included in the river forcing. This parameter is optional. If not provided,
         the river indices will be automatically determined based on the grid and the source dataset. If provided, it allows for explicit specification
         of river locations. The dictionary structure consists of river names as keys, and each value is a list of tuples. Each tuple represents
@@ -101,7 +105,7 @@ class RiverForcing:
     """Start time of the desired river forcing data."""
     end_time: datetime
     """End time of the desired river forcing data."""
-    source: dict[str, str | Path | list[str | Path]] = None
+    source: RawDataSource | None = None
     """Dictionary specifying the source of the river forcing data."""
     convert_to_climatology: str = "if_any_missing"
     """Determines when to compute climatology for river forcing."""
@@ -110,7 +114,7 @@ class RiverForcing:
     model_reference_date: datetime = datetime(2000, 1, 1)
     """Reference date for the ROMS simulation."""
 
-    indices: dict[str, dict[str, int | list[int]]] | None = None
+    indices: dict[str, list[tuple[int, int]]] | None = None
     """A dictionary of river indices.
 
     If not provided during initialization, it will be automatically determined based on
@@ -462,7 +466,7 @@ class RiverForcing:
 
         return ds_updated
 
-    def _get_overlapping_rivers(self) -> dict[tuple[int, int], list[str]]:
+    def _get_overlapping_rivers(self) -> TRiverIndex:
         """Identify grid cells shared by multiple rivers.
 
         Scans through the river indices and finds all grid cell indices
@@ -474,7 +478,10 @@ class RiverForcing:
             A dictionary mapping grid cell indices (eta_rho, xi_rho) to a list
             of river names that overlap at that grid cell.
         """
-        index_to_rivers = defaultdict(list)
+        if self.indices is None:
+            return {}
+
+        index_to_rivers: TRiverIndex = defaultdict(list)
 
         # Collect all index pairs used by multiple rivers
         for river_name, index_list in self.indices.items():
@@ -520,6 +527,9 @@ class RiverForcing:
             The volume-weighted tracer concentration at the overlapping grid cell,
             as a new 1-entry DataArray with updated coordinates.
         """
+        if self.indices is None:
+            self.indices = {}
+
         new_name = f"overlap_{i}"
         self.indices[new_name] = [idx_pair]
 
@@ -578,7 +588,7 @@ class RiverForcing:
         return combined_river_volume, combined_river_tracer
 
     def _reduce_river_volumes(
-        self, ds: xr.Dataset, overlapping_rivers: dict[tuple[int, int], list[str]]
+        self, ds: xr.Dataset, overlapping_rivers: TRiverIndex
     ) -> xr.Dataset:
         """Reduce river volumes for rivers contributing to overlapping grid cells.
 
@@ -595,8 +605,14 @@ class RiverForcing:
         ds : xr.Dataset
             Updated dataset with reduced river volumes.
         """
+        if self.indices is None:
+            raise ValueError(
+                "`self.indices` must be set before calling _reduce_river_volumes"
+            )
+
         # Count number of overlaps for each river
-        river_overlap_count = defaultdict(int)
+        river_overlap_count: dict[str, int] = defaultdict(int)
+
         for rivers in overlapping_rivers.values():
             for name in rivers:
                 river_overlap_count[name] += 1
@@ -671,13 +687,15 @@ class RiverForcing:
         Warning
             If NaN values are found in any of the dataset variables, a warning message is logged.
         """
-        for var_name in ds.data_vars:
-            da = ds[var_name]
-            if da.isnull().any().values:
+        var_name = "river_volume"
+        da = ds[var_name]
+        if da.isnull().any().values:
+            logging.warning(
+                f"NaNs detected in '{var_name}' and set to zero. This may indicate missing river data and affect model accuracy. "
+            )
+            if not self.climatology:
                 logging.warning(
-                    f"NaN values detected in the '{var_name}' field. These values are being set to zero. "
-                    "This may indicate missing river data, which could affect model accuracy. Consider setting "
-                    "`convert_to_climatology = 'if_any_missing'` to automatically fill missing values with climatological data."
+                    "Consider `convert_to_climatology='if_any_missing'` to fill missing values with climatological data."
                 )
 
     def plot_locations(self, river_names: list[str] | str = INCLUDE_ALL_RIVER_NAMES):
@@ -691,7 +709,8 @@ class RiverForcing:
             Defaults to "all".
 
         """
-        valid_river_names = list(self.indices.keys())
+        valid_river_names = list(self.indices or [])
+
         river_names = _validate_river_names(river_names, valid_river_names)
         if len(valid_river_names) > MAX_DISTINCT_COLORS:
             colors = assign_category_colors(river_names)
@@ -806,7 +825,8 @@ class RiverForcing:
             Defaults to "all".
 
         """
-        valid_river_names = list(self.indices.keys())
+        valid_river_names = list(self.indices or [])
+
         river_names = _validate_river_names(river_names, valid_river_names)
         if len(valid_river_names) > MAX_DISTINCT_COLORS:
             colors = assign_category_colors(river_names)
@@ -900,11 +920,17 @@ class RiverForcing:
         """
         forcing_dict = to_dict(self, exclude=["climatology"])
 
-        # Convert indices format
-        indices_data = forcing_dict["RiverForcing"]["indices"]
-        serialized_indices = {"_convention": "eta_rho, xi_rho"}
+        indices_data = forcing_dict.get("RiverForcing", {}).get("indices")
+        if not indices_data:
+            # If no indices, just write the dict as is
+            write_to_yaml(forcing_dict, filepath)
+            return
+
+        # Convert tuple indices to string format for YAML
+        serialized_indices: dict[str, str | list[str]] = {}
         for key, value in indices_data.items():
             serialized_indices[key] = [f"{tup[0]}, {tup[1]}" for tup in value]
+        serialized_indices["_convention"] = "eta_rho, xi_rho"
 
         # Remove keys starting with "overlap_"
         filtered_indices = {

roms_tools/setup/surface_forcing.py

@@ -16,6 +16,7 @@ from roms_tools.setup.datasets import (
     ERA5ARCODataset,
     ERA5Correction,
     ERA5Dataset,
+    RawDataSource,
     UnifiedBGCSurfaceDataset,
 )
 from roms_tools.setup.utils import (
@@ -56,7 +57,7 @@ class SurfaceForcing:
         The end time of the desired surface forcing data. This time is used to filter the dataset
         to include only records on or before this time, with a single record at or after this time.
         If no time filtering is desired, set it to None. Default is None.
-    source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool]
+    source : RawDataSource
         Dictionary specifying the source of the surface forcing data. Keys include:
 
           - "name" (str): Name of the data source. Currently supported: "ERA5"
@@ -116,7 +117,7 @@ class SurfaceForcing:
     """The start time of the desired surface forcing data."""
     end_time: datetime | None = None
     """The end time of the desired surface forcing data."""
-    source: dict[str, str | Path | list[str | Path]]
+    source: RawDataSource
     """Dictionary specifying the source of the surface forcing data."""
     type: str = "physics"
     """Specifies the type of forcing data ("physics", "bgc")."""
@@ -169,7 +170,6 @@ class SurfaceForcing:
 
         data.choose_subdomain(
             target_coords,
-            buffer_points=20,  # lateral fill needs some buffer from data margin
         )
         # Enforce double precision to ensure reproducibility
         data.convert_to_float64()
@@ -447,9 +447,7 @@ class SurfaceForcing:
             "lat": data.ds[data.dim_names["latitude"]],
             "lon": data.ds[data.dim_names["longitude"]],
         }
-        correction_data.choose_subdomain(
-            coords_correction, straddle=self.target_coords["straddle"]
-        )
+        correction_data.match_subdomain(coords_correction)
         correction_data.ds["mask"] = data.ds["mask"]  # use mask from ERA5 data
         correction_data.ds["time"] = correction_data.ds["time"].dt.days
 

roms_tools/setup/tides.py

@@ -80,7 +80,7 @@ class TidalForcing:
 
     grid: Grid
     """Object representing the grid information."""
-    source: dict[str, str | Path | list[str | Path]]
+    source: dict[str, str | Path | dict[str, str | Path]]
     """Dictionary specifying the source of the tidal data."""
     ntides: int = 10
     """Number of constituents to consider."""
@@ -105,7 +105,6 @@ class TidalForcing:
             if key != "omega":
                 data.choose_subdomain(
                     target_coords,
-                    buffer_points=20,
                 )
                 # Enforce double precision to ensure reproducibility
                 data.convert_to_float64()

roms_tools/setup/topography.py

@@ -12,7 +12,7 @@ from roms_tools.setup.datasets import ETOPO5Dataset, SRTM15Dataset
 from roms_tools.setup.utils import handle_boundaries
 
 
-def _add_topography(
+def add_topography(
     ds,
     target_coords,
     topography_source,
@@ -241,7 +241,7 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
         rmax_log = 0.0
 
     # Apply hmin threshold
-    h = _clip_depth(h, hmin)
+    h = clip_depth(h, hmin)
 
     # Perform logarithmic transformation of the height field
     h_log = np.log(h / hmin)
@@ -324,7 +324,7 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
         h = hmin * np.exp(h_log)
 
         # Apply hmin threshold again
-        h = _clip_depth(h, hmin)
+        h = clip_depth(h, hmin)
 
         # Compute maximum slope parameter r
         r_eta, r_xi = _compute_rfactor(h)
@@ -335,7 +335,7 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
     return h
 
 
-def _clip_depth(h: xr.DataArray, hmin: float) -> xr.DataArray:
+def clip_depth(h: xr.DataArray, hmin: float) -> xr.DataArray:
     """Ensures that depth values do not fall below a minimum threshold.
 
     This function replaces all depth values in `h` that are less than `hmin` with `hmin`,

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 enum import StrEnum
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 import cftime
 import numba as nb
@@ -18,11 +20,53 @@ 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 = "="
+
+
+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 +481,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 +800,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 +815,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 +866,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 +905,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 +1132,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: