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

roms_tools/tests/test_setup/test_utils.py

@@ -7,11 +7,11 @@ import pytest
 import xarray as xr
 
 from roms_tools import BoundaryForcing, Grid
-from roms_tools.download import download_test_data
-from roms_tools.setup.lat_lon_datasets import ERA5Correction
+from roms_tools.datasets.download import download_test_data
 from roms_tools.setup.utils import (
+    _infer_valid_boundaries_from_mask,
+    check_and_set_boundaries,
     get_target_coords,
-    interpolate_from_climatology,
     validate_names,
 )
 
@@ -101,18 +101,6 @@ class TestGetTargetCoords:
         assert "mask" in result
 
 
-def test_interpolate_from_climatology(use_dask):
-    fname = download_test_data("ERA5_regional_test_data.nc")
-    era5_times = xr.open_dataset(fname).time
-
-    climatology = ERA5Correction(use_dask=use_dask)
-    field = climatology.ds["ssr_corr"]
-    field["time"] = field["time"].dt.days
-
-    interpolated_field = interpolate_from_climatology(field, "time", era5_times)
-    assert len(interpolated_field.time) == len(era5_times)
-
-
 # Test yaml roundtrip with multiple source files
 @pytest.fixture()
 def boundary_forcing_from_multiple_source_files(request, use_dask):
@@ -210,3 +198,100 @@ def test_non_string_elements_in_list_raises():
 def test_custom_label_in_errors():
     with pytest.raises(ValueError, match="Invalid foozs: z"):
         validate_names(["z"], VALID_NAMES, SENTINEL, MAX_TO_PLOT, label="fooz")
+
+
+# test _infer_valid_boundaries_from_mask
+@pytest.fixture
+def simple_mask():
+    data = np.array(
+        [
+            [0, 0, 0, 0],  # south
+            [1, 0, 0, 0],
+            [1, 0, 0, 0],
+            [1, 1, 1, 1],  # north
+        ]
+    )
+    return xr.DataArray(data, dims=("eta_rho", "xi_rho"))
+
+
+def test_infer_valid_boundaries_partial(simple_mask):
+    out = _infer_valid_boundaries_from_mask(simple_mask)
+
+    assert out == {
+        "south": False,
+        "north": True,
+        "west": True,
+        "east": True,
+    }
+
+
+def test_infer_valid_boundaries_all_ocean():
+    mask = xr.DataArray(np.ones((3, 3)), dims=("eta_rho", "xi_rho"))
+
+    out = _infer_valid_boundaries_from_mask(mask)
+    assert all(out.values())
+
+
+# test check_and_set_boundaries
+
+
+def test_check_and_set_default_boundaries(simple_mask, monkeypatch):
+    monkeypatch.setattr(
+        "roms_tools.setup.utils._infer_valid_boundaries_from_mask",
+        lambda mask: {
+            "south": True,
+            "north": False,
+            "west": True,
+            "east": False,
+        },
+    )
+
+    result = check_and_set_boundaries(None, simple_mask)
+
+    assert result == {
+        "south": True,
+        "north": False,
+        "west": True,
+        "east": False,
+    }
+
+
+def test_check_and_set_partial_boundaries(simple_mask, monkeypatch):
+    monkeypatch.setattr(
+        "roms_tools.setup.utils._infer_valid_boundaries_from_mask",
+        lambda mask: {
+            "south": True,
+            "north": False,
+            "west": True,
+            "east": False,
+        },
+    )
+
+    user = {"south": False}  # user overrides south → False
+
+    result = check_and_set_boundaries(user, simple_mask)
+
+    assert result == {
+        "south": False,  # user-preserved
+        "north": False,  # inferred
+        "west": True,  # inferred
+        "east": False,  # inferred
+    }
+
+
+def test_check_and_set_type_error(simple_mask):
+    with pytest.raises(TypeError):
+        check_and_set_boundaries({"south": "yes"}, simple_mask)
+
+
+def test_check_and_set_invalid_key(simple_mask):
+    with pytest.raises(ValueError):
+        check_and_set_boundaries({"northeast": True}, simple_mask)
+
+
+def test_check_and_set_full_user_boundaries(simple_mask):
+    boundaries = {"south": False, "north": True, "east": False, "west": True}
+
+    result = check_and_set_boundaries(boundaries, simple_mask)
+
+    assert result == boundaries  # unchanged

roms_tools/tests/test_setup/test_validation.py

@@ -5,6 +5,15 @@ from pathlib import Path
 import pytest
 import xarray as xr
 
