roms-tools 2.2.1__py3-none-any.whl → 2.4.0__py3-none-any.whl

roms_tools/tests/test_setup/test_initial_conditions.py

@@ -4,17 +4,20 @@ from roms_tools import InitialConditions, Grid
 import xarray as xr
 import numpy as np
 import textwrap
-from roms_tools.setup.download import download_test_data
+import logging
+from roms_tools.download import download_test_data
 from roms_tools.setup.datasets import CESMBGCDataset
 from pathlib import Path
-from conftest import calculate_file_hash
+from conftest import calculate_data_hash
 
 
 @pytest.mark.parametrize(
     "ic_fixture",
     [
         "initial_conditions",
+        "initial_conditions_adjusted_for_zeta",
         "initial_conditions_with_bgc",
+        "initial_conditions_with_bgc_adjusted_for_zeta",
         "initial_conditions_with_bgc_from_climatology",
     ],
 )
@@ -29,6 +32,7 @@ def test_initial_conditions_creation(ic_fixture, request):
         "path": Path(download_test_data("GLORYS_coarse_test_data.nc")),
         "climatology": False,
     }
+    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
@@ -144,6 +148,128 @@ def test_initial_conditions_default_bgc_climatology(example_grid, use_dask):
     assert initial_conditions.bgc_source["climatology"] is False
 
 
+def test_info_depth(caplog, use_dask):
+
+    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
+    )
+
+    fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
+
+    with caplog.at_level(logging.INFO):
+
+        InitialConditions(
+            grid=grid,
+            ini_time=datetime(2021, 6, 29),
+            source={"path": fname, "name": "GLORYS"},
+            adjust_depth_for_sea_surface_height=True,
+            use_dask=use_dask,
+        )
+    # Verify the warning message in the log
+    assert "Sea surface height will be used to adjust depth coordinates." in caplog.text
+
+    # Clear the log before the next test
+    caplog.clear()
+
+    with caplog.at_level(logging.INFO):
+
+        InitialConditions(
+            grid=grid,
+            ini_time=datetime(2021, 6, 29),
+            source={"path": fname, "name": "GLORYS"},
+            adjust_depth_for_sea_surface_height=False,
+            use_dask=use_dask,
+        )
+    # Verify the warning message in the log
+    assert (
+        "Sea surface height will NOT be used to adjust depth coordinates."
+        in caplog.text
+    )
+
+
+@pytest.mark.parametrize(
+    "initial_conditions_fixture",
+    [
+        "initial_conditions_adjusted_for_zeta",
+        "initial_conditions_with_bgc_adjusted_for_zeta",
+    ],
+)
+def test_correct_depth_coords_adjusted_for_zeta(
+    initial_conditions_fixture, request, use_dask
+):
+
+    initial_conditions = request.getfixturevalue(initial_conditions_fixture)
+
+    # compute interface depth at rho-points and write it into .ds_depth_coords
+    zeta = initial_conditions.ds.zeta
+    initial_conditions._get_depth_coordinates(
+        zeta, location="rho", depth_type="interface"
+    )
+    # Test that lowermost interface coincides with topography
+    assert np.allclose(
+        initial_conditions.ds_depth_coords["interface_depth_rho"]
+        .isel(s_w=0)
+        .squeeze()
+        .values,  # Extract raw NumPy array
+        initial_conditions.grid.ds.h.values,
+        atol=1e-6,  # Adjust tolerance as needed
+    )
+
+    # Test that uppermost interface coincides with sea surface height
+    assert np.allclose(
+        initial_conditions.ds_depth_coords["interface_depth_rho"]
+        .isel(s_w=-1)
+        .squeeze()
+        .values,
+        -zeta.values,
+        atol=1e-6,
+    )
+
+
+@pytest.mark.parametrize(
+    "initial_conditions_fixture",
+    [
+        "initial_conditions",
+        "initial_conditions",
+    ],
+)
+def test_correct_depth_coords_zero_zeta(initial_conditions_fixture, request, use_dask):
+
+    initial_conditions = request.getfixturevalue(initial_conditions_fixture)
+
+    # compute interface depth at rho-points and write it into .ds_depth_coords
+    initial_conditions._get_depth_coordinates(0, location="rho", depth_type="interface")
+    # Test that lowermost interface coincides with topography
+    assert np.allclose(
+        initial_conditions.ds_depth_coords["interface_depth_rho"]
+        .isel(s_w=0)
+        .squeeze()
+        .values,  # Extract raw NumPy array
+        initial_conditions.grid.ds.h.values,
+        atol=1e-6,  # Adjust tolerance as needed
+    )
+
+    # Test that uppermost interface coincides with sea surface height
+    assert np.allclose(
+        initial_conditions.ds_depth_coords["interface_depth_rho"]
+        .isel(s_w=-1)
+        .squeeze()
+        .values,
+        0 * initial_conditions.grid.ds.h.values,
+        atol=1e-6,
+    )
+
+
 def test_interpolation_from_climatology(
     initial_conditions_with_bgc_from_climatology, use_dask
 ):
