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

roms_tools/tests/test_setup/test_nesting.py

@@ -89,7 +89,6 @@ class TestInterpolateIndices:
     )
     def test_correct_indices_of_same_grid(self, grid, caplog, request):
         """Verify boundary indices are correctly interpolated for the same grid."""
-
         grid = request.getfixturevalue(grid)
 
         bdry_coords_dict = get_boundary_coords()
@@ -140,7 +139,6 @@ class TestInterpolateIndices:
         self, big_grid_fixture, small_grid_fixture, request
     ):
         """Ensure interpolated indices fall within the parent grid's bounds."""
-
         big_grid = request.getfixturevalue(big_grid_fixture)
         small_grid = request.getfixturevalue(small_grid_fixture)
 
@@ -168,7 +166,6 @@ class TestInterpolateIndices:
 class TestMapChildBoundaries:
     def test_update_indices_does_nothing_if_no_parent_land(self, small_grid, baby_grid):
         """Verify no change in indices when parent grid has no land at boundaries."""
-
         ds_without_updated_indices = map_child_boundaries_onto_parent_grid_indices(
             small_grid.ds, baby_grid.ds, update_land_indices=False
         )
@@ -189,7 +186,6 @@ class TestMapChildBoundaries:
         self, big_grid_fixture, small_grid_fixture, request
     ):
         """Check updated indices map to wet points on the parent grid."""
-
         big_grid = request.getfixturevalue(big_grid_fixture)
         small_grid = request.getfixturevalue(small_grid_fixture)
 
@@ -223,8 +219,8 @@ class TestMapChildBoundaries:
         self, big_grid_fixture, small_grid_fixture, request
     ):
         """Test that child boundary indices are monotonically increasing or decreasing
-        in both the xi and eta directions, for all boundaries and locations."""
-
+        in both the xi and eta directions, for all boundaries and locations.
+        """
         big_grid = request.getfixturevalue(big_grid_fixture)
         small_grid = request.getfixturevalue(small_grid_fixture)
 
@@ -262,7 +258,6 @@ class TestBoundaryDistance:
         self, grid_fixture, request
     ):
         """Ensure boundary distance is zero for grids without land along boundaries."""
-
         grid = request.getfixturevalue(grid_fixture)
         alpha = compute_boundary_distance(grid.ds.mask_rho)
 
@@ -282,7 +277,8 @@ class TestBoundaryDistance:
 
     def test_boundary_distance_for_grid_with_land_along_boundary(self, big_grid):
         """Test that there are 1s along the boundary of alpha if the grid has land along
-        the boundary."""
+        the boundary.
+        """
         alpha = compute_boundary_distance(big_grid.ds.mask_rho)
         assert (alpha.isel(eta_rho=0) == 1).any()
         assert (alpha.isel(eta_rho=-1) == 1).any()
@@ -295,7 +291,6 @@ class TestModifyChid:
         self, small_grid, baby_grid
     ):
         """Confirm child mask remains unchanged if no parent land is at boundaries."""
-
         mask_original = baby_grid.ds.mask_rho.copy()
         modified_baby_grid_ds = modify_child_topography_and_mask(
             small_grid.ds, baby_grid.ds
@@ -312,7 +307,6 @@ class TestModifyChid:
     )
     def test_no_modification_if_parent_and_child_coincide(self, grid_fixture, request):
         """Ensure no changes occur when parent and child grids coincide."""
-
         grid = request.getfixturevalue(grid_fixture)
 
         h_original = grid.ds.h.copy()
