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

roms_tools/setup/river_forcing.py

@@ -1,32 +1,43 @@
 import logging
+from collections import defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Dict, List, Optional, Union
 
-import cartopy.crs as ccrs
-import matplotlib.cm as cm
 import matplotlib.pyplot as plt
 import numpy as np
 import xarray as xr
 
 from roms_tools import Grid
-from roms_tools.plot import _get_projection, _plot
-from roms_tools.setup.datasets import DaiRiverDataset
+from roms_tools.constants import MAX_DISTINCT_COLORS
+from roms_tools.plot import (
+    assign_category_colors,
+    get_projection,
+    plot_2d_horizontal_field,
+    plot_location,
+)
+from roms_tools.setup.datasets import (
+    DaiRiverDataset,
+    get_indices_of_nearest_grid_cell_for_rivers,
+)
 from roms_tools.setup.utils import (
-    _from_yaml,
-    _to_dict,
-    _write_to_yaml,
     add_time_info_to_ds,
     add_tracer_metadata_to_ds,
+    from_yaml,
     gc_dist,
     get_target_coords,
     get_tracer_defaults,
     get_variable_metadata,
     substitute_nans_by_fillvalue,
+    to_dict,
+    validate_names,
+    write_to_yaml,
 )
 from roms_tools.utils import save_datasets
 
+INCLUDE_ALL_RIVER_NAMES = "all"
+MAX_RIVERS_TO_PLOT = 20  # must be <= MAX_DISTINCT_COLORS
+
 
 @dataclass(kw_only=True)
 class RiverForcing:
@@ -63,7 +74,7 @@ class RiverForcing:
     include_bgc : bool, optional
         Whether to include BGC tracers. Defaults to `False`.
     model_reference_date : datetime, optional
-        Reference date for the model. Default is January 1, 2000.
+        Reference date for the ROMS simulation. Default is January 1, 2000.
     indices : dict[str, list[tuple]], 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
@@ -82,29 +93,34 @@ class RiverForcing:
 
         In the example, the dictionary provides the river names as keys, and the values are lists of tuples, where each tuple represents the
         `(eta_rho, xi_rho)` indices for a river location.
-
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The xarray Dataset containing the river forcing data.
-    climatology : bool
-        Indicates whether the final river forcing is climatological.
-    Dict[str, Union[int, List[int]]]
-        A dictionary of river indices. If not provided during initialization, it will be automatically determined
-        based on the grid and the source dataset. The dictionary structure is the same as described in the `indices` parameter docstring.
     """
 
     grid: Grid
+    """Object representing the grid information."""
     start_time: datetime
+    """Start time of the desired river forcing data."""
     end_time: datetime
-    source: Dict[str, Union[str, Path, List[Union[str, Path]]]] = None
+    """End time of the desired river forcing data."""
+    source: dict[str, str | Path | list[str | Path]] = 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."""
     include_bgc: bool = False
+    """Whether to include BGC tracers."""
     model_reference_date: datetime = datetime(2000, 1, 1)
-    indices: Optional[Dict[str, Dict[str, Union[int, List[int]]]]] = None
+    """Reference date for the ROMS simulation."""
+
+    indices: dict[str, dict[str, int | list[int]]] | None = None
+    """A dictionary of river indices.
 
+    If not provided during initialization, it will be automatically determined based on
+    the grid and the source dataset.
+    """
     ds: xr.Dataset = field(init=False, repr=False)
+    """An xarray Dataset containing post-processed variables ready for input into
+    ROMS."""
     climatology: xr.Dataset = field(init=False, repr=False)
+    """Indicates whether the final river forcing is climatological."""
 
     def __post_init__(self):
         self._input_checks()
@@ -135,6 +151,7 @@ class RiverForcing:
             data.extract_named_rivers(self.indices)
 
         ds = self._create_river_forcing(data)
+        ds = self._handle_overlapping_rivers(ds)
         ds = self._write_indices_into_dataset(ds)
         self._validate(ds)
 
@@ -179,8 +196,8 @@ class RiverForcing:
                         f"Data for river `{river_name}` must be a list of tuples."
                     )
 
-                # Ensure each element in the list is a tuple of length 2
                 seen_tuples = set()
