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

roms_tools/tests/test_setup/test_tides.py

@@ -0,0 +1,298 @@
+import pytest
+from roms_tools import Grid, TidalForcing
+import xarray as xr
+from roms_tools.setup.download import download_test_data
+import textwrap
+from pathlib import Path
+from conftest import calculate_file_hash
+
+
+@pytest.fixture
+def grid_that_lies_within_bounds_of_regional_tpxo_data():
+    grid = Grid(
+        nx=3, ny=3, size_x=1500, size_y=1500, center_lon=235, center_lat=25, rot=-20
+    )
+    return grid
+
+
+@pytest.fixture
+def grid_that_is_out_of_bounds_of_regional_tpxo_data():
+    grid = Grid(
+        nx=3, ny=3, size_x=1800, size_y=1500, center_lon=235, center_lat=25, rot=-20
+    )
+    return grid
+
+
+@pytest.fixture
+def grid_that_straddles_dateline():
+    """
+    Fixture for creating a domain that straddles the dateline.
+    """
+    grid = Grid(
+        nx=5,
+        ny=5,
+        size_x=1800,
+        size_y=2400,
+        center_lon=-10,
+        center_lat=30,
+        rot=20,
+    )
+
+    return grid
+
+
+@pytest.fixture
+def grid_that_straddles_180_degree_meridian():
+    """
+    Fixture for creating a domain that straddles 180 degree meridian.
+    """
+
+    grid = Grid(
+        nx=5,
+        ny=5,
+        size_x=1800,
+        size_y=2400,
+        center_lon=180,
+        center_lat=30,
+        rot=20,
+    )
+
+    return grid
+
+
+@pytest.mark.parametrize(
+    "grid_fixture",
+    [
+        "grid_that_lies_within_bounds_of_regional_tpxo_data",
+        "grid_that_is_out_of_bounds_of_regional_tpxo_data",
+        "grid_that_straddles_dateline",
+        "grid_that_straddles_180_degree_meridian",
+    ],
+)
+def test_successful_initialization_with_global_data(grid_fixture, request, use_dask):
+
+    fname = download_test_data("TPXO_global_test_data.nc")
+
+    grid = request.getfixturevalue(grid_fixture)
+
+    tidal_forcing = TidalForcing(
+        grid=grid, source={"name": "TPXO", "path": fname}, ntides=2, use_dask=use_dask
+    )
+
+    assert isinstance(tidal_forcing.ds, xr.Dataset)
+    assert "omega" in tidal_forcing.ds
+    assert "ssh_Re" in tidal_forcing.ds
+    assert "ssh_Im" in tidal_forcing.ds
+    assert "pot_Re" in tidal_forcing.ds
+    assert "pot_Im" in tidal_forcing.ds
+    assert "u_Re" in tidal_forcing.ds
+    assert "u_Im" in tidal_forcing.ds
+    assert "v_Re" in tidal_forcing.ds
+    assert "v_Im" in tidal_forcing.ds
+
+    assert tidal_forcing.source == {"name": "TPXO", "path": fname}
+    assert tidal_forcing.ntides == 2
+
+
+def test_successful_initialization_with_regional_data(
+    grid_that_lies_within_bounds_of_regional_tpxo_data, use_dask
+):
+
+    fname = download_test_data("TPXO_regional_test_data.nc")
+
+    tidal_forcing = TidalForcing(
+        grid=grid_that_lies_within_bounds_of_regional_tpxo_data,
+        source={"name": "TPXO", "path": fname},
+        ntides=10,
+        use_dask=use_dask,
+    )
+
+    assert isinstance(tidal_forcing.ds, xr.Dataset)
+    assert "omega" in tidal_forcing.ds
+    assert "ssh_Re" in tidal_forcing.ds
+    assert "ssh_Im" in tidal_forcing.ds
+    assert "pot_Re" in tidal_forcing.ds
+    assert "pot_Im" in tidal_forcing.ds
+    assert "u_Re" in tidal_forcing.ds
+    assert "u_Im" in tidal_forcing.ds
+    assert "v_Re" in tidal_forcing.ds
+    assert "v_Im" in tidal_forcing.ds
+
+    assert tidal_forcing.source == {"name": "TPXO", "path": fname}
+    assert tidal_forcing.ntides == 10
+
+
+def test_unsuccessful_initialization_with_regional_data_due_to_nans(
+    grid_that_is_out_of_bounds_of_regional_tpxo_data, use_dask
+):
+
+    fname = download_test_data("TPXO_regional_test_data.nc")
+
+    with pytest.raises(ValueError, match="NaN values found"):
+        TidalForcing(
+            grid=grid_that_is_out_of_bounds_of_regional_tpxo_data,
+            source={"name": "TPXO", "path": fname},
+            ntides=10,
+            use_dask=use_dask,
+        )
+
+
+@pytest.mark.parametrize(
+    "grid_fixture",
+    ["grid_that_straddles_dateline", "grid_that_straddles_180_degree_meridian"],
+)
+def test_unsuccessful_initialization_with_regional_data_due_to_no_overlap(
+    grid_fixture, request, use_dask
+):
+
+    fname = download_test_data("TPXO_regional_test_data.nc")
+
+    grid = request.getfixturevalue(grid_fixture)
+
+    with pytest.raises(
+        ValueError, match="Selected longitude range does not intersect with dataset"
+    ):
+        TidalForcing(
+            grid=grid,
+            source={"name": "TPXO", "path": fname},
+            ntides=10,
+            use_dask=use_dask,
+        )
+
+
+def test_insufficient_number_of_consituents(grid_that_straddles_dateline, use_dask):
+
+    fname = download_test_data("TPXO_global_test_data.nc")
+
+    with pytest.raises(ValueError, match="The dataset contains fewer"):
+        TidalForcing(
+            grid=grid_that_straddles_dateline,
+            source={"name": "TPXO", "path": fname},
+            ntides=10,
+            use_dask=use_dask,
+        )
+
+
+def test_tidal_forcing_plot_save(tidal_forcing, tmp_path):
+    """
+    Test plot and save methods in the same test since we dask arrays are already computed.
+    """
+    tidal_forcing.ds.load()
+
+    tidal_forcing.plot(varname="ssh_Re", ntides=0)
+
+    for file_str in ["test_tides", "test_tides.nc"]:
+        # Create a temporary filepath using the tmp_path fixture
+        for filepath in [
+            tmp_path / file_str,
+            str(tmp_path / file_str),
+        ]:  # test for Path object and str
+
+            # Test saving without partitioning
+            saved_filenames = tidal_forcing.save(filepath)
+            # Check if the .nc file was created
+            filepath = Path(filepath).with_suffix(".nc")
+            assert saved_filenames == [filepath]
+            assert filepath.exists()
+            # Clean up the .nc file
+            filepath.unlink()
+
+            # Test saving with partitioning
+            saved_filenames = tidal_forcing.save(filepath, np_eta=3, np_xi=3)
+
+            filepath_str = str(filepath.with_suffix(""))
+            expected_filepath_list = [
+                Path(filepath_str + f".{index}.nc") for index in range(9)
+            ]
+            assert saved_filenames == expected_filepath_list
+            for expected_filepath in expected_filepath_list:
+                assert expected_filepath.exists()
+                expected_filepath.unlink()
+
+
+def test_roundtrip_yaml(tidal_forcing, tmp_path, use_dask):
+    """Test that creating a TidalForcing object, saving its parameters to yaml file, and re-opening yaml file creates the same object."""
+
+    # Create a temporary filepath using the tmp_path fixture
+    file_str = "test_yaml"
+    for filepath in [
+        tmp_path / file_str,
+        str(tmp_path / file_str),
+    ]:  # test for Path object and str
+
+        tidal_forcing.to_yaml(filepath)
+
+        tidal_forcing_from_file = TidalForcing.from_yaml(filepath, use_dask=use_dask)
+
+        assert tidal_forcing == tidal_forcing_from_file
+
+        filepath = Path(filepath)
+        filepath.unlink()
+
+
+def test_files_have_same_hash(tidal_forcing, tmp_path, use_dask):
+
+    yaml_filepath = tmp_path / "test_yaml.yaml"
+    filepath1 = tmp_path / "test1.nc"
+    filepath2 = tmp_path / "test2.nc"
+
+    print(yaml_filepath)
+    print(filepath1)
+    print(filepath2)
+    tidal_forcing.to_yaml(yaml_filepath)
+    tidal_forcing.save(filepath1)
+    tidal_forcing_from_file = TidalForcing.from_yaml(yaml_filepath, use_dask=use_dask)
+    tidal_forcing_from_file.save(filepath2)
+
+    hash1 = calculate_file_hash(filepath1)
+    hash2 = calculate_file_hash(filepath2)
+
+    assert hash1 == hash2, f"Hashes do not match: {hash1} != {hash2}"
+
+    yaml_filepath.unlink()
+    filepath1.unlink()
+    filepath2.unlink()
+
+
+def test_from_yaml_missing_tidal_forcing(tmp_path, use_dask):
+    yaml_content = textwrap.dedent(
+        """\
+    ---
+    roms_tools_version: 0.0.0
+    ---
+    Grid:
+      nx: 100
+      ny: 100
+      size_x: 1800
+      size_y: 2400
+      center_lon: -10
+      center_lat: 61
+      rot: -20
+      topography_source: ETOPO5
+      smooth_factor: 8
+      hmin: 5.0
+      rmax: 0.2
+    """
+    )
+
+    # Create a temporary filepath using the tmp_path fixture
+    file_str = "test_yaml"
+    for yaml_filepath in [
+        tmp_path / file_str,
+        str(tmp_path / file_str),
+    ]:  # test for Path object and str
+
+        # Write YAML content to file
+        if isinstance(yaml_filepath, Path):
+            yaml_filepath.write_text(yaml_content)
+        else:
+            with open(yaml_filepath, "w") as f:
+                f.write(yaml_content)
+
+        with pytest.raises(
+            ValueError, match="No TidalForcing configuration found in the YAML file."
+        ):
+            TidalForcing.from_yaml(yaml_filepath, use_dask=use_dask)
+
+        yaml_filepath = Path(yaml_filepath)
+        yaml_filepath.unlink()