@@ -176,47 +302,61 @@ def test_interpolation_from_climatology(
     )
 
 
-def test_initial_conditions_plot(initial_conditions_with_bgc_from_climatology):
+@pytest.mark.parametrize(
+    "initial_conditions_fixture",
+    [
+        "initial_conditions_with_bgc_adjusted_for_zeta",
+        "initial_conditions_with_bgc_from_climatology",
+    ],
+)
+def test_initial_conditions_plot(initial_conditions_fixture, request):
     """Test plot method."""
 
-    initial_conditions_with_bgc_from_climatology.plot(var_name="temp", s=0)
-    initial_conditions_with_bgc_from_climatology.plot(
-        var_name="temp", s=0, depth_contours=True
-    )
-    initial_conditions_with_bgc_from_climatology.plot(
-        var_name="temp", eta=0, layer_contours=True
-    )
-    initial_conditions_with_bgc_from_climatology.plot(
-        var_name="temp", xi=0, layer_contours=True
-    )
-    initial_conditions_with_bgc_from_climatology.plot(var_name="temp", eta=0)
-    initial_conditions_with_bgc_from_climatology.plot(var_name="temp", xi=0)
-    initial_conditions_with_bgc_from_climatology.plot(var_name="temp", s=0, xi=0)
-    initial_conditions_with_bgc_from_climatology.plot(var_name="temp", eta=0, xi=0)
-    initial_conditions_with_bgc_from_climatology.plot(
-        var_name="u", s=0, layer_contours=True
-    )
-    initial_conditions_with_bgc_from_climatology.plot(var_name="u", s=0)
-    initial_conditions_with_bgc_from_climatology.plot(var_name="u", eta=0)
-    initial_conditions_with_bgc_from_climatology.plot(var_name="u", xi=0)
-    initial_conditions_with_bgc_from_climatology.plot(
-        var_name="v", s=0, layer_contours=True
-    )
-    initial_conditions_with_bgc_from_climatology.plot(var_name="v", s=0)
-    initial_conditions_with_bgc_from_climatology.plot(var_name="v", eta=0)
-    initial_conditions_with_bgc_from_climatology.plot(var_name="v", xi=0)
-    initial_conditions_with_bgc_from_climatology.plot(var_name="zeta")
-    initial_conditions_with_bgc_from_climatology.plot(var_name="ubar")
-    initial_conditions_with_bgc_from_climatology.plot(var_name="vbar")
-    initial_conditions_with_bgc_from_climatology.plot(var_name="ALK", s=0, xi=0)
-    initial_conditions_with_bgc_from_climatology.plot(var_name="ALK", eta=0, xi=0)
-
-
-def test_initial_conditions_save(
-    initial_conditions_with_bgc_from_climatology, tmp_path
-):
+    initial_conditions = request.getfixturevalue(initial_conditions_fixture)
+
+    # horizontal slices plots
+    for depth_contours in [True, False]:
+        initial_conditions.plot(var_name="temp", s=0, depth_contours=depth_contours)
+        initial_conditions.plot(var_name="u", s=0, depth_contours=depth_contours)
+        initial_conditions.plot(var_name="v", s=0, depth_contours=depth_contours)
+        initial_conditions.plot(var_name="ALK", s=0, depth_contours=depth_contours)
+
+    # section plots
+    for layer_contours in [True, False]:
+        initial_conditions.plot(var_name="temp", eta=0, layer_contours=layer_contours)
+        initial_conditions.plot(var_name="temp", xi=0, layer_contours=layer_contours)
+        initial_conditions.plot(var_name="ALK", eta=0, layer_contours=layer_contours)
+        initial_conditions.plot(var_name="ALK", xi=0, layer_contours=layer_contours)
+        initial_conditions.plot(var_name="u", eta=0, layer_contours=layer_contours)
+        initial_conditions.plot(var_name="u", xi=0, layer_contours=layer_contours)
+        initial_conditions.plot(var_name="v", eta=0, layer_contours=layer_contours)
+        initial_conditions.plot(var_name="v", xi=0, layer_contours=layer_contours)
+
+    # 1D plot in horizontal
+    initial_conditions.plot(var_name="temp", s=0, xi=0)
+    initial_conditions.plot(var_name="ALK", s=0, xi=0)
+
+    # 1D plot in vertical
+    initial_conditions.plot(var_name="temp", eta=0, xi=0)
+    initial_conditions.plot(var_name="ALK", eta=0, xi=0)
+
+    initial_conditions.plot(var_name="zeta")
+    initial_conditions.plot(var_name="ubar")
+    initial_conditions.plot(var_name="vbar")
+
+
+@pytest.mark.parametrize(
+    "initial_conditions_fixture",
+    [
+        "initial_conditions",
+        "initial_conditions_adjusted_for_zeta",
+    ],
+)
+def test_initial_conditions_save(initial_conditions_fixture, request, tmp_path):
     """Test save method."""
 
