roms-tools 0.20__py3-none-any.whl → 1.0.0__py3-none-any.whl

roms_tools/tests/test_datasets.py

@@ -2,7 +2,7 @@ import pytest
 from datetime import datetime
 import numpy as np
 import xarray as xr
-from roms_tools.setup.datasets import Dataset
+from roms_tools.setup.datasets import Dataset, ERA5Correction
 import tempfile
 import os
 
@@ -123,7 +123,7 @@ def test_select_times(data_fixture, expected_time_values, request):
         # Instantiate Dataset object using the temporary file
         dataset = Dataset(
             filename=filepath,
-            var_names=["var"],
+            var_names={"var": "var"},
             start_time=start_time,
             end_time=end_time,
         )
@@ -158,7 +158,9 @@ def test_select_times_no_end_time(data_fixture, expected_time_values, request):
         dataset.to_netcdf(filepath)
     try:
         # Instantiate Dataset object using the temporary file
-        dataset = Dataset(filename=filepath, var_names=["var"], start_time=start_time)
+        dataset = Dataset(
+            filename=filepath, var_names={"var": "var"}, start_time=start_time
+        )
 
         assert dataset.ds is not None
         assert len(dataset.ds.time) == len(expected_time_values)
@@ -184,7 +186,7 @@ def test_multiple_matching_times(global_dataset_with_multiple_times_per_day):
             ValueError,
             match="There must be exactly one time matching the start_time. Found 2 matching times.",
         ):
-            Dataset(filename=filepath, var_names=["var"], start_time=start_time)
+            Dataset(filename=filepath, var_names={"var": "var"}, start_time=start_time)
     finally:
         os.remove(filepath)
 
@@ -205,7 +207,7 @@ def test_no_matching_times(global_dataset):
         with pytest.raises(ValueError, match="No matching times found."):
             Dataset(
                 filename=filepath,
-                var_names=["var"],
+                var_names={"var": "var"},
                 start_time=start_time,
                 end_time=end_time,
             )
