roms-tools 1.3.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

roms_tools/tests/test_setup/test_datasets.py

@@ -2,9 +2,14 @@ import pytest
 from datetime import datetime
 import numpy as np
 import xarray as xr
-from roms_tools.setup.datasets import Dataset, ERA5Correction
-import tempfile
-import os
+from roms_tools.setup.datasets import (
+    Dataset,
+    GLORYSDataset,
+    ERA5Correction,
+    CESMBGCDataset,
+)
+from roms_tools.setup.download import download_test_data
+from pathlib import Path
 
 
 @pytest.fixture
@@ -119,7 +124,7 @@ def non_global_dataset():
         ),
     ],
 )
-def test_select_times(data_fixture, expected_time_values, request, use_dask):
+def test_select_times(data_fixture, expected_time_values, request, tmp_path, use_dask):
     """
     Test selecting times with different datasets.
     """
@@ -129,26 +134,20 @@ def test_select_times(data_fixture, expected_time_values, request, use_dask):
     # Get the fixture dynamically based on the parameter
     dataset = request.getfixturevalue(data_fixture)
 
-    # Create a temporary file
-    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
-        filepath = tmpfile.name
-        dataset.to_netcdf(filepath)
-    try:
-        # Instantiate Dataset object using the temporary file
-        dataset = Dataset(
-            filename=filepath,
-            var_names={"var": "var"},
-            start_time=start_time,
-            end_time=end_time,
-            use_dask=use_dask,
-        )
+    filepath = tmp_path / "test.nc"
+    dataset.to_netcdf(filepath)
+    dataset = Dataset(
+        filename=filepath,
+        var_names={"var": "var"},
+        start_time=start_time,
+        end_time=end_time,
+        use_dask=use_dask,
+    )
 
-        assert dataset.ds is not None
-        assert len(dataset.ds.time) == len(expected_time_values)
-        for expected_time in expected_time_values:
-            assert expected_time in dataset.ds.time.values
-    finally:
-        os.remove(filepath)
+    assert dataset.ds is not None
+    assert len(dataset.ds.time) == len(expected_time_values)
+    for expected_time in expected_time_values:
+        assert expected_time in dataset.ds.time.values
 
 
 @pytest.mark.parametrize(
@@ -158,8 +157,8 @@ def test_select_times(data_fixture, expected_time_values, request, use_dask):
         ("global_dataset_with_noon_times", [np.datetime64("2022-02-01T12:00:00")]),
     ],
 )
-def test_select_times_no_end_time(
-    data_fixture, expected_time_values, request, use_dask
+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.
@@ -170,241 +169,224 @@ def test_select_times_no_end_time(
     dataset = request.getfixturevalue(data_fixture)
 
     # Create a temporary file
-    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
-        filepath = tmpfile.name
-        dataset.to_netcdf(filepath)
-    try:
-        # Instantiate Dataset object using the temporary file
+    filepath = tmp_path / "test.nc"
+    dataset.to_netcdf(filepath)
+
+    # Instantiate Dataset object using the temporary file
+    dataset = Dataset(
+        filename=filepath,
+        var_names={"var": "var"},
+        start_time=start_time,
+        use_dask=use_dask,
+    )
+
+    assert dataset.ds is not None
+    assert len(dataset.ds.time) == len(expected_time_values)
+    for expected_time in expected_time_values:
+        assert expected_time in dataset.ds.time.values
+
+
+@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")]),
+    ],
+)
+def test_select_times_invalid_start_no_end_time(
+    data_fixture, expected_time_values, request, tmp_path, use_dask
+):
+    """
+    Test selecting times with only start_time specified.
+    """
+    # Get the fixture dynamically based on the parameter
+    dataset = request.getfixturevalue(data_fixture)
+
+    # Create a temporary file
+    filepath = tmp_path / "test.nc"
+    dataset.to_netcdf(filepath)
+
+    with pytest.raises(
+        ValueError,
+        match="The dataset does not contain any time entries between the specified start_time",
+    ):
         dataset = Dataset(
             filename=filepath,
             var_names={"var": "var"},
-            start_time=start_time,
+            start_time=datetime(2022, 5, 1),
             use_dask=use_dask,
         )
 
