roms-tools 2.6.0__py3-none-any.whl → 2.6.2__py3-none-any.whl

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_index/.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_index/.zattrs

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

roms_tools/tests/test_setup/test_datasets.py

@@ -532,6 +532,8 @@ def test_climatology_error(use_dask):
         ("tpxo_data", 1 / 6),
         ("cesm_bgc_data", 1.0),
         ("cesm_surface_bgc_data", 1.0),
+        ("unified_bgc_data", 2.0),
+        ("unified_surface_bgc_data", 2.0),
     ],
 )
 def test_horizontal_resolution(data_fixture, expected_resolution, request):

roms_tools/tests/test_setup/test_initial_conditions.py

@@ -6,7 +6,7 @@ import numpy as np
 import textwrap
 import logging
 from roms_tools.download import download_test_data
-from roms_tools.setup.datasets import CESMBGCDataset
+from roms_tools.setup.datasets import CESMBGCDataset, UnifiedBGCDataset
 from pathlib import Path
 from conftest import calculate_data_hash
 
@@ -19,6 +19,7 @@ from conftest import calculate_data_hash
         "initial_conditions_with_bgc",
         "initial_conditions_with_bgc_adjusted_for_zeta",
         "initial_conditions_with_bgc_from_climatology",
+        "initial_conditions_with_unified_bgc_from_climatology",
     ],
 )
 def test_initial_conditions_creation(ic_fixture, request):
@@ -42,6 +43,59 @@ def test_initial_conditions_creation(ic_fixture, request):
     assert ic.ds.coords["ocean_time"].attrs["units"] == "seconds"
 
 
+@pytest.mark.parametrize(
+    "ic_fixture",
+    [
+        "initial_conditions_with_bgc",
+        "initial_conditions_with_bgc_adjusted_for_zeta",
+        "initial_conditions_with_bgc_from_climatology",
+        "initial_conditions_with_unified_bgc_from_climatology",
+    ],
+)
+def test_initial_condition_creation_with_bgc(ic_fixture, request):
+    """Test the creation of the BoundaryForcing object."""
+
+    ic = request.getfixturevalue(ic_fixture)
+
+    expected_bgc_variables = [
+        "PO4",
+        "NO3",
+        "SiO3",
+        "NH4",
+        "Fe",
+        "Lig",
+        "O2",
+        "DIC",
+        "DIC_ALT_CO2",
+        "ALK",
+        "ALK_ALT_CO2",
+        "DOC",
+        "DON",
+        "DOP",
+        "DOCr",
+        "DONr",
+        "DOPr",
+        "zooC",
+        "spChl",
+        "spC",
+        "spP",
+        "spFe",
+        "spCaCO3",
+        "diatChl",
+        "diatC",
+        "diatP",
+        "diatFe",
+        "diatSi",
+        "diazChl",
+        "diazC",
+        "diazP",
+        "diazFe",
+    ]
+
+    for var in expected_bgc_variables:
+        assert var in ic.ds
+
+
 @pytest.fixture
 def example_grid():
 
@@ -256,7 +310,8 @@ def test_correct_depth_coords_adjusted_for_zeta(
     "initial_conditions_fixture",
     [
         "initial_conditions",
-        "initial_conditions",
+        "initial_conditions_with_bgc",
+        "initial_conditions_with_unified_bgc_from_climatology",
     ],
 )
 def test_correct_depth_coords_zero_zeta(initial_conditions_fixture, request, use_dask):
@@ -286,10 +341,9 @@ def test_correct_depth_coords_zero_zeta(initial_conditions_fixture, request, use
     )
 
 
-def test_interpolation_from_climatology(
-    initial_conditions_with_bgc_from_climatology, use_dask
-):
+def test_interpolation_from_climatology(use_dask):
 
+    # CESM climatology
     fname_bgc = download_test_data("CESM_regional_coarse_test_data_climatology.nc")
     ds = xr.open_dataset(fname_bgc)
 
@@ -317,12 +371,42 @@ def test_interpolation_from_climatology(
         equal_nan=True,
     )
 
