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

roms_tools/tests/test_utils.py

@@ -3,17 +3,29 @@ from pathlib import Path
 from unittest import mock
 
 import numpy as np
+import pandas as pd
 import pytest
 import xarray as xr
 
+from roms_tools.datasets.download import download_test_data
+from roms_tools.datasets.lat_lon_datasets import ERA5Correction
 from roms_tools.utils import (
+    _interpolate_generic,
     _path_list_from_input,
     generate_focused_coordinate_range,
     get_dask_chunks,
     has_copernicus,
     has_dask,
     has_gcsfs,
+    interpolate_cyclic_time,
+    interpolate_from_climatology,
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
+    interpolate_from_u_to_rho,
+    interpolate_from_v_to_rho,
     load_data,
+    rotate_velocities,
+    wrap_longitudes,
 )
 
 
@@ -241,3 +253,356 @@ def test_time_chunking_false_roms():
     dim_names = {"time": "ocean_time"}
     result = get_dask_chunks(dim_names, time_chunking=False)
     assert "ocean_time" not in result
+
+
+# test interpolate_from_climatology
+
+
+@pytest.fixture
+def climatology_data():
+    """Create a simple annual cycle dataset with 12 time points."""
+    time_coord = np.arange(1, 13)  # months as day_of_year approximation
+    da = xr.DataArray(np.arange(12), dims=("time",), coords={"time": time_coord})
+    ds = xr.Dataset({"var1": da, "var2": da * 2})
+    return da, ds, "time", "time"
+
+
+def test_interpolate_dataarray_single_time(climatology_data):
+    da, _, time_dim, time_coord = climatology_data
+    target_time = pd.Timestamp("2000-03-15")  # day_of_year ~ 75
+    interpolated = interpolate_from_climatology(da, time_dim, time_coord, target_time)
+    assert isinstance(interpolated, xr.DataArray)
+    assert interpolated.sizes[time_dim] == 1
+
+
+def test_interpolate_dataset_multiple_times(climatology_data):
+    _, ds, time_dim, time_coord = climatology_data
+    target_times = pd.date_range("2000-01-01", periods=3, freq="ME")
+    interpolated = interpolate_from_climatology(ds, time_dim, time_coord, target_times)
+    assert isinstance(interpolated, xr.Dataset)
+    assert all(interpolated[var].sizes[time_dim] == 3 for var in interpolated.data_vars)
+
+
+def test_interpolate_dataarray_time_dim_not_equal_time_coord():
+    time_values = np.arange(1, 13)
+    da = xr.DataArray(
+        np.arange(12),
+        dims=("time_dim",),
+        coords={"time_coord": ("time_dim", time_values)},
+    )
+    target_time = pd.Timestamp("2000-06-15")
+    interpolated = interpolate_from_climatology(
+        da, time_dim="time_dim", time_coord="time_coord", time=target_time
+    )
+    assert interpolated.sizes["time_dim"] == 1
+    assert np.issubdtype(interpolated.dtype, np.number)
+
+
+def test_interpolate_cyclic_time_basic():
+    time_values = np.arange(1, 13)
+    da = xr.DataArray(np.arange(12), dims=("time",), coords={"time": time_values})
+    target_days = [0.5, 6.5, 12.5]  # fractional days, include cyclic behavior
+    interpolated = interpolate_cyclic_time(
+        da, time_dim="time", time_coord="time", day_of_year=target_days
+    )
+    assert isinstance(interpolated, xr.DataArray)
+    assert interpolated.sizes["time"] == len(target_days)
+
+
+def test_interpolate_from_climatology_invalid_input():
+    with pytest.raises(TypeError):
+        interpolate_from_climatology(
+            "not a dataset", "time", "time", pd.Timestamp("2000-01-01")
+        )
+
+
+def test_interpolate_from_real_climatology(use_dask):
+    fname = download_test_data("ERA5_regional_test_data.nc")
+    era5_times = xr.open_dataset(fname).time
+
+    climatology = ERA5Correction(use_dask=use_dask)
+    field = climatology.ds["ssr_corr"]
+    field["time"] = field["time"].dt.days
+
+    interpolated_field = interpolate_from_climatology(field, "time", "time", era5_times)
+    assert len(interpolated_field.time) == len(era5_times)
+
+
+def test_wrap_longitudes_staggered():
+    # Dimensions
+    eta_rho, xi_rho, xi_u = 3, 4, 5
+    eta_v = 2
+
+    # Create 2D coordinates
+    lon_rho = xr.DataArray(
+        np.linspace(0, 360, eta_rho * xi_rho).reshape(eta_rho, xi_rho),
+        dims=("eta_rho", "xi_rho"),
+        attrs={"units": "degrees_east"},
+    )
+    lon_u = xr.DataArray(
+        np.linspace(-190, 190, eta_rho * xi_u).reshape(eta_rho, xi_u),
+        dims=("eta_rho", "xi_u"),
+        attrs={"units": "degrees_east"},
+    )
+    lon_v = xr.DataArray(
+        np.linspace(-180, 180, eta_v * xi_rho).reshape(eta_v, xi_rho),
+        dims=("eta_v", "xi_rho"),
+        attrs={"units": "degrees_east"},
+    )
+
+    # Dummy variables
+    ds = xr.Dataset(
+        {
+            "dummy_rho": (("eta_rho", "xi_rho"), np.zeros((eta_rho, xi_rho))),
+            "dummy_u": (("eta_rho", "xi_u"), np.zeros((eta_rho, xi_u))),
+            "dummy_v": (("eta_v", "xi_rho"), np.zeros((eta_v, xi_rho))),
+        },
+        coords={"lon_rho": lon_rho, "lon_u": lon_u, "lon_v": lon_v},
+    )
+
+    # Wrap to [-180, 180]
+    ds_wrapped = wrap_longitudes(ds, straddle=True)
+
+    # Check values: all >180 should be shifted
+    assert ds_wrapped.lon_rho.max().values <= 180
+    assert ds_wrapped.lon_u.max().values <= 180
+    assert ds_wrapped.lon_v.max().values <= 180
+
+    # Wrap to [0, 360]
+    ds_wrapped2 = wrap_longitudes(ds, straddle=False)
+    assert ds_wrapped2.lon_rho.min().values >= 0
+    assert ds_wrapped2.lon_u.min().values >= 0
+    assert ds_wrapped2.lon_v.min().values >= 0
+
+    # Check attributes preserved
+    for name in ["lon_rho", "lon_u", "lon_v"]:
+        assert ds.coords[name].attrs == ds_wrapped.coords[name].attrs
+
+
+# test _interpolate_generic and its wrappers
+
+# -------------------------
+# Fixtures
+# -------------------------
+
+
+@pytest.fixture
+def sample_rho_field() -> xr.DataArray:
+    """Create a simple rho-point field for testing."""
+    data = np.arange(12, dtype=float).reshape(3, 4)
+    eta = np.arange(3)
+    xi = np.arange(4)
+
+    return xr.DataArray(
+        data,
+        dims=("eta_rho", "xi_rho"),
+        coords={
+            "lat_rho": (("eta_rho", "xi_rho"), eta[:, None] * np.ones((1, 4))),
+            "lon_rho": (("eta_rho", "xi_rho"), np.ones((3, 1)) * xi[None, :]),
+        },
+    )
+
+
+@pytest.fixture
+def sample_u_field() -> xr.DataArray:
+    """Create a simple u-point field for testing."""
+    data = np.arange(9, dtype=float).reshape(3, 3)
+    eta = np.arange(3)
+    xi = np.arange(3)
+
+    return xr.DataArray(
+        data,
+        dims=("eta_rho", "xi_u"),
+        coords={
+            "lat_u": (("eta_rho", "xi_u"), eta[:, None] * np.ones((1, 3))),
+            "lon_u": (("eta_rho", "xi_u"), np.ones((3, 1)) * xi[None, :]),
+        },
+    )
+
+
+@pytest.fixture
+def sample_v_field() -> xr.DataArray:
+    """Create a simple v-point field for testing."""
+    data = np.arange(8, dtype=float).reshape(2, 4)
+    eta = np.arange(2)
+    xi = np.arange(4)
+
+    return xr.DataArray(
+        data,
+        dims=("eta_v", "xi_rho"),
+        coords={
+            "lat_v": (("eta_v", "xi_rho"), eta[:, None] * np.ones((1, 4))),
+            "lon_v": (("eta_v", "xi_rho"), np.ones((2, 1)) * xi[None, :]),
+        },
+    )
+
+
+# -------------------------
+# Generic interpolation tests
+# -------------------------
+
+
+def test_interpolate_from_rho_to_u_additive(sample_rho_field: xr.DataArray):
+    result = _interpolate_generic(
+        sample_rho_field, dim_in="xi_rho", dim_out="xi_u", method="additive"
+    )
+
+    # One fewer point along xi
+    assert result.shape[1] == sample_rho_field.shape[1] - 1
+
+    expected = 0.5 * (sample_rho_field.values[:, 1:] + sample_rho_field.values[:, :-1])
+    np.testing.assert_allclose(result.values, expected)
+
+
+def test_interpolate_from_rho_to_u_multiplicative(sample_rho_field: xr.DataArray):
+    result = _interpolate_generic(
+        sample_rho_field, dim_in="xi_rho", dim_out="xi_u", method="multiplicative"
+    )
+
+    expected = sample_rho_field.values[:, 1:] * sample_rho_field.values[:, :-1]
+    np.testing.assert_allclose(result.values, expected)
+
+
+# -------------------------
+# Wrapper tests
+# -------------------------
+
+
+def test_rho_to_u_wrapper_additive(sample_rho_field: xr.DataArray):
+    result = interpolate_from_rho_to_u(sample_rho_field, method="additive")
+
+    # Dimension swap
+    assert "xi_u" in result.dims
+    assert "xi_rho" not in result.dims
+
+    # Coordinates dropped
+    for coord in ("lat_rho", "lon_rho"):
+        assert coord not in result.coords
+
+    # Shape check
+    assert result.sizes["xi_u"] == sample_rho_field.sizes["xi_rho"] - 1
+
+
+def test_rho_to_v_wrapper_additive(sample_rho_field: xr.DataArray):
+    result = interpolate_from_rho_to_v(sample_rho_field, method="additive")
+
+    # Dimension swap
+    assert "eta_v" in result.dims
+    assert "eta_rho" not in result.dims
+
+    # Coordinates dropped
+    for coord in ("lat_rho", "lon_rho"):
+        assert coord not in result.coords
+
+    # Shape check
+    assert result.sizes["eta_v"] == sample_rho_field.sizes["eta_rho"] - 1
+
+
+def test_u_to_rho_wrapper_additive(sample_u_field: xr.DataArray):
+    result = interpolate_from_u_to_rho(sample_u_field, method="additive")
+
+    # Dimension swap
+    assert "xi_rho" in result.dims
+    assert "xi_u" not in result.dims
+
+    # Coordinates dropped
+    for coord in ("lat_u", "lon_u"):
+        assert coord not in result.coords
+
+    # Shape: one more along xi due to padding
+    assert result.sizes["xi_rho"] == sample_u_field.sizes["xi_u"] + 1
+
+
+def test_v_to_rho_wrapper_additive(sample_v_field: xr.DataArray):
+    result = interpolate_from_v_to_rho(sample_v_field, method="additive")
+
+    # Dimension swap
+    assert "eta_rho" in result.dims
+    assert "eta_v" not in result.dims
+
+    # Coordinates dropped
+    for coord in ("lat_v", "lon_v"):
+        assert coord not in result.coords
+
+    # Shape: one more along eta due to padding
+    assert result.sizes["eta_rho"] == sample_v_field.sizes["eta_v"] + 1
+
+
+# -------------------------
+# Error handling
+# -------------------------
+
+
+def test_invalid_method_raises(
+    sample_rho_field: xr.DataArray,
+    sample_u_field: xr.DataArray,
+    sample_v_field: xr.DataArray,
+):
+    with pytest.raises(NotImplementedError):
+        interpolate_from_rho_to_u(sample_rho_field, method="unsupported")
+
+    with pytest.raises(NotImplementedError):
+        interpolate_from_rho_to_v(sample_rho_field, method="unsupported")
+
+    with pytest.raises(NotImplementedError):
+        interpolate_from_u_to_rho(sample_u_field, method="unsupported")
+
+    with pytest.raises(NotImplementedError):
+        interpolate_from_v_to_rho(sample_v_field, method="unsupported")
+
+
+# Test rotate_velocities
+@pytest.fixture
+def sample_velocities_centered():
+    """Create a centered-grid velocity field with random values and grid angle."""
+    np.random.seed(42)  # For reproducibility
+
+    eta_rho, xi_rho = 10, 15
+
+    u = xr.DataArray(
+        np.random.rand(eta_rho, xi_rho),
+        dims=("eta_rho", "xi_rho"),
+        coords={
+            "eta_rho": np.arange(eta_rho),
+            "xi_rho": np.arange(xi_rho),
+        },
+    )
+
+    v = xr.DataArray(
+        np.random.rand(eta_rho, xi_rho),
+        dims=("eta_rho", "xi_rho"),
+        coords={
+            "eta_rho": np.arange(eta_rho),
+            "xi_rho": np.arange(xi_rho),
+        },
+    )
+
+    angle = xr.DataArray(
+        np.random.rand(eta_rho, xi_rho) * np.pi / 2
+        - np.pi / 4,  # random angles in [-45°, 45°]
+        dims=("eta_rho", "xi_rho"),
+        coords={
+            "eta_rho": np.arange(eta_rho),
+            "xi_rho": np.arange(xi_rho),
+        },
+    )
+
+    return u, v, angle
+
+
+def test_rotate_velocities_roundtrip(sample_velocities_centered):
+    """Test rotation to grid and back recovers original velocities."""
+    u, v, angle = sample_velocities_centered
+
+    # Rotate forward: lat-lon → model grid
+    u_rot, v_rot = rotate_velocities(
+        u, v, angle, interpolate_before=False, interpolate_after=False
+    )
+
+    # Rotate backward: model grid → lat-lon
+    u_back, v_back = rotate_velocities(
+        u_rot, v_rot, -angle, interpolate_before=False, interpolate_after=False
+    )
+
+    np.testing.assert_allclose(u.values, u_back.values)
+    np.testing.assert_allclose(v.values, v_back.values)