+try:
+    import xesmf  # type: ignore
+except ImportError:
+    xesmf = None
+
+skip_xesmf = pytest.mark.skipif(
+    xesmf is None, reason="xesmf required for ROMS regridding"
+)
+
 
 @pytest.mark.parametrize(
     "forcing_fixture",
@@ -14,6 +23,7 @@ import xarray as xr
         "tidal_forcing",
         "initial_conditions_with_bgc_from_climatology",
         "initial_conditions_with_unified_bgc_from_climatology",
+        pytest.param("initial_conditions_from_roms", marks=skip_xesmf),
         "surface_forcing",
         "coarse_surface_forcing",
         "corrected_surface_forcing",
@@ -59,6 +69,7 @@ def test_save_results(
         "tidal_forcing",
         "initial_conditions_with_bgc_from_climatology",
         "initial_conditions_with_unified_bgc_from_climatology",
+        pytest.param("initial_conditions_from_roms", marks=skip_xesmf),
         "surface_forcing",
         "coarse_surface_forcing",
         "corrected_surface_forcing",
@@ -93,13 +104,17 @@ def test_check_results(
     [
         "tidal_forcing",
         "initial_conditions_with_bgc_from_climatology",
+        "initial_conditions_with_unified_bgc_from_climatology",
+        pytest.param("initial_conditions_from_roms", marks=skip_xesmf),
         "surface_forcing",
         "coarse_surface_forcing",
         "corrected_surface_forcing",
         "bgc_surface_forcing",
         "bgc_surface_forcing_from_climatology",
+        "bgc_surface_forcing_from_unified_climatology",
         "boundary_forcing",
         "bgc_boundary_forcing_from_climatology",
+        "bgc_boundary_forcing_from_unified_climatology",
     ],
 )
 def test_dask_vs_no_dask(

roms_tools/tests/test_tiling/test_partition.py

@@ -1,6 +1,8 @@
 from pathlib import Path
 
+import numpy as np
 import pytest
+import xarray as xr
 import xarray.testing as xrt
 
 from roms_tools import Grid
@@ -18,7 +20,7 @@ class TestPartitionGrid:
     def test_partition_grid_along_x(self, grid):
         _, [ds1, ds2, ds3] = partition(grid.ds, np_eta=3, np_xi=1)
 
-        assert ds1.sizes == {
+        assert dict(ds1.sizes) == {
             "eta_rho": 11,
             "xi_rho": 32,
             "xi_u": 31,
@@ -30,7 +32,7 @@ class TestPartitionGrid:
             "s_rho": 100,
             "s_w": 101,
         }
-        assert ds2.sizes == {
+        assert dict(ds2.sizes) == {
             "eta_rho": 10,
             "xi_rho": 32,
             "xi_u": 31,
@@ -42,7 +44,7 @@ class TestPartitionGrid:
             "s_rho": 100,
             "s_w": 101,
         }
-        assert ds3.sizes == {
+        assert dict(ds3.sizes) == {
             "eta_rho": 11,
             "xi_rho": 32,
             "xi_u": 31,
@@ -58,7 +60,7 @@ class TestPartitionGrid:
     def test_partition_grid_along_y(self, grid):
         _, [ds1, ds2, ds3] = partition(grid.ds, np_eta=1, np_xi=3)
 
-        assert ds1.sizes == {
+        assert dict(ds1.sizes) == {
             "eta_rho": 32,
             "xi_rho": 11,
             "xi_u": 10,
@@ -70,7 +72,7 @@ class TestPartitionGrid:
             "s_rho": 100,
             "s_w": 101,
         }
-        assert ds2.sizes == {
+        assert dict(ds2.sizes) == {
             "eta_rho": 32,
             "xi_rho": 10,
             "xi_u": 10,
@@ -82,7 +84,7 @@ class TestPartitionGrid:
             "s_rho": 100,
             "s_w": 101,
         }
-        assert ds3.sizes == {
+        assert dict(ds3.sizes) == {
             "eta_rho": 32,
             "xi_rho": 11,
             "xi_u": 11,
@@ -103,7 +105,7 @@ class TestPartitionGrid:
          ds7, ds8, ds9] = partition(grid.ds, np_eta=3, np_xi=3)
         # fmt: on
 
-        assert ds1.sizes == {
+        assert dict(ds1.sizes) == {
             "eta_rho": 11,
             "xi_rho": 11,
             "xi_u": 10,
@@ -115,7 +117,7 @@ class TestPartitionGrid:
             "s_rho": 100,
             "s_w": 101,
         }
-        assert ds4.sizes == {
+        assert dict(ds4.sizes) == {
             "eta_rho": 10,
             "xi_rho": 11,
             "xi_u": 10,
@@ -127,7 +129,7 @@ class TestPartitionGrid:
             "s_rho": 100,
             "s_w": 101,
         }
-        assert ds7.sizes == {
+        assert dict(ds7.sizes) == {
             "eta_rho": 11,
             "xi_rho": 11,
             "xi_u": 10,
@@ -139,7 +141,7 @@ class TestPartitionGrid:
             "s_rho": 100,
             "s_w": 101,
         }
-        assert ds2.sizes == {
+        assert dict(ds2.sizes) == {
             "eta_rho": 11,
             "xi_rho": 10,
             "xi_u": 10,
@@ -151,7 +153,7 @@ class TestPartitionGrid:
             "s_rho": 100,
             "s_w": 101,
         }
-        assert ds5.sizes == {
+        assert dict(ds5.sizes) == {
             "eta_rho": 10,
             "xi_rho": 10,
             "xi_u": 10,
@@ -163,7 +165,7 @@ class TestPartitionGrid:
             "s_rho": 100,
             "s_w": 101,
         }
-        assert ds8.sizes == {
+        assert dict(ds8.sizes) == {
             "eta_rho": 11,
             "xi_rho": 10,
             "xi_u": 10,
@@ -175,7 +177,7 @@ class TestPartitionGrid:
             "s_rho": 100,
             "s_w": 101,
         }
-        assert ds3.sizes == {
+        assert dict(ds3.sizes) == {
             "eta_rho": 11,
             "xi_rho": 11,
             "xi_u": 11,
@@ -187,7 +189,7 @@ class TestPartitionGrid:
             "s_rho": 100,
             "s_w": 101,
         }
-        assert ds6.sizes == {
+        assert dict(ds6.sizes) == {
             "eta_rho": 10,
             "xi_rho": 11,
             "xi_u": 11,
@@ -199,7 +201,7 @@ class TestPartitionGrid:
             "s_rho": 100,
             "s_w": 101,
         }
-        assert ds9.sizes == {
+        assert dict(ds9.sizes) == {
             "eta_rho": 11,
             "xi_rho": 11,
             "xi_u": 11,
@@ -246,6 +248,52 @@ class TestPartitionGrid:
             assert ds.sizes["xi_coarse"] == grid.ds.sizes["xi_coarse"]
 
 
+class TestPartitionGridWithExtraDims:
+    @pytest.fixture
+    def ds_with_extra_dims(self):
+        # Base dims
+        eta_rho = 10
+        xi_rho = 12
+        s_rho = 5
+        eta_v = eta_rho - 1
+        xi_u = xi_rho - 1
+
+        # Extra dims
+        eta_u = eta_rho
+        xi_v = xi_rho
+
+        ds = xr.Dataset(
+            {
+                "zeta": (("eta_rho", "xi_rho"), np.zeros((eta_rho, xi_rho))),
+                "u": (("s_rho", "eta_u", "xi_u"), np.zeros((s_rho, eta_u, xi_u))),
+                "v": (("s_rho", "eta_v", "xi_v"), np.zeros((s_rho, eta_v, xi_v))),
+            },
+            coords={
+                "eta_rho": np.arange(eta_rho),
+                "eta_u": np.arange(eta_u),
+                "eta_v": np.arange(eta_v),
+                "xi_rho": np.arange(xi_rho),
+                "xi_u": np.arange(xi_u),
+                "xi_v": np.arange(xi_v),
+                "s_rho": np.arange(s_rho),
+            },
+        )
+
+        return ds
+
+    def test_partition_with_extra_dims(self, ds_with_extra_dims):
+        file_numbers, parts = partition(ds_with_extra_dims, np_eta=2, np_xi=2)
+
+        # Test that partitioned datasets contain eta_u and xi_v
+        for ds_part in parts:
+            assert "eta_u" in ds_part.dims
+            assert "xi_v" in ds_part.dims
+            assert ds_part.sizes["eta_u"] < ds_with_extra_dims.sizes["eta_u"]
+            assert ds_part.sizes["xi_v"] < ds_with_extra_dims.sizes["xi_v"]
+            assert ds_part.sizes["eta_u"] > 0
+            assert ds_part.sizes["xi_v"] > 0
+
+
 class TestPartitionMissingDims:
     def test_partition_missing_dims(self, grid):
         dims_to_drop = ["xi_u", "eta_v", "eta_coarse", "xi_coarse"]