@@ -324,8 +318,8 @@ class TestModifyChid:
 
     def test_modification_only_along_boundaries(self, big_grid, small_grid):
         """Test that modifications to the child grid's topography and mask occur only
-        along the boundaries, leaving the interior unchanged."""
-
+        along the boundaries, leaving the interior unchanged.
+        """
         # Make copies of original data for comparison
         h_original = small_grid.ds.h.copy()
         mask_original = small_grid.ds.mask_rho.copy()
@@ -340,15 +334,15 @@ class TestModifyChid:
         # Assert that the center values remain the same
         assert mask_original.isel(
             eta_rho=eta_center, xi_rho=xi_center
-        ) == modified_ds.mask_rho.isel(
-            eta_rho=eta_center, xi_rho=xi_center
-        ), "Mask at the grid center was modified."
+        ) == modified_ds.mask_rho.isel(eta_rho=eta_center, xi_rho=xi_center), (
+            "Mask at the grid center was modified."
+        )
 
         assert h_original.isel(
             eta_rho=eta_center, xi_rho=xi_center
-        ) == modified_ds.h.isel(
-            eta_rho=eta_center, xi_rho=xi_center
-        ), "Topography at the grid center was modified."
+        ) == modified_ds.h.isel(eta_rho=eta_center, xi_rho=xi_center), (
+            "Topography at the grid center was modified."
+        )
 
 
 class TestNesting:
@@ -428,11 +422,9 @@ class TestNesting:
 
     def test_save(self, child_grid, tmp_path):
         """Test save methods."""
-
         for file_str in ["test_grid", "test_grid.nc"]:
             # Create a temporary filepath using the tmp_path fixture
             for filepath in [tmp_path / file_str, str(tmp_path / file_str)]:
-
                 saved_filenames = child_grid.save(filepath)
                 # Check if the .nc file was created
                 filepath = Path(filepath).with_suffix(".nc")
@@ -444,7 +436,6 @@ class TestNesting:
         for file_str in ["test_nesting", "test_nesting.nc"]:
             # Create a temporary filepath using the tmp_path fixture
             for filepath in [tmp_path / file_str, str(tmp_path / file_str)]:
-
                 saved_filenames = child_grid.save_nesting(filepath)
                 # Check if the .nc file was created
                 filepath = Path(filepath).with_suffix(".nc")
@@ -463,15 +454,14 @@ class TestNesting:
 
     def test_roundtrip_yaml(self, child_grid, tmp_path):
         """Test that creating a ChildGrid object, saving its parameters to yaml file,
-        and re-opening yaml file creates the same object."""
-
+        and re-opening yaml file creates the same object.
+        """
         # Create a temporary filepath using the tmp_path fixture
         file_str = "test_yaml"
         for filepath in [
             tmp_path / file_str,
             str(tmp_path / file_str),
         ]:  # test for Path object and str
-
             child_grid.to_yaml(filepath)
 
             child_grid_from_file = ChildGrid.from_yaml(filepath)
@@ -482,7 +472,6 @@ class TestNesting:
             filepath.unlink()
 
     def test_files_have_same_hash(self, child_grid, tmp_path):
-
         yaml_filepath = tmp_path / "test_yaml.yaml"
         filepath1 = tmp_path / "test1.nc"
         filepath2 = tmp_path / "test2.nc"

roms_tools/tests/test_setup/test_river_forcing.py

@@ -9,6 +9,7 @@ import xarray as xr
 
 from conftest import calculate_file_hash
 from roms_tools import Grid, RiverForcing
+from roms_tools.constants import MAX_DISTINCT_COLORS
 
 
 @pytest.fixture
@@ -21,7 +22,6 @@ def iceland_test_grid():
 @pytest.fixture
 def river_forcing_climatology(iceland_test_grid):
     """Fixture for creating a RiverForcing object from the global Dai river dataset."""
-
     start_time = datetime(1998, 1, 1)
     end_time = datetime(1998, 3, 1)
 
@@ -36,8 +36,8 @@ def river_forcing_climatology(iceland_test_grid):
 @pytest.fixture
 def river_forcing_for_grid_that_straddles_dateline():
     """Fixture for creating a RiverForcing object from the global Dai river dataset for
-    a grid that straddles the dateline."""
-
+    a grid that straddles the dateline.
+    """
     grid = Grid(
         nx=18,
         ny=18,
@@ -58,6 +58,29 @@ def river_forcing_for_grid_that_straddles_dateline():
     )
 
 
+@pytest.fixture
+def river_forcing_for_gulf_of_mexico():
+    """Fixture for creating a RiverForcing object for the Gulf of Mexico with 45 rivers."""
+    grid = Grid(
+        nx=20,
+        ny=15,
+        size_x=2000,
+        size_y=1500,
+        center_lon=-89,
+        center_lat=24,
+        rot=0,
+        N=3,
+    )
+    start_time = datetime(2012, 1, 1)
+    end_time = datetime(2012, 1, 31)
+
+    return RiverForcing(
+        grid=grid,
+        start_time=start_time,
+        end_time=end_time,
+    )
+
+
 @pytest.fixture
 def single_cell_indices():
     # These are the indices that the `river_forcing` fixture generates automatically.
@@ -90,8 +113,8 @@ def river_forcing_with_prescribed_single_cell_indices(
 ):
     """Fixture for creating a RiverForcing object based on the global Dai river dataset,
     using manually specified single-cell river indices instead of relying on automatic
