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

roms_tools/tests/test_setup/test_nesting.py

@@ -11,9 +11,10 @@ from roms_tools.setup.nesting import (
     compute_boundary_distance,
     interpolate_indices,
     map_child_boundaries_onto_parent_grid_indices,
-    modify_child_topography_and_mask,
+    modify_child_mask,
+    modify_child_topography,
 )
-from roms_tools.setup.utils import get_boundary_coords
+from roms_tools.setup.utils import get_boundary_coords, wrap_longitudes
 
 
 @pytest.fixture()
@@ -46,9 +47,7 @@ def big_grid_that_straddles():
 
 @pytest.fixture()
 def small_grid_that_straddles():
-    return Grid(
-        nx=10, ny=10, center_lon=10, center_lat=61, rot=-20, size_x=500, size_y=500
-    )
+    return Grid(nx=10, ny=10, center_lon=0, center_lat=61, rot=0, size_x=50, size_y=200)
 
 
 @pytest.fixture()
@@ -71,11 +70,11 @@ def child_grid_that_straddles(big_grid_that_straddles):
         parent_grid=big_grid_that_straddles,
         nx=10,
         ny=10,
-        center_lon=10,
+        center_lon=0,
         center_lat=61,
-        rot=-20,
-        size_x=500,
-        size_y=500,
+        rot=0,
+        size_x=50,
+        size_y=200,
     )
 
 
@@ -100,13 +99,14 @@ class TestInterpolateIndices:
             mask = grid.ds[f"mask_{location}"].isel(**bdry_coords)
 
             with caplog.at_level(logging.WARNING):
-                i_eta, i_xi = interpolate_indices(grid.ds, lon, lat, mask)
+                i_eta, i_xi = interpolate_indices(grid.ds, lon, lat, mask, direction)
 
-            # Verify the warning message in the log
-            assert (
-                "Some boundary points of the child grid are very close to the boundary of the parent grid."
-                in caplog.text
-            )
+            if mask.sum() > 0:
+                # Verify the warning message in the log
+                assert (
+                    "boundary of the child grid lie very close to the edges of the parent grid"
+                    in caplog.text
+                )
 
             if direction == "south":
                 expected_i_eta = -0.5 * xr.ones_like(grid.ds.xi_rho)
@@ -142,20 +142,27 @@ class TestInterpolateIndices:
         big_grid = request.getfixturevalue(big_grid_fixture)
         small_grid = request.getfixturevalue(small_grid_fixture)
 
+        big_grid_ds = wrap_longitudes(big_grid.ds.copy(), straddle=big_grid.straddle)
+        small_grid_ds = wrap_longitudes(
+            small_grid.ds.copy(), straddle=big_grid.straddle
+        )
+
         bdry_coords_dict = get_boundary_coords()
         for location in ["rho", "u", "v"]:
             for direction in ["south", "east", "north", "west"]:
                 bdry_coords = bdry_coords_dict[location][direction]
-                lon = small_grid.ds[f"lon_{location}"].isel(**bdry_coords)
-                lat = small_grid.ds[f"lat_{location}"].isel(**bdry_coords)
-                mask = small_grid.ds[f"mask_{location}"].isel(**bdry_coords)
+                lon = small_grid_ds[f"lon_{location}"].isel(**bdry_coords)
+                lat = small_grid_ds[f"lat_{location}"].isel(**bdry_coords)
+                mask = small_grid_ds[f"mask_{location}"].isel(**bdry_coords)
 
-                i_eta, i_xi = interpolate_indices(big_grid.ds, lon, lat, mask)
+                i_eta, i_xi = interpolate_indices(
+                    big_grid_ds, lon, lat, mask, direction
+                )
 
                 expected_i_eta_min = -0.5
-                expected_i_eta_max = big_grid.ds.eta_rho[-1] - 0.5
+                expected_i_eta_max = big_grid_ds.eta_rho[-1] - 0.5
                 expected_i_xi_min = -0.5
-                expected_i_xi_max = big_grid.ds.xi_rho[-1] - 0.5
+                expected_i_xi_max = big_grid_ds.xi_rho[-1] - 0.5
 
                 assert (i_eta >= expected_i_eta_min).all()
                 assert (i_eta <= expected_i_eta_max).all()