+    # Unified BGC climatology
+    fname_bgc = download_test_data("coarsened_UNIFIED_bgc_dataset.nc")
+    ds = xr.open_dataset(fname_bgc)
+
+    # check if interpolated value for Feb 14 is indeed February value from climatology
+    bgc_data = UnifiedBGCDataset(
+        filename=fname_bgc,
+        start_time=datetime(2012, 2, 14),
+        climatology=True,
+        use_dask=use_dask,
+        apply_post_processing=False,
+    )
+    assert np.allclose(ds["Alk"].isel(month=1), bgc_data.ds["Alk"], equal_nan=True)
+
+    # check if interpolated value for Jan 30.25 is indeed average of January and February value from climatology
+    bgc_data = UnifiedBGCDataset(
+        filename=fname_bgc,
+        start_time=datetime(2012, 1, 30, 6),  # time: 6 am, Jan 30
+        climatology=True,
+        use_dask=use_dask,
+        apply_post_processing=False,
+    )
+
+    assert np.allclose(
+        0.5 * (ds["Alk"].isel(month=0) + ds["Alk"].isel(month=1)),
+        bgc_data.ds["Alk"],
+        equal_nan=True,
+    )
+
 
 @pytest.mark.parametrize(
     "initial_conditions_fixture",
     [
         "initial_conditions_with_bgc_adjusted_for_zeta",
         "initial_conditions_with_bgc_from_climatology",
+        "initial_conditions_with_unified_bgc_from_climatology",
     ],
 )
 def test_initial_conditions_plot(initial_conditions_fixture, request):
@@ -332,29 +416,26 @@ def test_initial_conditions_plot(initial_conditions_fixture, request):
 
     # 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)
+        for var_name in ["temp", "u", "v", "ALK", "DOC"]:
+            initial_conditions.plot(
+                var_name=var_name, 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)
+        for var_name in ["temp", "u", "v", "ALK", "DOC"]:
+            initial_conditions.plot(
+                var_name=var_name, eta=0, layer_contours=layer_contours
+            )
+            initial_conditions.plot(
+                var_name=var_name, xi=0, layer_contours=layer_contours
+            )
+
+    for var_name in ["temp", "ALK", "DOC"]:
+        # 1D plot in horizontal
+        initial_conditions.plot(var_name=var_name, s=0, xi=0)
+        # 1D plot in vertical
+        initial_conditions.plot(var_name=var_name, eta=0, xi=0)
 
     initial_conditions.plot(var_name="zeta")
     initial_conditions.plot(var_name="ubar")
@@ -366,6 +447,8 @@ def test_initial_conditions_plot(initial_conditions_fixture, request):
     [
         "initial_conditions",
         "initial_conditions_adjusted_for_zeta",
+        "initial_conditions_with_bgc_from_climatology",
+        "initial_conditions_with_unified_bgc_from_climatology",
     ],
 )
 def test_initial_conditions_save(initial_conditions_fixture, request, tmp_path):
@@ -394,6 +477,8 @@ def test_initial_conditions_save(initial_conditions_fixture, request, tmp_path):
     [
         "initial_conditions",
         "initial_conditions_adjusted_for_zeta",
+        "initial_conditions_with_bgc_from_climatology",
+        "initial_conditions_with_unified_bgc_from_climatology",
     ],
 )
 def test_roundtrip_yaml(initial_conditions_fixture, request, tmp_path, use_dask):
@@ -426,6 +511,8 @@ def test_roundtrip_yaml(initial_conditions_fixture, request, tmp_path, use_dask)
     [
         "initial_conditions",
         "initial_conditions_adjusted_for_zeta",
+        "initial_conditions_with_bgc_from_climatology",
+        "initial_conditions_with_unified_bgc_from_climatology",
     ],
 )
 def test_files_have_same_hash(initial_conditions_fixture, request, tmp_path, use_dask):

roms_tools/tests/test_setup/test_river_forcing.py

@@ -236,7 +236,8 @@ class TestRiverForcingGeneral:
         for name in indices.keys():
             for (eta_rho, xi_rho) in indices[name]:
                 assert coast[eta_rho, xi_rho]