-        assert dataset.ds is not None
-        assert len(dataset.ds.time) == len(expected_time_values)
-        for expected_time in expected_time_values:
-            assert expected_time in dataset.ds.time.values
-    finally:
-        os.remove(filepath)
-
 
-def test_multiple_matching_times(global_dataset_with_multiple_times_per_day, use_dask):
+def test_multiple_matching_times(
+    global_dataset_with_multiple_times_per_day, tmp_path, use_dask
+):
     """
     Test handling when multiple matching times are found when end_time is not specified.
     """
-    start_time = datetime(2022, 1, 1)
+    filepath = tmp_path / "test.nc"
+    global_dataset_with_multiple_times_per_day.to_netcdf(filepath)
+    dataset = Dataset(
+        filename=filepath,
+        var_names={"var": "var"},
+        start_time=datetime(2022, 1, 31, 22, 0),
+        use_dask=use_dask,
+    )
 
-    # Create a temporary file
-    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
-        filepath = tmpfile.name
-        global_dataset_with_multiple_times_per_day.to_netcdf(filepath)
-    try:
-        # Instantiate Dataset object using the temporary file
-        with pytest.raises(
-            ValueError,
-            match="There must be exactly one time matching the start_time. Found 2 matching times.",
-        ):
-            Dataset(
-                filename=filepath,
-                var_names={"var": "var"},
-                start_time=start_time,
-                use_dask=use_dask,
-            )
-    finally:
-        os.remove(filepath)
-
-
-def test_warnings_times(global_dataset, use_dask):
+    assert dataset.ds["time"].values == np.datetime64(datetime(2022, 2, 1, 0, 0))
+
+
+def test_warnings_times(global_dataset, tmp_path, use_dask):
     """
     Test handling when no matching times are found.
     """
     # Create a temporary file
-    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
-        filepath = tmpfile.name
-        global_dataset.to_netcdf(filepath)
-    try:
-        # Instantiate Dataset object using the temporary file
-        with pytest.warns(
-            Warning, match="No records found at or before the start_time."
-        ):
-            start_time = datetime(2021, 1, 1)
-            end_time = datetime(2021, 2, 1)
-
-            Dataset(
-                filename=filepath,
-                var_names={"var": "var"},
-                start_time=start_time,
-                end_time=end_time,
-                use_dask=use_dask,
-            )
-
-        with pytest.warns(Warning, match="No records found at or after the end_time."):
-            start_time = datetime(2024, 1, 1)
-            end_time = datetime(2024, 2, 1)
-
-            Dataset(
-                filename=filepath,
-                var_names={"var": "var"},
-                start_time=start_time,
-                end_time=end_time,
-                use_dask=use_dask,
-            )
-    finally:
-        os.remove(filepath)
-
-
-def test_reverse_latitude_choose_subdomain_negative_depth(global_dataset, use_dask):
-    """
-    Test reversing latitude when it is not ascending, the choose_subdomain method, and the convert_to_negative_depth method of the Dataset class.
-    """
-    start_time = datetime(2022, 1, 1)
+    filepath = tmp_path / "test.nc"
+    global_dataset.to_netcdf(filepath)
+    with pytest.warns(Warning, match="No records found at or before the start_time."):
+        start_time = datetime(2021, 1, 1)
+        end_time = datetime(2021, 2, 1)
 