@@ -166,11 +173,16 @@ 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."""
+        small_grid_ds = wrap_longitudes(
+            small_grid.ds.copy(), straddle=small_grid.straddle
+        )
+        baby_grid_ds = wrap_longitudes(baby_grid.ds.copy(), straddle=baby_grid.straddle)
+
         ds_without_updated_indices = map_child_boundaries_onto_parent_grid_indices(
-            small_grid.ds, baby_grid.ds, update_land_indices=False
+            small_grid_ds, baby_grid_ds, update_land_indices=False
         )
         ds_with_updated_indices = map_child_boundaries_onto_parent_grid_indices(
-            small_grid.ds, baby_grid.ds, update_land_indices=True
+            small_grid_ds, baby_grid_ds, update_land_indices=True
         )
 
         xr.testing.assert_allclose(ds_without_updated_indices, ds_with_updated_indices)
@@ -189,7 +201,12 @@ class TestMapChildBoundaries:
         big_grid = request.getfixturevalue(big_grid_fixture)
         small_grid = request.getfixturevalue(small_grid_fixture)
 
-        ds = map_child_boundaries_onto_parent_grid_indices(big_grid.ds, small_grid.ds)
+        big_grid_ds = wrap_longitudes(big_grid.ds.copy(), straddle=big_grid.straddle)
+        small_grid_ds = wrap_longitudes(
+            small_grid.ds.copy(), straddle=big_grid.straddle
+        )
+
+        ds = map_child_boundaries_onto_parent_grid_indices(big_grid_ds, small_grid_ds)
         for direction in ["south", "east", "north", "west"]:
             for location in ["rho", "u", "v"]:
                 if location == "rho":
@@ -224,9 +241,14 @@ class TestMapChildBoundaries:
         big_grid = request.getfixturevalue(big_grid_fixture)
         small_grid = request.getfixturevalue(small_grid_fixture)
 
+        big_grid_ds = wrap_longitudes(big_grid.ds.copy(), straddle=big_grid.straddle)
+        small_grid_ds = wrap_longitudes(
+            small_grid.ds.copy(), straddle=big_grid.straddle
+        )
+
         for update_land_indices in [False, True]:
             ds = map_child_boundaries_onto_parent_grid_indices(
-                big_grid.ds, small_grid.ds, update_land_indices=update_land_indices
+                big_grid_ds, small_grid_ds, update_land_indices=update_land_indices
             )
 
             for direction in ["south", "east", "north", "west"]:
@@ -292,9 +314,7 @@ class TestModifyChid:
     ):
         """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