-                assert river_forcing.ds["river_location"][eta_rho, xi_rho] > 0
+                assert river_forcing.ds["river_index"][eta_rho, xi_rho] > 0
+                assert river_forcing.ds["river_fraction"][eta_rho, xi_rho] > 0
 
     def test_missing_source_name(self, iceland_test_grid):
         with pytest.raises(ValueError, match="`source` must include a 'name'."):
@@ -475,14 +476,14 @@ class TestRiverForcingWithPrescribedIndices:
 
         # check that all values are integers for single cell rivers
         non_zero_values = river_forcing_with_prescribed_single_cell_indices.ds[
-            "river_location"
+            "river_fraction"
         ]
         is_integer = non_zero_values == np.floor(non_zero_values)
         assert (is_integer).all()
 
         # check that not all values are integers for multi cell rivers
         non_zero_values = river_forcing_with_prescribed_multi_cell_indices.ds[
-            "river_location"
+            "river_fraction"
         ]
         is_integer = non_zero_values == np.floor(non_zero_values)
         assert not (is_integer).all()

roms_tools/tests/test_setup/test_surface_forcing.py

@@ -398,7 +398,17 @@ def test_start_time_end_time_warning(grid_that_straddles_dateline, use_dask, cap
     )
 
 
-def test_nans_filled_in(grid_that_straddles_dateline, use_dask):
+@pytest.mark.parametrize(
+    "name, fname, type, climatology",
+    [
+        ("ERA5", "ERA5_regional_test_data.nc", "physics", False),
+        ("CESM_REGRIDDED", "CESM_surface_global_test_data_climatology.nc", "bgc", True),
+        ("UNIFIED", "coarsened_UNIFIED_bgc_dataset.nc", "bgc", True),
+    ],
+)
+def test_nans_filled_in(
+    grid_that_straddles_dateline, name, fname, type, climatology, use_dask
+):
     """Test that the surface forcing fields contain no NaNs.
 
     The test is performed twice:
@@ -409,8 +419,10 @@ def test_nans_filled_in(grid_that_straddles_dateline, use_dask):
     start_time = datetime(2020, 1, 31)
     end_time = datetime(2020, 2, 2)
 
-    fname = Path(download_test_data("ERA5_regional_test_data.nc"))
-    fname_bgc = Path(download_test_data("CESM_surface_global_test_data_climatology.nc"))
+    fname = Path(download_test_data(fname))
+    print(fname)
+    print(type)
+    print(name)
 
     for coarse_grid_mode in ["always", "never"]:
         sfc_forcing = SurfaceForcing(
@@ -418,42 +430,36 @@ def test_nans_filled_in(grid_that_straddles_dateline, use_dask):
             coarse_grid_mode=coarse_grid_mode,
             start_time=start_time,
             end_time=end_time,
-            source={"name": "ERA5", "path": fname},
+            source={"name": name, "path": fname, "climatology": climatology},
+            type=type,
             use_dask=use_dask,
         )
 
         # Check that no NaNs are in surface forcing fields (they could make ROMS blow up)
         # Note that ROMS-Tools should replace NaNs with a fill value after the nan_check has successfully
         # completed; the nan_check passes if there are NaNs only over land
-        assert not sfc_forcing.ds["uwnd"].isnull().any().values.item()
-        assert not sfc_forcing.ds["vwnd"].isnull().any().values.item()
-        assert not sfc_forcing.ds["rain"].isnull().any().values.item()
+        for var in sfc_forcing.ds.data_vars:
+            assert not sfc_forcing.ds[var].isnull().any().values.item()
 
-        sfc_forcing = SurfaceForcing(
-            grid=grid_that_straddles_dateline,
-            coarse_grid_mode=coarse_grid_mode,
-            start_time=start_time,
-            end_time=end_time,
-            source={"name": "CESM_REGRIDDED", "path": fname_bgc, "climatology": True},
-            type="bgc",
-            use_dask=use_dask,
-        )
-
-        # Check that no NaNs are in surface forcing fields (they could make ROMS blow up)
-        # Note that ROMS-Tools should replace NaNs with a fill value after the nan_check has successfully
-        # completed; the nan_check passes if there are NaNs only over land
-        assert not sfc_forcing.ds["pco2_air"].isnull().any().values.item()
 
-
-def test_time_attr_climatology(bgc_surface_forcing_from_climatology):
+@pytest.mark.parametrize(
+    "bgc_surface_forcing_fixture",
+    [
+        "bgc_surface_forcing_from_climatology",
+        "bgc_surface_forcing_from_unified_climatology",
+    ],
+)
+def test_time_attr_climatology(bgc_surface_forcing_fixture, request):
     """Test that the 'cycle_length' attribute is present in the time coordinate of the
     BGC dataset when using climatology data."""
+
+    bgc_surface_forcing = request.getfixturevalue(bgc_surface_forcing_fixture)
     for time_coord in ["pco2_time", "iron_time", "dust_time", "nox_time", "nhy_time"]:
         assert hasattr(
-            bgc_surface_forcing_from_climatology.ds[time_coord],
+            bgc_surface_forcing.ds[time_coord],
             "cycle_length",
         )
-    assert hasattr(bgc_surface_forcing_from_climatology.ds, "climatology")
+    assert hasattr(bgc_surface_forcing.ds, "climatology")
 
 
 def test_time_attr(bgc_surface_forcing):
@@ -468,22 +474,30 @@ def test_time_attr(bgc_surface_forcing):
 
 
 @pytest.mark.parametrize(
-    "sfc_forcing_fixture, expected_climatology, expected_fname",
+    "sfc_forcing_fixture, expected_name, expected_climatology, expected_fname",
     [
         (
             "bgc_surface_forcing",
+            "CESM_REGRIDDED",
             False,
             Path(download_test_data("CESM_surface_global_test_data.nc")),
         ),
         (
             "bgc_surface_forcing_from_climatology",
+            "CESM_REGRIDDED",
             True,
             Path(download_test_data("CESM_surface_global_test_data_climatology.nc")),
         ),
+        (
+            "bgc_surface_forcing_from_unified_climatology",
+            "UNIFIED",
+            True,
+            Path(download_test_data("coarsened_UNIFIED_bgc_dataset.nc")),
+        ),
     ],
 )
 def test_surface_forcing_creation(
-    sfc_forcing_fixture, expected_climatology, expected_fname, request
+    sfc_forcing_fixture, expected_name, expected_climatology, expected_fname, request
 ):
     """Test the creation and initialization of the SurfaceForcing object with BGC.
 
