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

roms_tools/tests/test_datasets/test_roms_dataset.py

@@ -308,36 +308,136 @@ def test_that_coordinates_and_masks_are_added(use_dask):
     assert "mask_v" in output.ds
 
 
-def test_apply_lateral_fill():
-    grid_parameters = {
-        "nx": 10,
-        "ny": 10,
-        "size_x": 2500,
-        "size_y": 3000,
-        "center_lon": -30,
-        "center_lat": 57,
-        "rot": -20,
-    }
+# Test applying lateral fill
+
+
+def make_roms_dataset(ds, grid):
+    roms = ROMSDataset.__new__(ROMSDataset)
+    roms.ds = ds
+    roms.grid = grid
+    return roms
+
 
-    grid = Grid(**grid_parameters)
+def test_apply_lateral_fill_rho_success():
+    grid = Grid(
+        nx=10, ny=10, size_x=2500, size_y=3000, center_lon=-30, center_lat=57, rot=-20
+    )
 
     ds = xr.Dataset()
     ds["field"] = 5 * grid.ds.mask_rho.copy()
     ds["mask_rho"] = grid.ds.mask_rho.copy()
+
+    roms = make_roms_dataset(ds, grid)
+
+    # land exists initially
+    assert (roms.ds["field"] == 0).any()
+
+    roms.apply_lateral_fill()
+
+    assert (roms.ds["field"] == 5).all()
+
+
+def test_apply_lateral_fill_rho_missing_mask_raises():
+    grid = Grid(
+        nx=10, ny=10, size_x=2500, size_y=3000, center_lon=-30, center_lat=57, rot=-20
+    )
+
+    ds = xr.Dataset()
+    ds["field"] = grid.ds.mask_rho.copy()
+
+    roms = make_roms_dataset(ds, grid)
+
+    with pytest.raises(ValueError, match="mask_rho"):
+        roms.apply_lateral_fill()
+
+
+def test_apply_lateral_fill_u_success():
+    grid = Grid(
+        nx=10, ny=10, size_x=2500, size_y=3000, center_lon=-30, center_lat=57, rot=-20
+    )
+
+    ds = xr.Dataset()
+    ds["u_field"] = 7 * grid.ds.mask_u.copy()
     ds["mask_u"] = grid.ds.mask_u.copy()
+
+    roms = make_roms_dataset(ds, grid)
+
+    assert (roms.ds["u_field"] == 0).any()
+
+    roms.apply_lateral_fill()
+
+    assert (roms.ds["u_field"] == 7).all()
+
+
+def test_apply_lateral_fill_u_missing_mask_raises():
+    grid = Grid(
+        nx=10, ny=10, size_x=2500, size_y=3000, center_lon=-30, center_lat=57, rot=-20
+    )
+
+    ds = xr.Dataset()
+    ds["u_field"] = grid.ds.mask_u.copy()
+    ds["mask_rho"] = grid.ds.mask_rho.copy()
+
+    roms = make_roms_dataset(ds, grid)
+
+    with pytest.raises(ValueError, match="mask_u"):
+        roms.apply_lateral_fill()
+
+
+def test_apply_lateral_fill_v_success():
+    grid = Grid(
+        nx=10, ny=10, size_x=2500, size_y=3000, center_lon=-30, center_lat=57, rot=-20
+    )
+
+    ds = xr.Dataset()
+    ds["v_field"] = 9 * grid.ds.mask_v.copy()
     ds["mask_v"] = grid.ds.mask_v.copy()
 
-    roms_dataset = ROMSDataset.__new__(ROMSDataset)
-    roms_dataset.ds = ds
-    roms_dataset.grid = grid
+    roms = make_roms_dataset(ds, grid)
 
-    # Check that initially some values are zero (land)
-    assert (roms_dataset.ds["field"].values == 0).any()
+    assert (roms.ds["v_field"] == 0).any()
 
-    roms_dataset.apply_lateral_fill()
+    roms.apply_lateral_fill()
 