+    initial_conditions = request.getfixturevalue(initial_conditions_fixture)
+
     for file_str in ["test_ic", "test_ic.nc"]:
         # Create a temporary filepath using the tmp_path fixture
         for filepath in [
@@ -224,10 +364,7 @@ def test_initial_conditions_save(
             str(tmp_path / file_str),
         ]:  # test for Path object and str
 
-            # Test saving without partitioning
-            saved_filenames = initial_conditions_with_bgc_from_climatology.save(
-                filepath
-            )
+            saved_filenames = initial_conditions.save(filepath)
             # Check if the .nc file was created
             filepath = Path(filepath).with_suffix(".nc")
             assert saved_filenames == [filepath]
@@ -235,27 +372,20 @@ def test_initial_conditions_save(
             # Clean up the .nc file
             filepath.unlink()
 
-            # Test saving with partitioning
-            saved_filenames = initial_conditions_with_bgc_from_climatology.save(
-                filepath, np_eta=2
-            )
-
-            filepath_str = str(filepath.with_suffix(""))
-            expected_filepath_list = [
-                Path(filepath_str + f".{index}.nc") for index in range(2)
-            ]
-            assert saved_filenames == expected_filepath_list
-            for expected_filepath in expected_filepath_list:
-                assert expected_filepath.exists()
-                expected_filepath.unlink()
 
-
-def test_roundtrip_yaml(
-    initial_conditions_with_bgc_from_climatology, tmp_path, use_dask
-):
+@pytest.mark.parametrize(
+    "initial_conditions_fixture",
+    [
+        "initial_conditions",
+        "initial_conditions_adjusted_for_zeta",
+    ],
+)
+def test_roundtrip_yaml(initial_conditions_fixture, request, tmp_path, use_dask):
     """Test that creating an InitialConditions object, saving its parameters to yaml
     file, and re-opening yaml file creates the same object."""
 
+    initial_conditions = request.getfixturevalue(initial_conditions_fixture)
+
     # Create a temporary filepath using the tmp_path fixture
     file_str = "test_yaml"
     for filepath in [
@@ -263,21 +393,28 @@ def test_roundtrip_yaml(
         str(tmp_path / file_str),
     ]:  # test for Path object and str
 
-        initial_conditions_with_bgc_from_climatology.to_yaml(filepath)
+        initial_conditions.to_yaml(filepath)
 
         initial_conditions_from_file = InitialConditions.from_yaml(
             filepath, use_dask=use_dask
         )
 
-        assert (
-            initial_conditions_with_bgc_from_climatology == initial_conditions_from_file
-        )
+        assert initial_conditions == initial_conditions_from_file
 
         filepath = Path(filepath)
         filepath.unlink()
 
 
-def test_files_have_same_hash(initial_conditions, tmp_path, use_dask):
+@pytest.mark.parametrize(
+    "initial_conditions_fixture",
+    [
+        "initial_conditions",
+        "initial_conditions_adjusted_for_zeta",
+    ],
+)
+def test_files_have_same_hash(initial_conditions_fixture, request, tmp_path, use_dask):
+
+    initial_conditions = request.getfixturevalue(initial_conditions_fixture)
 
     yaml_filepath = tmp_path / "test_yaml.yaml"
     filepath1 = tmp_path / "test1.nc"
@@ -288,8 +425,9 @@ def test_files_have_same_hash(initial_conditions, tmp_path, use_dask):
     ic_from_file = InitialConditions.from_yaml(yaml_filepath, use_dask)
     ic_from_file.save(filepath2)
 
-    hash1 = calculate_file_hash(filepath1)
-    hash2 = calculate_file_hash(filepath2)
+    # Only compare hash of datasets because metadata is non-deterministic with dask
+    hash1 = calculate_data_hash(filepath1)
+    hash2 = calculate_data_hash(filepath2)
 
     assert hash1 == hash2, f"Hashes do not match: {hash1} != {hash2}"
 

roms_tools/tests/test_setup/test_nesting.py

@@ -417,7 +417,6 @@ class TestNesting:
                 [tmp_path / grid_file_str, str(tmp_path / grid_file_str)],
             ):  # test for Path object and str
 
-                # Test saving without partitioning
                 saved_filenames = nesting.save(filepath, grid_filepath)
                 # Check if the .nc file was created
                 filepath = Path(filepath).with_suffix(".nc")
@@ -429,21 +428,6 @@ class TestNesting:
                 filepath.unlink()
                 grid_filepath.unlink()
 
-                # Test saving with partitioning
-                saved_filenames = nesting.save(
-                    filepath, grid_filepath, np_eta=5, np_xi=5
-                )
-
-                filepath_str = str(filepath.with_suffix(""))
-                grid_filepath_str = str(grid_filepath.with_suffix(""))
-                expected_filepath_list = [
-                    Path(filepath_str + f".{index}.nc") for index in range(25)
-                ] + [Path(grid_filepath_str + f".{index}.nc") for index in range(25)]
-                assert saved_filenames == expected_filepath_list
-                for expected_filepath in expected_filepath_list:
-                    assert expected_filepath.exists()
-                    expected_filepath.unlink()
-
     def test_roundtrip_yaml(self, nesting, tmp_path):
         """Test that creating a Nesting object, saving its parameters to yaml file, and
         re-opening yaml file creates the same object."""

roms_tools/tests/test_setup/test_river_forcing.py

@@ -6,7 +6,7 @@ import textwrap
 from pathlib import Path
 import pytest
 from conftest import calculate_file_hash
-from roms_tools.setup.download import download_river_data
+from roms_tools.download import download_river_data
 
 
 @pytest.fixture
@@ -170,7 +170,7 @@ def test_river_locations_are_along_coast(river_forcing_fixture, request):
         eta_rho = indices["eta_rho"][i]
         xi_rho = indices["xi_rho"][i]
         assert coast[eta_rho, xi_rho]
-        assert river_forcing.grid.ds["river_flux"][eta_rho, xi_rho] > 0
+        assert river_forcing.ds["river_location"][eta_rho, xi_rho] > 0
 
 
 def test_missing_source_name():
@@ -224,7 +224,7 @@ def test_reproducibility_same_grid(river_forcing):
     assert river_forcing == the_same_river_forcing
 
 
-def test_update_river_flux_variable_without_conflicts(river_forcing, tmp_path):
+def test_update_river_location_variable_without_conflicts(river_forcing, tmp_path):
 
     fname = download_river_data("dai_trenberth_may2019.nc")
     ds = xr.open_dataset(fname, decode_times=False)
@@ -257,41 +257,17 @@ def test_river_forcing_plot(river_forcing_with_bgc):
 def test_river_forcing_save(river_forcing_with_bgc, tmp_path):
     """Test save method."""
 
-    for file_str, grid_file_str in zip(
-        ["test_rivers", "test_rivers.nc"], ["test_grid", "test_grid.nc"]
-    ):
+    for file_str in ["test_rivers", "test_rivers.nc"]:
         # Create a temporary filepath using the tmp_path fixture
-        for filepath, grid_filepath in zip(
-            [tmp_path / file_str, str(tmp_path / file_str)],
-            [tmp_path / grid_file_str, str(tmp_path / grid_file_str)],
-        ):  # test for Path object and str
+        for filepath in [tmp_path / file_str, str(tmp_path / file_str)]:
 
-            # Test saving without partitioning
-            saved_filenames = river_forcing_with_bgc.save(filepath, grid_filepath)
+            saved_filenames = river_forcing_with_bgc.save(filepath)
             # Check if the .nc file was created
             filepath = Path(filepath).with_suffix(".nc")
-            grid_filepath = Path(grid_filepath).with_suffix(".nc")
-            assert saved_filenames == [filepath, grid_filepath]
+            assert saved_filenames == [filepath]
             assert filepath.exists()
-            assert grid_filepath.exists()
             # Clean up the .nc file
             filepath.unlink()
-            grid_filepath.unlink()
-
-            # Test saving with partitioning
-            saved_filenames = river_forcing_with_bgc.save(
-                filepath, grid_filepath, np_eta=3, np_xi=3
-            )
-
-            filepath_str = str(filepath.with_suffix(""))
-            grid_filepath_str = str(grid_filepath.with_suffix(""))
-            expected_filepath_list = [
-                Path(filepath_str + f".{index}.nc") for index in range(9)
-            ] + [Path(grid_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()
 
 
 @pytest.mark.parametrize(
@@ -340,13 +316,11 @@ def test_files_have_same_hash(river_forcing_fixture, request, tmp_path):
     yaml_filepath = tmp_path / "test_yaml.yaml"
     filepath1 = tmp_path / "test1.nc"
     filepath2 = tmp_path / "test2.nc"
-    grid_filepath1 = tmp_path / "grid_test1.nc"
-    grid_filepath2 = tmp_path / "grid_test2.nc"
 
     river_forcing.to_yaml(yaml_filepath)
-    river_forcing.save(filepath1, grid_filepath1)
+    river_forcing.save(filepath1)
     rf_from_file = RiverForcing.from_yaml(yaml_filepath)
-    rf_from_file.save(filepath2, grid_filepath2)
+    rf_from_file.save(filepath2)
 
     hash1 = calculate_file_hash(filepath1)
     hash2 = calculate_file_hash(filepath2)
@@ -356,8 +330,6 @@ def test_files_have_same_hash(river_forcing_fixture, request, tmp_path):
     yaml_filepath.unlink()
     filepath1.unlink()
     filepath2.unlink()
-    grid_filepath1.unlink()
-    grid_filepath2.unlink()
 
 
 def test_from_yaml_missing_river_forcing(tmp_path):