roms-tools 3.3.0__py3-none-any.whl → 3.5.0__py3-none-any.whl

roms_tools/tests/test_regrid.py

@@ -2,7 +2,7 @@ import numpy as np
 import pytest
 import xarray as xr
 
-from roms_tools.regrid import VerticalRegridToROMS
+from roms_tools.regrid import VerticalRegrid, VerticalRegridToROMS
 
 try:
     import xesmf  # type: ignore
@@ -134,3 +134,122 @@ def test_vertical_regrid(request, depth_values, layer_depth_rho_values, temp_dat
     regridded = vertical_regrid.apply(data.temp_data, fill_nans=True)
     expected = np.interp(layer_depth_rho_values, depth_values, temp_data)
     assert np.allclose(expected, regridded.data, equal_nan=True)
+
+
+# Test VerticalRegrid
+def test_vertical_regrid_2d_depths_different_vertical_levels():
+    """
+    Vertical regridding with 2D (eta, xi) depth coordinates where
+    source and target have different numbers of vertical levels.
+    """
+    # --- dimensions ---
+    time = xr.DataArray(
+        np.array(["2000-01-01", "2000-01-02"], dtype="datetime64[ns]"),
+        dims="time",
+    )
+
+    s_rho_src = xr.DataArray(np.arange(6), dims="s_rho")  # source: 6 levels
+    s_rho_tgt = xr.DataArray(np.arange(10), dims="s_rho")  # target: 10 levels
+
+    eta = xr.DataArray(np.arange(3), dims="eta")
+    xi = xr.DataArray(np.arange(4), dims="xi")
+
+    # --- source depth coords: 2D + s_rho ---
+    base_depth = -(100 + 10 * eta.values[:, None] + 5 * xi.values[None, :])
+
+    source_depth = xr.DataArray(
+        base_depth,
+        dims=("eta", "xi"),
+        coords={"eta": eta, "xi": xi},
+    ).expand_dims(time=time, s_rho=s_rho_src)
+
+    source_depth = source_depth + 20 * s_rho_src
+
+    # --- target depth coords: 2D + different s_rho ---
+    target_depth = xr.DataArray(
+        base_depth,
+        dims=("eta", "xi"),
+        coords={"eta": eta, "xi": xi},
+    ).expand_dims(s_rho=s_rho_tgt)
+
+    target_depth = target_depth + 20 * s_rho_tgt
+
+    # --- synthetic temperature field: varies linearly with depth ---
+    temp_data = xr.DataArray(
+        np.broadcast_to(
+            s_rho_src.values[None, :, None, None] * 2.0,
+            (len(time), len(s_rho_src), len(eta), len(xi)),
+        ),
+        dims=("time", "s_rho", "eta", "xi"),
+        coords={
+            "time": time,
+            "s_rho": s_rho_src,
+            "eta": eta,
+            "xi": xi,
+        },
+        name="temp",
+    )
+
+    ds = xr.Dataset(
+        {
+            "temp": temp_data,
+        },
+        coords={
+            "time": time,
+            "s_rho": s_rho_src,
+            "eta": eta,
+            "xi": xi,
+        },
+    )
+
+    # --- regrid ---
+    regridder = VerticalRegrid(ds)
+    out = regridder.apply(
+        temp_data,
+        source_depth_coords=source_depth,
+        target_depth_coords=target_depth,
+    )
+
+    # --- assertions ---
+    assert isinstance(out, xr.DataArray)
+
+    # vertical dimension changed to target resolution
+    assert out.sizes["s_rho"] == s_rho_tgt.size
+
+    # horizontal + time preserved
+    assert out.sizes["time"] == time.size
+    assert out.sizes["eta"] == eta.size
+    assert out.sizes["xi"] == xi.size
+
+    # output contains finite values
+    assert np.isfinite(out).any()
+
+    # interpolation stays within source bounds
+    assert out.min() >= temp_data.min()
+    assert out.max() <= temp_data.max()
+
+
+def test_vertical_regrid_mask_edges():
+    """Values outside source depth range should be masked when mask_edges=True."""
+    s_rho = xr.DataArray(np.linspace(-1, 0, 5), dims="s_rho")
+    source_depth = xr.DataArray([-100, -75, -50, -25, 0], dims="s_rho")
+
+    target_depth = xr.DataArray([-200, -50, 10], dims="s_rho")
+
+    data = xr.DataArray(
+        source_depth.values,
+        dims="s_rho",
+        coords={"s_rho": s_rho},
+    )
+
+    ds = xr.Dataset(
+        {"data": data, "depth": source_depth},
+        coords={"s_rho": s_rho},
+    )
+
+    regridder = VerticalRegrid(ds)
+    out = regridder.apply(data, source_depth, target_depth, mask_edges=True)
+
+    assert np.isnan(out[0])
+    assert np.isnan(out[-1])
+    assert not np.isnan(out[1])

roms_tools/tests/test_setup/test_boundary_forcing.py

@@ -12,7 +12,7 @@ import xarray as xr
 
 from conftest import calculate_data_hash
 from roms_tools import BoundaryForcing, Grid
-from roms_tools.download import download_test_data
+from roms_tools.datasets.download import download_test_data
 from roms_tools.tests.test_setup.utils import download_regional_and_bigger
 
 try:
@@ -25,9 +25,7 @@ except ImportError:
     "boundary_forcing_fixture",
     [
         "boundary_forcing",
-        "boundary_forcing_adjusted_for_zeta",
         "boundary_forcing_with_2d_fill",
-        "boundary_forcing_with_2d_fill_adjusted_for_zeta",
     ],
 )
 def test_boundary_forcing_creation(boundary_forcing_fixture, request):
