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

roms_tools/tests/test_setup/test_grid.py

@@ -18,7 +18,7 @@ from roms_tools.constants import (
     UPPER_BOUND_THETA_B,
     UPPER_BOUND_THETA_S,
 )
-from roms_tools.download import download_test_data
+from roms_tools.datasets.download import download_test_data
 from roms_tools.setup.topography import _compute_rfactor
 
 try:
@@ -121,6 +121,25 @@ def grid_that_straddles_180_degree_meridian_with_global_srtm15_data():
     return grid
 
 
+@pytest.fixture()
+def grid_with_emod_data():
+    grid = Grid(
+        nx=2,
+        ny=2,
+        size_x=32,
+        size_y=19.2,
+        center_lon=-21.68,
+        center_lat=64.325,
+        rot=0,
+        topography_source={
+            "name": "EMOD",
+            "path": download_test_data("EMODnet_C2_coarse100.nc"),
+        },
+    )
+
+    return grid
+
+
 @pytest.fixture()
 def grid_with_gshhs_coastlines():
     iceland_fjord_kwargs = {
@@ -203,6 +222,7 @@ def test_coords_relation(grid_fixture, request):
         "grid_that_straddles_dateline_with_global_srtm15_data",
         "grid_that_straddles_180_degree_meridian_with_global_srtm15_data",
         "grid_with_gshhs_coastlines",
+        "grid_with_emod_data",
     ],
 )
 def test_successful_initialization_with_topography(grid_fixture, request):
@@ -344,6 +364,7 @@ def test_grid_straddle_crosses_meridian():
         "grid_that_straddles_dateline_with_shifted_global_etopo_data",
         "grid_that_straddles_dateline_with_global_srtm15_data",
         "grid_with_gshhs_coastlines",
+        "grid_with_emod_data",
     ],
 )
 def test_roundtrip_netcdf(grid_fixture, tmp_path, request):
@@ -379,6 +400,7 @@ def test_roundtrip_netcdf(grid_fixture, tmp_path, request):
         "grid_that_straddles_dateline_with_shifted_global_etopo_data",
         "grid_that_straddles_dateline_with_global_srtm15_data",
         "grid_with_gshhs_coastlines",
+        "grid_with_emod_data",
     ],
 )
 def test_roundtrip_yaml(grid_fixture, tmp_path, request):
@@ -411,6 +433,7 @@ def test_roundtrip_yaml(grid_fixture, tmp_path, request):
         "grid_that_straddles_dateline_with_shifted_global_etopo_data",
         "grid_that_straddles_dateline_with_global_srtm15_data",
         "grid_with_gshhs_coastlines",
+        "grid_with_emod_data",
     ],
 )
 def test_roundtrip_from_file_yaml(grid_fixture, tmp_path, request):
@@ -441,6 +464,7 @@ def test_roundtrip_from_file_yaml(grid_fixture, tmp_path, request):
         "grid_that_straddles_dateline_with_shifted_global_etopo_data",
         "grid_that_straddles_dateline_with_global_srtm15_data",
         "grid_with_gshhs_coastlines",
+        "grid_with_emod_data",
     ],
 )
 def test_files_have_same_hash(grid_fixture, tmp_path, request):
@@ -741,6 +765,23 @@ def test_hmin_criterion_and_update_topography():
     assert np.less_equal(grid.hmin, grid.ds.h.min())
 
 
+def test_update_topography_raises_if_grid_loaded_from_file_has_no_source_info():
+    fname = download_test_data("grid_created_with_matlab.nc")
+    grid = Grid.from_file(fname)
+
+    with pytest.raises(
+        ValueError,
+        match="Topography source information is not available",
+    ):
+        grid.update_topography(hmin=15)
+
+    with pytest.raises(
+        ValueError,
+        match="Minimal ocean depth is not available",
+    ):
+        grid.update_topography(topography_source={"name": "ETOPO5"})
+
+
 # Mask tests
 
 

roms_tools/tests/test_setup/test_initial_conditions.py

@@ -1,4 +1,3 @@
-import logging
 import textwrap
 from datetime import datetime
 from pathlib import Path
@@ -10,8 +9,8 @@ import xarray as xr
 
 from conftest import calculate_data_hash
 from roms_tools import Grid, InitialConditions