-    detection."""
-
+    detection.
+    """
     start_time = datetime(1998, 1, 1)
     end_time = datetime(1998, 3, 1)
 
@@ -109,8 +132,8 @@ def river_forcing_with_prescribed_multi_cell_indices(
 ):
     """Fixture for creating a RiverForcing object based on the global Dai river dataset,
     using manually specified multi-cell river indices instead of relying on automatic
-    detection."""
-
+    detection.
+    """
     start_time = datetime(1998, 1, 1)
     end_time = datetime(1998, 3, 1)
 
@@ -201,7 +224,6 @@ class TestRiverForcingGeneral:
         )
 
     def test_reproducibility_same_grid(self, river_forcing):
-
         the_same_river_forcing = RiverForcing(
             grid=river_forcing.grid,
             start_time=datetime(1998, 1, 1),
@@ -234,7 +256,7 @@ class TestRiverForcingGeneral:
 
         indices = river_forcing.indices
         for name in indices.keys():
-            for (eta_rho, xi_rho) in indices[name]:
+            for eta_rho, xi_rho in indices[name]:
                 assert coast[eta_rho, xi_rho]
                 assert river_forcing.ds["river_index"][eta_rho, xi_rho] > 0
                 assert river_forcing.ds["river_fraction"][eta_rho, xi_rho] > 0
@@ -249,14 +271,46 @@ class TestRiverForcingGeneral:
             )
 
     def test_river_forcing_plot(self, river_forcing_with_bgc):
-        """Test plot method."""
+        """Test plot methods with and without specifying river_names."""
+        river_names = list(river_forcing_with_bgc.indices.keys())[0:2]
 
+        # Test plot_locations
         river_forcing_with_bgc.plot_locations()
-        river_forcing_with_bgc.plot("river_volume")
-        river_forcing_with_bgc.plot("river_temp")
-        river_forcing_with_bgc.plot("river_salt")
-        river_forcing_with_bgc.plot("river_ALK")
-        river_forcing_with_bgc.plot("river_PO4")
+        river_forcing_with_bgc.plot_locations(river_names=river_names)
+
+        # Fields to test
+        variables = [
+            "river_volume",
+            "river_temp",
+            "river_salt",
+            "river_ALK",
+            "river_PO4",
+        ]
+
+        for var in variables:
+            river_forcing_with_bgc.plot(var)
+            river_forcing_with_bgc.plot(var, river_names=river_names)
+
+    def test_plot_max_releases(self, caplog, river_forcing_for_gulf_of_mexico):
+        river_names = list(river_forcing_for_gulf_of_mexico.indices.keys())
+
+        caplog.clear()
+        with caplog.at_level("WARNING"):
+            river_forcing_for_gulf_of_mexico.plot_locations()
+        assert any(
+            f"Only the first {MAX_DISTINCT_COLORS} rivers will be plotted" in message
+            for message in caplog.messages
+        )
+
+        with caplog.at_level("WARNING"):
+            river_forcing_for_gulf_of_mexico.plot(
+                "river_volume", river_names=river_names
+            )
+
+        assert any(
+            f"Only the first {MAX_DISTINCT_COLORS} rivers will be plotted" in message
+            for message in caplog.messages
+        )
 
     @pytest.mark.parametrize(
         "river_forcing_fixture",
@@ -267,12 +321,10 @@ class TestRiverForcingGeneral:
     )
     def test_river_forcing_save(self, river_forcing_fixture, tmp_path, request):
         """Test save method."""
-
         river_forcing = request.getfixturevalue(river_forcing_fixture)
         for file_str in ["test_rivers", "test_rivers.nc"]:
             # Create a temporary filepath using the tmp_path fixture
             for filepath in [tmp_path / file_str, str(tmp_path / file_str)]:
-
                 saved_filenames = river_forcing.save(filepath)
                 # Check if the .nc file was created
                 filepath = Path(filepath).with_suffix(".nc")
@@ -293,8 +345,8 @@ class TestRiverForcingGeneral:
     )
     def test_roundtrip_yaml(self, river_forcing_fixture, request, tmp_path, caplog):
         """Test that creating an RiverForcing object, saving its parameters to yaml