roms_tools/tests/{test_topography.py → test_setup/test_topography.py}

@@ -56,25 +56,3 @@ def test_hmin_criterion():
 
     assert grid.hmin == 10.0
     assert np.less_equal(grid.hmin, grid.ds.h.min())
-
-
-def test_data_consistency():
-    """
-    Test that the topography generation remains consistent.
-    """
-
-    grid = Grid(
-        nx=3, ny=3, size_x=1500, size_y=1500, center_lon=235, center_lat=25, rot=-20
-    )
-
-    expected_h = np.array(
-        [
-            [4505.16995868, 4505.16995868, 4407.37986032, 4306.51226663, 4306.51226663],
-            [4505.16995868, 4505.16995868, 4407.37986032, 4306.51226663, 4306.51226663],
-            [4400.69482254, 4400.69482254, 3940.84931344, 3060.19573878, 3060.19573878],
-            [4234.97356606, 4234.97356606, 2880.90226836, 2067.46801754, 2067.46801754],
-            [4234.97356606, 4234.97356606, 2880.90226836, 2067.46801754, 2067.46801754],
-        ]
-    )
-
-    assert np.allclose(grid.ds["h"].values, expected_h)

roms_tools/tests/test_setup/test_utils.py

@@ -0,0 +1,16 @@
+from roms_tools.setup.utils import interpolate_from_climatology
+from roms_tools.setup.datasets import ERA5Correction
+from roms_tools.setup.download import download_test_data
+import xarray as xr
+
+
+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"]
+
+    interpolated_field = interpolate_from_climatology(field, "time", era5_times)
+    assert len(interpolated_field.time) == len(era5_times)

