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

roms_tools/tests/test_setup/test_datasets.py

@@ -1,7 +1,8 @@
 import logging
 from collections import OrderedDict
-from datetime import datetime
+from datetime import datetime, timedelta
 from pathlib import Path
+from unittest import mock
 
 import numpy as np
 import pytest
@@ -9,13 +10,27 @@ import xarray as xr
 
 from roms_tools.download import download_test_data
 from roms_tools.setup.datasets import (
+    GLORYS_GLOBAL_GRID_PATH,
     CESMBGCDataset,
     Dataset,
+    ERA5ARCODataset,
     ERA5Correction,
     GLORYSDataset,
+    GLORYSDefaultDataset,
     RiverDataset,
     TPXODataset,
+    _concatenate_longitudes,
+    choose_subdomain,
+    get_glorys_bounds,
 )
+from roms_tools.setup.surface_forcing import DEFAULT_ERA5_ARCO_PATH
+from roms_tools.setup.utils import get_target_coords
+from roms_tools.tests.test_setup.utils import download_regional_and_bigger
+
+try:
+    import copernicusmarine  # type: ignore
+except ImportError:
+    copernicusmarine = None
 
 
 @pytest.fixture
@@ -25,9 +40,9 @@ def global_dataset():
     depth = np.linspace(0, 2000, 10)
     time = [
         np.datetime64("2022-01-01T00:00:00"),
-        np.datetime64("2022-02-01T00:00:00"),
-        np.datetime64("2022-03-01T00:00:00"),
-        np.datetime64("2022-04-01T00:00:00"),
+        np.datetime64("2022-01-02T00:00:00"),
+        np.datetime64("2022-01-03T00:00:00"),
+        np.datetime64("2022-01-04T00:00:00"),
     ]
     data = np.random.rand(4, 10, 180, 360)
     ds = xr.Dataset(
@@ -48,9 +63,9 @@ def global_dataset_with_noon_times():
     lat = np.linspace(-90, 90, 180)
     time = [
         np.datetime64("2022-01-01T12:00:00"),
-        np.datetime64("2022-02-01T12:00:00"),
-        np.datetime64("2022-03-01T12:00:00"),
-        np.datetime64("2022-04-01T12:00:00"),
+        np.datetime64("2022-01-02T12:00:00"),
+        np.datetime64("2022-01-03T12:00:00"),
+        np.datetime64("2022-01-04T12:00:00"),
     ]
     data = np.random.rand(4, 180, 360)
     ds = xr.Dataset(
@@ -71,12 +86,12 @@ def global_dataset_with_multiple_times_per_day():
     time = [
         np.datetime64("2022-01-01T00:00:00"),
         np.datetime64("2022-01-01T12:00:00"),
-        np.datetime64("2022-02-01T00:00:00"),
-        np.datetime64("2022-02-01T12:00:00"),
-        np.datetime64("2022-03-01T00:00:00"),
-        np.datetime64("2022-03-01T12:00:00"),
-        np.datetime64("2022-04-01T00:00:00"),
-        np.datetime64("2022-04-01T12:00:00"),
+        np.datetime64("2022-01-02T00:00:00"),
+        np.datetime64("2022-01-02T12:00:00"),
+        np.datetime64("2022-01-03T00:00:00"),
+        np.datetime64("2022-01-03T12:00:00"),
+        np.datetime64("2022-01-04T00:00:00"),
+        np.datetime64("2022-01-04T12:00:00"),
     ]
     data = np.random.rand(8, 180, 360)
     ds = xr.Dataset(
@@ -108,34 +123,37 @@ def non_global_dataset():
         (
             "global_dataset",
             [
-                np.datetime64("2022-02-01T00:00:00"),
-                np.datetime64("2022-03-01T00:00:00"),
+                np.datetime64("2022-01-02T00:00:00"),
+                np.datetime64("2022-01-03T00:00:00"),
+                np.datetime64("2022-01-04T00:00:00"),
             ],
         ),
         (
             "global_dataset_with_noon_times",
             [
                 np.datetime64("2022-01-01T12:00:00"),
-                np.datetime64("2022-02-01T12:00:00"),
-                np.datetime64("2022-03-01T12:00:00"),
+                np.datetime64("2022-01-02T12:00:00"),
+                np.datetime64("2022-01-03T12:00:00"),
+                np.datetime64("2022-01-04T12:00:00"),
             ],
         ),
         (
             "global_dataset_with_multiple_times_per_day",
             [
-                np.datetime64("2022-02-01T00:00:00"),
-                np.datetime64("2022-02-01T12:00:00"),
-                np.datetime64("2022-03-01T00:00:00"),
+                np.datetime64("2022-01-02T00:00:00"),
+                np.datetime64("2022-01-02T12:00:00"),
+                np.datetime64("2022-01-03T00:00:00"),
+                np.datetime64("2022-01-03T12:00:00"),
+                np.datetime64("2022-01-04T00:00:00"),
             ],
         ),
     ],
 )
 def test_select_times(data_fixture, expected_time_values, request, tmp_path, use_dask):
     """Test selecting times with different datasets."""
-    start_time = datetime(2022, 2, 1)
-    end_time = datetime(2022, 3, 1)
+    start_time = datetime(2022, 1, 2)
+    end_time = datetime(2022, 1, 4)
 
-    # Get the fixture dynamically based on the parameter
     dataset = request.getfixturevalue(data_fixture)
 
     filepath = tmp_path / "test.nc"
@@ -157,16 +175,15 @@ def test_select_times(data_fixture, expected_time_values, request, tmp_path, use
 @pytest.mark.parametrize(
     "data_fixture, expected_time_values",
     [
-        ("global_dataset", [np.datetime64("2022-02-01T00:00:00")]),
-        ("global_dataset_with_noon_times", [np.datetime64("2022-02-01T12:00:00")]),
+        ("global_dataset", [np.datetime64("2022-01-02T00:00:00")]),
+        ("global_dataset_with_noon_times", [np.datetime64("2022-01-02T12:00:00")]),
     ],
 )
 def test_select_times_valid_start_no_end_time(
     data_fixture, expected_time_values, request, tmp_path, use_dask
 ):
     """Test selecting times with only start_time specified."""
-    start_time = datetime(2022, 2, 1)
-
+    start_time = datetime(2022, 1, 2)
     # Get the fixture dynamically based on the parameter
     dataset = request.getfixturevalue(data_fixture)
 
@@ -180,9 +197,11 @@ def test_select_times_valid_start_no_end_time(
         var_names={"var": "var"},
         start_time=start_time,
         use_dask=use_dask,
+        allow_flex_time=True,
     )
 
     assert dataset.ds is not None
+    assert "time" in dataset.ds.dims
     assert len(dataset.ds.time) == len(expected_time_values)
     for expected_time in expected_time_values:
         assert expected_time in dataset.ds.time.values
@@ -208,7 +227,7 @@ def test_select_times_invalid_start_no_end_time(
 
     with pytest.raises(
         ValueError,
-        match="The dataset does not contain any time entries between the specified start_time",
+        match="No exact match found ",
     ):
         dataset = Dataset(
             filename=filepath,
@@ -229,11 +248,12 @@ def test_multiple_matching_times(
     dataset = Dataset(
         filename=filepath,
         var_names={"var": "var"},
-        start_time=datetime(2022, 1, 31, 22, 0),
+        start_time=datetime(2021, 12, 31, 22, 0),
         use_dask=use_dask,
+        allow_flex_time=True,
     )
 
-    assert dataset.ds["time"].values == np.datetime64(datetime(2022, 2, 1, 0, 0))
+    assert dataset.ds["time"].values == np.datetime64(datetime(2022, 1, 1, 0, 0))
 
 
 def test_warnings_times(global_dataset, tmp_path, caplog, use_dask):
@@ -253,7 +273,7 @@ def test_warnings_times(global_dataset, tmp_path, caplog, use_dask):
             use_dask=use_dask,
         )
     # Verify the warning message in the log
-    assert "No records found at or before the start_time." in caplog.text
+    assert "No records found at or before the start_time" in caplog.text
 
     with caplog.at_level(logging.WARNING):
         start_time = datetime(2024, 1, 1)
@@ -267,7 +287,7 @@ def test_warnings_times(global_dataset, tmp_path, caplog, use_dask):
             use_dask=use_dask,
         )
     # Verify the warning message in the log
-    assert "No records found at or after the end_time." in caplog.text
+    assert "No records found at or after the end_time" in caplog.text
 
 
 def test_from_ds(global_dataset, global_dataset_with_noon_times, use_dask, tmp_path):
@@ -288,6 +308,7 @@ def test_from_ds(global_dataset, global_dataset_with_noon_times, use_dask, tmp_p
         },
         start_time=start_time,
         use_dask=use_dask,
+        allow_flex_time=True,
     )
 
     new_dataset = Dataset.from_ds(dataset, global_dataset_with_noon_times)
@@ -329,6 +350,7 @@ def test_reverse_latitude_reverse_depth_choose_subdomain(
         },
         start_time=start_time,
         use_dask=use_dask,
+        allow_flex_time=True,
     )
 
     assert np.all(np.diff(dataset.ds["latitude"]) > 0)
@@ -427,16 +449,122 @@ def test_check_dataset(global_dataset, tmp_path, use_dask):
         )
 
 
