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

roms_tools/tests/test_setup/test_datasets.py

@@ -1,18 +1,21 @@
-import pytest
 import logging
+from collections import OrderedDict
 from datetime import datetime
+from pathlib import Path
+
 import numpy as np
+import pytest
 import xarray as xr
+
+from roms_tools.download import download_test_data
 from roms_tools.setup.datasets import (
+    CESMBGCDataset,
     Dataset,
-    GLORYSDataset,
     ERA5Correction,
-    CESMBGCDataset,
+    GLORYSDataset,
+    RiverDataset,
     TPXODataset,
 )
-from roms_tools.download import download_test_data
-from pathlib import Path
-from collections import OrderedDict
 
 
 @pytest.fixture
@@ -219,7 +222,8 @@ def test_multiple_matching_times(
     global_dataset_with_multiple_times_per_day, tmp_path, use_dask
 ):
     """Test handling when multiple matching times are found when end_time is not
-    specified."""
+    specified.
+    """
     filepath = tmp_path / "test.nc"
     global_dataset_with_multiple_times_per_day.to_netcdf(filepath)
     dataset = Dataset(
@@ -268,7 +272,6 @@ def test_warnings_times(global_dataset, tmp_path, caplog, use_dask):
 
 def test_from_ds(global_dataset, global_dataset_with_noon_times, use_dask, tmp_path):
     """Test the from_ds method of the Dataset class."""
-
     start_time = datetime(2022, 1, 1)
 
     filepath = tmp_path / "test.nc"
@@ -297,16 +300,17 @@ def test_from_ds(global_dataset, global_dataset_with_noon_times, use_dask, tmp_p
     # Ensure other attributes are copied correctly
     for attr in vars(dataset):
         if attr != "ds":
-            assert getattr(new_dataset, attr) == getattr(
-                dataset, attr
-            ), f"Attribute {attr} not copied correctly."
+            assert getattr(new_dataset, attr) == getattr(dataset, attr), (
+                f"Attribute {attr} not copied correctly."
+            )
 
 
 def test_reverse_latitude_reverse_depth_choose_subdomain(
     global_dataset, tmp_path, use_dask
 ):
     """Test reversing latitude when it is not ascending, the choose_subdomain method,
-    and the convert_to_negative_depth method of the Dataset class."""
+    and the convert_to_negative_depth method of the Dataset class.
+    """
     start_time = datetime(2022, 1, 1)
 
     filepath = tmp_path / "test.nc"
