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

roms_tools/setup/river_forcing.py

@@ -4,14 +4,18 @@ from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
 
-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_2d_horizontal_field
+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,
@@ -26,10 +30,14 @@ from roms_tools.setup.utils import (
     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:
@@ -672,8 +680,24 @@ 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
         lat_deg = self.grid.ds.lat_rho
@@ -695,53 +719,37 @@ class RiverForcing:
         for ax in axs:
             plot_2d_horizontal_field(field, kwargs=kwargs, ax=ax, 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],
-                    )
+        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.
 
@@ -791,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
@@ -814,15 +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)
@@ -965,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",
+    )

roms_tools/setup/surface_forcing.py

@@ -150,12 +150,20 @@ class SurfaceForcing:
             use_coarse_grid = False
         elif self.coarse_grid_mode == "auto":
             use_coarse_grid = self._determine_coarse_grid_usage(data)
-        if use_coarse_grid:
-            logging.info("Data will be interpolated onto grid coarsened by factor 2.")
-        else:
-            logging.info("Data will be interpolated onto fine grid.")
         self.use_coarse_grid = use_coarse_grid
 
+        opt_file = "bulk_frc.opt" if self.type == "physics" else "bgc.opt"
+        grid_desc = "grid coarsened by factor 2" if use_coarse_grid else "fine grid"
+        interp_flag = 1 if use_coarse_grid else 0
+
+        logging.info(
+            "Data will be interpolated onto the %s. "
+            "Remember to set `interp_frc = %d` in your `%s` ROMS option file.",
+            grid_desc,
+            interp_flag,
+            opt_file,
+        )
+
         target_coords = get_target_coords(self.grid, self.use_coarse_grid)
         self.target_coords = target_coords
 

roms_tools/setup/utils.py

@@ -1,4 +1,5 @@
 import importlib.metadata
+import logging
 from collections.abc import Sequence
 from dataclasses import asdict, fields, is_dataclass
 from datetime import datetime
@@ -1794,3 +1795,59 @@ def to_float(val):
     if isinstance(val, list):
         return [float(v) for v in val]
     return float(val)
+
+
+def validate_names(
+    names: list[str] | str,
+    valid_names: list[str],
+    include_all_sentinel: str,
+    max_to_plot: int,
+    label: str = "item",
+) -> list[str]:
+    """
+    Generic validation and filtering for a list of names.
+
+    Parameters
+    ----------
+    names : list of str or sentinel
+        Names to validate, or sentinel value to include all valid names.
+    valid_names : list of str
+        List of valid names to check against.
+    include_all_sentinel : str
+        Sentinel value to indicate all names should be included.
+    max_to_plot : int
+        Maximum number of names to return.
+    label : str, default "item"
+        Label to use in error and warning messages.
+
+    Returns
+    -------
+    list of str
+        Validated and possibly truncated list of names.
+
+    Raises
+    ------
+    ValueError
+        If any names are invalid or input is not a list of strings.
+    """
+    if names == include_all_sentinel:
+        names = valid_names
+
+    if isinstance(names, list):
+        if not all(isinstance(n, str) for n in names):
+            raise ValueError(f"All elements in `{label}_names` must be strings.")
+    else:
+        raise ValueError(f"`{label}_names` should be a list of strings.")
+
+    invalid = [n for n in names if n not in valid_names]
+    if invalid:
+        raise ValueError(f"Invalid {label}s: {', '.join(invalid)}")
+
+    if len(names) > max_to_plot:
+        logging.warning(
+            f"Only the first {max_to_plot} {label}s will be plotted "
+            f"(received {len(names)})."
+        )
+        names = names[:max_to_plot]
+
+    return names

roms_tools/tests/test_setup/test_boundary_forcing.py

@@ -1,7 +1,9 @@
 import logging
+import os
 import textwrap
 from datetime import datetime
 from pathlib import Path