-        file, and re-opening yaml file creates the same object."""
-
+        file, and re-opening yaml file creates the same object.
+        """
         river_forcing = request.getfixturevalue(river_forcing_fixture)
 
         # Create a temporary filepath using the tmp_path fixture
@@ -303,7 +355,6 @@ class TestRiverForcingGeneral:
             tmp_path / file_str,
             str(tmp_path / file_str),
         ]:  # test for Path object and str
-
             river_forcing.to_yaml(filepath)
 
             # Clear caplog before running the test
@@ -329,7 +380,6 @@ class TestRiverForcingGeneral:
         ],
     )
     def test_files_have_same_hash(self, river_forcing_fixture, request, tmp_path):
-
         river_forcing = request.getfixturevalue(river_forcing_fixture)
 
         yaml_filepath = tmp_path / "test_yaml.yaml"
@@ -376,7 +426,6 @@ class TestRiverForcingGeneral:
             tmp_path / file_str,
             str(tmp_path / file_str),
         ]:  # test for Path object and str
-
             # Write YAML content to file
             if isinstance(yaml_filepath, Path):
                 yaml_filepath.write_text(yaml_content)
@@ -399,7 +448,6 @@ class TestRiverForcingWithoutPrescribedIndices:
     end_time = datetime(1998, 3, 1)
 
     def test_logging_message(self, iceland_test_grid, caplog):