-from roms_tools.download import download_test_data
-from roms_tools.setup.lat_lon_datasets import (
+from roms_tools.datasets.download import download_test_data
+from roms_tools.datasets.lat_lon_datasets import (
     CESMBGCDataset,
     UnifiedBGCDataset,
 )
@@ -46,9 +45,7 @@ def example_grid():
     "ic_fixture",
     [
         "initial_conditions",
-        "initial_conditions_adjusted_for_zeta",
         "initial_conditions_with_bgc",
-        "initial_conditions_with_bgc_adjusted_for_zeta",
         "initial_conditions_with_bgc_from_climatology",
         "initial_conditions_with_unified_bgc_from_climatology",
     ],
@@ -66,6 +63,7 @@ def test_initial_conditions_creation_with_nondefault_glorys_dataset(
         "climatology": False,
     }
     assert hasattr(ic.ds, "adjust_depth_for_sea_surface_height")
+    assert ic.ds.attrs["adjust_depth_for_sea_surface_height"] == "False"
     assert isinstance(ic.ds, xr.Dataset)
     assert ic.ds.coords["ocean_time"].attrs["units"] == "seconds"
     expected_vars = {"temp", "salt", "u", "v", "zeta", "ubar", "vbar"}
@@ -171,7 +169,6 @@ def test_initial_conditions_creation_with_duplicates(use_dask: bool) -> None:
     "ic_fixture",
     [
         "initial_conditions_with_bgc",
-        "initial_conditions_with_bgc_adjusted_for_zeta",
         "initial_conditions_with_bgc_from_climatology",
         "initial_conditions_with_unified_bgc_from_climatology",
     ],
@@ -303,90 +300,6 @@ def test_initial_conditions_default_bgc_climatology(example_grid, use_dask):
     assert initial_conditions.bgc_source["climatology"] is False
 
 
-def test_info_depth(caplog, use_dask):
-    grid = Grid(
-        nx=2,
-        ny=2,
-        size_x=500,
-        size_y=1000,
-        center_lon=0,
-        center_lat=55,
-        rot=10,
-        N=3,  # number of vertical levels
-        theta_s=5.0,  # surface control parameter
-        theta_b=2.0,  # bottom control parameter
-        hc=250.0,  # critical depth
-    )
-
-    fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
-
-    with caplog.at_level(logging.INFO):
-        InitialConditions(
-            grid=grid,
-            ini_time=datetime(2021, 6, 29),
-            source={"path": fname, "name": "GLORYS"},
-            adjust_depth_for_sea_surface_height=True,
-            use_dask=use_dask,
-        )
-    # Verify the warning message in the log
-    assert "Sea surface height will be used to adjust depth coordinates." in caplog.text
-
-    # Clear the log before the next test
-    caplog.clear()
-
-    with caplog.at_level(logging.INFO):
-        InitialConditions(
-            grid=grid,
-            ini_time=datetime(2021, 6, 29),
-            source={"path": fname, "name": "GLORYS"},
-            adjust_depth_for_sea_surface_height=False,
-            use_dask=use_dask,
-        )
-    # Verify the warning message in the log
-    assert (
-        "Sea surface height will NOT be used to adjust depth coordinates."
-        in caplog.text
-    )
-
-
-@pytest.mark.parametrize(
-    "initial_conditions_fixture",
-    [
-        "initial_conditions_adjusted_for_zeta",
-        "initial_conditions_with_bgc_adjusted_for_zeta",
-    ],
-)
-def test_correct_depth_coords_adjusted_for_zeta(
-    initial_conditions_fixture, request, use_dask
-):
-    initial_conditions = request.getfixturevalue(initial_conditions_fixture)
-
-    # compute interface depth at rho-points and write it into .ds_depth_coords
-    zeta = initial_conditions.ds.zeta
-    initial_conditions._get_depth_coordinates(
-        zeta, location="rho", depth_type="interface"
-    )
-    # Test that lowermost interface coincides with topography
-    assert np.allclose(
-        initial_conditions.ds_depth_coords["interface_depth_rho"]
-        .isel(s_w=0)
-        .squeeze()
-        .values,  # Extract raw NumPy array
-        initial_conditions.grid.ds.h.values,
-        atol=1e-6,  # Adjust tolerance as needed
-    )
-
-    # Test that uppermost interface coincides with sea surface height
-    assert np.allclose(
-        initial_conditions.ds_depth_coords["interface_depth_rho"]
-        .isel(s_w=-1)
-        .squeeze()
-        .values,
-        -zeta.values,
-        atol=1e-6,
-    )
-
-
 @pytest.mark.parametrize(
     "initial_conditions_fixture",
     [
@@ -498,7 +411,6 @@ def test_computed_missing_optional_fields(
 @pytest.mark.parametrize(
     "initial_conditions_fixture",
     [
-        "initial_conditions_with_bgc_adjusted_for_zeta",
         "initial_conditions_with_bgc_from_climatology",
         "initial_conditions_with_unified_bgc_from_climatology",
     ],
@@ -546,7 +458,6 @@ def test_initial_conditions_plot(initial_conditions_fixture, request):
     "initial_conditions_fixture",
     [
         "initial_conditions",
-        "initial_conditions_adjusted_for_zeta",
         "initial_conditions_with_bgc_from_climatology",
         "initial_conditions_with_unified_bgc_from_climatology",
     ],
@@ -574,7 +485,6 @@ def test_initial_conditions_save(initial_conditions_fixture, request, tmp_path):
     "initial_conditions_fixture",
     [
         "initial_conditions",
-        "initial_conditions_adjusted_for_zeta",
         "initial_conditions_with_bgc_from_climatology",
         "initial_conditions_with_unified_bgc_from_climatology",
     ],
@@ -607,7 +517,6 @@ def test_roundtrip_yaml(initial_conditions_fixture, request, tmp_path, use_dask)
     "initial_conditions_fixture",
     [
         "initial_conditions",
-        "initial_conditions_adjusted_for_zeta",
         "initial_conditions_with_bgc_from_climatology",
         "initial_conditions_with_unified_bgc_from_climatology",
     ],

roms_tools/tests/test_setup/test_nesting.py

@@ -14,7 +14,8 @@ from roms_tools.setup.nesting import (
     modify_child_mask,
     modify_child_topography,
 )
-from roms_tools.setup.utils import get_boundary_coords, wrap_longitudes
+from roms_tools.setup.utils import get_boundary_coords
+from roms_tools.utils import wrap_longitudes
 
 
 @pytest.fixture()

roms_tools/tests/test_setup/test_surface_forcing.py

@@ -9,7 +9,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.datasets.download import download_test_data
 from roms_tools.setup.utils import RawDataSource
 
 

roms_tools/tests/test_setup/test_tides.py

@@ -6,7 +6,7 @@ import xarray as xr
 
 from conftest import calculate_data_hash
 from roms_tools import Grid, TidalForcing
-from roms_tools.download import download_test_data
+from roms_tools.datasets.download import download_test_data
 
 
 @pytest.fixture(scope="session")

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_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"]