@@ -228,7 +230,7 @@ def test_reverse_latitude_choose_subdomain_negative_depth(global_dataset):
         # Instantiate Dataset object using the temporary file
         dataset = Dataset(
             filename=filepath,
-            var_names=["var"],
+            var_names={"var": "var"},
             dim_names={
                 "latitude": "latitude",
                 "longitude": "longitude",
@@ -253,7 +255,7 @@ def test_reverse_latitude_choose_subdomain_negative_depth(global_dataset):
         # test choosing subdomain for domain that does not straddle the dateline
         dataset = Dataset(
             filename=filepath,
-            var_names=["var"],
+            var_names={"var": "var"},
             dim_names={
                 "latitude": "latitude",
                 "longitude": "longitude",
@@ -285,7 +287,7 @@ def test_check_if_global_with_global_dataset(global_dataset):
         filepath = tmpfile.name
         global_dataset.to_netcdf(filepath)
     try:
-        dataset = Dataset(filename=filepath, var_names=["var"])
+        dataset = Dataset(filename=filepath, var_names={"var": "var"})
         is_global = dataset.check_if_global(dataset.ds)
         assert is_global
     finally:
@@ -298,9 +300,71 @@ def test_check_if_global_with_non_global_dataset(non_global_dataset):
         filepath = tmpfile.name
         non_global_dataset.to_netcdf(filepath)
     try:
-        dataset = Dataset(filename=filepath, var_names=["var"])
+        dataset = Dataset(filename=filepath, var_names={"var": "var"})
         is_global = dataset.check_if_global(dataset.ds)
 
         assert not is_global
     finally:
         os.remove(filepath)
+
+
+def test_check_dataset(global_dataset):
+
+    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,
+            )
+    finally:
+        os.remove(filepath)
+
+    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,
+            )
+    finally:
+        os.remove(filepath)
+
+
+def test_era5_correction_choose_subdomain():
+
+    data = ERA5Correction()
+    lats = data.ds.latitude[10:20]
+    lons = data.ds.longitude[10:20]
+    coords = {"latitude": lats, "longitude": lons}
+    data.choose_subdomain(coords, straddle=False)
+    assert (data.ds["latitude"] == lats).all()
+    assert (data.ds["longitude"] == lons).all()

roms_tools/tests/test_initial_conditions.py

@@ -6,7 +6,8 @@ import numpy as np
 import tempfile
 import os
 import textwrap
-from roms_tools.setup.datasets import download_test_data
+from roms_tools.setup.download import download_test_data
+from roms_tools.setup.datasets import CESMBGCDataset
 
 
 @pytest.fixture
@@ -49,37 +50,205 @@ def initial_conditions(example_grid, example_vertical_coordinate):
         grid=example_grid,
         vertical_coordinate=example_vertical_coordinate,
         ini_time=datetime(2021, 6, 29),
-        filename=fname,
+        physics_source={"path": fname, "name": "GLORYS"},
     )
 
 
-def test_initial_conditions_creation(initial_conditions):
+@pytest.fixture
+def initial_conditions_with_bgc(example_grid, example_vertical_coordinate):
     """
-    Test the creation of the InitialConditions object.
+    Fixture for creating a dummy InitialConditions object.
     """
-    assert initial_conditions.ini_time == datetime(2021, 6, 29)
-    assert initial_conditions.filename == download_test_data("GLORYS_test_data.nc")
-    assert initial_conditions.source == "GLORYS"
+
+    fname = download_test_data("GLORYS_test_data.nc")
+    fname_bgc = download_test_data("CESM_regional_test_data_one_time_slice.nc")
+
+    return InitialConditions(
+        grid=example_grid,
+        vertical_coordinate=example_vertical_coordinate,
+        ini_time=datetime(2021, 6, 29),
+        physics_source={"path": fname, "name": "GLORYS"},
+        bgc_source={"path": fname_bgc, "name": "CESM_REGRIDDED"},
+    )
 
 
-def test_initial_conditions_ds_attribute(initial_conditions):
+@pytest.fixture
+def initial_conditions_with_bgc_from_climatology(
+    example_grid, example_vertical_coordinate
+):
     """
-    Test the ds attribute of the InitialConditions object.
+    Fixture for creating a dummy InitialConditions object.
     """
-    assert isinstance(initial_conditions.ds, xr.Dataset)
-    assert "temp" in initial_conditions.ds
-    assert "salt" in initial_conditions.ds
-    assert "u" in initial_conditions.ds
-    assert "v" in initial_conditions.ds
-    assert "zeta" in initial_conditions.ds
+
+    fname = download_test_data("GLORYS_test_data.nc")
+    fname_bgc = download_test_data("CESM_regional_test_data_climatology.nc")
+
+    return InitialConditions(
+        grid=example_grid,
+        vertical_coordinate=example_vertical_coordinate,
+        ini_time=datetime(2021, 6, 29),
+        physics_source={"path": fname, "name": "GLORYS"},
+        bgc_source={
+            "path": fname_bgc,
+            "name": "CESM_REGRIDDED",
+            "climatology": True,
+        },
+    )
 
 
-def test_initial_conditions_data_consistency_plot_save(initial_conditions, tmp_path):
+@pytest.mark.parametrize(
+    "ic_fixture",
+    [
+        "initial_conditions",
+        "initial_conditions_with_bgc",
+        "initial_conditions_with_bgc_from_climatology",
+    ],
+)
+def test_initial_conditions_creation(ic_fixture, request):
+    """
+    Test the creation of the InitialConditions object.
+    """
+
+    ic = request.getfixturevalue(ic_fixture)
+
+    assert ic.ini_time == datetime(2021, 6, 29)
+    assert ic.physics_source == {
+        "name": "GLORYS",
+        "path": download_test_data("GLORYS_test_data.nc"),
+        "climatology": False,
+    }
+    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
+
+
+# Test initialization with missing 'name' in physics_source
+def test_initial_conditions_missing_physics_name(
+    example_grid, example_vertical_coordinate
+):
+    with pytest.raises(ValueError, match="`physics_source` must include a 'name'."):
+        InitialConditions(
+            grid=example_grid,
+            vertical_coordinate=example_vertical_coordinate,
+            ini_time=datetime(2021, 6, 29),
+            physics_source={"path": "physics_data.nc"},
+        )
+
+
+# Test initialization with missing 'path' in physics_source
+def test_initial_conditions_missing_physics_path(
+    example_grid, example_vertical_coordinate
+):
+    with pytest.raises(ValueError, match="`physics_source` must include a 'path'."):
+        InitialConditions(
+            grid=example_grid,
+            vertical_coordinate=example_vertical_coordinate,
+            ini_time=datetime(2021, 6, 29),
+            physics_source={"name": "GLORYS"},
+        )
+
+
+# Test initialization with missing 'name' in bgc_source
+def test_initial_conditions_missing_bgc_name(example_grid, example_vertical_coordinate):
+
+    fname = download_test_data("GLORYS_test_data.nc")
+    with pytest.raises(
+        ValueError, match="`bgc_source` must include a 'name' if it is provided."
+    ):
+        InitialConditions(
+            grid=example_grid,
+            vertical_coordinate=example_vertical_coordinate,
+            ini_time=datetime(2021, 6, 29),
+            physics_source={"name": "GLORYS", "path": fname},
+            bgc_source={"path": "bgc_data.nc"},
+        )
+
+
+# Test initialization with missing 'path' in bgc_source
+def test_initial_conditions_missing_bgc_path(example_grid, example_vertical_coordinate):
+
+    fname = download_test_data("GLORYS_test_data.nc")
+    with pytest.raises(
+        ValueError, match="`bgc_source` must include a 'path' if it is provided."
+    ):
+        InitialConditions(
+            grid=example_grid,
+            vertical_coordinate=example_vertical_coordinate,
+            ini_time=datetime(2021, 6, 29),
+            physics_source={"name": "GLORYS", "path": fname},
+            bgc_source={"name": "CESM_REGRIDDED"},
+        )
+
+
+# Test default climatology value
+def test_initial_conditions_default_climatology(
+    example_grid, example_vertical_coordinate
+):
+
+    fname = download_test_data("GLORYS_test_data.nc")
+
+    initial_conditions = InitialConditions(
+        grid=example_grid,
+        vertical_coordinate=example_vertical_coordinate,
+        ini_time=datetime(2021, 6, 29),
+        physics_source={"name": "GLORYS", "path": fname},
+    )
+
+    assert initial_conditions.physics_source["climatology"] is False
+    assert initial_conditions.bgc_source is None
+
+
+def test_initial_conditions_default_bgc_climatology(
+    example_grid, example_vertical_coordinate
+):
+
+    fname = download_test_data("GLORYS_test_data.nc")
+    fname_bgc = download_test_data("CESM_regional_test_data_one_time_slice.nc")
+
+    initial_conditions = InitialConditions(
+        grid=example_grid,
+        vertical_coordinate=example_vertical_coordinate,
+        ini_time=datetime(2021, 6, 29),
+        physics_source={"name": "GLORYS", "path": fname},
+        bgc_source={"name": "CESM_REGRIDDED", "path": fname_bgc},
+    )
+
+    assert initial_conditions.bgc_source["climatology"] is True
+
+
+def test_interpolation_from_climatology(initial_conditions_with_bgc_from_climatology):
+
+    fname_bgc = download_test_data("CESM_regional_test_data_climatology.nc")
+    ds = xr.open_dataset(fname_bgc)
+
+    # check if interpolated value for Jan 15 is indeed January value from climatology
+    bgc_data = CESMBGCDataset(
+        filename=fname_bgc, start_time=datetime(2012, 1, 15), climatology=True
+    )
+    assert np.allclose(ds["ALK"].sel(month=1), bgc_data.ds["ALK"], equal_nan=True)
+
+    # check if interpolated value for Jan 30 is indeed average of January and February value from climatology
+    bgc_data = CESMBGCDataset(
+        filename=fname_bgc, start_time=datetime(2012, 1, 30), climatology=True
+    )
+    assert np.allclose(
+        0.5 * (ds["ALK"].sel(month=1) + ds["ALK"].sel(month=2)),
+        bgc_data.ds["ALK"],
+        equal_nan=True,
+    )
+
+
+def test_initial_conditions_data_consistency_plot_save(
+    initial_conditions_with_bgc_from_climatology, tmp_path
+):
     """
     Test that the data within the InitialConditions object remains consistent.
     Also test plot and save methods in the same test since we dask arrays are already computed.
     """
-    initial_conditions.ds.load()
+    initial_conditions_with_bgc_from_climatology.ds.load()
 
     # Define the expected data
     expected_temp = np.array(
@@ -226,24 +395,75 @@ def test_initial_conditions_data_consistency_plot_save(initial_conditions, tmp_p
         dtype=np.float32,
     )
 
+    expected_alk = np.array(
+        [
+            [
+                [
+                    [2341.926, 2340.8894, 2340.557],
+                    [2317.8875, 2315.86, 2315.2148],
+                    [2297.689, 2285.8933, 2284.404],
+                    [2276.4216, 2258.4436, 2256.1062],
+                ],
+                [
+                    [2330.5837, 2329.8225, 2329.5264],
+                    [2317.878, 2316.6787, 2316.3088],
+                    [2278.9314, 2269.464, 2268.904],
+                    [2259.975, 2247.7456, 2246.0632],
+                ],
+                [
+                    [2376.7534, 2373.4402, 2372.9192],
+                    [2362.5308, 2360.5066, 2360.2224],
+                    [2350.3384, 2344.3135, 2343.6768],
+                    [2310.4275, 2287.6785, 2281.5872],
+                ],
+                [
+                    [2384.8064, 2386.2126, 2386.632],
+                    [2383.737, 2385.1553, 2385.6685],
+                    [2380.2297, 2381.4849, 2381.8616],
+                    [2350.0762, 2342.5403, 2339.2244],
+                ],
+            ]
+        ],
+        dtype=np.float32,
+    )
+
     # Check the values in the dataset
-    assert np.allclose(initial_conditions.ds["temp"].values, expected_temp)
-    assert np.allclose(initial_conditions.ds["salt"].values, expected_salt)
-    assert np.allclose(initial_conditions.ds["zeta"].values, expected_zeta)
-    assert np.allclose(initial_conditions.ds["u"].values, expected_u)
-    assert np.allclose(initial_conditions.ds["v"].values, expected_v)
-    assert np.allclose(initial_conditions.ds["ubar"].values, expected_ubar)
-    assert np.allclose(initial_conditions.ds["vbar"].values, expected_vbar)
-
-    initial_conditions.plot(varname="temp", s=0)
-    initial_conditions.plot(varname="temp", eta=0)
-    initial_conditions.plot(varname="temp", xi=0)
-    initial_conditions.plot(varname="temp", s=0, xi=0)
-    initial_conditions.plot(varname="temp", eta=0, xi=0)
-    initial_conditions.plot(varname="zeta")
+    assert np.allclose(
+        initial_conditions_with_bgc_from_climatology.ds["temp"].values, expected_temp
+    )
+    assert np.allclose(
+        initial_conditions_with_bgc_from_climatology.ds["salt"].values, expected_salt
+    )
+    assert np.allclose(
+        initial_conditions_with_bgc_from_climatology.ds["zeta"].values, expected_zeta
+    )
+    assert np.allclose(
+        initial_conditions_with_bgc_from_climatology.ds["u"].values, expected_u
+    )
+    assert np.allclose(
+        initial_conditions_with_bgc_from_climatology.ds["v"].values, expected_v
+    )
+    assert np.allclose(
+        initial_conditions_with_bgc_from_climatology.ds["ubar"].values, expected_ubar
+    )
+    assert np.allclose(
+        initial_conditions_with_bgc_from_climatology.ds["vbar"].values, expected_vbar
+    )
+    assert np.allclose(
+        initial_conditions_with_bgc_from_climatology.ds["ALK"].values, expected_alk
+    )
+
+    initial_conditions_with_bgc_from_climatology.plot(varname="temp", s=0)
+    initial_conditions_with_bgc_from_climatology.plot(varname="temp", eta=0)
+    initial_conditions_with_bgc_from_climatology.plot(varname="temp", xi=0)
+    initial_conditions_with_bgc_from_climatology.plot(varname="temp", s=0, xi=0)
+    initial_conditions_with_bgc_from_climatology.plot(varname="temp", eta=0, xi=0)
+    initial_conditions_with_bgc_from_climatology.plot(varname="zeta")
+    initial_conditions_with_bgc_from_climatology.plot(varname="ALK", s=0, xi=0)
+    initial_conditions_with_bgc_from_climatology.plot(varname="ALK", eta=0, xi=0)
 
     filepath = tmp_path / "initial_conditions.nc"
-    initial_conditions.save(filepath)
+    initial_conditions_with_bgc_from_climatology.save(filepath)
     assert filepath.exists()