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

roms_tools/tests/test_setup/test_grid.py

@@ -121,6 +121,33 @@ def grid_that_straddles_180_degree_meridian_with_global_srtm15_data():
     return grid
 
 
+@pytest.fixture()
+def grid_with_gshhs_coastlines():
+    iceland_fjord_kwargs = {
+        "nx": 80,
+        "ny": 40,
+        "size_x": 40,
+        "size_y": 20,
+        "center_lon": -21.76,
+        "center_lat": 64.325,
+        "rot": 0,
+        "N": 3,
+    }
+
+    # Make sure all 4 L1 files are downloaded
+    _ = download_test_data("GSHHS_l_L1.dbf")
+    _ = download_test_data("GSHHS_l_L1.prj")
+    _ = download_test_data("GSHHS_l_L1.shx")
+    shapefile = download_test_data("GSHHS_l_L1.shp")
+
+    grid = Grid(
+        **iceland_fjord_kwargs,
+        mask_shapefile=shapefile,
+    )
+
+    return grid
+
+
 def test_grid_creation(grid):
     assert grid.nx == 1
     assert grid.ny == 1
@@ -175,11 +202,33 @@ def test_coords_relation(grid_fixture, request):
         "grid_that_straddles_180_degree_meridian_with_shifted_global_etopo_data",
         "grid_that_straddles_dateline_with_global_srtm15_data",
         "grid_that_straddles_180_degree_meridian_with_global_srtm15_data",
+        "grid_with_gshhs_coastlines",
     ],
 )
 def test_successful_initialization_with_topography(grid_fixture, request):
     grid = request.getfixturevalue(grid_fixture)
-    assert grid is not None
+
+    expected_attrs = [
+        "nx",
+        "ny",
+        "size_x",
+        "size_y",
+        "center_lon",
+        "center_lat",
+        "rot",
+        "N",
+        "theta_s",
+        "theta_b",
+        "hc",
+        "topography_source",
+        "hmin",
+        "mask_shapefile",
+        "verbose",
+        "straddle",
+    ]
+
+    for attr in expected_attrs:
+        assert hasattr(grid, attr), f"Missing attribute: {attr}"
 
 
 def test_plot(grid_that_straddles_180_degree_meridian):
@@ -187,6 +236,20 @@ def test_plot(grid_that_straddles_180_degree_meridian):
     grid_that_straddles_180_degree_meridian.plot(with_dim_names=True)
 
 
+def test_plot_wide_grid():
+    # This grid should be handled via different cartopy projection
+    grid = Grid(
+        nx=10,
+        ny=10,
+        size_x=15000,
+        size_y=15000,
+        center_lon=-161,
+        center_lat=14.4,
+        rot=-3,
+    )
+    grid.plot()
+
+
 @pytest.mark.skipif(xesmf is None, reason="xesmf required")
 def test_plot_along_lat_lon(grid_that_straddles_180_degree_meridian):
     grid_that_straddles_180_degree_meridian.plot(lat=61)
@@ -280,6 +343,7 @@ def test_grid_straddle_crosses_meridian():
         "grid",
         "grid_that_straddles_dateline_with_shifted_global_etopo_data",
         "grid_that_straddles_dateline_with_global_srtm15_data",
+        "grid_with_gshhs_coastlines",
     ],
 )
 def test_roundtrip_netcdf(grid_fixture, tmp_path, request):
@@ -301,8 +365,8 @@ def test_roundtrip_netcdf(grid_fixture, tmp_path, request):
             # Load the grid from the file
             grid_from_file = Grid.from_file(filepath.with_suffix(".nc"))
 
-            # Assert that the initial grid and the loaded grid are equivalent (including the 'ds' attribute)
             assert grid == grid_from_file
+            xr.testing.assert_equal(grid.ds, grid_from_file.ds)
 
             # Clean up the .nc file
             (filepath.with_suffix(".nc")).unlink()
@@ -314,6 +378,7 @@ def test_roundtrip_netcdf(grid_fixture, tmp_path, request):
         "grid",
         "grid_that_straddles_dateline_with_shifted_global_etopo_data",
         "grid_that_straddles_dateline_with_global_srtm15_data",
+        "grid_with_gshhs_coastlines",
     ],
 )
 def test_roundtrip_yaml(grid_fixture, tmp_path, request):
@@ -332,8 +397,8 @@ def test_roundtrip_yaml(grid_fixture, tmp_path, request):
 
         grid_from_file = Grid.from_yaml(filepath)
 