@@ -44,25 +42,24 @@ def test_boundary_forcing_creation(boundary_forcing_fixture, request):
         "climatology": False,
     }
     assert boundary_forcing.model_reference_date == datetime(2000, 1, 1)
-    assert boundary_forcing.boundaries == {
-        "south": True,
-        "east": True,
-        "north": True,
-        "west": True,
-    }
+    assert all(
+        k in boundary_forcing.boundaries for k in ["south", "east", "north", "west"]
+    )
 
     assert boundary_forcing.ds.source == "GLORYS"
     for direction in ["south", "east", "north", "west"]:
-        assert f"temp_{direction}" in boundary_forcing.ds
-        assert f"salt_{direction}" in boundary_forcing.ds
-        assert f"u_{direction}" in boundary_forcing.ds
-        assert f"v_{direction}" in boundary_forcing.ds
-        assert f"zeta_{direction}" in boundary_forcing.ds
+        if boundary_forcing.boundaries[direction]:
+            assert f"temp_{direction}" in boundary_forcing.ds
+            assert f"salt_{direction}" in boundary_forcing.ds
+            assert f"u_{direction}" in boundary_forcing.ds
+            assert f"v_{direction}" in boundary_forcing.ds
+            assert f"zeta_{direction}" in boundary_forcing.ds
 
     assert len(boundary_forcing.ds.bry_time) == 2
     assert boundary_forcing.ds.coords["bry_time"].attrs["units"] == "days"
     assert not hasattr(boundary_forcing.ds, "climatology")
     assert hasattr(boundary_forcing.ds, "adjust_depth_for_sea_surface_height")
+    assert boundary_forcing.ds.attrs["adjust_depth_for_sea_surface_height"] == "False"
     assert hasattr(boundary_forcing.ds, "apply_2d_horizontal_fill")
 
 
@@ -104,12 +101,9 @@ def test_bgc_boundary_forcing_creation(boundary_forcing_fixture, request):
     assert boundary_forcing.end_time == datetime(2021, 6, 30)
     assert boundary_forcing.source["climatology"]
     assert boundary_forcing.model_reference_date == datetime(2000, 1, 1)
-    assert boundary_forcing.boundaries == {
-        "south": True,
-        "east": True,
-        "north": True,
-        "west": True,
-    }
+    assert all(
+        k in boundary_forcing.boundaries for k in ["south", "east", "north", "west"]
+    )
 
     expected_bgc_variables = [
         "PO4",
@@ -147,8 +141,9 @@ def test_bgc_boundary_forcing_creation(boundary_forcing_fixture, request):
     ]
 
     for direction in ["south", "east", "north", "west"]:
-        for var in expected_bgc_variables:
-            assert f"{var}_{direction}" in boundary_forcing.ds
+        if boundary_forcing.boundaries[direction]:
+            for var in expected_bgc_variables:
+                assert f"{var}_{direction}" in boundary_forcing.ds
 
     assert len(boundary_forcing.ds.bry_time) == 12
     assert boundary_forcing.ds.coords["bry_time"].attrs["units"] == "days"
@@ -280,56 +275,6 @@ def test_boundary_divided_by_land_warning(caplog, use_dask):
     assert "divided by land" in caplog.text
 
 