-    # After filling the ocean values (all 5) should have propagated into land
-    assert (roms_dataset.ds["field"].values == 5).all()
+    assert (roms.ds["v_field"] == 9).all()
+
+
+def test_apply_lateral_fill_v_missing_mask_raises():
+    grid = Grid(
+        nx=10, ny=10, size_x=2500, size_y=3000, center_lon=-30, center_lat=57, rot=-20
+    )
+
+    ds = xr.Dataset()
+    ds["v_field"] = grid.ds.mask_v.copy()
+    ds["mask_rho"] = grid.ds.mask_rho.copy()
+
+    roms = make_roms_dataset(ds, grid)
+
+    with pytest.raises(ValueError, match="mask_v"):
+        roms.apply_lateral_fill()
+
+
+def test_apply_lateral_fill_mixed_grids():
+    grid = Grid(
+        nx=10, ny=10, size_x=2500, size_y=3000, center_lon=-30, center_lat=57, rot=-20
+    )
+
+    ds = xr.Dataset()
+    ds["rho"] = 1 * grid.ds.mask_rho.copy()
+    ds["u"] = 2 * grid.ds.mask_u.copy()
+    ds["v"] = 3 * grid.ds.mask_v.copy()
+
+    ds["mask_rho"] = grid.ds.mask_rho.copy()
+    ds["mask_u"] = grid.ds.mask_u.copy()
+    ds["mask_v"] = grid.ds.mask_v.copy()
+
+    roms = make_roms_dataset(ds, grid)
+
+    roms.apply_lateral_fill()
+
+    assert (roms.ds["rho"] == 1).all()
+    assert (roms.ds["u"] == 2).all()
+    assert (roms.ds["v"] == 3).all()
 
 
 # Test choose_subdomain
@@ -513,8 +613,7 @@ def test_choose_subdomain_with(big_params, small_params):
     target_coords = get_target_coords(small)
 
     # --- apply function ---
-    sub = choose_subdomain(ds, big.ds, target_coords, buffer_points=1)
-
+    sub = choose_subdomain(ds, big.ds, target_coords, buffer_points=10)
     # --- rho tests ---
     assert sub.lat_rho.shape[0] <= ds.lat_rho.shape[0]
     assert sub.lat_rho.shape[1] <= ds.lat_rho.shape[1]
@@ -537,3 +636,108 @@ def test_choose_subdomain_with(big_params, small_params):
         assert float(sub.lat_v.min()) >= float(ds.lat_v.min()) - 1e-6
         assert float(sub.lat_v.max()) <= float(ds.lat_v.max()) + 1e-6
         assert not sub.field_v.isnull().any()