@@ -495,23 +509,19 @@ def test_surface_forcing_creation(
     sfc_forcing = request.getfixturevalue(sfc_forcing_fixture)
 
     assert sfc_forcing.ds is not None
-    assert "pco2_air" in sfc_forcing.ds
-    assert "pco2_air_alt" in sfc_forcing.ds
-    assert "iron" in sfc_forcing.ds
-    assert "dust" in sfc_forcing.ds
-    assert "nox" in sfc_forcing.ds
-    assert "nhy" in sfc_forcing.ds
+    for var_name in ["pco2_air", "pco2_air_alt", "iron", "dust", "nox", "nhy"]:
+        assert var_name in sfc_forcing.ds
 
     assert sfc_forcing.start_time == datetime(2020, 2, 1)
     assert sfc_forcing.end_time == datetime(2020, 2, 1)
     assert sfc_forcing.type == "bgc"
     assert sfc_forcing.source == {
-        "name": "CESM_REGRIDDED",
+        "name": expected_name,
         "path": expected_fname,
         "climatology": expected_climatology,
     }
     assert not sfc_forcing.use_coarse_grid
-    assert sfc_forcing.ds.attrs["source"] == "CESM_REGRIDDED"
+    assert sfc_forcing.ds.attrs["source"] == expected_name
     for time_coord in ["pco2_time", "iron_time", "dust_time", "nox_time", "nhy_time"]:
         assert sfc_forcing.ds.coords[time_coord].attrs["units"] == "days"
 
@@ -523,6 +533,7 @@ def test_surface_forcing_creation(
     [
         "bgc_surface_forcing",
         "bgc_surface_forcing_from_climatology",
+        "bgc_surface_forcing_from_unified_climatology",
     ],
 )
 def test_surface_forcing_pco2_replication(sfc_forcing_fixture, request):
@@ -656,11 +667,18 @@ def test_surface_forcing_bgc_save(bgc_surface_forcing, tmp_path):
             expected_filepath.unlink()
 
 
-def test_surface_forcing_bgc_from_clim_save(
-    bgc_surface_forcing_from_climatology, tmp_path
-):
+@pytest.mark.parametrize(
+    "sfc_forcing_fixture",
+    [
+        "bgc_surface_forcing_from_climatology",
+        "bgc_surface_forcing_from_unified_climatology",
+    ],
+)
+def test_surface_forcing_bgc_from_clim_save(sfc_forcing_fixture, tmp_path, request):
     """Test save method."""
 
+    bgc_surface_forcing_from_climatology = request.getfixturevalue(sfc_forcing_fixture)
+
     for file_str in ["test_sf", "test_sf.nc"]:
         # Create a temporary filepath using the tmp_path fixture
         for filepath in [
@@ -697,6 +715,7 @@ def test_surface_forcing_bgc_from_clim_save(
         "corrected_surface_forcing",
         "bgc_surface_forcing",
         "bgc_surface_forcing_from_climatology",
+        "bgc_surface_forcing_from_unified_climatology",
     ],
 )
 def test_roundtrip_yaml(sfc_forcing_fixture, request, tmp_path, use_dask):
@@ -760,9 +779,16 @@ def test_files_have_same_hash(sfc_forcing_fixture, request, tmp_path, use_dask):
     Path(expected_filepath2).unlink()
 
 
-def test_files_have_same_hash_clim(
-    bgc_surface_forcing_from_climatology, tmp_path, use_dask
-):
+@pytest.mark.parametrize(
+    "sfc_forcing_fixture",
+    [
+        "bgc_surface_forcing_from_climatology",
+        "bgc_surface_forcing_from_unified_climatology",
+    ],
+)
+def test_files_have_same_hash_clim(sfc_forcing_fixture, tmp_path, use_dask, request):
+
+    bgc_surface_forcing_from_climatology = request.getfixturevalue(sfc_forcing_fixture)
 
     yaml_filepath = tmp_path / "test_yaml"
     filepath1 = tmp_path / "test1.nc"

roms_tools/tests/test_setup/test_validation.py

@@ -16,13 +16,16 @@ def _get_fname(name):
         "grid_that_straddles_dateline",
         "tidal_forcing",
         "initial_conditions_with_bgc_from_climatology",
+        "initial_conditions_with_unified_bgc_from_climatology",
         "surface_forcing",
         "coarse_surface_forcing",
         "corrected_surface_forcing",
         "bgc_surface_forcing",
         "bgc_surface_forcing_from_climatology",
+        "bgc_surface_forcing_from_unified_climatology",
         "boundary_forcing",
         "bgc_boundary_forcing_from_climatology",
+        "bgc_boundary_forcing_from_unified_climatology",
         "river_forcing_with_bgc",
         "river_forcing_no_climatology",
     ],