+from unittest import mock
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -758,3 +760,58 @@ def test_from_yaml_missing_boundary_forcing(tmp_path, use_dask):
 
         yaml_filepath = Path(yaml_filepath)
         yaml_filepath.unlink()
+
+
+@pytest.mark.stream
+@pytest.mark.use_dask
+@pytest.mark.use_copernicus
+def test_default_glorys_dataset_loading(tiny_grid: Grid) -> None:
+    """Verify the default GLORYS dataset is loaded when a path is not provided."""
+    start_time = datetime(2010, 2, 1)
+    end_time = datetime(2010, 3, 1)
+
+    with mock.patch.dict(
+        os.environ, {"PYDEVD_WARN_EVALUATION_TIMEOUT": "90"}, clear=True
+    ):
+        bf = BoundaryForcing(
+            grid=tiny_grid,
+            source={"name": "GLORYS"},
+            type="physics",
+            start_time=start_time,
+            end_time=end_time,
+            use_dask=True,
+            bypass_validation=True,
+        )
+
+        expected_vars = {"u_south", "v_south", "temp_south", "salt_south"}
+        assert set(bf.ds.data_vars).issuperset(expected_vars)
+
+
+@pytest.mark.parametrize(
+    "use_dask",
+    [pytest.param(True, marks=pytest.mark.use_dask), False],
+)
+def test_nondefault_glorys_dataset_loading(small_grid: Grid, use_dask: bool) -> None:
+    """Verify a non-default GLORYS dataset is loaded when a path is provided."""
+    start_time = datetime(2012, 1, 1)
+    end_time = datetime(2012, 12, 31)
+
+    local_path = Path(download_test_data("GLORYS_NA_20120101.nc"))
+
+    with mock.patch.dict(
+        os.environ, {"PYDEVD_WARN_EVALUATION_TIMEOUT": "90"}, clear=True
+    ):
+        bf = BoundaryForcing(
+            grid=small_grid,
+            source={
+                "name": "GLORYS",
+                "path": local_path,
+            },
+            type="physics",
+            start_time=start_time,
+            end_time=end_time,
+            use_dask=use_dask,
+        )
+
+        expected_vars = {"u_south", "v_south", "temp_south", "salt_south"}
+        assert set(bf.ds.data_vars).issuperset(expected_vars)

roms_tools/tests/test_setup/test_cdr_forcing.py

@@ -9,7 +9,7 @@ from pydantic import ValidationError
 
 from conftest import calculate_file_hash
 from roms_tools import CDRForcing, Grid, TracerPerturbation, VolumeRelease
-from roms_tools.constants import NUM_TRACERS
+from roms_tools.constants import MAX_DISTINCT_COLORS, NUM_TRACERS
 from roms_tools.setup.cdr_forcing import (
     CDRForcingDatasetBuilder,
     ReleaseCollector,
@@ -725,6 +725,8 @@ class TestCDRForcing:
             rot=0,
             N=3,
         )
+        self.grid = grid
+
         grid_that_straddles = Grid(
             nx=18,
             ny=18,
@@ -817,8 +819,13 @@ class TestCDRForcing:
             self.volume_release_cdr_forcing_with_straddling_grid,
         ]:
             cdr.plot_volume_flux()
+            cdr.plot_volume_flux(release_names=["first_release"])
+
             cdr.plot_tracer_concentration("ALK")
+            cdr.plot_tracer_concentration("ALK", release_names=["first_release"])
+
             cdr.plot_tracer_concentration("DIC")
+            cdr.plot_tracer_concentration("DIC", release_names=["first_release"])
 
         self.volume_release_cdr_forcing.plot_locations()
         self.volume_release_cdr_forcing.plot_locations(release_names=["first_release"])
@@ -830,13 +837,56 @@ class TestCDRForcing:
             self.tracer_perturbation_cdr_forcing_with_straddling_grid,
         ]:
             cdr.plot_tracer_flux("ALK")
+            cdr.plot_tracer_flux("ALK", release_names=["first_release"])
+
             cdr.plot_tracer_flux("DIC")
