roms-tools 3.0.0__py3-none-any.whl → 3.1.1__py3-none-any.whl

roms_tools/setup/datasets.py

@@ -1,9 +1,9 @@
 import logging
 import time
+from collections import Counter, defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 from pathlib import Path
-from typing import Dict, List, Optional, Union
 
 import numpy as np
 import xarray as xr
@@ -79,23 +79,23 @@ class Dataset:
     ... )
     """
 
-    filename: Union[str, Path, List[Union[str, Path]]]
-    start_time: Optional[datetime] = None
-    end_time: Optional[datetime] = None
-    dim_names: Dict[str, str] = field(
+    filename: str | Path | list[str | Path]
+    start_time: datetime | None = None
+    end_time: datetime | None = None
+    dim_names: dict[str, str] = field(
         default_factory=lambda: {
             "longitude": "longitude",
             "latitude": "latitude",
             "time": "time",
         }
     )
-    var_names: Dict[str, str]
-    opt_var_names: Optional[Dict[str, str]] = field(default_factory=dict)
-    climatology: Optional[bool] = False
-    needs_lateral_fill: Optional[bool] = True
-    use_dask: Optional[bool] = False
-    apply_post_processing: Optional[bool] = True
-    read_zarr: Optional[bool] = False
+    var_names: dict[str, str]
+    opt_var_names: dict[str, str] | None = field(default_factory=dict)
+    climatology: bool | None = False
+    needs_lateral_fill: bool | None = True
+    use_dask: bool | None = False
+    apply_post_processing: bool | None = True
+    read_zarr: bool | None = False
 
     is_global: bool = field(init=False, repr=False)
     ds: xr.Dataset = field(init=False, repr=False)
@@ -109,7 +109,6 @@ class Dataset:
         4. Ensures latitude values and depth values are in ascending order.
         5. Checks if the dataset covers the entire globe and adjusts if necessary.
         """