-
         with caplog.at_level(logging.INFO):
             RiverForcing(
                 grid=iceland_test_grid,
@@ -424,7 +472,6 @@ class TestRiverForcingWithoutPrescribedIndices:
         compare_dictionaries(river_forcing.indices, river_forcing_climatology.indices)
 
     def test_no_rivers_found(self):
-
         # Create a grid over open ocean
         grid = Grid(
             nx=2, ny=2, size_x=50, size_y=50, center_lon=0, center_lat=55, rot=10, N=3
@@ -439,7 +486,6 @@ class TestRiverForcingWithPrescribedIndices:
     end_time = datetime(1998, 3, 1)
 
     def test_logging_message(self, single_cell_indices, caplog, iceland_test_grid):
-
         with caplog.at_level(logging.INFO):
             RiverForcing(
                 grid=iceland_test_grid,
@@ -689,3 +735,257 @@ class TestRiverForcingWithPrescribedIndices:
                 end_time=self.end_time,
                 indices=indices,
             )
+
+
+class TestRiverForcingWithOverlappingIndices:
+    def test_get_overlapping_rivers(self):
+        indices = {
+            "RiverA": [(1, 2), (3, 4)],
+            "RiverB": [(1, 2), (5, 6)],
+            "RiverC": [(7, 8)],
+            "RiverD": [(3, 4)],
+        }
+
+        # Create a dummy RiverForcing instance with these indices
+        rf = RiverForcing.__new__(RiverForcing)
+        rf.indices = indices
+
+        overlapping = rf._get_overlapping_rivers()
+
+        expected = {
+            (1, 2): ["RiverA", "RiverB"],
+            (3, 4): ["RiverA", "RiverD"],
+        }
+
+        # Check keys
+        assert set(overlapping.keys()) == set(expected.keys())
+
+        # Check that all expected rivers are present for each overlapping cell
+        for idx, rivers in expected.items():
+            assert set(overlapping[idx]) == set(rivers)
+
+        # Confirm non-overlapping indices are not present
+        assert (5, 6) not in overlapping
+        assert (7, 8) not in overlapping
+
+    @pytest.fixture
+    def mock_river_dataset_with_simple_overlap(self):
+        """Creates a small mock dataset with overlapping rivers."""
+        ntracers = 1
+        river_names = ["RiverA", "RiverB", "RiverC"]
+        tracer_names = ["temp"]
+        nriver = len(river_names)
+
+        ds = xr.Dataset(
+            {
+                "river_volume": (("time", "nriver"), [[2.0, 3.0, 4.0]]),
+                "river_tracer": (
+                    ("time", "ntracers", "nriver"),
+                    [[[10.0, 20.0, 30.0]]],
+                ),
+            },
+            coords={
+                "time": [0],
+                "ntracers": range(ntracers),
+                "nriver": range(nriver),
+                "river_name": ("nriver", river_names),
+                "tracer_name": ("ntracers", tracer_names),
+            },
+        )
+
+        # Assign overlapping indices to two rivers
+        indices = {
+            "RiverA": [(10, 20), (10, 21)],
+            "RiverB": [(10, 20)],
+            "RiverC": [(15, 25)],
+        }
+
+        return ds, indices
+
+    @pytest.fixture
+    def mock_river_dataset_with_complex_overlap(self):
+        """Creates a small mock dataset with overlapping rivers."""
+        ntracers = 1
+        river_names = ["RiverA", "RiverB", "RiverC"]
+        tracer_names = ["temp"]
+        nriver = len(river_names)
+
+        ds = xr.Dataset(
+            {
+                "river_volume": (("time", "nriver"), [[1.0, 2.0, 3.0]]),
+                "river_tracer": (
+                    ("time", "ntracers", "nriver"),
+                    [[[10.0, 20.0, 30.0]]],
+                ),
+            },
+            coords={
+                "time": [0],
+                "ntracers": range(ntracers),
+                "nriver": range(nriver),
+                "river_name": ("nriver", river_names),
+                "tracer_name": ("ntracers", tracer_names),
+            },
+        )
+
+        # Assign overlapping indices to two rivers
+        indices = {
+            "RiverA": [(10, 20), (10, 22), (11, 20)],
+            "RiverB": [(10, 20), (10, 21), (10, 22)],
+            "RiverC": [(15, 25), (10, 21), (11, 23)],
+        }
+
+        return ds, indices
+
+    @pytest.mark.parametrize(
+        "fixture_name, river_list, idx_pair, expected_volume, expected_tracer",
+        [
+            # Simple overlap: RiverA and RiverB both map to (10, 20)
+            (
+                "mock_river_dataset_with_simple_overlap",
+                ["RiverA", "RiverB"],
+                (10, 20),
+                2.0 / 2 + 3.0 / 1,
+                (10 * (2.0 / 2) + 20 * (3.0 / 1)) / (2.0 / 2 + 3.0 / 1),
+            ),
+            # Complex overlap: RiverA and RiverB at (10, 20)
+            (
+                "mock_river_dataset_with_complex_overlap",
+                ["RiverA", "RiverB"],
+                (10, 20),
+                1.0 / 3 + 2.0 / 3,
+                (10 * (1 / 3) + 20 * (2 / 3)) / (1.0 / 3 + 2.0 / 3),
+            ),
+            # Complex overlap: RiverA and RiverB at (10, 22)
+            (
+                "mock_river_dataset_with_complex_overlap",
+                ["RiverA", "RiverB"],
+                (10, 22),
+                1.0 / 3 + 2.0 / 3,
+                (10 * (1 / 3) + 20 * (2 / 3)) / (1.0 / 3 + 2.0 / 3),
+            ),
+            # Complex overlap: RiverB and RiverC at (10, 21)
+            (
+                "mock_river_dataset_with_complex_overlap",
+                ["RiverB", "RiverC"],
+                (10, 21),
+                2.0 / 3 + 3.0 / 3,
+                (20 * (2 / 3) + 30 * (3 / 3)) / (2.0 / 3 + 3.0 / 3),
+            ),
+        ],
+    )
+    def test_create_combined_river(
+        self,
+        request,
+        fixture_name,
+        river_list,
+        idx_pair,
+        expected_volume,
+        expected_tracer,
+    ):
+        ds, indices = request.getfixturevalue(fixture_name)
+
+        rf = RiverForcing.__new__(RiverForcing)
+        rf.indices = indices
+
+        combined_volume, combined_tracer = rf._create_combined_river(
+            ds=ds,
+            i=1,
+            idx_pair=idx_pair,
+            river_list=river_list,
+        )
+
+        assert combined_volume.sizes["nriver"] == 1
+        assert combined_tracer.sizes["nriver"] == 1
+        assert combined_volume.coords["river_name"].item() == "overlap_1"
+        assert combined_tracer.coords["river_name"].item() == "overlap_1"
+
+        np.testing.assert_allclose(combined_volume.values, expected_volume, rtol=1e-6)
+        np.testing.assert_allclose(
+            combined_tracer.values, [[[expected_tracer]]], rtol=1e-6
+        )
+
+    @pytest.mark.parametrize(
+        "fixture_name, expected_volumes",
+        [
+            # RiverA has 2 cells, 1 overlap -> scale = 0.5
+            # RiverB has 1 cell, 1 overlap -> scale = 0.0
+            # RiverC has 1 cell, 0 overlap -> stays the same
+            ("mock_river_dataset_with_simple_overlap", [1.0, 0.0, 4.0]),
+            # RiverA has 3 cells, 2 overlaps -> scale = 1/3
+            # RiverB has 3 cells, 3 overlaps -> scale = 0.0
+            # RiverC has 3 cells, 1 overlap -> scale = 2/3
+            (
+                "mock_river_dataset_with_complex_overlap",
+                [1.0 * (1 / 3), 0.0, 3.0 * (2 / 3)],
+            ),
+        ],
+    )
+    def test_reduce_river_volumes(self, request, fixture_name, expected_volumes):
+        ds, indices = request.getfixturevalue(fixture_name)
+
+        # Manually create RiverForcing instance with mock indices
+        rf = RiverForcing.__new__(RiverForcing)
+        rf.indices = indices
+
+        # Compute overlapping rivers
+        overlapping = rf._get_overlapping_rivers()
+
+        # Apply volume reduction
+        ds_out = rf._reduce_river_volumes(ds.copy(), overlapping)
+
+        # Validate reduced volumes
+        actual = ds_out["river_volume"].sel(time=0).values
+        np.testing.assert_allclose(actual, expected_volumes, rtol=1e-6)
+
+    @pytest.mark.parametrize(
+        "mock_river_dataset_fixture, expected_synthetic_count",
+        [
+            ("mock_river_dataset_with_simple_overlap", 1),
+            ("mock_river_dataset_with_complex_overlap", 3),
+        ],
+    )
+    def test_handle_overlapping_rivers(
+        self, mock_river_dataset_fixture, expected_synthetic_count, request
+    ):
+        ds, indices = request.getfixturevalue(mock_river_dataset_fixture)
+
+        rf = RiverForcing.__new__(RiverForcing)
+        rf.indices = indices
+        rf.original_indices = indices.copy()
+        rf.grid = None  # Not needed for this test
+
+        ds_out = rf._handle_overlapping_rivers(ds)
+
+        # Assert number of synthetic rivers added
+        expected_nriver = ds.dims["nriver"] + expected_synthetic_count
+        assert ds_out.dims["nriver"] == expected_nriver
+
+        # Check no river's volume increased
+        for name in ds["river_name"].values:
+            river_idx = np.where(ds["river_name"].values == name)[0].item()
+            assert float(ds_out["river_volume"].isel(nriver=river_idx)) <= float(
+                ds["river_volume"].isel(nriver=river_idx)
+            )
+
+        # Check rivers that contribute to overlapping indices have strictly decreased volume
+        overlapping = rf._get_overlapping_rivers()
+        for name, idx_list in indices.items():
+            if name in ds["river_name"].values and any(
+                idx in overlapping for idx in idx_list
+            ):
+                river_idx = np.where(ds["river_name"].values == name)[0].item()
+                assert float(ds_out["river_volume"].isel(nriver=river_idx)) < float(
+                    ds["river_volume"].isel(nriver=river_idx)
+                )
+
+        # Check that total volume is preserved (within tolerance)
+        np.testing.assert_allclose(
+            ds_out["river_volume"].sum(dim="nriver").values,
+            ds["river_volume"].sum(dim="nriver").values,
+            rtol=1e-6,
+        )
+
+        # Check that the synthetic river volume is positive and tracer is not NaN
+        synthetic_idx = ds_out.dims["nriver"] - 1
+        assert float(ds_out["river_volume"].isel(nriver=synthetic_idx)) > 0
+        assert not np.isnan(ds_out["river_tracer"].isel(nriver=synthetic_idx).item())