roms_tools/tests/test_setup/test_validation.py

@@ -0,0 +1,72 @@
+import pytest
+import os
+import shutil
+import xarray as xr
+
+
+def _get_fname(name):
+    dirname = os.path.dirname(__file__)
+    return os.path.join(dirname, "test_data", f"{name}.zarr")
+
+
+@pytest.mark.parametrize(
+    "forcing_fixture",
+    [
+        "grid",
+        "grid_that_straddles_dateline",
+        "tidal_forcing",
+        "initial_conditions_with_bgc_from_climatology",
+        "surface_forcing",
+        "coarse_surface_forcing",
+        "corrected_surface_forcing",
+        "bgc_surface_forcing",
+        "bgc_surface_forcing_from_climatology",
+        "boundary_forcing",
+        "bgc_boundary_forcing_from_climatology",
+    ],
+)
+# this test will not be run by default
+# to run it and overwrite the test data, invoke pytest as follows
+# pytest --overwrite=tidal_forcing --overwrite=boundary_forcing
+def test_save_results(forcing_fixture, request):
+
+    overwrite = request.config.getoption("--overwrite")
+
+    # Skip the test if the fixture isn't marked for overwriting, unless 'all' is specified
+    if "all" not in overwrite and forcing_fixture not in overwrite:
+        pytest.skip(f"Skipping overwrite for {forcing_fixture}")
+
+    forcing = request.getfixturevalue(forcing_fixture)
+    fname = _get_fname(forcing_fixture)
+
+    # Check if the Zarr directory exists and delete it if it does
+    if os.path.exists(fname):
+        shutil.rmtree(fname)
+
+    forcing.ds.to_zarr(fname)
+
+
+@pytest.mark.parametrize(
+    "forcing_fixture",
+    [
+        "grid",
+        "grid_that_straddles_dateline",
+        "tidal_forcing",
+        "initial_conditions_with_bgc_from_climatology",
+        "surface_forcing",
+        "coarse_surface_forcing",
+        "corrected_surface_forcing",
+        "bgc_surface_forcing",
+        "bgc_surface_forcing_from_climatology",
+        "boundary_forcing",
+        "bgc_boundary_forcing_from_climatology",
+    ],
+)
+def test_check_results(forcing_fixture, request):
+
+    fname = _get_fname(forcing_fixture)
+    expected_forcing_ds = xr.open_zarr(fname, decode_timedelta=False)
+
+    forcing = request.getfixturevalue(forcing_fixture)
+
+    xr.testing.assert_allclose(forcing.ds, expected_forcing_ds, rtol=1.0e-5)