+
+
+def test_choose_subdomain_dataset_and_grid(use_dask):
+    fname_grid = Path(download_test_data("epac25km_grd.nc"))
+    grid = Grid.from_file(fname_grid)
+
+    roms_dataset = ROMSDataset(
+        grid=grid,
+        path=Path(download_test_data("eastpac25km_rst.19980106000000.nc")),
+        use_dask=use_dask,
+    )
+
+    # Save original sizes
+    orig_grid_sizes = dict(roms_dataset.grid.ds.sizes)
+    orig_ds_sizes = dict(roms_dataset.ds.sizes)
+
+    # Create a small child grid and target coordinates
+    child_grid = Grid(
+        nx=5,
+        ny=5,
+        size_x=100,
+        size_y=100,
+        center_lon=-128.0,
+        center_lat=9.0,
+        rot=32.0,
+    )
+    target_coords = get_target_coords(child_grid)
+
+    # Apply subdomain selection
+    roms_dataset.choose_subdomain(target_coords, buffer_points=1)
+
+    new_grid_sizes = dict(roms_dataset.grid.ds.sizes)
+    new_ds_sizes = dict(roms_dataset.ds.sizes)
+
+    # Grid and dataset should both be reduced
+    assert new_grid_sizes != orig_grid_sizes, (
+        "Grid sizes did not change after choose_subdomain"
+    )
+    assert new_ds_sizes != orig_ds_sizes, (
+        "Dataset sizes did not change after choose_subdomain"
+    )
+
+    # Grid and dataset should remain consistent with each other
+    joint_dims = ["eta_rho", "xi_rho", "xi_u", "eta_v", "s_rho"]
+
+    for dim in joint_dims:
+        assert new_grid_sizes.get(dim) == new_ds_sizes.get(dim), (
+            f"Mismatch in dimension '{dim}': "
+            f"grid={new_grid_sizes.get(dim)}, ds={new_ds_sizes.get(dim)}"
+        )
+
+
+def test_choose_subdomain_does_not_mutate_shared_grid(use_dask):
+    fname_grid = Path(download_test_data("epac25km_grd.nc"))
+    grid = Grid.from_file(fname_grid)
+
+    rd1 = ROMSDataset(
+        grid=grid,
+        path=Path(download_test_data("eastpac25km_rst.19980106000000.nc")),
+        use_dask=use_dask,
+    )
+    rd2 = ROMSDataset(
+        grid=grid,
+        path=Path(download_test_data("eastpac25km_rst.19980106000000.nc")),
+        use_dask=use_dask,
+    )
+
+    child_grid = Grid(
+        nx=5,
+        ny=5,
+        size_x=100,
+        size_y=100,
+        center_lon=-128.0,
+        center_lat=9.0,
+        rot=32.0,
+    )
+    target_coords = get_target_coords(child_grid)
+    rd1.choose_subdomain(target_coords, buffer_points=1)
+
+    assert rd2.grid.ds.sizes == grid.ds.sizes
+
+
+def test_choose_subdomain_then_compute_depth_coordinates(use_dask):
+    fname_grid = Path(download_test_data("epac25km_grd.nc"))
+    grid = Grid.from_file(fname_grid)
+
+    rd = ROMSDataset(
+        grid=grid,
+        path=Path(download_test_data("eastpac25km_rst.19980106000000.nc")),
+        use_dask=use_dask,
+    )
+
+    child_grid = Grid(
+        nx=5,
+        ny=5,
+        size_x=100,
+        size_y=100,
+        center_lon=-128.0,
+        center_lat=9.0,
+        rot=32.0,
+    )
+    target_coords = get_target_coords(child_grid)
+    rd.choose_subdomain(target_coords, buffer_points=1)
+
+    rd._get_depth_coordinates(locations=["rho", "u", "v"])

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_data/initial_conditions_from_roms.zarr/ALK/zarr.json

@@ -0,0 +1,57 @@
+{
+  "shape": [
+    1,
+    3,
+    7,
+    7
+  ],
+  "data_type": "float32",
+  "chunk_grid": {
+    "name": "regular",
+    "configuration": {
+      "chunk_shape": [
+        1,
+        3,
+        7,
+        7
+      ]
+    }
+  },
+  "chunk_key_encoding": {
+    "name": "default",
+    "configuration": {
+      "separator": "/"
+    }
+  },
+  "fill_value": 0.0,
+  "codecs": [
+    {
+      "name": "bytes",
+      "configuration": {
+        "endian": "little"
+      }
+    },
+    {
+      "name": "zstd",
+      "configuration": {
+        "level": 0,
+        "checksum": false
+      }
+    }
+  ],
+  "attributes": {
+    "long_name": "alkalinity",
+    "units": "meq/m^3",
+    "coordinates": "abs_time",
+    "_FillValue": "AAAAAAAA+H8="
+  },
+  "dimension_names": [
+    "ocean_time",
+    "s_rho",
+    "eta_rho",
+    "xi_rho"
+  ],
+  "zarr_format": 3,
+  "node_type": "array",
+  "storage_transformers": []
+}

roms_tools/tests/test_setup/test_data/initial_conditions_from_roms.zarr/ALK_ALT_CO2/zarr.json

@@ -0,0 +1,57 @@
+{
+  "shape": [
+    1,
+    3,
+    7,
+    7
+  ],
+  "data_type": "float32",
+  "chunk_grid": {
+    "name": "regular",
+    "configuration": {
+      "chunk_shape": [
+        1,
+        3,
+        7,
+        7
+      ]
+    }
+  },
+  "chunk_key_encoding": {
+    "name": "default",
+    "configuration": {
+      "separator": "/"
+    }
+  },
+  "fill_value": 0.0,
+  "codecs": [
+    {
+      "name": "bytes",
+      "configuration": {
+        "endian": "little"
+      }
+    },
+    {
+      "name": "zstd",
+      "configuration": {
+        "level": 0,
+        "checksum": false
+      }
+    }
+  ],
+  "attributes": {
+    "long_name": "alkalinity, alternative CO2",
+    "units": "meq/m^3",
+    "coordinates": "abs_time",
+    "_FillValue": "AAAAAAAA+H8="
+  },
+  "dimension_names": [
+    "ocean_time",
+    "s_rho",
+    "eta_rho",
+    "xi_rho"
+  ],
+  "zarr_format": 3,
+  "node_type": "array",
+  "storage_transformers": []
+}