-
         # Validate start_time and end_time
         if self.start_time is not None and not isinstance(self.start_time, datetime):
             raise TypeError(
@@ -168,7 +167,6 @@ class Dataset:
         ValueError
             If a list of files is provided but self.dim_names["time"] is not available or use_dask=False.
         """
-
         ds = _load_data(
             self.filename, self.dim_names, self.use_dask, read_zarr=self.read_zarr
         )
@@ -206,7 +204,6 @@ class Dataset:
         ValueError
             If the dataset does not contain the specified variables or dimensions.
         """
-
         _check_dataset(ds, self.dim_names, self.var_names)
 
     def select_relevant_fields(self, ds) -> xr.Dataset:
@@ -223,7 +220,6 @@ class Dataset:
         xr.Dataset
             A dataset containing only the variables specified in `self.var_names`.
         """
-
         for var in ds.data_vars:
             if (
                 var not in self.var_names.values()
@@ -298,7 +294,6 @@ class Dataset:
         - If the dataset uses `cftime` datetime objects, these will be converted to standard
           `np.datetime64` objects before filtering.
         """
-
         time_dim = self.dim_names["time"]
 
         ds = _select_relevant_times(
@@ -392,7 +387,6 @@ class Dataset:
             The smallest horizontal grid spacing derived from the latitude
             and longitude differences, in meters.
         """
-
         lat_dim = self.dim_names["latitude"]
         lon_dim = self.dim_names["longitude"]
 
@@ -463,7 +457,6 @@ class Dataset:
         ds_concatenated : xr.Dataset
             The concatenated dataset.
         """
-
         if verbose:
             start_time = time.time()
 
@@ -594,7 +587,6 @@ class Dataset:
         ValueError
             If the selected latitude or longitude range does not intersect with the dataset.
         """
-
         lat_min = target_coords["lat"].min().values
         lat_max = target_coords["lat"].max().values
         lon_min = target_coords["lon"].min().values
@@ -713,7 +705,6 @@ class Dataset:
         dataset variable is filled only once, even if multiple entries in `self.var_names`
         point to the same variable in the dataset.
         """
-
         if self.needs_lateral_fill:
             logging.info(
                 "Applying 2D horizontal fill to the source data before regridding."
@@ -761,7 +752,6 @@ class Dataset:
         For each variable with a depth dimension, fills missing values at the bottom by
         propagating the deepest available data downward.
         """
-
         if "depth" in self.dim_names:
             for var_name in self.ds.data_vars:
                 if self.dim_names["depth"] in self.ds[var_name].dims:
@@ -836,8 +826,8 @@ class TPXODataset(Dataset):
     filename: str
     grid_filename: str
     location: str
-    var_names: Dict[str, str]
-    dim_names: Dict[str, str] = field(
+    var_names: dict[str, str]
+    dim_names: dict[str, str] = field(
         default_factory=lambda: {"longitude": "nx", "latitude": "ny", "ntides": "nc"}
     )
     tolerate_coord_mismatch: bool = False
@@ -872,7 +862,6 @@ class TPXODataset(Dataset):
             ValueError
                 If longitude or latitude values do not match the grid.
         """
-
         ds_grid = _load_data(self.grid_filename, self.dim_names, self.use_dask)
 
         # Define mask and coordinate names based on location
@@ -974,7 +963,7 @@ class TPXODataset(Dataset):
 
         return ds
 
-    def select_constituents(self, ntides: int, omega: Dict[str, float]):
+    def select_constituents(self, ntides: int, omega: dict[str, float]):
         """Selects the first `ntides` tidal constituents based on the provided omega
         values.
 
@@ -1002,7 +991,6 @@ class TPXODataset(Dataset):
             selected from the `omega` dictionary, a `ValueError` is raised, indicating the mismatch
             between the expected and present constituents in the dataset.
         """
-
         # Expected constituents based on the first 'ntides' from the omega dictionary
         expected_constituents = list(omega.keys())[:ntides]
 
@@ -1035,7 +1023,7 @@ class TPXODataset(Dataset):
 class GLORYSDataset(Dataset):
     """Represents GLORYS data on original grid."""
 
-    var_names: Dict[str, str] = field(
+    var_names: dict[str, str] = field(
         default_factory=lambda: {
             "temp": "thetao",
             "salt": "so",
@@ -1045,7 +1033,7 @@ class GLORYSDataset(Dataset):
         }
     )
 
-    dim_names: Dict[str, str] = field(
+    dim_names: dict[str, str] = field(
         default_factory=lambda: {
             "longitude": "longitude",
             "latitude": "latitude",
@@ -1054,7 +1042,7 @@ class GLORYSDataset(Dataset):
         }
     )
 
-    climatology: Optional[bool] = False
+    climatology: bool | None = False
 
     def post_process(self):
         """Apply a mask to the dataset based on the 'zeta' variable, with 0 where 'zeta'
@@ -1070,7 +1058,6 @@ class GLORYSDataset(Dataset):
         None
             The dataset is modified in-place by applying the mask to each variable.
         """
-
         mask = xr.where(
             self.ds[self.var_names["zeta"]].isel({self.dim_names["time"]: 0}).isnull(),
             0,
@@ -1099,7 +1086,7 @@ class UnifiedDataset(Dataset):
     attribute is set to `False`.
     """
 
-    needs_lateral_fill: Optional[bool] = False
+    needs_lateral_fill: bool | None = False
 
     # overwrite clean_up method from parent class
     def clean_up(self, ds: xr.Dataset) -> xr.Dataset:
@@ -1112,7 +1099,6 @@ class UnifiedDataset(Dataset):
         ds : xr.Dataset
             The xarray Dataset with the correct time dimension assigned or added.
         """
-
         ds = ds.rename({"lon": "longitude", "lat": "latitude", "dep": "depth"})
         ds = ds.assign_coords(
             {
@@ -1164,14 +1150,14 @@ class UnifiedDataset(Dataset):
 
 @dataclass(kw_only=True)
 class UnifiedBGCDataset(UnifiedDataset):
-    dim_names: Dict[str, str] = field(
+    dim_names: dict[str, str] = field(
         default_factory=lambda: {
             "longitude": "lon",
             "latitude": "lat",
             "depth": "dep",
         }
     )
-    var_names: Dict[str, str] = field(
+    var_names: dict[str, str] = field(
         default_factory=lambda: {
             "PO4": "PO4",
             "NO3": "NO3",
@@ -1182,7 +1168,7 @@ class UnifiedBGCDataset(UnifiedDataset):
             "ALK": "Alk",
         }
     )
-    opt_var_names: Dict[str, str] = field(
+    opt_var_names: dict[str, str] = field(
         default_factory=lambda: {
             "NH4": "NH4",
             "Lig": "Lig",
@@ -1213,21 +1199,21 @@ class UnifiedBGCDataset(UnifiedDataset):
         }
     )
 
-    climatology: Optional[bool] = True
+    climatology: bool | None = True
 
 
 @dataclass(kw_only=True)
 class UnifiedBGCSurfaceDataset(UnifiedDataset):
-    dim_names: Dict[str, str] = field(
+    dim_names: dict[str, str] = field(
         default_factory=lambda: {
             "longitude": "lon",
             "latitude": "lat",
         }
     )
-    var_names: Dict[str, str] = field(
+    var_names: dict[str, str] = field(
         default_factory=lambda: {"pco2_air": "pco2_air", "dust": "dust", "iron": "iron"}
     )
-    opt_var_names: Dict[str, str] = field(
+    opt_var_names: dict[str, str] = field(
         default_factory=lambda: {
             "pco2_air_alt": "pco2_air_alt",
             "nox": "nox",
@@ -1235,7 +1221,7 @@ class UnifiedBGCSurfaceDataset(UnifiedDataset):
         }
     )
 
-    climatology: Optional[bool] = True
+    climatology: bool | None = True
 
 
 @dataclass(kw_only=True)
@@ -1253,7 +1239,6 @@ class CESMDataset(Dataset):
         ds : xr.Dataset
             The xarray Dataset with the correct time dimension assigned or added.
         """
-
         if "time" not in self.dim_names:
             if "time" in ds.dims:
                 self.dim_names["time"] = "time"
@@ -1306,7 +1291,7 @@ class CESMDataset(Dataset):
 class CESMBGCDataset(CESMDataset):
     """Represents CESM BGC data on original grid."""
 
-    var_names: Dict[str, str] = field(
+    var_names: dict[str, str] = field(
         default_factory=lambda: {
             "PO4": "PO4",
             "NO3": "NO3",
@@ -1343,7 +1328,7 @@ class CESMBGCDataset(CESMDataset):
         }
     )
 
-    dim_names: Dict[str, str] = field(
+    dim_names: dict[str, str] = field(
         default_factory=lambda: {
             "longitude": "lon",
             "latitude": "lat",
@@ -1351,7 +1336,7 @@ class CESMBGCDataset(CESMDataset):
         }
     )
 
-    climatology: Optional[bool] = False
+    climatology: bool | None = False
 
     def post_process(self):
         """
@@ -1404,7 +1389,7 @@ class CESMBGCDataset(CESMDataset):
 class CESMBGCSurfaceForcingDataset(CESMDataset):
     """Represents CESM BGC surface forcing data on original grid."""
 
-    var_names: Dict[str, str] = field(
+    var_names: dict[str, str] = field(
         default_factory=lambda: {
             "pco2_air": "pCO2SURF",
             "pco2_air_alt": "pCO2SURF",
@@ -1415,14 +1400,14 @@ class CESMBGCSurfaceForcingDataset(CESMDataset):
         }
     )
 
-    dim_names: Dict[str, str] = field(
+    dim_names: dict[str, str] = field(
         default_factory=lambda: {
             "longitude": "lon",
             "latitude": "lat",
         }
     )
 
-    climatology: Optional[bool] = False
+    climatology: bool | None = False
 
     def post_process(self):
         """Perform post-processing on the dataset to remove specific variables.
@@ -1431,7 +1416,6 @@ class CESMBGCSurfaceForcingDataset(CESMDataset):
         the variable is removed from the dataset. The modified dataset is then
         reassigned to the `ds` attribute of the object.
         """
-
         if "z_t" in self.ds.variables:
             ds = self.ds.drop_vars("z_t")
             self.ds = ds
@@ -1451,7 +1435,7 @@ class CESMBGCSurfaceForcingDataset(CESMDataset):
 class ERA5Dataset(Dataset):
     """Represents ERA5 data on original grid."""
 
-    var_names: Dict[str, str] = field(
+    var_names: dict[str, str] = field(
         default_factory=lambda: {
             "uwnd": "u10",
             "vwnd": "v10",
@@ -1464,7 +1448,7 @@ class ERA5Dataset(Dataset):
         }
     )
 
-    dim_names: Dict[str, str] = field(
+    dim_names: dict[str, str] = field(
         default_factory=lambda: {
             "longitude": "longitude",
             "latitude": "latitude",
@@ -1472,7 +1456,7 @@ class ERA5Dataset(Dataset):
         }
     )
 
-    climatology: Optional[bool] = False
+    climatology: bool | None = False
 
     def post_process(self):
         """
@@ -1549,7 +1533,7 @@ class ERA5Dataset(Dataset):
 class ERA5ARCODataset(ERA5Dataset):
     # The GCP ARCO version of the dataset has different variable names
 
-    var_names: Dict[str, str] = field(
+    var_names: dict[str, str] = field(
         default_factory=lambda: {
             "uwnd": "10m_u_component_of_wind",
             "vwnd": "10m_v_component_of_wind",
@@ -1586,22 +1570,21 @@ class ERA5Correction(Dataset):
     filename: str = field(
         default_factory=lambda: download_correction_data("SSR_correction.nc")
     )
-    var_names: Dict[str, str] = field(
+    var_names: dict[str, str] = field(
         default_factory=lambda: {
             "swr_corr": "ssr_corr",  # multiplicative correction factor for ERA5 shortwave radiation
         }
     )
-    dim_names: Dict[str, str] = field(
+    dim_names: dict[str, str] = field(
         default_factory=lambda: {
             "longitude": "longitude",
             "latitude": "latitude",
             "time": "time",
         }
     )
-    climatology: Optional[bool] = True
+    climatology: bool | None = True
 
     def __post_init__(self):
-
         if not self.climatology:
             raise NotImplementedError(
                 "Correction data must be a climatology. Set climatology to True."
@@ -1636,7 +1619,6 @@ class ERA5Correction(Dataset):
         -----
         - The dataset (`self.ds`) is updated in place to reflect the chosen subdomain.
         """
-
         # Select the subdomain in latitude direction (so that we have to concatenate fewer latitudes below if concatenation is performed)
         subdomain = self.ds.sel({self.dim_names["latitude"]: target_coords["lat"]})
 
@@ -1666,12 +1648,12 @@ class ETOPO5Dataset(Dataset):
     """Represents topography data on the original grid from the ETOPO5 dataset."""
 
     filename: str = field(default_factory=lambda: download_topo("etopo5.nc"))
-    var_names: Dict[str, str] = field(
+    var_names: dict[str, str] = field(
         default_factory=lambda: {
             "topo": "topo",
         }
     )
-    dim_names: Dict[str, str] = field(
+    dim_names: dict[str, str] = field(
         default_factory=lambda: {"longitude": "lon", "latitude": "lat"}
     )
 
@@ -1701,12 +1683,12 @@ class ETOPO5Dataset(Dataset):
 class SRTM15Dataset(Dataset):
     """Represents topography data on the original grid from the SRTM15 dataset."""
 
-    var_names: Dict[str, str] = field(
+    var_names: dict[str, str] = field(
         default_factory=lambda: {
             "topo": "z",
         }
     )
-    dim_names: Dict[str, str] = field(
+    dim_names: dict[str, str] = field(
         default_factory=lambda: {"longitude": "lon", "latitude": "lat"}
     )
 
@@ -1743,17 +1725,16 @@ class RiverDataset:
         The xarray Dataset containing the forcing data on its original grid.
     """
 
-    filename: Union[str, Path, List[Union[str, Path]]]
+    filename: str | Path | list[str | Path]
     start_time: datetime
     end_time: datetime
-    dim_names: Dict[str, str]
-    var_names: Dict[str, str]
-    opt_var_names: Optional[Dict[str, str]] = field(default_factory=dict)
-    climatology: Optional[bool] = False
+    dim_names: dict[str, str]
+    var_names: dict[str, str]
+    opt_var_names: dict[str, str] | None = field(default_factory=dict)
+    climatology: bool | None = False
     ds: xr.Dataset = field(init=False, repr=False)
 
     def __post_init__(self):
-
         # Validate start_time and end_time
         if not isinstance(self.start_time, datetime):
             raise TypeError(
@@ -1767,6 +1748,9 @@ class RiverDataset:
         ds = self.load_data()
         ds = self.clean_up(ds)
         self.check_dataset(ds)
+        ds = _deduplicate_river_names(
+            ds, self.var_names["name"], self.dim_names["station"]
+        )
 
         # Select relevant times
         ds = self.add_time_info(ds)
@@ -1804,7 +1788,6 @@ class RiverDataset:
         ds : xr.Dataset
             The dataset with the decoded 'name' variable.
         """
-
         if ds[self.var_names["name"]].dtype == "object":
             names = []
             for i in range(len(ds[self.dim_names["station"]])):
@@ -1828,7 +1811,7 @@ class RiverDataset:
         return ds
 
     def check_dataset(self, ds: xr.Dataset) -> None:
-        """Check if the dataset contains the specified variables and dimensions.
+        """Validate required variables, dimensions, and uniqueness of river names.
 
         Parameters
         ----------
@@ -1840,7 +1823,6 @@ class RiverDataset:
         ValueError
             If the dataset does not contain the specified variables or dimensions.
         """
-
         _check_dataset(ds, self.dim_names, self.var_names, self.opt_var_names)
 
     def add_time_info(self, ds: xr.Dataset) -> xr.Dataset:
@@ -1891,7 +1873,6 @@ class RiverDataset:
         UserWarning
             If the dataset does not contain any time dimension or the time dimension is incorrectly named.
         """
-
         time_dim = self.dim_names["time"]
 
         ds = _select_relevant_times(ds, time_dim, self.start_time, self.end_time, False)
@@ -1935,7 +1916,6 @@ class RiverDataset:
             The dataset with rivers sorted by their volume in descending order.
             If the volume variable is not available, the original dataset is returned.
         """
-
         if "vol" in self.opt_var_names:
             volume_values = ds[self.opt_var_names["vol"]].values
             if isinstance(volume_values, np.ndarray):
@@ -1991,7 +1971,6 @@ class RiverDataset:
             the target coordinates. The dictionary structure consists of river names as keys, and each value is a list of tuples. Each tuple represents
             a pair of indices corresponding to the `eta_rho` and `xi_rho` grid coordinates of the river.
         """
-
         # Retrieve longitude and latitude of river mouths
         river_lon = self.ds[self.var_names["longitude"]]
         river_lat = self.ds[self.var_names["latitude"]]
@@ -2012,24 +1991,8 @@ class RiverDataset:
             dist = dist.where(dist_min < dx, drop=True).transpose(
                 self.dim_names["station"], "eta_rho", "xi_rho"
             )
-            dist_min = dist_min.where(dist_min < dx, drop=True)
-
-            # Find the indices of the closest grid cell to the river mouth
-            indices = np.where(dist == dist_min)
-            stations = indices[0]
-            eta_rho_values = indices[1]
-            xi_rho_values = indices[2]
-            names = (
-                ds[self.var_names["name"]]
-                .isel({self.dim_names["station"]: stations})
-                .values
-            )
-            river_indices = {}
-            for i in range(len(stations)):
-                river_name = names[i]
-                river_indices[river_name] = [
-                    (int(eta_rho_values[i]), int(xi_rho_values[i]))
-                ]  # list of tuples
+
+            river_indices = get_indices_of_nearest_grid_cell_for_rivers(dist, self)
         else:
             ds = xr.Dataset()
             river_indices = {}
@@ -2063,7 +2026,6 @@ class RiverDataset:
             - If `indices` is not a dictionary.
             - If any of the requested river names are not found in the dataset.
         """
-
         if not isinstance(indices, dict):
             raise ValueError("`indices` must be a dictionary.")
 
@@ -2091,16 +2053,16 @@ class RiverDataset:
 class DaiRiverDataset(RiverDataset):
     """Represents river data from the Dai river dataset."""
 
-    filename: Union[str, Path, List[Union[str, Path]]] = field(
+    filename: str | Path | list[str | Path] = field(
         default_factory=lambda: download_river_data("dai_trenberth_may2019.nc")
     )
-    dim_names: Dict[str, str] = field(
+    dim_names: dict[str, str] = field(
         default_factory=lambda: {
             "station": "station",
             "time": "time",
         }
     )
-    var_names: Dict[str, str] = field(
+    var_names: dict[str, str] = field(
         default_factory=lambda: {
             "latitude": "lat_mou",
             "longitude": "lon_mou",
@@ -2109,12 +2071,12 @@ class DaiRiverDataset(RiverDataset):
             "name": "riv_name",
         }
     )
-    opt_var_names: Dict[str, str] = field(
+    opt_var_names: dict[str, str] = field(
         default_factory=lambda: {
             "vol": "vol_stn",
         }
     )
-    climatology: Optional[bool] = False
+    climatology: bool | None = False
 
     def add_time_info(self, ds: xr.Dataset) -> xr.Dataset:
         """Adds time information to the dataset based on the climatology flag and
@@ -2193,10 +2155,9 @@ class TPXOManager:
     ntides: int
     reference_date: datetime = datetime(1992, 1, 1)
     allan_factor: float = 2.0
-    use_dask: Optional[bool] = False
+    use_dask: bool | None = False
 
     def __post_init__(self):
-
         fname_sal = download_sal_data("sal_tpxo9.v2a.nc")
 
         # Initialize the data_dict with TPXODataset instances
@@ -2307,7 +2268,6 @@ class TPXOManager:
 
         The amplitudes and elasticity factors are sourced from the `constit.h` file in the OTPSnc package.
         """
-
         # Amplitudes for 15 tidal constituents (from variable amp_d in constit.h of OTPSnc package)
         A = xr.DataArray(
             data=np.array(
@@ -2429,7 +2389,6 @@ class TPXOManager:
         - Egbert, G.D., and S.Y. Erofeeva. "Efficient inverse modeling of barotropic ocean
           tides." Journal of Atmospheric and Oceanic Technology 19, no. 2 (2002): 183-204.
         """
-
         year = date.year
         month = date.month
         day = date.day
@@ -2582,7 +2541,6 @@ class TPXOManager:
             The dataset is modified in-place with corrected real and imaginary components for ssh, u, v, and the
             potential field ('pot_Re', 'pot_Im').
         """
-
         datasets = self.datasets
         omega = self.datasets["omega"].isel(ntides=slice(None, self.ntides))
 
@@ -2649,9 +2607,9 @@ class TPXOManager:
 
 def _check_dataset(
     ds: xr.Dataset,
-    dim_names: Dict[str, str],
-    var_names: Dict[str, str],
-    opt_var_names: Optional[Dict[str, str]] = None,
+    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.
 
@@ -2751,7 +2709,6 @@ def _select_relevant_times(
     - If the dataset uses `cftime` datetime objects, these will be converted to standard
       `np.datetime64` objects before filtering.
     """
-
     if time_dim in ds.variables:
         if climatology:
             if len(ds[time_dim]) != 12:
@@ -2835,7 +2792,6 @@ def _select_relevant_times(
 
 
 def decode_string(byte_array):
-
     # Decode each byte and handle errors with 'ignore'
     decoded_string = "".join(
         [
@@ -2890,7 +2846,6 @@ def modified_julian_days(year, month, day, hour=0):
     >>> modified_julian_days(1582, 10, 4)
     -141428.5
     """
-
     if month < 3:
         year -= 1
         month += 12
@@ -2904,3 +2859,79 @@ def modified_julian_days(year, month, day, hour=0):
     mjd = jd - 2400000.5
 
     return mjd
+
+
+def get_indices_of_nearest_grid_cell_for_rivers(
+    dist: xr.DataArray, data: RiverDataset
+) -> dict[str, list[tuple[int, int]]]:
+    """Get the indices of the nearest grid cell for each river based on distance.
+
+    Parameters
+    ----------
+    dist : xr.DataArray
+        A 2D or 3D array representing distances from each river to coastal grid cells,
+        with dimensions including "eta_rho" and "xi_rho".
+    data : RiverDataset
+        An instance of RiverDataset containing river names and dimension metadata.
+
+    Returns
+    -------
+    dict[str, list[tuple[int, int]]]
+        Dictionary mapping each river name to a list containing the (eta_rho, xi_rho) index
+        of the closest coastal grid cell.
+    """
+    # Find indices of the nearest coastal grid cell for each river
+    indices = dist.argmin(dim=["eta_rho", "xi_rho"])
+
+    eta_rho_values = indices["eta_rho"].values
+    xi_rho_values = indices["xi_rho"].values
+
+    # Get the corresponding station indices and river names
+    stations = indices["eta_rho"][data.dim_names["station"]].values
+    names = (
+        data.ds[data.var_names["name"]]
+        .sel({data.dim_names["station"]: stations})
+        .values
+    )
+
+    # Build dictionary of river name to grid index
+    river_indices = {
+        str(names[i]): [(int(eta_rho_values[i]), int(xi_rho_values[i]))]
+        for i in range(len(stations))
+    }
+
+    return river_indices
+
+
+def _deduplicate_river_names(
+    ds: xr.Dataset, name_var: str, station_dim: str
+) -> xr.Dataset:
+    """Ensure river names are unique by appending _1, _2 to duplicates, excluding non-
+    duplicates.
+    """
+    original = ds[name_var]
+
+    # Force cast to plain Python strings
+    names = [str(name) for name in original.values]
+
+    # Count all names
+    name_counts = Counter(names)
+    seen = defaultdict(int)
+
+    unique_names = []
+    for name in names:
+        if name_counts[name] > 1:
+            seen[name] += 1
+            unique_names.append(f"{name}_{seen[name]}")
+        else:
+            unique_names.append(name)
+
+    # Replace with updated names while preserving dtype, dims, attrs
+    updated_array = xr.DataArray(
+        data=np.array(unique_names, dtype=f"<U{max(len(n) for n in unique_names)}"),
+        dims=original.dims,
+        attrs=original.attrs,
+    )
+    ds[name_var] = updated_array
+
+    return ds