roms_tools/tests/test_setup/test_vertical_coordinate.py

@@ -0,0 +1,95 @@
+import pytest
+from roms_tools import Grid
+
+
+def test_invalid_theta_s_value():
+    """
+    Test the validation of the theta_s value.
+    """
+    with pytest.raises(ValueError):
+
+        Grid(
+            nx=2,
+            ny=2,
+            size_x=500,
+            size_y=1000,
+            center_lon=0,
+            center_lat=55,
+            rot=10,
+            N=3,
+            theta_s=11.0,  # Invalid value, should be 0 < theta_s <= 10
+            theta_b=2.0,
+            hc=250.0,
+        )
+
+
+def test_invalid_theta_b_value():
+    """
+    Test the validation of the theta_b value.
+    """
+    with pytest.raises(ValueError):
+        Grid(
+            nx=2,
+            ny=2,
+            size_x=500,
+            size_y=1000,
+            center_lon=0,
+            center_lat=55,
+            rot=10,
+            N=3,
+            theta_s=5.0,
+            theta_b=5.0,  # Invalid value, should be 0 < theta_b <= 4
+            hc=250.0,
+        )
+
+
+def test_update_vertical_coordinate():
+
+    grid = Grid(
+        nx=2, ny=2, size_x=500, size_y=1000, center_lon=0, center_lat=55, rot=10
+    )
+
+    assert grid.N == 100
+    assert grid.theta_s == 5.0
+    assert grid.theta_b == 2.0
+    assert grid.hc == 300.0
+    assert len(grid.ds.s_rho) == 100
+
+    grid.update_vertical_coordinate(N=3, theta_s=10.0, theta_b=1.0, hc=400.0)
+
+    assert grid.N == 3
+    assert grid.theta_s == 10.0
+    assert grid.theta_b == 1.0
+    assert grid.hc == 400.0
+    assert len(grid.ds.s_rho) == 3
+
+
+def test_plot():
+    grid = Grid(
+        nx=2,
+        ny=2,
+        size_x=500,
+        size_y=1000,
+        center_lon=0,
+        center_lat=55,
+        rot=10,
+        N=3,
+        theta_s=5.0,
+        theta_b=2.0,
+        hc=250.0,
+    )
+    grid.plot_vertical_coordinate("layer_depth_u", s=0)
+    grid.plot_vertical_coordinate("layer_depth_rho", s=-1)
+    grid.plot_vertical_coordinate("interface_depth_v", s=-1)
+    grid.plot_vertical_coordinate("layer_depth_rho", eta=0)
+    grid.plot_vertical_coordinate("layer_depth_u", eta=0)
+    grid.plot_vertical_coordinate("layer_depth_v", eta=0)
+    grid.plot_vertical_coordinate("interface_depth_rho", eta=0)
+    grid.plot_vertical_coordinate("interface_depth_u", eta=0)
+    grid.plot_vertical_coordinate("interface_depth_v", eta=0)
+    grid.plot_vertical_coordinate("layer_depth_rho", xi=0)
+    grid.plot_vertical_coordinate("layer_depth_u", xi=0)
+    grid.plot_vertical_coordinate("layer_depth_v", xi=0)
+    grid.plot_vertical_coordinate("interface_depth_rho", xi=0)
+    grid.plot_vertical_coordinate("interface_depth_u", xi=0)
+    grid.plot_vertical_coordinate("interface_depth_v", xi=0)