roms_tools/tests/test_setup/test_data/initial_conditions_from_roms.zarr/Cs_r/zarr.json

@@ -0,0 +1,47 @@
+{
+  "shape": [
+    3
+  ],
+  "data_type": "float32",
+  "chunk_grid": {
+    "name": "regular",
+    "configuration": {
+      "chunk_shape": [
+        3
+      ]
+    }
+  },
+  "chunk_key_encoding": {
+    "name": "default",
+    "configuration": {
+      "separator": "/"
+    }
+  },
+  "fill_value": 0.0,
+  "codecs": [
+    {
+      "name": "bytes",
+      "configuration": {
+        "endian": "little"
+      }
+    },
+    {
+      "name": "zstd",
+      "configuration": {
+        "level": 0,
+        "checksum": false
+      }
+    }
+  ],
+  "attributes": {
+    "long_name": "Vertical stretching function at rho-points",
+    "units": "nondimensional",
+    "_FillValue": "AAAAAAAA+H8="
+  },
+  "dimension_names": [
+    "s_rho"
+  ],
+  "zarr_format": 3,
+  "node_type": "array",
+  "storage_transformers": []
+}

roms_tools/tests/test_setup/test_data/initial_conditions_from_roms.zarr/Cs_w/zarr.json

@@ -0,0 +1,47 @@
+{
+  "shape": [
+    4
+  ],
+  "data_type": "float32",
+  "chunk_grid": {
+    "name": "regular",
+    "configuration": {
+      "chunk_shape": [
+        4
+      ]
+    }
+  },
+  "chunk_key_encoding": {
+    "name": "default",
+    "configuration": {
+      "separator": "/"
+    }
+  },
+  "fill_value": 0.0,
+  "codecs": [
+    {
+      "name": "bytes",
+      "configuration": {
+        "endian": "little"
+      }
+    },
+    {
+      "name": "zstd",
+      "configuration": {
+        "level": 0,
+        "checksum": false
+      }
+    }
+  ],
+  "attributes": {
+    "long_name": "Vertical stretching function at w-points",
+    "units": "nondimensional",
+    "_FillValue": "AAAAAAAA+H8="
+  },
+  "dimension_names": [
+    "s_w"
+  ],
+  "zarr_format": 3,
+  "node_type": "array",
+  "storage_transformers": []
+}

roms_tools/tests/test_setup/test_data/initial_conditions_from_roms.zarr/DIC/zarr.json

@@ -0,0 +1,57 @@
+{
+  "shape": [
+    1,
+    3,
+    7,
+    7
+  ],
+  "data_type": "float32",
+  "chunk_grid": {
+    "name": "regular",
+    "configuration": {
+      "chunk_shape": [
+        1,
+        3,
+        7,
+        7
+      ]
+    }
+  },
+  "chunk_key_encoding": {
+    "name": "default",
+    "configuration": {
+      "separator": "/"
+    }
+  },
+  "fill_value": 0.0,
+  "codecs": [
+    {
+      "name": "bytes",
+      "configuration": {
+        "endian": "little"
+      }
+    },
+    {
+      "name": "zstd",
+      "configuration": {
+        "level": 0,
+        "checksum": false
+      }
+    }
+  ],
+  "attributes": {
+    "long_name": "dissolved inorganic carbon",
+    "units": "mmol/m^3",
+    "coordinates": "abs_time",
+    "_FillValue": "AAAAAAAA+H8="
+  },
+  "dimension_names": [
+    "ocean_time",
+    "s_rho",
+    "eta_rho",
+    "xi_rho"
+  ],
+  "zarr_format": 3,
+  "node_type": "array",
+  "storage_transformers": []
+}