-        # Assert that the initial grid and the loaded grid are equivalent (including the 'ds' attribute)
         assert grid == grid_from_file
+        xr.testing.assert_equal(grid.ds, grid_from_file.ds)
 
         filepath = Path(filepath)
         filepath.unlink()
@@ -345,6 +410,7 @@ def test_roundtrip_yaml(grid_fixture, tmp_path, request):
         "grid",
         "grid_that_straddles_dateline_with_shifted_global_etopo_data",
         "grid_that_straddles_dateline_with_global_srtm15_data",
+        "grid_with_gshhs_coastlines",
     ],
 )
 def test_roundtrip_from_file_yaml(grid_fixture, tmp_path, request):
@@ -359,30 +425,33 @@ def test_roundtrip_from_file_yaml(grid_fixture, tmp_path, request):
     filepath_yaml = Path(tmp_path / "test.yaml")
     grid_from_file.to_yaml(filepath_yaml)
 
+    grid_from_yaml = Grid.from_yaml(filepath_yaml)
+
+    assert grid_from_yaml == grid
+    xr.testing.assert_equal(grid.ds, grid_from_yaml.ds)
+
     filepath.unlink()
     filepath_yaml.unlink()
 
 
-def test_files_have_same_hash(tmp_path):
-    # Initialize a Grid object using the initializer
-    grid_init = Grid(
-        nx=10,
-        ny=15,
-        size_x=100.0,
-        size_y=150.0,
-        center_lon=0.0,
-        center_lat=0.0,
-        rot=0.0,
-        topography_source={"name": "ETOPO5"},
-        hmin=5.0,
-    )
+@pytest.mark.parametrize(
+    "grid_fixture",
+    [
+        "grid",
+        "grid_that_straddles_dateline_with_shifted_global_etopo_data",
+        "grid_that_straddles_dateline_with_global_srtm15_data",
+        "grid_with_gshhs_coastlines",
+    ],
+)
+def test_files_have_same_hash(grid_fixture, tmp_path, request):
+    grid = request.getfixturevalue(grid_fixture)
 
     yaml_filepath = tmp_path / "test_yaml"
     filepath1 = tmp_path / "test1.nc"
     filepath2 = tmp_path / "test2.nc"
 
-    grid_init.to_yaml(yaml_filepath)
-    grid_init.save(filepath1)
+    grid.to_yaml(yaml_filepath)
+    grid.save(filepath1)
 
     grid_from_file = Grid.from_yaml(yaml_filepath)
     grid_from_file.save(filepath2)
@@ -505,6 +574,9 @@ def test_from_yaml_version_mismatch(tmp_path, caplog):
         yaml_filepath.unlink()
 
 
+# Vertical coordinate tests
+
+
 def test_invalid_theta_s_value():
     """Test the validation of the theta_s value."""
     with pytest.raises(ValueError):
@@ -608,6 +680,152 @@ def test_plot_vertical_coordinate():
         grid.plot_vertical_coordinate(eta=-1, xi=0, s=-1)
 
 
