roms-tools 2.3.0__py3-none-any.whl → 2.5.0__py3-none-any.whl

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/nriver/.zarray

@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        6
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<i8",
+    "fill_value": null,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        6
+    ],
+    "zarr_format": 2
+}

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/nriver/.zattrs

@@ -0,0 +1,6 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "nriver"
+    ],
+    "long_name": "River ID (1-based Fortran indexing)"
+}

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_location/.zarray

@@ -0,0 +1,22 @@
+{
+    "chunks": [
+        20,
+        20
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<f8",
+    "fill_value": "NaN",
+    "filters": null,
+    "order": "C",
+    "shape": [
+        20,
+        20
+    ],
+    "zarr_format": 2
+}

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_location/.zattrs

@@ -0,0 +1,8 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "eta_rho",
+        "xi_rho"
+    ],
+    "long_name": "River ID plus local volume fraction",
+    "units": "none"
+}

roms_tools/tests/test_setup/test_grid.py

@@ -90,7 +90,6 @@ def test_plot_save_methods(tmp_path):
             str(tmp_path / file_str),
         ]:  # test for Path object and str
 
-            # Test saving without partitioning
             saved_filenames = grid.save(filepath)
             # Check if the .nc file was created
             filepath = Path(filepath).with_suffix(".nc")
@@ -100,18 +99,6 @@ def test_plot_save_methods(tmp_path):
             # Clean up the .nc file
             filepath.unlink()
 
-            # Test saving with partitioning
-            saved_filenames = grid.save(filepath, np_eta=2, np_xi=5)
-
-            filepath_str = str(filepath.with_suffix(""))
-            expected_filepath_list = [
-                Path(filepath_str + f".{index}.nc") for index in range(10)
-            ]
-            assert saved_filenames == expected_filepath_list
-            for expected_filepath in expected_filepath_list:
-                assert expected_filepath.exists()
-                expected_filepath.unlink()
-
 
 def test_raise_if_domain_too_large():
     with pytest.raises(ValueError, match="Domain size exceeds"):

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
+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
@@ -112,6 +116,22 @@ def test_initial_conditions_missing_bgc_path(example_grid, use_dask):
         )
 
 
+# Test initialization with missing ini_time
+def test_initial_conditions_missing_ini_time(example_grid, use_dask):
+
+    fname = Path(download_test_data("GLORYS_coarse_test_data.nc"))
+    with pytest.raises(
+        ValueError,
+        match="`ini_time` must be a valid datetime object and cannot be None.",
+    ):
+        InitialConditions(
+            grid=example_grid,
+            ini_time=None,
+            source={"name": "GLORYS", "path": fname},
+            use_dask=use_dask,
+        )
+
+
 # Test default climatology value
 def test_initial_conditions_default_climatology(example_grid, use_dask):
 
@@ -144,6 +164,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 +318,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 +380,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 +388,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 +409,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 +441,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}"