+            cdr.plot_tracer_flux("DIC", release_names=["first_release"])
 
         self.tracer_perturbation_cdr_forcing.plot_locations()
         self.tracer_perturbation_cdr_forcing.plot_locations(
             release_names=["first_release"]
         )
 
+    def test_plot_max_releases(self, caplog):
+        # Prepare releases with more than MAX_DISTINCT_COLORS unique names
+        releases = []
+        for i in range(MAX_DISTINCT_COLORS + 1):
+            release = self.first_volume_release.__replace__(name=f"release_{i}")
+            releases.append(release)
+
+        # Construct a CDRForcing object with too many releases to plot
+        cdr_forcing = CDRForcing(
+            grid=self.grid,
+            start_time=self.start_time,
+            end_time=self.end_time,
+            releases=releases,
+        )
+
+        release_names = [r.name for r in releases]
+
+        plot_methods_with_release_names = [
+            cdr_forcing.plot_locations,
+            cdr_forcing.plot_volume_flux,
+        ]
+
+        for plot_func in plot_methods_with_release_names:
+            caplog.clear()
+            with caplog.at_level("WARNING"):
+                plot_func(release_names=release_names)
+            assert any(
+                f"Only the first {MAX_DISTINCT_COLORS} releases will be plotted"
+                in message
+                for message in caplog.messages
+            ), f"Warning not raised by {plot_func.__name__}"
+
+        with caplog.at_level("WARNING"):
+            cdr_forcing.plot_locations(release_names=release_names)
+
+        assert any(
+            f"Only the first {MAX_DISTINCT_COLORS} releases will be plotted" in message
+            for message in caplog.messages
+        )
+
     @pytest.mark.skipif(xesmf is None, reason="xesmf required")
     def test_plot_distribution(self):
         self.volume_release_cdr_forcing.plot_distribution("first_release")
@@ -856,10 +906,10 @@ class TestCDRForcing:
         with pytest.raises(ValueError, match="Invalid releases"):
             self.volume_release_cdr_forcing.plot_locations(release_names=["fake"])
 
-        with pytest.raises(ValueError, match="should be a string"):
+        with pytest.raises(ValueError, match="should be a list"):
             self.volume_release_cdr_forcing.plot_locations(release_names=4)
 
-        with pytest.raises(ValueError, match="list must be strings"):
+        with pytest.raises(ValueError, match="must be strings"):
             self.volume_release_cdr_forcing.plot_locations(release_names=[4])
 
     def test_cdr_forcing_save(self, tmp_path):

roms_tools/tests/test_setup/test_datasets.py

@@ -2,6 +2,7 @@ import logging
 from collections import OrderedDict
 from datetime import datetime
 from pathlib import Path
+from unittest import mock
 
 import numpy as np
 import pytest
@@ -11,11 +12,14 @@ from roms_tools.download import download_test_data
 from roms_tools.setup.datasets import (
     CESMBGCDataset,
     Dataset,
+    ERA5ARCODataset,
     ERA5Correction,
     GLORYSDataset,
+    GLORYSDefaultDataset,
     RiverDataset,
     TPXODataset,
 )
+from roms_tools.setup.surface_forcing import DEFAULT_ERA5_ARCO_PATH
 
 
 @pytest.fixture
@@ -437,6 +441,78 @@ def test_era5_correction_choose_subdomain(use_dask):
     assert (data.ds["longitude"] == lons).all()
 
 