+                # Ensure each element in the list is a tuple of length 2
                 for idx_pair in river_data:
                     if not isinstance(idx_pair, tuple) or len(idx_pair) != 2:
                         raise ValueError(
@@ -202,14 +219,14 @@ class RiverForcing:
                     # Check that eta_rho and xi_rho are within the valid range
                     if not (0 <= eta_rho < len(self.grid.ds.eta_rho)):
                         raise ValueError(
-                            f"Value of eta_rho for river `{river_name}` ({eta_rho}) is out of valid range [0, {len(self.grid.ds.eta_rho)-1}]."
+                            f"Value of eta_rho for river `{river_name}` ({eta_rho}) is out of valid range [0, {len(self.grid.ds.eta_rho) - 1}]."
                         )
                     if not (0 <= xi_rho < len(self.grid.ds.xi_rho)):
                         raise ValueError(
-                            f"Value of xi_rho for river `{river_name}` ({xi_rho}) is out of valid range [0, {len(self.grid.ds.xi_rho)-1}]."
+                            f"Value of xi_rho for river `{river_name}` ({xi_rho}) is out of valid range [0, {len(self.grid.ds.xi_rho) - 1}]."
                         )
 
-                    # Check for duplicate tuples
+                    # Check for duplicate tuples for a single river
                     if idx_pair in seen_tuples:
                         raise ValueError(
                             f"Duplicate location {idx_pair} found for river `{river_name}`."
@@ -217,7 +234,6 @@ class RiverForcing:
                     seen_tuples.add(idx_pair)
 
     def _get_data(self):
-
         data_dict = {
             "start_time": self.start_time,
             "end_time": self.end_time,
@@ -233,6 +249,64 @@ class RiverForcing:
 
         return data
 
+    def _move_rivers_to_closest_coast(self, target_coords, data):
+        """Move river mouths to the closest coastal grid cell.
+
+        This method computes the closest coastal grid point to each river mouth
+        based on geographical distance. It identifies the nearest grid point on the coast and returns the updated river mouth indices.
+
+        Parameters:
+        -----------
+        target_coords : dict
+            A dictionary containing the following keys:
+            - "lon" (xarray.DataArray): Longitude coordinates of the target grid points.
+            - "lat" (xarray.DataArray): Latitude coordinates of the target grid points.
+            - "straddle" (bool): A flag indicating whether the river mouth crosses the International Date Line.
+
+        data : object
+            An object that contains the dataset and related variables. It must have the following attributes:
+            - `ds`: The dataset containing river information.
+            - `var_names`: A dictionary of variable names in the dataset (e.g., longitude, latitude, station names).
+            - `dim_names`: A dictionary containing dimension names for the dataset (e.g., "station", "eta_rho", "xi_rho").
+
+        Returns
+        -------
+        indices : dict[str, list[tuple]]
+            A dictionary consisting 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 = data.ds[data.var_names["longitude"]]
+        river_lat = data.ds[data.var_names["latitude"]]
+
+        # Adjust longitude based on whether it crosses the International Date Line (straddle case)
+        if target_coords["straddle"]:
+            river_lon = xr.where(river_lon > 180, river_lon - 360, river_lon)
+        else:
+            river_lon = xr.where(river_lon < 0, river_lon + 360, river_lon)
+
+        mask = self.grid.ds.mask_rho
+        faces = (
+            mask.shift(eta_rho=1)
+            + mask.shift(eta_rho=-1)
+            + mask.shift(xi_rho=1)
+            + mask.shift(xi_rho=-1)
+        )
+
+        # We want all grid points on land that are adjacent to the ocean
+        coast = (1 - mask) * (faces > 0)
+        dist_coast = gc_dist(
+            target_coords["lon"].where(coast),
+            target_coords["lat"].where(coast),
+            river_lon,
+            river_lat,
+        ).transpose(data.dim_names["station"], "eta_rho", "xi_rho")
+
+        # Find the indices of the closest coastal grid cell to the river mouth
+        river_indices = get_indices_of_nearest_grid_cell_for_rivers(dist_coast, data)
+
+        return river_indices
+
     def _create_river_forcing(self, data):
         """Create river forcing data for volume flux and tracers (temperature, salinity,
         BGC tracers).
@@ -328,80 +402,219 @@ class RiverForcing:
 
         return ds
 
-    def _move_rivers_to_closest_coast(self, target_coords, data):
-        """Move river mouths to the closest coastal grid cell.
+    def _handle_overlapping_rivers(self, ds: xr.Dataset) -> xr.Dataset:
+        """Detect and resolve overlapping river grid cell assignments.
 
-        This method computes the closest coastal grid point to each river mouth
-        based on geographical distance. It identifies the nearest grid point on the coast and returns the updated river mouth indices.
+        If multiple rivers are assigned to the same grid cell (i.e., overlapping index pairs),
+        this method creates new uniquely named rivers ('overlap_1', 'overlap_2', ...)
+        that consolidate the contributions from all original rivers sharing that location.
 
-        Parameters:
-        -----------
-        target_coords : dict
-            A dictionary containing the following keys:
-            - "lon" (xarray.DataArray): Longitude coordinates of the target grid points.
-            - "lat" (xarray.DataArray): Latitude coordinates of the target grid points.
-            - "straddle" (bool): A flag indicating whether the river mouth crosses the International Date Line.
+        For each overlapping grid cell:
+        - A new river is created using a volume-weighted sum of the original rivers' volume and tracer.
+        - The volume fraction of the original river is reduced proportionally to reflect the removal of the shared grid cell.
+        - The new river is appended to the dataset and indexed in `self.indices`.
 
-        data : object
-            An object that contains the dataset and related variables. It must have the following attributes:
-            - `ds`: The dataset containing river information.
-            - `var_names`: A dictionary of variable names in the dataset (e.g., longitude, latitude, station names).
-            - `dim_names`: A dictionary containing dimension names for the dataset (e.g., "station", "eta_rho", "xi_rho").
+        Parameters
+        ----------
+        ds : xr.Dataset
+            Dataset containing the existing river forcing fields, including:
+            - `river_volume` : volume transport per river [river_time, nriver]
+            - `river_tracer` : tracer concentration per river [river_time, ntracers, nriver]
+            - `river_name`   : river name labels [nriver]
 
         Returns
         -------
-        indices : dict[str, list[tuple]]
-            A dictionary consisting 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.
+        xr.Dataset
+            A new dataset with overlapping rivers resolved and new entries added.
         """
+        overlapping_rivers = self._get_overlapping_rivers()
 
-        # Retrieve longitude and latitude of river mouths
-        river_lon = data.ds[data.var_names["longitude"]]
-        river_lat = data.ds[data.var_names["latitude"]]
+        if len(overlapping_rivers) > 0:
+            logging.info(
+                f"Creating {len(overlapping_rivers)} synthetic river(s) to handle overlapping entries."
+            )
 
-        # Adjust longitude based on whether it crosses the International Date Line (straddle case)
-        if target_coords["straddle"]:
-            river_lon = xr.where(river_lon > 180, river_lon - 360, river_lon)
+            # Add new unique river for each overlapping index
+            combined_river_volumes = []
+            combined_river_tracers = []
+
+            for i, (idx_pair, river_list) in enumerate(overlapping_rivers.items()):
+                (
+                    combined_river_volume,
+                    combined_river_tracer,
+                ) = self._create_combined_river(ds, i + 1, idx_pair, river_list)
+                combined_river_volumes.append(combined_river_volume)
+                combined_river_tracers.append(combined_river_tracer)
+
+            ds_updated = xr.Dataset()
+            ds_updated["river_volume"] = xr.concat(
+                [ds["river_volume"], *combined_river_volumes], dim="nriver"
+            )
+            ds_updated["river_tracer"] = xr.concat(
+                [ds["river_tracer"], *combined_river_tracers], dim="nriver"
+            )
+            ds_updated.attrs = ds.attrs
         else:
-            river_lon = xr.where(river_lon < 0, river_lon + 360, river_lon)
+            ds_updated = ds
 
-        mask = self.grid.ds.mask_rho
-        faces = (
-            mask.shift(eta_rho=1)
-            + mask.shift(eta_rho=-1)
-            + mask.shift(xi_rho=1)
-            + mask.shift(xi_rho=-1)
+        # Reduce volume fraction of original rivers by appropriate amount
+        ds_updated = self._reduce_river_volumes(ds_updated, overlapping_rivers)
+
+        return ds_updated
+
+    def _get_overlapping_rivers(self) -> dict[tuple[int, int], list[str]]:
+        """Identify grid cells shared by multiple rivers.
+
+        Scans through the river indices and finds all grid cell indices
+        (as tuples) that are associated with more than one river.
+
+        Returns
+        -------
+        overlapping_rivers : dict[tuple[int, int], list[str]]
+            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)
+
+        # Collect all index pairs used by multiple rivers
+        for river_name, index_list in self.indices.items():
+            for idx_pair in index_list:
+                index_to_rivers[idx_pair].append(river_name)
+
+        overlapping_rivers = {
+            idx: names for idx, names in index_to_rivers.items() if len(names) > 1
+        }
+
+        return overlapping_rivers
+
+    def _create_combined_river(
+        self,
+        ds: xr.Dataset,
+        i: int,
+        idx_pair: tuple[int, int],
+        river_list: list[str],
+    ) -> tuple[xr.DataArray, xr.DataArray]:
+        """Create a combined river entry from multiple overlapping rivers.
+
+        This method generates a synthetic river by merging contributions from several rivers
+        that map to the same grid cell. It performs a weighted sum of river volumes and
+        computes volume-weighted averages of tracer concentrations. The new synthetic river is also
+        registered in `self.indices`.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            Dataset containing river data, including `river_volume`, `river_tracer`, and `river_name`.
+        i : int
+            Index used for naming the new synthetic river (e.g., "overlap_1", "overlap_2", ...).
+        idx_pair : tuple[int, int]
+            Grid cell index (eta_rho, xi_rho) that the combined river will occupy.
+        river_list : list[str]
+            List of river names contributing to the overlapping grid cell.
+
+        Returns
+        -------
+        combined_river_volume : xr.DataArray
+            The total river volume at the overlapping grid cell, as a new 1-entry DataArray with updated coordinates.
+        combined_river_tracer : xr.DataArray
+            The volume-weighted tracer concentration at the overlapping grid cell,
+            as a new 1-entry DataArray with updated coordinates.
+        """
+        new_name = f"overlap_{i}"
+        self.indices[new_name] = [idx_pair]
+
+        # Get index of all rivers contributing to this overlapping cell
+        contributing_indices = [
+            np.where(ds["river_name"].values == name)[0].item() for name in river_list
+        ]
+
+        # Get the number of grid points each river originally contributed to
+        num_cells_per_river = [len(self.indices[name]) for name in river_list]
+
+        # Weighted sum of river volume contributions at the overlapping location
+        weighted_volumes = xr.concat(
+            [
+                (ds["river_volume"].isel(nriver=i) / n_cells).astype("float64")
+                for i, n_cells in zip(contributing_indices, num_cells_per_river)
+            ],
+            dim="tmp",
         )
+        combined_river_volume = weighted_volumes.sum(dim="tmp", skipna=True)
 
-        # We want all grid points on land that are adjacent to the ocean
-        coast = (1 - mask) * (faces > 0)
-        dist_coast = gc_dist(
-            target_coords["lon"].where(coast),
-            target_coords["lat"].where(coast),
-            river_lon,
-            river_lat,
-        ).transpose(data.dim_names["station"], "eta_rho", "xi_rho")
-        dist_coast_min = dist_coast.min(dim=["eta_rho", "xi_rho"])
+        # Volume-weighted sum of river tracer contributions at the overlapping location
+        weighted_tracers = xr.concat(
+            [
+                (ds["river_tracer"].isel(nriver=i) * weight.fillna(0.0)).astype(
+                    "float64"
+                )
+                for i, weight in zip(contributing_indices, weighted_volumes)
+            ],
+            dim="tmp",
+        )
+        combined_river_tracer = (
+            weighted_tracers.sum(dim="tmp", skipna=True) / combined_river_volume
+        )
+        # If combined_river_volume is 0.0, the result will be NaN. Replace with default for clarity.
+        # This is mainly for plotting and to avoid confusing users—ROMS will ignore tracers with zero volume anyway.
+        defaults = get_tracer_defaults()
+        for ntracer in range(combined_river_tracer.sizes["ntracers"]):
+            tracer_name = combined_river_tracer.tracer_name[ntracer].item()
+            default = defaults[tracer_name]
+            combined_river_tracer.loc[{"ntracers": ntracer}] = (
+                combined_river_tracer.loc[{"ntracers": ntracer}].fillna(default)
+            )
 
-        # Find the indices of the closest coastal grid cell to the river mouth
-        indices = np.where(dist_coast == dist_coast_min)
-        stations = indices[0]
-        eta_rho_values = indices[1]
-        xi_rho_values = indices[2]
-        names = (
-            data.ds[data.var_names["name"]]
-            .isel({data.dim_names["station"]: stations})
-            .values
+        # Expand, assign coordinates, and name for both volume and tracer
+        new_nriver = ds.sizes["nriver"] + i
+        combined_river_volume = combined_river_volume.expand_dims(nriver=1)
+        combined_river_volume = combined_river_volume.assign_coords(
+            nriver=[new_nriver], river_name=new_name
+        )
+        combined_river_tracer = combined_river_tracer.expand_dims(nriver=1)
+        combined_river_tracer = combined_river_tracer.assign_coords(
+            nriver=[new_nriver], river_name=new_name
         )
-        # Return the indices in a dictionary format
-        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
 
-        return river_indices
+        return combined_river_volume, combined_river_tracer
+
+    def _reduce_river_volumes(
+        self, ds: xr.Dataset, overlapping_rivers: dict[tuple[int, int], list[str]]
+    ) -> xr.Dataset:
+        """Reduce river volumes for rivers contributing to overlapping grid cells.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            Dataset containing `river_name` and `river_volume` variables.
+        overlapping_rivers : dict
+            Mapping from (eta_rho, xi_rho) grid cell indices to lists of river names
+            contributing to that cell.
+
+        Returns
+        -------
+        ds : xr.Dataset
+            Updated dataset with reduced river volumes.
+        """
+        # Count number of overlaps for each river
+        river_overlap_count = defaultdict(int)
+        for rivers in overlapping_rivers.values():
+            for name in rivers:
+                river_overlap_count[name] += 1
+
+        for name in ds["river_name"].values:
+            n_cells = len(self.indices[name])
+            n_overlaps = river_overlap_count.get(name, 0)
+
+            if n_cells == 0:
+                continue  # Avoid division by zero
+
+            # Scale river volume based on non-overlapping cells
+            scale_factor = (n_cells - n_overlaps) / n_cells
+            river_idx = np.where(ds["river_name"].values == name)[0].item()
+            nriver_val = ds["nriver"].values[river_idx]
+            ds["river_volume"].loc[{"nriver": nriver_val}] *= scale_factor
+
+        return ds
 
     def _write_indices_into_dataset(self, ds):
         """Adds river location indices to the dataset as the "river_index" and
@@ -420,7 +633,6 @@ class RiverForcing:
         xarray.Dataset
             The modified dataset with the "river_index" and "river_fraction" variables added.
         """
-
         river_index = xr.zeros_like(self.grid.ds.h, dtype=np.float32)
         river_fraction = xr.zeros_like(self.grid.ds.h, dtype=np.float32)
 
@@ -428,9 +640,7 @@ class RiverForcing:
             river_name = str(ds.river_name.sel(nriver=nriver).values)
             indices = self.indices[river_name]
             fraction = 1.0 / len(indices)
-
             for eta_index, xi_index in indices:
-
                 # Assign unique nriver ID (Fortran-based indexing)
                 river_index[eta_index, xi_index] = nriver
                 # Fractional contribution for multiple grid points
@@ -461,7 +671,6 @@ 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:
@@ -471,8 +680,23 @@ class RiverForcing:
                     "`convert_to_climatology = 'if_any_missing'` to automatically fill missing values with climatological data."
                 )
 
-    def plot_locations(self):
-        """Plots the original and updated river locations on a map projection."""
+    def plot_locations(self, river_names: list[str] | str = INCLUDE_ALL_RIVER_NAMES):
+        """Plots the original and updated river locations on a map projection.
+
+        Parameters
+        ----------
+        river_names : list[str], or str, optional
+            A list of release names to plot.
+            If a string equal to "all", all rivers will be plotted.
+            Defaults to "all".
+
+        """
+        valid_river_names = list(self.indices.keys())
+        river_names = _validate_river_names(river_names, valid_river_names)
+        if len(valid_river_names) > MAX_DISTINCT_COLORS:
+            colors = assign_category_colors(river_names)
+        else:
+            colors = assign_category_colors(valid_river_names)
 
         field = self.grid.ds.mask_rho
         lon_deg = self.grid.ds.lon_rho
@@ -486,62 +710,46 @@ class RiverForcing:
         cmap = plt.colormaps.get_cmap("Blues")
         kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
-        trans = _get_projection(lon_deg, lat_deg)
+        trans = get_projection(lon_deg, lat_deg)
 
         fig, axs = plt.subplots(
             1, 2, figsize=(13, 13), subplot_kw={"projection": trans}
         )
 
         for ax in axs:
-            _plot(field, kwargs=kwargs, ax=ax, c=None, add_colorbar=False)
-
-        proj = ccrs.PlateCarree()
-
-        if len(self.indices) <= 10:
-            color_map = cm.get_cmap("tab10")
-        elif len(self.indices) <= 20:
-            color_map = cm.get_cmap("tab20")
-        else:
-            color_map = cm.get_cmap("tab20b")
-        # Create a dictionary of colors
-        colors = {name: color_map(i) for i, name in enumerate(self.indices.keys())}
-
-        for ax, indices in zip(axs, [self.original_indices, self.indices]):
-            added_labels = set()
-            for name in indices.keys():
-                for tuple in indices[name]:
-                    eta_index = tuple[0]
-                    xi_index = tuple[1]
-
-                    # transform coordinates to projected space
-                    transformed_lon, transformed_lat = trans.transform_point(
-                        self.grid.ds.lon_rho[eta_index, xi_index],
-                        self.grid.ds.lat_rho[eta_index, xi_index],
-                        proj,
-                    )
-
-                    if name not in added_labels:
-                        added_labels.add(name)
-                        label = name
-                    else:
-                        label = "_None"
-
-                    ax.plot(
-                        transformed_lon,
-                        transformed_lat,
-                        marker="x",
-                        markersize=8,
-                        markeredgewidth=2,
-                        label=label,
-                        color=colors[name],
-                    )
+            plot_2d_horizontal_field(field, kwargs=kwargs, ax=ax, add_colorbar=False)
+
+        points = {}
+        for j, (ax, indices) in enumerate(
+            [(ax, ind) for ax, ind in zip(axs, [self.original_indices, self.indices])]
+        ):
+            for name in river_names:
+                if name in indices:
+                    for i, (eta_index, xi_index) in enumerate(indices[name]):
+                        lon = self.grid.ds.lon_rho[eta_index, xi_index].item()
+                        lat = self.grid.ds.lat_rho[eta_index, xi_index].item()
+                        key = name if i == 0 else f"_{name}_{i}"
+                        points[key] = {
+                            "lon": lon,
+                            "lat": lat,
+                            "color": colors[name],
+                        }
+
+            plot_location(
+                grid_ds=self.grid.ds,
+                points=points,
+                ax=ax,
+                include_legend=(j == 1),
+            )
 
         axs[0].set_title("Original river locations")
         axs[1].set_title("Updated river locations")
 
-        axs[1].legend(loc="center left", bbox_to_anchor=(1.1, 0.5))
-
-    def plot(self, var_name="river_volume"):
+    def plot(
+        self,
+        var_name: str = "river_volume",
+        river_names: list[str] | str = INCLUDE_ALL_RIVER_NAMES,
+    ):
         """Plots the river flux (e.g., volume, temperature, or salinity) over time for
         all rivers.
 
@@ -591,8 +799,19 @@ class RiverForcing:
             - 'river_diazFe' : river diazFe (from river_tracer).
 
             The default is 'river_volume'.
+
+        river_names : list[str], or str, optional
+            A list of release names to plot.
+            If a string equal to "all", all rivers will be plotted.
+            Defaults to "all".
+
         """
-        fig, ax = plt.subplots(1, 1, figsize=(9, 5))
+        valid_river_names = list(self.indices.keys())
+        river_names = _validate_river_names(river_names, valid_river_names)
+        if len(valid_river_names) > MAX_DISTINCT_COLORS:
+            colors = assign_category_colors(river_names)
+        else:
+            colors = assign_category_colors(valid_river_names)
 
         if self.climatology:
             xticks = self.ds.month.values
@@ -614,16 +833,19 @@ class RiverForcing:
             units = d[var_name_wo_river]["units"]
             long_name = f"River {d[var_name_wo_river]['long_name']}"
 
-        for i in range(len(self.ds.nriver)):
+        fig, ax = plt.subplots(1, 1, figsize=(9, 5))
+        for name in river_names:
+            nriver = np.where(self.ds["river_name"].values == name)[0].item()
 
             ax.plot(
                 xticks,
-                field.isel(nriver=i),
+                field.isel(nriver=nriver),
                 marker="x",
                 markersize=8,
                 markeredgewidth=2,
                 lw=2,
-                label=self.ds.isel(nriver=i).river_name.values,
+                label=name,
+                color=colors[name],
             )
 
         ax.set_xticks(xticks)
@@ -639,7 +861,7 @@ class RiverForcing:
 
     def save(
         self,
-        filepath: Union[str, Path],
+        filepath: str | Path,
     ) -> None:
         """Save the river forcing to netCDF4 file.
 
@@ -653,7 +875,6 @@ class RiverForcing:
         List[Path]
             A list of `Path` objects for the saved files. Each element in the list corresponds to a file that was saved.
         """
-
         # Ensure filepath is a Path object
         filepath = Path(filepath)
 
@@ -668,7 +889,7 @@ class RiverForcing:
 
         return saved_filenames
 
-    def to_yaml(self, filepath: Union[str, Path]) -> None:
+    def to_yaml(self, filepath: str | Path) -> None:
         """Export the parameters of the class to a YAML file, including the version of
         roms-tools.
 
@@ -677,20 +898,27 @@ class RiverForcing:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-
-        forcing_dict = _to_dict(self)
+        forcing_dict = to_dict(self, exclude=["climatology"])
 
         # Convert indices format
         indices_data = forcing_dict["RiverForcing"]["indices"]
         serialized_indices = {"_convention": "eta_rho, xi_rho"}
         for key, value in indices_data.items():
             serialized_indices[key] = [f"{tup[0]}, {tup[1]}" for tup in value]
-        forcing_dict["RiverForcing"]["indices"] = serialized_indices
 
-        _write_to_yaml(forcing_dict, filepath)
+        # Remove keys starting with "overlap_"
+        filtered_indices = {
+            key: value
+            for key, value in serialized_indices.items()
+            if not key.startswith("overlap_")
+        }
+
+        forcing_dict["RiverForcing"]["indices"] = filtered_indices
+
+        write_to_yaml(forcing_dict, filepath)
 
     @classmethod
-    def from_yaml(cls, filepath: Union[str, Path]) -> "RiverForcing":
+    def from_yaml(cls, filepath: str | Path) -> "RiverForcing":
         """Create an instance of the RiverForcing class from a YAML file.
 
         Parameters
@@ -706,7 +934,7 @@ class RiverForcing:
         filepath = Path(filepath)
 
         grid = Grid.from_yaml(filepath)
-        params = _from_yaml(cls, filepath)
+        params = from_yaml(cls, filepath)
 
         def convert_indices_format(indices):
             indices = {
@@ -743,7 +971,6 @@ def check_river_locations_are_along_coast(mask, indices):
     ValueError
         If any river is not located on the coast.
     """
-
     faces = (
         mask.shift(eta_rho=1)
         + mask.shift(eta_rho=-1)
@@ -761,3 +988,38 @@ def check_river_locations_are_along_coast(mask, indices):
                 raise ValueError(
                     f"River `{key}` is not located on the coast at grid cell ({eta_rho}, {xi_rho})."
                 )
+
+
+def _validate_river_names(
+    river_names: list[str] | str, valid_river_names: list[str]
+) -> list[str]:
+    """
+    Validate and filter a list of river names.
+
+    Ensures that each river name exists in `valid_river_names` and limits the list
+    to `MAX_RIVERS_TO_PLOT` entries with a warning if truncated.
+
+    Parameters
+    ----------
+    river_names : list of str or INCLUDE_ALL_RIVER_NAMES
+        Names of rivers to plot, or sentinel to include all.
+    valid_river_names : list of str
+        List of valid river names.
+
+    Returns
+    -------
+    list of str
+        Validated and truncated list of river names.
+
+    Raises
+    ------
+    ValueError
+        If any names are invalid.
+    """
+    return validate_names(
+        river_names,
+        valid_river_names,
+        INCLUDE_ALL_RIVER_NAMES,
+        MAX_RIVERS_TO_PLOT,
+        label="river",
+    )