-    # Create a temporary file
-    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
-        filepath = tmpfile.name
-        global_dataset["latitude"] = global_dataset["latitude"][::-1]
-        global_dataset.to_netcdf(filepath)
-    try:
-        # Instantiate Dataset object using the temporary file
-        dataset = Dataset(
+        Dataset(
             filename=filepath,
             var_names={"var": "var"},
-            dim_names={
-                "latitude": "latitude",
-                "longitude": "longitude",
-                "time": "time",
-                "depth": "depth",
-            },
             start_time=start_time,
+            end_time=end_time,
             use_dask=use_dask,
         )
 
-        assert np.all(np.diff(dataset.ds["latitude"]) > 0)
+    with pytest.warns(Warning, match="No records found at or after the end_time."):
+        start_time = datetime(2024, 1, 1)
+        end_time = datetime(2024, 2, 1)
 
-        # test choosing subdomain for domain that straddles the dateline
-        dataset.choose_subdomain(
-            latitude_range=(-10, 10), longitude_range=(-10, 10), margin=1, straddle=True
-        )
-
-        assert -11 <= dataset.ds["latitude"].min() <= 11
-        assert -11 <= dataset.ds["latitude"].max() <= 11
-        assert -11 <= dataset.ds["longitude"].min() <= 11
-        assert -11 <= dataset.ds["longitude"].max() <= 11
-
-        # test choosing subdomain for domain that does not straddle the dateline
-        dataset = Dataset(
+        Dataset(
             filename=filepath,
             var_names={"var": "var"},
-            dim_names={
-                "latitude": "latitude",
-                "longitude": "longitude",
-                "time": "time",
-                "depth": "depth",
-            },
             start_time=start_time,
+            end_time=end_time,
             use_dask=use_dask,
         )
-        dataset.choose_subdomain(
-            latitude_range=(-10, 10), longitude_range=(10, 20), margin=1, straddle=False
-        )
 
-        assert -11 <= dataset.ds["latitude"].min() <= 11
-        assert -11 <= dataset.ds["latitude"].max() <= 11
-        assert 9 <= dataset.ds["longitude"].min() <= 21
-        assert 9 <= dataset.ds["longitude"].max() <= 21
 
-        dataset.convert_to_negative_depth()
+def test_reverse_latitude_reverse_depth_choose_subdomain(
+    global_dataset, tmp_path, use_dask
+):
+    """
+    Test reversing latitude when it is not ascending, the choose_subdomain method, and the convert_to_negative_depth method of the Dataset class.
+    """
+    start_time = datetime(2022, 1, 1)
 
-        assert (dataset.ds["depth"] <= 0).all()
+    filepath = tmp_path / "test.nc"
+    global_dataset["latitude"] = global_dataset["latitude"][::-1]
+    global_dataset["depth"] = global_dataset["depth"][::-1]
+    global_dataset.to_netcdf(filepath)
+
+    dataset = Dataset(
+        filename=filepath,
+        var_names={"var": "var"},
+        dim_names={
+            "latitude": "latitude",
+            "longitude": "longitude",
+            "time": "time",
+            "depth": "depth",
+        },
+        start_time=start_time,
+        use_dask=use_dask,
+    )
 
-    finally:
-        os.remove(filepath)
+    assert np.all(np.diff(dataset.ds["latitude"]) > 0)
+    assert np.all(np.diff(dataset.ds["depth"]) > 0)
 
+    # test choosing subdomain for domain that straddles the dateline
+    ds = dataset.choose_subdomain(
+        latitude_range=(-10, 10),
+        longitude_range=(-10, 10),
+        margin=1,
+        straddle=True,
+        return_subdomain=True,
+    )
 
-def test_check_if_global_with_global_dataset(global_dataset, use_dask):
+    assert -11 <= ds["latitude"].min() <= 11
+    assert -11 <= ds["latitude"].max() <= 11
+    assert -11 <= ds["longitude"].min() <= 11
+    assert -11 <= ds["longitude"].max() <= 11
+
+    ds = dataset.choose_subdomain(
+        latitude_range=(-10, 10),
+        longitude_range=(10, 20),
+        margin=1,
+        straddle=False,
+        return_subdomain=True,
+    )
 
-    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
-        filepath = tmpfile.name
-        global_dataset.to_netcdf(filepath)
-    try:
-        dataset = Dataset(
-            filename=filepath, var_names={"var": "var"}, use_dask=use_dask
-        )
-        is_global = dataset.check_if_global(dataset.ds)
-        assert is_global
-    finally:
-        os.remove(filepath)
+    assert -11 <= ds["latitude"].min() <= 11
+    assert -11 <= ds["latitude"].max() <= 11
+    assert 9 <= ds["longitude"].min() <= 21
+    assert 9 <= ds["longitude"].max() <= 21
 
 
-def test_check_if_global_with_non_global_dataset(non_global_dataset, use_dask):
+def test_check_if_global_with_global_dataset(global_dataset, tmp_path, use_dask):
 
-    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
-        filepath = tmpfile.name
-        non_global_dataset.to_netcdf(filepath)
-    try:
-        dataset = Dataset(
-            filename=filepath, var_names={"var": "var"}, use_dask=use_dask
-        )
-        is_global = dataset.check_if_global(dataset.ds)
+    filepath = tmp_path / "test.nc"
+    global_dataset.to_netcdf(filepath)
+    dataset = Dataset(filename=filepath, var_names={"var": "var"}, use_dask=use_dask)
+    is_global = dataset.check_if_global(dataset.ds)
+    assert is_global
+
+
+def test_check_if_global_with_non_global_dataset(
+    non_global_dataset, tmp_path, use_dask
+):
 
-        assert not is_global
-    finally:
-        os.remove(filepath)
+    filepath = tmp_path / "test.nc"
+    non_global_dataset.to_netcdf(filepath)
+    dataset = Dataset(filename=filepath, var_names={"var": "var"}, use_dask=use_dask)
+    is_global = dataset.check_if_global(dataset.ds)
 
+    assert not is_global
 
-def test_check_dataset(global_dataset, use_dask):
+
+def test_check_dataset(global_dataset, tmp_path, use_dask):
 
     ds = global_dataset.copy()
     ds = ds.drop_vars("var")
 
-    # Create a temporary file
-    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
-        filepath = tmpfile.name
-        ds.to_netcdf(filepath)
-    try:
-        # Instantiate Dataset object using the temporary file
-        start_time = datetime(2022, 2, 1)
-        end_time = datetime(2022, 3, 1)
-        with pytest.raises(
-            ValueError, match="Dataset does not contain all required variables."
-        ):
-
-            Dataset(
-                filename=filepath,
-                var_names={"var": "var"},
-                start_time=start_time,
-                end_time=end_time,
-                use_dask=use_dask,
-            )
-    finally:
-        os.remove(filepath)
+    filepath = tmp_path / "test.nc"
+    ds.to_netcdf(filepath)
+
+    start_time = datetime(2022, 2, 1)
+    end_time = datetime(2022, 3, 1)
+    with pytest.raises(
+        ValueError, match="Dataset does not contain all required variables."
+    ):
+
+        Dataset(
+            filename=filepath,
+            var_names={"var": "var"},
+            start_time=start_time,
+            end_time=end_time,
+            use_dask=use_dask,
+        )
 
     ds = global_dataset.copy()
     ds = ds.rename({"latitude": "lat", "longitude": "long"})
 
-    # Create a temporary file
-    with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
-        filepath = tmpfile.name
-        ds.to_netcdf(filepath)
-    try:
-        # Instantiate Dataset object using the temporary file
-        start_time = datetime(2022, 2, 1)
-        end_time = datetime(2022, 3, 1)
-        with pytest.raises(
-            ValueError, match="Dataset does not contain all required dimensions."
-        ):
-
-            Dataset(
-                filename=filepath,
-                var_names={"var": "var"},
-                start_time=start_time,
-                end_time=end_time,
-                use_dask=use_dask,
-            )
-    finally:
-        os.remove(filepath)
+    filepath = tmp_path / "test2.nc"
+    ds.to_netcdf(filepath)
+
+    start_time = datetime(2022, 2, 1)
+    end_time = datetime(2022, 3, 1)
+    with pytest.raises(
+        ValueError, match="Dataset does not contain all required dimensions."
+    ):
+
+        Dataset(
+            filename=filepath,
+            var_names={"var": "var"},
+            start_time=start_time,
+            end_time=end_time,
+            use_dask=use_dask,
+        )
 
 
 def test_era5_correction_choose_subdomain(use_dask):
@@ -416,3 +398,90 @@ def test_era5_correction_choose_subdomain(use_dask):
     data.choose_subdomain(coords, straddle=False)
     assert (data.ds["latitude"] == lats).all()
     assert (data.ds["longitude"] == lons).all()
+
+
+def test_data_concatenation(use_dask):
+
+    fname = download_test_data("GLORYS_NA_2012.nc")
+    data = GLORYSDataset(
+        filename=fname,
+        start_time=datetime(2012, 1, 1),
+        end_time=datetime(2013, 1, 1),
+        use_dask=use_dask,
+    )
+
+    # Concatenating the datasets at fname0 and fname1 should result in the dataset at fname
+    fname0 = download_test_data("GLORYS_NA_20120101.nc")
+    fname1 = download_test_data("GLORYS_NA_20121231.nc")
+
+    # Test concatenation based on wildcards
+    directory_path = Path(fname0).parent
+    data_concatenated = GLORYSDataset(
+        filename=str(directory_path) + "/GLORYS_NA_2012????.nc",
+        start_time=datetime(2012, 1, 1),
+        end_time=datetime(2013, 1, 1),
+        use_dask=use_dask,
+    )
+    assert data.ds.equals(data_concatenated.ds)
+
+    # Test concatenation based on lists
+    data_concatenated = GLORYSDataset(
+        filename=[fname0, fname1],
+        start_time=datetime(2012, 1, 1),
+        end_time=datetime(2013, 1, 1),
+        use_dask=use_dask,
+    )
+    assert data.ds.equals(data_concatenated.ds)
+
+
+def test_time_validation(use_dask):
+
+    fname = download_test_data("GLORYS_NA_2012.nc")
+
+    with pytest.raises(TypeError, match="start_time must be a datetime object"):
+        GLORYSDataset(
+            filename=fname,
+            start_time="dummy",
+            end_time=datetime(2013, 1, 1),
+            use_dask=use_dask,
+        )
+    with pytest.raises(TypeError, match="end_time must be a datetime object"):
+
+        GLORYSDataset(
+            filename=fname,
+            start_time=datetime(2012, 1, 1),
+            end_time="dummy",
+            use_dask=use_dask,
+        )
+
+
+def test_climatology_error(use_dask):
+
+    fname = download_test_data("GLORYS_NA_2012.nc")
+
+    with pytest.raises(
+        ValueError,
+        match="The dataset contains 2 time steps, but the climatology flag is set to True, which requires exactly 12 time steps.",
+    ):
+        GLORYSDataset(
+            filename=fname,
+            start_time=datetime(2012, 1, 1),
+            end_time=datetime(2013, 1, 1),
+            climatology=True,
+            use_dask=use_dask,
+        )
+
+    fname_bgc = download_test_data("CESM_regional_coarse_test_data_climatology.nc")
+
+    with pytest.raises(
+        ValueError,
+        match="The dataset contains integer time values, which are only supported when the climatology flag is set to True. However, your climatology flag is set to False.",
+    ):
+
+        CESMBGCDataset(
+            filename=fname_bgc,
+            start_time=datetime(2012, 1, 1),
+            end_time=datetime(2013, 1, 1),
+            climatology=False,
+            use_dask=use_dask,
+        )