+# Topography tests
+
+
+def test_enclosed_regions():
+    """Test that there are only two connected regions, one dry and one wet."""
+    grid = Grid(
+        nx=100,
+        ny=100,
+        size_x=1800,
+        size_y=2400,
+        center_lon=30,
+        center_lat=61,
+        rot=20,
+    )
+
+    reg, nreg = label(grid.ds.mask_rho)
+    npt.assert_equal(nreg, 2)
+
+
+def test_rmax_criterion():
+    grid = Grid(
+        nx=100,
+        ny=100,
+        size_x=1800,
+        size_y=2400,
+        center_lon=30,
+        center_lat=61,
+        rot=20,
+    )
+    r_eta, r_xi = _compute_rfactor(grid.ds.h)
+    rmax0 = np.max([r_eta.max(), r_xi.max()])
+    npt.assert_array_less(rmax0, 0.2)
+
+
+def test_hmin_criterion_and_update_topography():
+    grid = Grid(
+        nx=100,
+        ny=100,
+        size_x=1800,
+        size_y=2400,
+        center_lon=30,
+        center_lat=61,
+        rot=20,
+        hmin=5.0,
+    )
+
+    assert grid.hmin == 5.0
+    assert np.less_equal(grid.hmin, grid.ds.h.min())
+
+    grid.update_topography(hmin=10.0)
+
+    assert grid.hmin == 10.0
+    assert np.less_equal(grid.hmin, grid.ds.h.min())
+
+    # this should not do anything
+    grid.update_topography()
+
+    assert grid.hmin == 10.0
+    assert np.less_equal(grid.hmin, grid.ds.h.min())
+
+
+# Mask tests
+
+
+def test_update_mask():
+    iceland_fjord_kwargs = {
+        "nx": 80,
+        "ny": 40,
+        "size_x": 40,
+        "size_y": 20,
+        "center_lon": -21.76,
+        "center_lat": 64.325,
+        "rot": 0,
+        "N": 3,
+    }
+
+    # Make sure all 4 L1 files are downloaded
+    _ = download_test_data("GSHHS_l_L1.dbf")
+    _ = download_test_data("GSHHS_l_L1.prj")
+    _ = download_test_data("GSHHS_l_L1.shx")
+    shapefile = download_test_data("GSHHS_l_L1.shp")
+
+    grid = Grid(
+        **iceland_fjord_kwargs,
+        mask_shapefile=shapefile,
+    )
+
+    assert grid.mask_shapefile == shapefile
+
+    # Save original mask
+    mask_orig = grid.ds.mask_rho.copy()
+
+    # Update mask (switches to Natural Earth)
+    grid.update_mask()
+
+    assert grid.mask_shapefile is None
+
+    # New mask after update
+    mask_new = grid.ds.mask_rho.copy()
+
+    assert abs(mask_new - mask_orig).max() == 1, (
+        "Mask should change after update_mask()"
+    )
+
+
+# Boundary tests
+
+
+def test_mask_topography_boundary():
+    """Test that the mask and topography along the grid boundaries (north, south, east,
+    west) are identical to the adjacent inland cells.
+    """
+    # Create a grid with some land along the northern boundary
+    grid = Grid(
+        nx=10, ny=10, size_x=1000, size_y=1000, center_lon=-20, center_lat=60, rot=0
+    )
+
+    # Toopography
+    np.testing.assert_array_equal(
+        grid.ds.h.isel(eta_rho=0).data, grid.ds.h.isel(eta_rho=1).data
+    )
+    np.testing.assert_array_equal(
+        grid.ds.h.isel(eta_rho=-1).data, grid.ds.h.isel(eta_rho=-2).data
+    )
+    np.testing.assert_array_equal(
+        grid.ds.h.isel(xi_rho=0).data, grid.ds.h.isel(xi_rho=1).data
+    )
+    np.testing.assert_array_equal(
+        grid.ds.h.isel(xi_rho=-1).data, grid.ds.h.isel(xi_rho=-2).data
+    )
+
+    # Mask
+    np.testing.assert_array_equal(
+        grid.ds.mask_rho.isel(eta_rho=0).data, grid.ds.mask_rho.isel(eta_rho=1).data
+    )
+    np.testing.assert_array_equal(
+        grid.ds.mask_rho.isel(eta_rho=-1).data, grid.ds.mask_rho.isel(eta_rho=-2).data
+    )
+    np.testing.assert_array_equal(
+        grid.ds.mask_rho.isel(xi_rho=0).data, grid.ds.mask_rho.isel(xi_rho=1).data
+    )
+    np.testing.assert_array_equal(
+        grid.ds.mask_rho.isel(xi_rho=-1).data, grid.ds.mask_rho.isel(xi_rho=-2).data
+    )
+
+
 # More Grid.from_file() tests
 
 
@@ -723,100 +941,3 @@ def test_from_file_partial_parameters_raises_error(grid, tmp_path):
 
     with pytest.raises(ValueError, match="must provide all of"):
         Grid.from_file(path, theta_s=5.0)