@@ -365,7 +369,6 @@ def test_reverse_latitude_reverse_depth_choose_subdomain(
 
 
 def test_check_if_global_with_global_dataset(global_dataset, tmp_path, use_dask):
-
     filepath = tmp_path / "test.nc"
     global_dataset.to_netcdf(filepath)
     dataset = Dataset(filename=filepath, var_names={"var": "var"}, use_dask=use_dask)
@@ -376,7 +379,6 @@ def test_check_if_global_with_global_dataset(global_dataset, tmp_path, use_dask)
 def test_check_if_global_with_non_global_dataset(
     non_global_dataset, tmp_path, use_dask
 ):
-
     filepath = tmp_path / "test.nc"
     non_global_dataset.to_netcdf(filepath)
     dataset = Dataset(filename=filepath, var_names={"var": "var"}, use_dask=use_dask)
@@ -386,7 +388,6 @@ def test_check_if_global_with_non_global_dataset(
 
 
 def test_check_dataset(global_dataset, tmp_path, use_dask):
-
     ds = global_dataset.copy()
     ds = ds.drop_vars("var")
 
@@ -398,7 +399,6 @@ def test_check_dataset(global_dataset, tmp_path, use_dask):
     with pytest.raises(
         ValueError, match="Dataset does not contain all required variables."
     ):
-
         Dataset(
             filename=filepath,
             var_names={"var": "var"},
@@ -418,7 +418,6 @@ def test_check_dataset(global_dataset, tmp_path, use_dask):
     with pytest.raises(
         ValueError, match="Dataset does not contain all required dimensions."
     ):
-
         Dataset(
             filename=filepath,
             var_names={"var": "var"},
@@ -429,7 +428,6 @@ def test_check_dataset(global_dataset, tmp_path, use_dask):
 
 
 def test_era5_correction_choose_subdomain(use_dask):
-
     data = ERA5Correction(use_dask=use_dask)
     lats = data.ds.latitude[10:20]
     lons = data.ds.longitude[10:20]
@@ -440,7 +438,6 @@ def test_era5_correction_choose_subdomain(use_dask):
 
 
 def test_data_concatenation(use_dask):
-
     fname = download_test_data("GLORYS_NA_2012.nc")
     data = GLORYSDataset(
         filename=fname,
@@ -474,7 +471,6 @@ def test_data_concatenation(use_dask):
 
 
 def test_time_validation(use_dask):
-
     fname = download_test_data("GLORYS_NA_2012.nc")
 
     with pytest.raises(TypeError, match="start_time must be a datetime object"):
@@ -485,7 +481,6 @@ def test_time_validation(use_dask):
             use_dask=use_dask,
         )
     with pytest.raises(TypeError, match="end_time must be a datetime object"):
-
         GLORYSDataset(
             filename=fname,
             start_time=datetime(2012, 1, 1),
@@ -495,7 +490,6 @@ def test_time_validation(use_dask):
 
 
 def test_climatology_error(use_dask):
-
     fname = download_test_data("GLORYS_NA_2012.nc")
 
     with pytest.raises(
@@ -516,7 +510,6 @@ def test_climatology_error(use_dask):
         ValueError,
         match="The dataset contains integer time values, which are only supported when the climatology flag is set to True. However, your climatology flag is set to False.",
     ):
-
         CESMBGCDataset(
             filename=fname_bgc,
             start_time=datetime(2012, 1, 1),
@@ -539,7 +532,6 @@ def test_climatology_error(use_dask):
     ],
 )
 def test_horizontal_resolution(data_fixture, expected_resolution, request):
-
     data = request.getfixturevalue(data_fixture)
     assert np.isclose(data.resolution, expected_resolution)
 
@@ -561,7 +553,7 @@ class TestTPXODataset:
 
     @pytest.fixture
     def global_tpxo_dataset(self, use_dask):
-        """TPXO dataset with global coverage and 1 constituent: M2"""
+        """TPXO dataset with global coverage and 1 constituent: M2."""
         fname_grid = Path(download_test_data("global_grid_tpxo10.v2.nc"))
         fname_h = Path(download_test_data("global_h_tpxo10.v2.nc"))
 
@@ -601,7 +593,6 @@ class TestTPXODataset:
         ],
     )
     def test_initialization_and_clean_up(self, tpxo_dataset_fixture, request):
-
         tpxo_dataset = request.getfixturevalue(tpxo_dataset_fixture)
         assert tpxo_dataset.location == "h"
 
@@ -614,14 +605,12 @@ class TestTPXODataset:
         assert "latitude" in tpxo_dataset.ds.variables
 
     def test_select_fewer_constituents(self, regional_tpxo_dataset, omega):
-
         regional_tpxo_dataset.select_constituents(2, omega)
         assert (
             regional_tpxo_dataset.ds["ntides"].values.astype("U3") == ["m2", "s2"]
         ).all()
 
     def test_select_constituents_with_reordering(self, regional_tpxo_dataset, omega):
-
         regional_tpxo_dataset.select_constituents(11, omega)
 
         assert len(regional_tpxo_dataset.ds["ntides"]) == 11
@@ -632,9 +621,8 @@ class TestTPXODataset:
         )
 
     def test_select_constituents_omega_mismatch(self, regional_tpxo_dataset, omega):
-
         omega = OrderedDict(
-            list(omega.items())[:3] + [("fake", 6.495854e-05)] + list(omega.items())[3:]
+            [*list(omega.items())[:3], ("fake", 6.495854e-05), *list(omega.items())[3:]]
         )
         with pytest.raises(ValueError, match="The dataset contains tidal constituents"):
             regional_tpxo_dataset.select_constituents(11, omega)
@@ -642,3 +630,45 @@ class TestTPXODataset:
     def test_select_constituents_too_few(self, global_tpxo_dataset, omega):
         with pytest.raises(ValueError, match="The dataset contains tidal constituents"):
             global_tpxo_dataset.select_constituents(11, omega)
+
+
+class TestRiverDataset:
+    def test_deduplicate_river_names(self, tmp_path):
+        sample_dim_and_var_names = {
+            "dim_names": {"station": "station", "time": "time"},
+            "var_names": {
+                "latitude": "lat",
+                "longitude": "lon",
+                "flux": "flux",
+                "ratio": "ratio",
+                "name": "name",
+            },
+        }
+
+        data = {
+            "lat": (["station"], [10.0, 20.0, 30.0]),
+            "lon": (["station"], [100.0, 110.0, 120.0]),
+            "flux": (["time", "station"], np.random.rand(1, 3)),
+            "ratio": (["time", "station"], np.random.rand(1, 3)),
+            "name": (["station"], ["Amazon", "Nile", "Amazon"]),  # duplicate
+        }
+        coords = {"station": [0, 1, 2], "time": [0]}
+        ds = xr.Dataset(data, coords=coords)
+
+        # Write to temporary NetCDF file
+        file_path = tmp_path / "rivers.nc"
+        ds.to_netcdf(file_path)
+
+        river_dataset = RiverDataset(
+            filename=file_path,
+            start_time=datetime(2000, 1, 1),
+            end_time=datetime(2000, 1, 2),
+            dim_names=sample_dim_and_var_names["dim_names"],
+            var_names=sample_dim_and_var_names["var_names"],
+        )
+
+        names = river_dataset.ds["name"].values
+        assert "Amazon_1" in names
+        assert "Amazon_2" in names
+        assert "Nile" in names
+        assert len(set(names)) == len(names)  # all names must be unique

roms_tools/tests/test_setup/test_fill.py

@@ -1,8 +1,9 @@
-import pytest
-from roms_tools.setup.fill import LateralFill
 import numpy as np
+import pytest
 import xarray as xr
 
+from roms_tools.setup.fill import LateralFill
+
 
 @pytest.mark.parametrize(
     "data_fixture",
@@ -37,7 +38,6 @@ def test_lateral_fill_no_nans(data_fixture, request):
 
 
 def test_lateral_fill_correct_order_of_magnitude(coarsened_cesm_bgc_data):
-
     lateral_fill = LateralFill(
         coarsened_cesm_bgc_data.ds["mask"],
         [
@@ -75,7 +75,6 @@ def test_lateral_fill_correct_order_of_magnitude(coarsened_cesm_bgc_data):
     ],
 )
 def test_lateral_fill_reproducibility(data_fixture, request):
-
     data = request.getfixturevalue(data_fixture)
     lateral_fill0 = LateralFill(
         data.ds["mask"],