+@pytest.mark.use_gcsfs
+def test_default_era5_dataset_loading_without_dask() -> None:
+    """Verify that loading the default ERA5 dataset fails if use_dask is not True."""
+    start_time = datetime(2020, 2, 1)
+    end_time = datetime(2020, 2, 2)
+
+    with pytest.raises(ValueError):
+        _ = ERA5ARCODataset(
+            filename=DEFAULT_ERA5_ARCO_PATH,
+            start_time=start_time,
+            end_time=end_time,
+            use_dask=False,
+        )
+
+
+@pytest.mark.skip("Temporary skip until memory consumption issue is addressed. # TODO")
+@pytest.mark.stream
+@pytest.mark.use_dask
+@pytest.mark.use_gcsfs
+def test_default_era5_dataset_loading() -> None:
+    """Verify the default ERA5 dataset is loaded correctly."""
+    start_time = datetime(2020, 2, 1)
+    end_time = datetime(2020, 2, 2)
+
+    ds = ERA5ARCODataset(
+        filename=DEFAULT_ERA5_ARCO_PATH,
+        start_time=start_time,
+        end_time=end_time,
+        use_dask=True,
+    )
+
+    expected_vars = {"uwnd", "vwnd", "swrad", "lwrad", "Tair", "rain"}
+    assert set(ds.var_names).issuperset(expected_vars)
+
+
+@pytest.mark.use_copernicus
+def test_default_glorys_dataset_loading_dask_not_installed() -> None:
+    """Verify that loading the default GLORYS dataset fails if dask is not available."""
+    start_time = datetime(2020, 2, 1)
+    end_time = datetime(2020, 2, 2)
+
+    with (
+        pytest.raises(RuntimeError),
+        mock.patch("roms_tools.utils._has_dask", return_value=False),
+    ):
+        _ = GLORYSDefaultDataset(
+            filename=GLORYSDefaultDataset.dataset_name,
+            start_time=start_time,
+            end_time=end_time,
+            use_dask=True,
+        )
+
+
+@pytest.mark.stream
+@pytest.mark.use_copernicus
+@pytest.mark.use_dask
+def test_default_glorys_dataset_loading() -> None:
+    """Verify the default GLORYS dataset is loaded correctly."""
+    start_time = datetime(2012, 1, 1)
+    end_time = datetime(2013, 1, 1)
+
+    ds = GLORYSDefaultDataset(
+        filename=GLORYSDefaultDataset.dataset_name,
+        start_time=start_time,
+        end_time=end_time,
+        use_dask=True,
+    )
+
+    expected_vars = {"temp", "salt", "u", "v", "zeta"}
+    assert set(ds.var_names).issuperset(expected_vars)
+
+
 def test_data_concatenation(use_dask):
     fname = download_test_data("GLORYS_NA_2012.nc")
     data = GLORYSDataset(

roms_tools/tests/test_setup/test_grid.py

@@ -21,6 +21,11 @@ from roms_tools.constants import (
 from roms_tools.download import download_test_data
 from roms_tools.setup.topography import _compute_rfactor
 
+try:
+    import xesmf  # type: ignore
+except ImportError:
+    xesmf = None
+
 
 @pytest.fixture()
 def counter_clockwise_rotated_grid():
@@ -177,13 +182,18 @@ def test_successful_initialization_with_topography(grid_fixture, request):
     assert grid is not None
 
 
-def test_plot():
-    grid = Grid(
-        nx=20, ny=20, size_x=100, size_y=100, center_lon=-20, center_lat=0, rot=0
-    )
+def test_plot(grid_that_straddles_180_degree_meridian):
+    grid_that_straddles_180_degree_meridian.plot(with_dim_names=False)
+    grid_that_straddles_180_degree_meridian.plot(with_dim_names=True)
+
+
+@pytest.mark.skipif(xesmf is None, reason="xesmf required")
+def test_plot_along_lat_lon(grid_that_straddles_180_degree_meridian):
+    grid_that_straddles_180_degree_meridian.plot(lat=61)
+    grid_that_straddles_180_degree_meridian.plot(lon=180)
 
-    grid.plot(with_dim_names=False)
-    grid.plot(with_dim_names=True)
+    with pytest.raises(ValueError, match="Specify either `lat` or `lon`, not both"):
+        grid_that_straddles_180_degree_meridian.plot(lat=61, lon=180)
 
 
 def test_save(tmp_path):