-        )
+        modified_baby_grid_ds = modify_child_mask(small_grid.ds, baby_grid.ds)
         xr.testing.assert_allclose(modified_baby_grid_ds.mask_rho, mask_original)
 
     @pytest.mark.parametrize(
@@ -311,7 +331,8 @@ class TestModifyChid:
 
         h_original = grid.ds.h.copy()
         mask_original = grid.ds.mask_rho.copy()
-        modified_grid_ds = modify_child_topography_and_mask(grid.ds, grid.ds)
+        modified_grid_ds = modify_child_mask(grid.ds, grid.ds)
+        modified_grid_ds = modify_child_topography(grid.ds, modified_grid_ds)
 
         xr.testing.assert_allclose(modified_grid_ds.h, h_original)
         xr.testing.assert_allclose(modified_grid_ds.mask_rho, mask_original)
@@ -325,7 +346,8 @@ class TestModifyChid:
         mask_original = small_grid.ds.mask_rho.copy()
 
         # Apply the modification function
-        modified_ds = modify_child_topography_and_mask(big_grid.ds, small_grid.ds)
+        modified_ds = modify_child_mask(big_grid.ds, small_grid.ds)
+        modified_ds = modify_child_topography(big_grid.ds, modified_ds)
 
         # Calculate the center indices for the grid
         eta_center = h_original.sizes["eta_rho"] // 2
@@ -405,9 +427,75 @@ class TestNesting:
         params = dataclasses.asdict(small_grid)
         del params["ds"], params["straddle"]
 
-        with pytest.raises(ValueError, match="Some points are outside the grid."):
+        with pytest.raises(
+            ValueError, match="boundary of the child grid lie outside the parent grid"
+        ):
             ChildGrid(parent_grid=big_grid, **params)
 
+    def test_no_error_if_land_child_points_beyond_parent_grid(self):
+        # coarse resolution Pacific domain
+        parent_grid = Grid(
+            nx=50,
+            ny=50,
+            size_x=23000,
+            size_y=12000,
+            center_lon=-161,
+            center_lat=14.4,
+            rot=-3,
+        )
+
+        # California Current System domain, where some land points extend beyond Pacific domain
+        child_grid_parameters = {
+            "nx": 50,
+            "ny": 50,
+            "size_x": 2688,
+            "size_y": 5280,
+            "center_lat": 39.6,
+            "center_lon": -134.5,
+            "rot": 33.3,
+        }
+
+        child_grid = ChildGrid(
+            **child_grid_parameters,
+            parent_grid=parent_grid,
+            boundaries={"north": True, "west": True, "south": True, "east": False},
+        )
+        assert isinstance(child_grid.ds, xr.Dataset)
+
+    def test_no_error_if_child_boundary_entirely_on_land(self):
+        parent_grid = Grid(
+            nx=7,
+            ny=7,
+            size_x=240,
+            size_y=240,
+            center_lon=-4.1,
+            center_lat=52.36,
+            rot=0,
+        )
+
+        child_grid_parameters = {
+            "nx": 5,
+            "ny": 5,
+            "size_x": 100,
+            "size_y": 100,
+            "center_lon": -4.1,
+            "center_lat": 52.36,
+        }
+
+        child_grid = ChildGrid(
+            **child_grid_parameters,
+            parent_grid=parent_grid,
+            boundaries={
+                "south": True,
+                "east": True,
+                "north": True,
+                "west": True,
+            },
+        )
+
+        assert isinstance(child_grid.ds, xr.Dataset)
+        assert isinstance(child_grid.ds_nesting, xr.Dataset)
+
     @pytest.mark.parametrize(
         "child_grid_fixture",
         ["child_grid", "child_grid_that_straddles"],

roms_tools/tests/test_setup/test_river_datasets.py

@@ -0,0 +1,48 @@
+from datetime import datetime
+
+import numpy as np
+import xarray as xr
+
+from roms_tools.setup.river_datasets import RiverDataset
+
+
+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_surface_forcing.py

@@ -10,6 +10,7 @@ import xarray as xr
 from conftest import calculate_data_hash
 from roms_tools import Grid, SurfaceForcing
 from roms_tools.download import download_test_data
+from roms_tools.setup.utils import RawDataSource
 
 
 @pytest.fixture
@@ -159,7 +160,7 @@ def _test_successful_initialization(
     grid: Grid,
     start_time: datetime,
     end_time: datetime,
-    source: dict[str, str],
+    source: RawDataSource,
     coarse_grid_mode: str,
     use_dask: bool,
     caplog,

roms_tools/tests/test_setup/test_utils.py

@@ -2,13 +2,103 @@ import logging
 from datetime import datetime
 from pathlib import Path
 
+import numpy as np
 import pytest
 import xarray as xr
 
 from roms_tools import BoundaryForcing, Grid
 from roms_tools.download import download_test_data
-from roms_tools.setup.datasets import ERA5Correction
-from roms_tools.setup.utils import interpolate_from_climatology, validate_names
+from roms_tools.setup.lat_lon_datasets import ERA5Correction
+from roms_tools.setup.utils import (
+    get_target_coords,
+    interpolate_from_climatology,
+    validate_names,
+)
+
+
+class DummyGrid:
+    """Lightweight grid wrapper mimicking the real Grid API for testing."""
+
+    def __init__(self, ds: xr.Dataset, straddle: bool):
+        """Initialize grid wrapper."""
+        self.ds = ds
+        self.straddle = straddle
+
+
+class TestGetTargetCoords:
+    def make_rho_grid(self, lons, lats, with_mask=False):
+        """Helper to create a minimal rho grid dataset."""
+        eta, xi = len(lats), len(lons)
+        lon_rho, lat_rho = np.meshgrid(lons, lats)
+        ds = xr.Dataset(
+            {
+                "lon_rho": (("eta_rho", "xi_rho"), lon_rho),
+                "lat_rho": (("eta_rho", "xi_rho"), lat_rho),
+                "angle": (("eta_rho", "xi_rho"), np.zeros_like(lon_rho)),
+            },
+            coords={"eta_rho": np.arange(eta), "xi_rho": np.arange(xi)},
+        )
+        if with_mask:
+            ds["mask_rho"] = (("eta_rho", "xi_rho"), np.ones_like(lon_rho))
+        return ds
+
+    def test_basic_rho_grid(self):
+        ds = self.make_rho_grid(lons=[-10, -5, 0, 5, 10], lats=[50, 55])
+        grid = DummyGrid(ds, straddle=True)
+        result = get_target_coords(grid)
+        assert "lat" in result and "lon" in result
+        assert np.allclose(result["lon"], ds.lon_rho)
+
+    def test_wrap_longitudes_to_minus180_180(self):
+        ds = self.make_rho_grid(lons=[190, 200], lats=[0, 1])
+        grid = DummyGrid(ds, straddle=True)
+        result = get_target_coords(grid)
+        # longitudes >180 should wrap to -170, -160
+        expected = np.array([[-170, -160], [-170, -160]])
+        assert np.allclose(result["lon"].values, expected)
+
+    def test_convert_to_0_360_if_far_from_greenwich(self):
+        ds = self.make_rho_grid(lons=[-170, -160], lats=[0, 1])
+        grid = DummyGrid(ds, straddle=False)
+        result = get_target_coords(grid)
+        # Should convert to 190, 200 since domain is far from Greenwich
+        expected = np.array([[190, 200], [190, 200]])
+        assert np.allclose(result["lon"].values, expected)
+        assert result["straddle"] is False
+
+    def test_close_to_greenwich_stays_minus180_180(self):
+        ds = self.make_rho_grid(lons=[-2, -1], lats=[0, 1])
+        grid = DummyGrid(ds, straddle=False)
+        result = get_target_coords(grid)
+        # Should remain unchanged (-2, -1), not converted to 358, 359
+        expected = np.array([[-2, -1], [-2, -1]])
+        assert np.allclose(result["lon"].values, expected)
+        assert result["straddle"] is True
+
+    def test_includes_optional_fields(self):
+        ds = self.make_rho_grid(lons=[-10, -5], lats=[0, 1], with_mask=True)
+        grid = DummyGrid(ds, straddle=True)
+        result = get_target_coords(grid)
+        assert result["mask"] is not None
+
+    def test_coarse_grid_selection(self):
+        lon = np.array([[190, 200]])
+        lat = np.array([[10, 10]])
+        ds = xr.Dataset(
+            {
+                "lon_coarse": (("eta_coarse", "xi_coarse"), lon),
+                "lat_coarse": (("eta_coarse", "xi_coarse"), lat),
+                "angle_coarse": (("eta_coarse", "xi_coarse"), np.zeros_like(lon)),
+                "mask_coarse": (("eta_coarse", "xi_coarse"), np.ones_like(lon)),
+            },
+            coords={"eta_coarse": [0], "xi_coarse": [0, 1]},
+        )
+        grid = DummyGrid(ds, straddle=True)
+        result = get_target_coords(grid, use_coarse_grid=True)
+        # Should wrap longitudes to -170, -160
+        expected = np.array([[-170, -160]])
+        assert np.allclose(result["lon"].values, expected)
+        assert "mask" in result
 
 
 def test_interpolate_from_climatology(use_dask):

roms_tools/tests/test_setup/utils.py

@@ -0,0 +1,71 @@
+from datetime import datetime
+from pathlib import Path
+
+from roms_tools import Grid, get_glorys_bounds
+
+try:
+    import copernicusmarine  # type: ignore
+except ImportError:
+    copernicusmarine = None
+
+
+def download_regional_and_bigger(
+    tmp_path: Path,
+    grid: Grid,
+    start_time: datetime,
+    variables: list[str] = ["thetao", "so", "uo", "vo", "zos"],
+) -> tuple[Path, Path]:
+    """
+    Helper: download minimal and slightly bigger GLORYS subsets.
+
+    Parameters
+    ----------
+    tmp_path : Path
+        Directory to store the downloaded NetCDF files.
+    grid : Grid
+        ROMS-Tools Grid object defining the target domain.
+    start_time : datetime
+        Start time of the requested subset.
+    variables : list[str]
+        What variables to download.
+
+    Returns
+    -------
+    Tuple[Path, Path]
+        Paths to the minimal and slightly bigger GLORYS subset files.
+    """
+    bounds = get_glorys_bounds(grid)
+
+    # minimal dataset
+    regional_file = tmp_path / "regional_GLORYS.nc"
+    copernicusmarine.subset(
+        dataset_id="cmems_mod_glo_phy_my_0.083deg_P1D-m",
+        variables=variables,
+        **bounds,
+        start_datetime=start_time,
+        end_datetime=start_time,
+        coordinates_selection_method="outside",
+        output_filename=str(regional_file),
+    )
+
+    # slightly bigger dataset
+    for key, delta in {
+        "minimum_latitude": -1,
+        "minimum_longitude": -1,
+        "maximum_latitude": +1,
+        "maximum_longitude": +1,
+    }.items():
+        bounds[key] += delta
+
+    bigger_regional_file = tmp_path / "bigger_regional_GLORYS.nc"
+    copernicusmarine.subset(
+        dataset_id="cmems_mod_glo_phy_my_0.083deg_P1D-m",
+        variables=variables,
+        **bounds,
+        start_datetime=start_time,
+        end_datetime=start_time,
+        coordinates_selection_method="outside",
+        output_filename=str(bigger_regional_file),
+    )
+
+    return regional_file, bigger_regional_file