-def test_info_depth(caplog, use_dask):
-    grid = Grid(
-        nx=3,
-        ny=3,
-        size_x=400,
-        size_y=400,
-        center_lon=-8,
-        center_lat=58,
-        rot=0,
-        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
-    )
-
-    fname1 = Path(download_test_data("GLORYS_NA_20120101.nc"))
-    fname2 = Path(download_test_data("GLORYS_NA_20121231.nc"))
-
-    with caplog.at_level(logging.INFO):
-        BoundaryForcing(
-            grid=grid,
-            start_time=datetime(2012, 1, 1),
-            end_time=datetime(2012, 12, 31),
-            source={"name": "GLORYS", "path": [fname1, fname2]},
-            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):
-        BoundaryForcing(
-            grid=grid,
-            start_time=datetime(2012, 1, 1),
-            end_time=datetime(2012, 12, 31),
-            source={"name": "GLORYS", "path": [fname1, fname2]},
-            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
-    )
-
-
 def test_1d_and_2d_fill_coincide_if_no_fill(use_dask):
     grid = Grid(
         nx=2,
@@ -395,29 +340,6 @@ def test_1d_and_2d_fill_coincide_if_no_land(use_dask):
     xr.testing.assert_allclose(bf_1d_fill.ds, bf_2d_fill.ds, rtol=1.0e-4)
 
 
-@pytest.mark.parametrize(
-    "boundary_forcing_fixture",
-    [
-        "boundary_forcing_adjusted_for_zeta",
-        "boundary_forcing_with_2d_fill_adjusted_for_zeta",
-    ],
-)
-def test_correct_depth_coords_adjusted_for_zeta(
-    boundary_forcing_fixture, request, use_dask
-):
-    boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
-
-    for direction in ["south", "east", "north", "west"]:
-        # Test that uppermost interface coincides with sea surface height
-        assert np.allclose(
-            boundary_forcing.ds_depth_coords[f"interface_depth_rho_{direction}"]
-            .isel(s_w=-1)
-            .values,
-            -boundary_forcing.ds[f"zeta_{direction}"].values,
-            atol=1e-6,
-        )
-
-
 @pytest.mark.parametrize(
     "boundary_forcing_fixture",
     [
@@ -429,14 +351,15 @@ def test_correct_depth_coords_zero_zeta(boundary_forcing_fixture, request, use_d
     boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
 
     for direction in ["south", "east", "north", "west"]:
-        # Test that uppermost interface coincides with sea surface height
-        assert np.allclose(
-            boundary_forcing.ds_depth_coords[f"interface_depth_rho_{direction}"]
-            .isel(s_w=-1)
-            .values,
-            0 * boundary_forcing.ds[f"zeta_{direction}"].values,
-            atol=1e-6,
-        )
+        if boundary_forcing.boundaries[direction]:
+            # Test that uppermost interface coincides with sea surface height
+            assert np.allclose(
+                boundary_forcing.ds_depth_coords[f"interface_depth_rho_{direction}"]
+                .isel(s_w=-1)
+                .values,
+                0 * boundary_forcing.ds[f"zeta_{direction}"].values,
+                atol=1e-6,
+            )
 
 
 def test_computed_missing_optional_fields(
@@ -446,16 +369,16 @@ def test_computed_missing_optional_fields(
 
     # Use tight tolerances because 'DOC' and 'DOCr' can have values order 1e-6
 
-    # 'DOCr' was missing in the source data and should have been filled with a constant default value
     for direction in ["south", "east", "north", "west"]:
-        assert np.allclose(
-            ds[f"DOCr_{direction}"].std(), 0.0, rtol=1e-10, atol=1e-10
-        ), "DOCr should be constant across space and time"
-    # 'DOC' was present in the source data and should show spatial or temporal variability
-    for direction in ["south", "east", "north", "west"]:
-        assert ds[f"DOC_{direction}"].std() > 1e-10, (
-            "DOC should vary across space and time"
-        )
+        if bgc_boundary_forcing_from_unified_climatology.boundaries[direction]:
+            # 'DOCr' was missing in the source data and should have been filled with a constant default value
+            assert np.allclose(
+                ds[f"DOCr_{direction}"].std(), 0.0, rtol=1e-10, atol=1e-10
+            ), "DOCr should be constant across space and time"
+            # 'DOC' was present in the source data and should show spatial or temporal variability
+            assert ds[f"DOC_{direction}"].std() > 1e-10, (
+                "DOC should vary across space and time"
+            )
 
 
 @pytest.mark.parametrize(
@@ -463,8 +386,6 @@ def test_computed_missing_optional_fields(
     [
         "boundary_forcing",
         "boundary_forcing_with_2d_fill",
-        "boundary_forcing_adjusted_for_zeta",
-        "boundary_forcing_with_2d_fill_adjusted_for_zeta",
     ],
 )
 def test_boundary_forcing_plot(boundary_forcing_fixture, request):
@@ -472,24 +393,25 @@ def test_boundary_forcing_plot(boundary_forcing_fixture, request):
     boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
 
     for direction in ["south", "east", "north", "west"]:
-        for layer_contours in [False, True]:
-            boundary_forcing.plot(
-                var_name=f"temp_{direction}", layer_contours=layer_contours
-            )
-            boundary_forcing.plot(
-                var_name=f"u_{direction}", layer_contours=layer_contours
-            )
-            boundary_forcing.plot(
-                var_name=f"v_{direction}", layer_contours=layer_contours
-            )
-        boundary_forcing.plot(var_name=f"zeta_{direction}")
-        boundary_forcing.plot(var_name=f"vbar_{direction}")
-        boundary_forcing.plot(var_name=f"ubar_{direction}")
-
-    # Test that passing a matplotlib.axes.Axes works
-    fig, ax = plt.subplots(1, 1)
-    boundary_forcing.plot(var_name=f"temp_{direction}", ax=ax)
-    boundary_forcing.plot(var_name=f"zeta_{direction}", ax=ax)
+        if boundary_forcing.boundaries[direction]:
+            for layer_contours in [False, True]:
+                boundary_forcing.plot(
+                    var_name=f"temp_{direction}", layer_contours=layer_contours
+                )
+                boundary_forcing.plot(
+                    var_name=f"u_{direction}", layer_contours=layer_contours
+                )
+                boundary_forcing.plot(
+                    var_name=f"v_{direction}", layer_contours=layer_contours
+                )
+            boundary_forcing.plot(var_name=f"zeta_{direction}")
+            boundary_forcing.plot(var_name=f"vbar_{direction}")
+            boundary_forcing.plot(var_name=f"ubar_{direction}")
+
+            # Test that passing a matplotlib.axes.Axes works
+            fig, ax = plt.subplots(1, 1)
+            boundary_forcing.plot(var_name=f"temp_{direction}", ax=ax)
+            boundary_forcing.plot(var_name=f"zeta_{direction}", ax=ax)
 
 
 @pytest.mark.parametrize(
@@ -497,8 +419,6 @@ def test_boundary_forcing_plot(boundary_forcing_fixture, request):
     [
         "boundary_forcing",
         "boundary_forcing_with_2d_fill",
-        "boundary_forcing_adjusted_for_zeta",
-        "boundary_forcing_with_2d_fill_adjusted_for_zeta",
     ],
 )
 def test_boundary_forcing_save(boundary_forcing_fixture, request, tmp_path):
@@ -543,10 +463,9 @@ def test_bgc_boundary_forcing_plot(boundary_forcing_fixture, request):
     """Test plot method."""
     bgc_boundary_forcing = request.getfixturevalue(boundary_forcing_fixture)
 
-    bgc_boundary_forcing.plot(var_name="ALK_south", layer_contours=True)
-    bgc_boundary_forcing.plot(var_name="ALK_east", layer_contours=True)
-    bgc_boundary_forcing.plot(var_name="ALK_north", layer_contours=True)
-    bgc_boundary_forcing.plot(var_name="ALK_west", layer_contours=True)
+    for direction in ["south", "east", "north", "west"]:
+        if bgc_boundary_forcing.boundaries[direction]:
+            bgc_boundary_forcing.plot(var_name=f"ALK_{direction}", layer_contours=True)
 
 
 @pytest.mark.parametrize(

roms_tools/tests/test_setup/test_cdr_release.py

@@ -328,11 +328,10 @@ class TestTracerPerturbation:
         with pytest.raises(ValidationError):
             TracerPerturbation(**self.params, tracer_fluxes={"ALK": ["not", "valid"]})
 
-    def test_tracer_fluxes_non_negative(self):
-        with pytest.raises(ValidationError):
-            TracerPerturbation(**self.params, tracer_fluxes={"ALK": -1})
-        with pytest.raises(ValidationError):
-            TracerPerturbation(**self.params, tracer_fluxes={"ALK": [-1, 15]})
+    def test_tracer_fluxes_negative(self):
+        tp = TracerPerturbation(**self.params, tracer_fluxes={"DIC": -1})
+
+        assert tp.tracer_fluxes["DIC"].values <= 0.0
 
     def test_zero_fill_strategy(self):
         alk_value = 100.0