@@ -55,13 +58,16 @@ def test_save_results(forcing_fixture, request):
         "grid_that_straddles_dateline",
         "tidal_forcing",
         "initial_conditions_with_bgc_from_climatology",
+        "initial_conditions_with_unified_bgc_from_climatology",
         "surface_forcing",
         "coarse_surface_forcing",
         "corrected_surface_forcing",
         "bgc_surface_forcing",
         "bgc_surface_forcing_from_climatology",
+        "bgc_surface_forcing_from_unified_climatology",
         "boundary_forcing",
         "bgc_boundary_forcing_from_climatology",
+        "bgc_boundary_forcing_from_unified_climatology",
         "river_forcing_with_bgc",
         "river_forcing_no_climatology",
     ],

roms_tools/tests/test_vertical_coordinate.py

@@ -0,0 +1,73 @@
+import pytest
+import numpy as np
+import xarray as xr
+
+from roms_tools.vertical_coordinate import (
+    compute_cs,
+    sigma_stretch,
+    compute_depth,
+    compute_depth_coordinates,
+)
+
+
+def test_compute_cs():
+    sigma = np.linspace(-1, 0, 10)
+    theta_s, theta_b = 5, 2
+    cs = compute_cs(sigma, theta_s, theta_b)
+    assert cs.shape == sigma.shape
+    assert np.all(cs <= 0) and np.all(cs >= -1)
+
+    with pytest.raises(ValueError, match="theta_s must be between 0 and 10"):
+        compute_cs(sigma, 15, 2)
+
+    with pytest.raises(ValueError, match="theta_b must be between 0 and 4"):
+        compute_cs(sigma, 5, 5)
+
+
+def test_sigma_stretch():
+    theta_s, theta_b, N = 5, 2, 10
+    cs, sigma = sigma_stretch(theta_s, theta_b, N, "r")
+    assert cs.shape == sigma.shape
+    assert isinstance(cs, xr.DataArray)
+    assert isinstance(sigma, xr.DataArray)
+
+    with pytest.raises(
+        ValueError,
+        match="Type must be either 'w' for vertical velocity points or 'r' for rho-points.",
+    ):
+        sigma_stretch(theta_s, theta_b, N, "invalid")
+
+
+def test_compute_depth():
+    zeta = xr.DataArray(0.5)
+    h = xr.DataArray(10.0)
+    hc = 5.0
+    cs = xr.DataArray(np.linspace(-1, 0, 10), dims="s_rho")
+    sigma = xr.DataArray(np.linspace(-1, 0, 10), dims="s_rho")
+
+    depth = compute_depth(zeta, h, hc, cs, sigma)
+    assert depth.shape == sigma.shape
+    assert isinstance(depth, xr.DataArray)
+
+
+def test_compute_depth_coordinates():
+    grid_ds = xr.Dataset(
+        {
+            "h": xr.DataArray([[10, 20], [30, 40]], dims=("eta_rho", "xi_rho")),
+            "Cs_r": xr.DataArray(np.linspace(-1, 0, 10), dims="s_rho"),
+            "sigma_r": xr.DataArray(np.linspace(-1, 0, 10), dims="s_rho"),
+            "Cs_w": xr.DataArray(np.linspace(-1, 0, 11), dims="s_w"),
+            "sigma_w": xr.DataArray(np.linspace(-1, 0, 11), dims="s_w"),
+        },
+        attrs={"hc": 5.0},
+    )
+
+    depth = compute_depth_coordinates(grid_ds, depth_type="layer", location="rho")
+    assert isinstance(depth, xr.DataArray)
+    assert "eta_rho" in depth.dims and "xi_rho" in depth.dims
+
+    with pytest.raises(ValueError, match="Invalid depth_type"):
+        compute_depth_coordinates(grid_ds, depth_type="invalid")
+
+    with pytest.raises(ValueError, match="Invalid location"):
+        compute_depth_coordinates(grid_ds, location="invalid")

