roms-tools 3.1.2__py3-none-any.whl → 3.3.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
@@ -16,11 +17,12 @@ from roms_tools.plot import (
     plot_2d_horizontal_field,
     plot_location,
 )
-from roms_tools.setup.datasets import (
+from roms_tools.setup.river_datasets import (
     DaiRiverDataset,
     get_indices_of_nearest_grid_cell_for_rivers,
 )
 from roms_tools.setup.utils import (
+    RawDataSource,
     add_time_info_to_ds,
     add_tracer_metadata_to_ds,
     from_yaml,
@@ -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

@@ -10,15 +10,16 @@ import xarray as xr
 from roms_tools import Grid
 from roms_tools.plot import plot
 from roms_tools.regrid import LateralRegridToROMS
-from roms_tools.setup.datasets import (
+from roms_tools.setup.lat_lon_datasets import (
     CESMBGCSurfaceForcingDataset,
-    Dataset,
     ERA5ARCODataset,
     ERA5Correction,
     ERA5Dataset,
+    LatLonDataset,
     UnifiedBGCSurfaceDataset,
 )
 from roms_tools.setup.utils import (
+    RawDataSource,
     add_time_info_to_ds,
     compute_missing_surface_bgc_variables,
     from_yaml,
@@ -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()
@@ -420,7 +420,7 @@ class SurfaceForcing:
         self.variable_info = variable_info
 
     def _apply_radiation_correction(
-        self, radiation: xr.DataArray, data: Dataset
+        self, radiation: xr.DataArray, data: LatLonDataset
     ) -> xr.DataArray:
         """Apply a climatological correction to shortwave radiation.
 
@@ -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

@@ -9,7 +9,7 @@ import xarray as xr
 from roms_tools import Grid
 from roms_tools.plot import plot
 from roms_tools.regrid import LateralRegridToROMS
-from roms_tools.setup.datasets import TPXOManager
+from roms_tools.setup.lat_lon_datasets import TPXOManager
 from roms_tools.setup.utils import (
     from_yaml,
     get_target_coords,
@@ -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

@@ -8,16 +8,16 @@ import numpy as np
 import xarray as xr
 
 from roms_tools.regrid import LateralRegridToROMS
-from roms_tools.setup.datasets import ETOPO5Dataset, SRTM15Dataset
+from roms_tools.setup.lat_lon_datasets import ETOPO5Dataset, SRTM15Dataset
 from roms_tools.setup.utils import handle_boundaries
 
 
-def _add_topography(
+def add_topography(
     ds,
     target_coords,
     topography_source,
     hmin,
-    smooth_factor=6.0,
+    smooth_factor=5.0,
     rmax=0.2,
     verbose=False,
 ) -> xr.Dataset:
@@ -77,10 +77,8 @@ def _add_topography(
     # smooth topography locally to satisfy r < rmax
     if verbose:
         start_time = time.time()
-    # inserting hraw * mask_rho into this function eliminates any inconsistencies between
-    # the land according to the topography and the land according to the mask; land points
-    # will always be set to hmin
-    ds["h"] = _smooth_topography_locally(hraw * ds["mask_rho"], hmin, rmax)
+
+    ds["h"] = _smooth_topography_locally(hraw, hmin, rmax)
     ds["h"].attrs = {
         "long_name": "Bathymetry at rho-points",
         "units": "meter",
@@ -241,7 +239,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 +322,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 +333,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`,