-
-
-# Topography tests
-def test_enclosed_regions():
-    """Test that there are only two connected regions, one dry and one wet."""
-    grid = Grid(
-        nx=100,
-        ny=100,
-        size_x=1800,
-        size_y=2400,
-        center_lon=30,
-        center_lat=61,
-        rot=20,
-    )
-
-    reg, nreg = label(grid.ds.mask_rho)
-    npt.assert_equal(nreg, 2)
-
-
-def test_rmax_criterion():
-    grid = Grid(
-        nx=100,
-        ny=100,
-        size_x=1800,
-        size_y=2400,
-        center_lon=30,
-        center_lat=61,
-        rot=20,
-    )
-    r_eta, r_xi = _compute_rfactor(grid.ds.h)
-    rmax0 = np.max([r_eta.max(), r_xi.max()])
-    npt.assert_array_less(rmax0, 0.2)
-
-
-def test_hmin_criterion():
-    grid = Grid(
-        nx=100,
-        ny=100,
-        size_x=1800,
-        size_y=2400,
-        center_lon=30,
-        center_lat=61,
-        rot=20,
-        hmin=5.0,
-    )
-
-    assert grid.hmin == 5.0
-    assert np.less_equal(grid.hmin, grid.ds.h.min())
-
-    grid.update_topography(hmin=10.0)
-
-    assert grid.hmin == 10.0
-    assert np.less_equal(grid.hmin, grid.ds.h.min())
-
-    # this should not do anything
-    grid.update_topography()
-
-    assert grid.hmin == 10.0
-    assert np.less_equal(grid.hmin, grid.ds.h.min())
-
-
-def test_mask_topography_boundary():
-    """Test that the mask and topography along the grid boundaries (north, south, east,
-    west) are identical to the adjacent inland cells.
-    """
-    # Create a grid with some land along the northern boundary
-    grid = Grid(
-        nx=10, ny=10, size_x=1000, size_y=1000, center_lon=-20, center_lat=60, rot=0
-    )
-
-    # Toopography
-    np.testing.assert_array_equal(
-        grid.ds.h.isel(eta_rho=0).data, grid.ds.h.isel(eta_rho=1).data
-    )
-    np.testing.assert_array_equal(
-        grid.ds.h.isel(eta_rho=-1).data, grid.ds.h.isel(eta_rho=-2).data
-    )
-    np.testing.assert_array_equal(
-        grid.ds.h.isel(xi_rho=0).data, grid.ds.h.isel(xi_rho=1).data
-    )
-    np.testing.assert_array_equal(
-        grid.ds.h.isel(xi_rho=-1).data, grid.ds.h.isel(xi_rho=-2).data
-    )
-
-    # Mask
-    np.testing.assert_array_equal(
-        grid.ds.mask_rho.isel(eta_rho=0).data, grid.ds.mask_rho.isel(eta_rho=1).data
-    )
-    np.testing.assert_array_equal(
-        grid.ds.mask_rho.isel(eta_rho=-1).data, grid.ds.mask_rho.isel(eta_rho=-2).data
-    )
-    np.testing.assert_array_equal(
-        grid.ds.mask_rho.isel(xi_rho=0).data, grid.ds.mask_rho.isel(xi_rho=1).data
-    )
-    np.testing.assert_array_equal(
-        grid.ds.mask_rho.isel(xi_rho=-1).data, grid.ds.mask_rho.isel(xi_rho=-2).data
-    )

roms_tools/tests/test_setup/test_initial_conditions.py

@@ -11,7 +11,35 @@ 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.datasets import CESMBGCDataset, UnifiedBGCDataset
+from roms_tools.setup.lat_lon_datasets import (
+    CESMBGCDataset,
+    UnifiedBGCDataset,
+)
+from roms_tools.tests.test_setup.utils import download_regional_and_bigger
+
+try:
+    import copernicusmarine  # type: ignore
+except ImportError:
+    copernicusmarine = None
+
+
+@pytest.fixture
+def example_grid():
+    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
+    )
+
+    return grid
 
 
 @pytest.mark.parametrize(
@@ -25,7 +53,9 @@ from roms_tools.setup.datasets import CESMBGCDataset, UnifiedBGCDataset
         "initial_conditions_with_unified_bgc_from_climatology",
     ],
 )
-def test_initial_conditions_creation(ic_fixture, request):
+def test_initial_conditions_creation_with_nondefault_glorys_dataset(
+    ic_fixture, request
+):
     """Test the creation of the InitialConditions object."""
     ic = request.getfixturevalue(ic_fixture)
 
@@ -37,12 +67,65 @@ def test_initial_conditions_creation(ic_fixture, request):
     }
     assert hasattr(ic.ds, "adjust_depth_for_sea_surface_height")
     assert isinstance(ic.ds, xr.Dataset)
-    assert "temp" in ic.ds
-    assert "salt" in ic.ds
-    assert "u" in ic.ds
-    assert "v" in ic.ds
-    assert "zeta" in ic.ds
     assert ic.ds.coords["ocean_time"].attrs["units"] == "seconds"