roms_tools/vertical_coordinate.py

@@ -171,6 +171,12 @@ def compute_depth_coordinates(
     - Spatial slicing (`eta`, `xi`) is performed before depth computation to optimize efficiency.
     - Depth calculations rely on the ROMS vertical stretching curves (`Cs`) and sigma-layers.
     """
+    # Validate location
+    valid_locations = {"rho", "u", "v"}
+    if location not in valid_locations:
+        raise ValueError(
+            f"Invalid location: {location}. Must be one of {valid_locations}."
+        )
 
     # Select the appropriate depth computation parameters
     if depth_type == "layer":

{roms_tools-2.6.0.dist-info → roms_tools-2.6.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: roms-tools
-Version: 2.6.0
+Version: 2.6.2
 Summary: Tools for running and analysing UCLA-ROMS simulations
 Author-email: Nora Loose <nora.loose@gmail.com>, Thomas Nicholas <tom@cworthy.org>
 License: Apache-2
@@ -117,7 +117,6 @@ install `ROMS-Tools` along with the additional dependency via:
 pip install -e ".[dask]"
 ```
 
-
 ### Run the tests
 
 Before running the tests, you can activate the conda environment created in the previous section:
@@ -142,7 +141,6 @@ To learn how to use `ROMS-Tools`, check out the [documentation](https://roms-too
 
 We also accept contributions in the form of Pull Requests.
 
-
 ## See also
 
 - [ROMS source code](https://github.com/CESR-lab/ucla-roms)