-def test_era5_correction_choose_subdomain(use_dask):
+def test_era5_correction_match_subdomain(use_dask):
     data = ERA5Correction(use_dask=use_dask)
     lats = data.ds.latitude[10:20]
     lons = data.ds.longitude[10:20]
     target_coords = {"lat": lats, "lon": lons}
-    data.choose_subdomain(target_coords, straddle=False)
+    data.match_subdomain(target_coords)
     assert (data.ds["latitude"] == lats).all()
     assert (data.ds["longitude"] == lons).all()
 
 
+@pytest.mark.use_gcsfs
+def test_default_era5_dataset_loading_without_dask() -> None:
+    """Verify that loading the default ERA5 dataset fails if use_dask is not True."""
+    start_time = datetime(2020, 2, 1)
+    end_time = datetime(2020, 2, 2)
+
+    with pytest.raises(ValueError):
+        _ = ERA5ARCODataset(
+            filename=DEFAULT_ERA5_ARCO_PATH,
+            start_time=start_time,
+            end_time=end_time,
+            use_dask=False,
+        )
+
+
+@pytest.mark.skip("Temporary skip until memory consumption issue is addressed. # TODO")
+@pytest.mark.stream
+@pytest.mark.use_dask
+@pytest.mark.use_gcsfs
+def test_default_era5_dataset_loading() -> None:
+    """Verify the default ERA5 dataset is loaded correctly."""
+    start_time = datetime(2020, 2, 1)
+    end_time = datetime(2020, 2, 2)
+
+    ds = ERA5ARCODataset(
+        filename=DEFAULT_ERA5_ARCO_PATH,
+        start_time=start_time,
+        end_time=end_time,
+        use_dask=True,
+    )
+
+    expected_vars = {"uwnd", "vwnd", "swrad", "lwrad", "Tair", "rain"}
+    assert set(ds.var_names).issuperset(expected_vars)
+
+
+@pytest.mark.use_copernicus
+def test_default_glorys_dataset_loading_dask_not_installed() -> None:
+    """Verify that loading the default GLORYS dataset fails if dask is not available."""
+    start_time = datetime(2020, 2, 1)
+    end_time = datetime(2020, 2, 2)
+
+    with (
+        pytest.raises(RuntimeError),
+        mock.patch("roms_tools.utils.has_dask", return_value=False),
+    ):
+        _ = GLORYSDefaultDataset(
+            filename=GLORYSDefaultDataset.dataset_name,
+            start_time=start_time,
+            end_time=end_time,
+            use_dask=True,
+        )
+
+
+@pytest.mark.stream
+@pytest.mark.use_copernicus
+@pytest.mark.use_dask
+def test_default_glorys_dataset_loading() -> None:
+    """Verify the default GLORYS dataset is loaded correctly."""
+    start_time = datetime(2012, 1, 1)
+
+    for end_time in [start_time, start_time + timedelta(days=0.5)]:
+        data = GLORYSDefaultDataset(
+            filename=GLORYSDefaultDataset.dataset_name,
+            start_time=start_time,
+            end_time=end_time,
+            use_dask=True,
+        )
+
+        expected_vars = {"temp", "salt", "u", "v", "zeta"}
+        assert set(data.var_names).issuperset(expected_vars)
+
+        expected_vars = {"thetao", "so", "uo", "vo", "zos"}
+        assert "time" in data.ds.dims
+        assert set(data.ds.data_vars).issuperset(expected_vars)
+
+
+@pytest.mark.parametrize(
+    "fname,start_time",
+    [
+        (download_test_data("GLORYS_NA_2012.nc"), datetime(2012, 1, 1, 12)),
+        (download_test_data("GLORYS_NA_20121231.nc"), datetime(2012, 12, 31, 12)),
+        (download_test_data("GLORYS_coarse_test_data.nc"), datetime(2021, 6, 29)),
+    ],
+)
+@pytest.mark.parametrize("allow_flex_time", [True, False])
+def test_non_default_glorys_dataset_loading(
+    fname, start_time, allow_flex_time, use_dask
+) -> None:
+    """Verify the default GLORYS dataset is loaded correctly."""
+    for end_time in [None, start_time, start_time]:
+        data = GLORYSDataset(
+            filename=fname,
+            start_time=start_time,
+            end_time=end_time,
+            use_dask=use_dask,
+            allow_flex_time=allow_flex_time,
+        )
+
+        expected_vars = {"temp", "salt", "u", "v", "zeta"}
+        assert set(data.var_names).issuperset(expected_vars)
+
+        expected_vars = {"thetao", "so", "uo", "vo", "zos"}
+        assert "time" in data.ds.dims
+        assert set(data.ds.data_vars).issuperset(expected_vars)
+
+
 def test_data_concatenation(use_dask):
     fname = download_test_data("GLORYS_NA_2012.nc")
     data = GLORYSDataset(
@@ -672,3 +800,299 @@ class TestRiverDataset:
         assert "Amazon_2" in names
         assert "Nile" in names
         assert len(set(names)) == len(names)  # all names must be unique
+
+
+# test _concatenate_longitudes
+
+
+@pytest.fixture
+def sample_ds(use_dask):
+    lon = xr.DataArray(np.array([0, 90, 180]), dims="lon", name="lon")
+    lat = xr.DataArray(np.array([-30, 0, 30]), dims="lat", name="lat")
+
+    var_with_lon = xr.DataArray(
+        np.arange(9).reshape(3, 3),
+        dims=("lat", "lon"),
+        coords={"lat": lat, "lon": lon},
+        name="var_with_lon",
+    )
+
+    var_no_lon = xr.DataArray(
+        np.array([1, 2, 3]),
+        dims="lat",
+        coords={"lat": lat},
+        name="var_no_lon",
+    )
+
+    ds = xr.Dataset({"var_with_lon": var_with_lon, "var_no_lon": var_no_lon})
+
+    if use_dask:
+        ds = ds.chunk({"lat": -1, "lon": -1})
+
+    return ds
+
+
+@pytest.mark.parametrize(
+    "end,expected_lons",
+    [
+        ("lower", [-360, -270, -180, 0, 90, 180]),
+        ("upper", [0, 90, 180, 360, 450, 540]),
+        ("both", [-360, -270, -180, 0, 90, 180, 360, 450, 540]),
+    ],
+)
+def test_concatenate_longitudes(sample_ds, end, expected_lons, use_dask):
+    dim_names = {"longitude": "lon"}
+
+    ds_concat = _concatenate_longitudes(
+        sample_ds, dim_names, end=end, use_dask=use_dask
+    )
+
+    # longitude should be extended as expected
+    np.testing.assert_array_equal(ds_concat.lon.values, expected_lons)
+
+    # variable with longitude should be extended in size
+    assert ds_concat.var_with_lon.shape[-1] == len(expected_lons)
+
+    # variable without longitude should remain untouched
+    np.testing.assert_array_equal(
+        ds_concat.var_no_lon.values,
+        sample_ds.var_no_lon.values,
+    )
+
+    if use_dask:
+        import dask
+
+        # Ensure dask array backing the data
+        assert isinstance(ds_concat.var_with_lon.data, dask.array.Array)
+        # Longitude dimension should be chunked (-1 → one chunk spanning the whole dim)
+        assert ds_concat.var_with_lon.chunks[-1] == (len(expected_lons),)
+    else:
+        # With use_dask=False, data should be a numpy array
+        assert isinstance(ds_concat.var_with_lon.data, np.ndarray)
+
+
+def test_invalid_end_raises(sample_ds):
+    dim_names = {"longitude": "lon"}
+    with pytest.raises(ValueError):
+        _concatenate_longitudes(sample_ds, dim_names, end="invalid")
+
+
+# test choose_subdomain
+
+
+def test_choose_subdomain_basic(global_dataset, use_dask):
+    target_coords = {
+        "lat": xr.DataArray([0, 10]),
+        "lon": xr.DataArray([30, 40]),
+        "straddle": False,
+    }
+    out = choose_subdomain(
+        global_dataset,
+        dim_names={"latitude": "latitude", "longitude": "longitude"},
+        resolution=1.0,
+        is_global=True,
+        target_coords=target_coords,
+        buffer_points=2,
+        use_dask=use_dask,
+    )
+    assert out.latitude.min() <= 0
+    assert out.latitude.max() >= 10
+    assert out.longitude.min() <= 30
+    assert out.longitude.max() >= 40
+
+
+def test_choose_subdomain_raises_on_empty_lon(non_global_dataset, use_dask):
+    target_coords = {
+        "lat": xr.DataArray([-10, 10]),
+        "lon": xr.DataArray([210, 215]),  # outside 0-180
+        "straddle": False,
+    }
+    with pytest.raises(ValueError, match="longitude range"):
+        choose_subdomain(
+            non_global_dataset,
+            dim_names={"latitude": "latitude", "longitude": "longitude"},
+            resolution=1.0,
+            is_global=False,
+            target_coords=target_coords,
+            use_dask=use_dask,
+        )
+
+
+def test_choose_subdomain_raises_on_empty_lat(global_dataset, use_dask):
+    target_coords = {
+        "lat": xr.DataArray([1000, 1010]),  # outside dataset range
+        "lon": xr.DataArray([30, 40]),
+        "straddle": False,
+    }
+    with pytest.raises(ValueError, match="latitude range"):
+        choose_subdomain(
+            global_dataset,
+            dim_names={"latitude": "latitude", "longitude": "longitude"},
+            resolution=1.0,
+            is_global=True,
+            target_coords=target_coords,
+            use_dask=use_dask,
+        )
+
+
+def test_choose_subdomain_straddle(global_dataset, use_dask):
+    target_coords = {
+        "lat": xr.DataArray([-10, 10]),
+        "lon": xr.DataArray([-170, 170]),  # cross the dateline
+        "straddle": True,
+    }
+    out = choose_subdomain(
+        global_dataset,
+        dim_names={"latitude": "latitude", "longitude": "longitude"},
+        resolution=1.0,
+        is_global=True,
+        target_coords=target_coords,
+        buffer_points=5,
+        use_dask=use_dask,
+    )
+    # Ensure output includes both sides of the dateline, mapped into -180 - 180
+    assert (out.longitude.min() < 0) and (out.longitude.max() > 0)
+
+
+def test_choose_subdomain_wraps_negative_lon(global_dataset, use_dask):
+    target_coords = {
+        "lat": xr.DataArray([0, 20]),
+        "lon": xr.DataArray([-20, -10]),  # negative longitudes
+        "straddle": False,
+    }
+    out = choose_subdomain(
+        global_dataset,
+        dim_names={"latitude": "latitude", "longitude": "longitude"},
+        resolution=1.0,
+        is_global=True,
+        target_coords=target_coords,
+        buffer_points=2,
+        use_dask=use_dask,
+    )
+    # Output longitudes should be shifted into [0, 360]
+    assert (out.longitude >= 0).all()
+
+
+def test_choose_subdomain_respects_buffer(global_dataset, use_dask):
+    target_coords = {
+        "lat": xr.DataArray([0, 0]),
+        "lon": xr.DataArray([50, 50]),
+        "straddle": False,
+    }
+    out = choose_subdomain(
+        global_dataset,
+        dim_names={"latitude": "latitude", "longitude": "longitude"},
+        resolution=1.0,
+        is_global=True,
+        target_coords=target_coords,
+        buffer_points=10,
+        use_dask=use_dask,
+    )
+    # Buffer should extend at least 10 degrees beyond target lon
+    assert out.longitude.min() <= 40
+    assert out.longitude.max() >= 60
+
+
+# test get_glorys_bounds
+
+
+@pytest.fixture
+def glorys_grid_0_360(tmp_path):
+    lats = np.linspace(-90, 90, 181)
+    lons = np.linspace(0, 360, 361)
+    ds = xr.Dataset(coords={"latitude": lats, "longitude": lons})
+    path = tmp_path / "GLORYS_0_360.nc"
+    ds.to_netcdf(path)
+    return path
+
+
+@pytest.fixture
+def glorys_grid_neg180_180(tmp_path):
+    lats = np.linspace(-90, 90, 181)
+    lons = np.linspace(-180, 180, 361)
+    ds = xr.Dataset(coords={"latitude": lats, "longitude": lons})
+    path = tmp_path / "GLORYS_neg180_180.nc"
+    ds.to_netcdf(path)
+    return path
+
+
+@pytest.fixture
+def glorys_grid_real():
+    return GLORYS_GLOBAL_GRID_PATH
+
+
+@pytest.mark.parametrize(
+    "grid_fixture",
+    [
+        "grid",
+        "grid_that_straddles_dateline",
+        "grid_that_straddles_180_degree_meridian",
+        "small_grid",
+        "tiny_grid",
+    ],
+)
+@pytest.mark.parametrize(
+    "glorys_grid_fixture",
+    ["glorys_grid_0_360", "glorys_grid_neg180_180", "glorys_grid_real"],
+)
+def test_get_glorys_bounds(tmp_path, grid_fixture, glorys_grid_fixture, request):
+    grid = request.getfixturevalue(grid_fixture)
+    glorys_grid_path = request.getfixturevalue(glorys_grid_fixture)
+
+    bounds = get_glorys_bounds(grid=grid, glorys_grid_path=glorys_grid_path)
+    assert set(bounds) == {
+        "minimum_latitude",
+        "maximum_latitude",
+        "minimum_longitude",
+        "maximum_longitude",
+    }
+
+
+@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):
+    start_time = datetime(2012, 1, 1)
+    grid = request.getfixturevalue(grid_fixture)
+    target_coords = get_target_coords(grid)
+
+    regional_file, bigger_regional_file = download_regional_and_bigger(
+        tmp_path, grid, start_time, variables=["thetao", "uo", "zos"]
+    )
+
+    # create datasets from regional and bigger regional data
+    regional_data = GLORYSDataset(
+        var_names={"temp": "thetao", "u": "uo", "zeta": "zos"},
+        filename=regional_file,
+        start_time=start_time,
+        climatology=False,
+        allow_flex_time=False,
+        use_dask=use_dask,
+    )
+    bigger_regional_data = GLORYSDataset(
+        var_names={"temp": "thetao", "u": "uo", "zeta": "zos"},
+        filename=bigger_regional_file,
+        start_time=start_time,
+        climatology=False,
+        allow_flex_time=False,
+        use_dask=use_dask,
+    )
+
+    # subset both datasets and check they are the same
+    regional_data.choose_subdomain(target_coords)
+    bigger_regional_data.choose_subdomain(target_coords)
+
+    # 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, so strict equality will fail
+    # even though the data are effectively identical.
+    # For grids that do not straddle the 180° meridian, strict equality still holds.
+    xr.testing.assert_allclose(bigger_regional_data.ds, regional_data.ds)