+    expected_vars = {"temp", "salt", "u", "v", "zeta", "ubar", "vbar"}
+    assert set(ic.ds.data_vars).issuperset(expected_vars)
+
+
+@pytest.mark.stream
+@pytest.mark.use_copernicus
+@pytest.mark.use_dask
+def test_initial_conditions_creation_with_default_glorys_dataset(example_grid: Grid):
+    """Verify the default GLORYS dataset is loaded when a path is not provided."""
+    ic = InitialConditions(
+        grid=example_grid,
+        ini_time=datetime(2021, 6, 29),
+        source={"name": "GLORYS"},
+        use_dask=True,
+        bypass_validation=True,
+    )
+    expected_vars = {"temp", "salt", "u", "v", "zeta", "ubar", "vbar"}
+    assert set(ic.ds.data_vars).issuperset(expected_vars)
+
+
+@pytest.mark.use_copernicus
+@pytest.mark.skipif(copernicusmarine is None, reason="copernicusmarine required")
+@pytest.mark.parametrize(
+    "grid_fixture",
+    [
+        "tiny_grid_that_straddles_dateline",
+        "tiny_grid_that_straddles_180_degree_meridian",
+        "tiny_rotated_grid",
+    ],
+)
+def test_invariance_to_get_glorys_bounds(tmp_path, grid_fixture, use_dask, request):
+    ini_time = datetime(2012, 1, 1)
+    grid = request.getfixturevalue(grid_fixture)
+    regional_file, bigger_regional_file = download_regional_and_bigger(
+        tmp_path, grid, ini_time
+    )
+
+    ic_from_regional = InitialConditions(
+        grid=grid,
+        source={"name": "GLORYS", "path": str(regional_file)},
+        ini_time=ini_time,
+        use_dask=use_dask,
+    )
+    ic_from_bigger_regional = InitialConditions(
+        grid=grid,
+        source={"name": "GLORYS", "path": str(bigger_regional_file)},
+        ini_time=ini_time,
+        use_dask=use_dask,
+    )
+
+    # Use assert_allclose instead of equals: necessary for grids that straddle the 180° meridian.
+    # Copernicus returns data on [-180, 180] by default, but if you request a range
+    # like [170, 190], it remaps longitudes. That remapping introduces tiny floating
+    # point differences in the longitude coordinate, which will then propagate into further differences once you do regridding.
+    # Need to adjust the tolerances for these grids that straddle the 180° meridian.
+    xr.testing.assert_allclose(
+        ic_from_bigger_regional.ds, ic_from_regional.ds, rtol=1e-4, atol=1e-5
+    )
 
 
 def test_initial_conditions_creation_with_duplicates(use_dask: bool) -> None:
@@ -67,6 +150,7 @@ def test_initial_conditions_creation_with_duplicates(use_dask: bool) -> None:
         grid=grid,
         ini_time=datetime(2012, 1, 1),
         source={"path": [fname1, fname2], "name": "GLORYS"},
+        allow_flex_time=True,
         use_dask=use_dask,
     )
 
@@ -74,6 +158,7 @@ def test_initial_conditions_creation_with_duplicates(use_dask: bool) -> None:
         grid=grid,
         ini_time=datetime(2012, 1, 1),
         source={"path": [fname1, fname1, fname2], "name": "GLORYS"},
+        allow_flex_time=True,
         use_dask=use_dask,
     )
 
@@ -134,28 +219,9 @@ def test_initial_condition_creation_with_bgc(ic_fixture, request):
         assert var in ic.ds
 
 
-@pytest.fixture
-def example_grid():
-    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
-    )
-
-    return grid
-
-
 # Test initialization with missing 'name' in source
 def test_initial_conditions_missing_physics_name(example_grid, use_dask):
-    with pytest.raises(ValueError, match="`source` must include a 'name'."):
+    with pytest.raises(ValueError, match="`source` must include a 'name'"):
         InitialConditions(
             grid=example_grid,
             ini_time=datetime(2021, 6, 29),
@@ -164,23 +230,10 @@ def test_initial_conditions_missing_physics_name(example_grid, use_dask):
         )
 
 
-# Test initialization with missing 'path' in source
-def test_initial_conditions_missing_physics_path(example_grid, use_dask):
-    with pytest.raises(ValueError, match="`source` must include a 'path'."):
-        InitialConditions(
-            grid=example_grid,
-            ini_time=datetime(2021, 6, 29),
-            source={"name": "GLORYS"},
-            use_dask=use_dask,
-        )
-
-
 # Test initialization with missing 'name' in bgc_source
 def test_initial_conditions_missing_bgc_name(example_grid, use_dask):
     fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
-    with pytest.raises(
-        ValueError, match="`bgc_source` must include a 'name' if it is provided."
-    ):
+    with pytest.raises(ValueError, match="`bgc_source` must include a 'name'"):
         InitialConditions(
             grid=example_grid,
             ini_time=datetime(2021, 6, 29),