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

roms_tools/tests/test_utils.py

@@ -10,6 +10,7 @@ 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,
@@ -18,7 +19,13 @@ from roms_tools.utils import (
     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,
 )
 
 
@@ -319,3 +326,283 @@ def test_interpolate_from_real_climatology(use_dask):
 
     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)

roms_tools/utils.py

@@ -477,96 +477,199 @@ def load_data(
     return ds
 
 
-def interpolate_from_rho_to_u(field, method="additive"):
-    """Interpolates the given field from rho points to u points.
-
-    This function performs an interpolation from the rho grid (cell centers) to the u grid
-    (cell edges in the xi direction). Depending on the chosen method, it either averages
-    (additive) or multiplies (multiplicative) the field values between adjacent rho points
-    along the xi dimension. It also handles the removal of unnecessary coordinate variables
-    and updates the dimensions accordingly.
+def _interpolate_generic(
+    field: xr.DataArray,
+    dim_in: str,
+    dim_out: str,
+    method: str = "additive",
+    drop_coords: Sequence[str] | None = None,
+    pad_end: bool = False,
+) -> xr.DataArray:
+    """
+    Generic interpolation along one horizontal dimension.
 
     Parameters
     ----------
     field : xr.DataArray
-        The input data array on the rho grid to be interpolated. It is assumed to have a dimension
-        named "xi_rho".
-
-    method : str, optional, default='additive'
-        The method to use for interpolation. Options are:
-        - 'additive': Average the field values between adjacent rho points.
-        - 'multiplicative': Multiply the field values between adjacent rho points. Appropriate for
-          binary masks.
+        Input array to interpolate.
+    dim_in : str
+        Dimension along which to interpolate (e.g., "xi_rho").
+    dim_out : str
+        New dimension name after interpolation (e.g., "xi_u").
+    method : str, default "additive"
+        Interpolation method:
+        - "additive": average adjacent points
+        - "multiplicative": multiply adjacent points (useful for masks)
+    drop_coords : Sequence[str] or None, optional
+        Coordinate variables to drop (e.g., ["lat_rho", "lon_rho"]).
+    pad_end : bool, default False
+        Whether to pad the last point with NaN (useful when interpolating back to rho grid).
 
     Returns
     -------
-    field_interpolated : xr.DataArray
-        The interpolated data array on the u grid with the dimension "xi_u".
+    xr.DataArray
+        Interpolated array with dimension `dim_out`.
     """
-    if method == "additive":
-        field_interpolated = 0.5 * (field + field.shift(xi_rho=1)).isel(
-            xi_rho=slice(1, None)
+    if not isinstance(field, xr.DataArray):
+        raise TypeError(
+            "_interpolate_generic expects an xarray.DataArray, "
+            f"got {type(field).__name__}"
         )
+
+    if drop_coords:
+        for coord in drop_coords:
+            if coord in field.coords:
+                field = field.drop_vars(coord)
+
+    if method == "additive":
+        interp = 0.5 * (field + field.shift(**{dim_in: 1}))
     elif method == "multiplicative":
-        field_interpolated = (field * field.shift(xi_rho=1)).isel(xi_rho=slice(1, None))
+        interp = field * field.shift(**{dim_in: 1})
     else:
         raise NotImplementedError(f"Unsupported method '{method}' specified.")
 
-    vars_to_drop = ["lat_rho", "lon_rho", "eta_rho", "xi_rho"]
-    for var in vars_to_drop:
-        if var in field_interpolated.coords:
-            field_interpolated = field_interpolated.drop_vars(var)
+    if pad_end:
+        pad_shape = {d: field.sizes[d] for d in field.dims}
+        pad_shape[dim_in] = 1
+        pad = xr.DataArray(
+            np.nan * np.ones(tuple(pad_shape[d] for d in field.dims)),
+            dims=field.dims,
+        )
+        interp = xr.concat([interp, pad], dim=dim_in)
+    else:
+        interp = interp.isel({dim_in: slice(1, None)})
+
+    interp = interp.swap_dims({dim_in: dim_out})
+
+    return interp
+
+
+def interpolate_from_rho_to_u(
+    field: xr.DataArray, method: str = "additive"
+) -> xr.DataArray:
+    """Interpolate a field from rho points to u points (xi direction)."""
+    return _interpolate_generic(
+        field,
+        dim_in="xi_rho",
+        dim_out="xi_u",
+        method=method,
+        drop_coords=["lat_rho", "lon_rho", "eta_rho", "xi_rho"],
+    )
+
+
+def interpolate_from_rho_to_v(
+    field: xr.DataArray, method: str = "additive"
+) -> xr.DataArray:
+    """Interpolate a field from rho points to v points (eta direction)."""
+    return _interpolate_generic(
+        field,
+        dim_in="eta_rho",
+        dim_out="eta_v",
+        method=method,
+        drop_coords=["lat_rho", "lon_rho", "eta_rho", "xi_rho"],
+    )
+
 
-    field_interpolated = field_interpolated.swap_dims({"xi_rho": "xi_u"})
+def interpolate_from_u_to_rho(
+    field: xr.DataArray, method: str = "additive"
+) -> xr.DataArray:
+    """Interpolate a field from u points back to rho points (xi direction)."""
+    return _interpolate_generic(
+        field,
+        dim_in="xi_u",
+        dim_out="xi_rho",
+        method=method,
+        drop_coords=["lat_u", "lon_u", "eta_rho", "xi_u"],
+        pad_end=True,
+    )
 
-    return field_interpolated
 
+def interpolate_from_v_to_rho(
+    field: xr.DataArray, method: str = "additive"
+) -> xr.DataArray:
+    """Interpolate a field from v points back to rho points (eta direction)."""
+    return _interpolate_generic(
+        field,
+        dim_in="eta_v",
+        dim_out="eta_rho",
+        method=method,
+        drop_coords=["lat_v", "lon_v", "eta_v", "xi_rho"],
+        pad_end=True,
+    )
 
-def interpolate_from_rho_to_v(field, method="additive"):
-    """Interpolates the given field from rho points to v points.
 
-    This function performs an interpolation from the rho grid (cell centers) to the v grid
-    (cell edges in the eta direction). Depending on the chosen method, it either averages
-    (additive) or multiplies (multiplicative) the field values between adjacent rho points
-    along the eta dimension. It also handles the removal of unnecessary coordinate variables
-    and updates the dimensions accordingly.
+def rotate_velocities(
+    u: xr.DataArray,
+    v: xr.DataArray,
+    angle: xr.DataArray,
+    interpolate_before: bool = False,
+    interpolate_after: bool = False,
+) -> tuple[xr.DataArray, xr.DataArray]:
+    """
+    Rotate horizontal velocity components to align with a rotated grid.
+
+    This function rotates zonal (u) and meridional (v) velocity components
+    using a grid angle field. It can be used to:
+
+    1. Rotate model velocities from the ROMS model grid to a lat-lon reference frame.
+    2. Rotate lat-lon velocities onto the ROMS model grid.
+
+    Optionally, velocities can be interpolated between staggered C-grid
+    locations (u-, v-, and rho-points) before and/or after rotation.
 
     Parameters
     ----------
-    field : xr.DataArray
-        The input data array on the rho grid to be interpolated. It is assumed to have a dimension
-        named "eta_rho".
-
-    method : str, optional, default='additive'
-        The method to use for interpolation. Options are:
-        - 'additive': Average the field values between adjacent rho points.
-        - 'multiplicative': Multiply the field values between adjacent rho points. Appropriate for
-          binary masks.
+    u : xarray.DataArray
+        Zonal (east-west) velocity component defined on u-points.
+    v : xarray.DataArray
+        Meridional (north-south) velocity component defined on v-points.
+    angle : xarray.DataArray
+        Grid orientation angle in radians, defined at rho-points. This is the
+        ROMS grid angle: the angle between the model xi-direction and true east.
+        Positive values indicate that the model grid is rotated counterclockwise
+        relative to east (which is mathematically equivalent to rotating velocity
+        vectors clockwise).
+        The rotation transforms velocity components between earth-relative
+        (east/north) and grid-relative (xi/eta) coordinates. To reverse the
+        transformation (e.g., model → lat-lon), provide ``-angle``.
+    interpolate_before : bool, optional
+        If True, interpolate ``u`` and ``v`` to rho-points before rotation.
+        Default is False.
+    interpolate_after : bool, optional
+        If True, interpolate the rotated velocities back to u- and v-points
+        after rotation. Default is True.
 
     Returns
     -------
-    field_interpolated : xr.DataArray
-        The interpolated data array on the v grid with the dimension "eta_v".
+    u_rot, v_rot : tuple of xarray.DataArray
+        Rotated velocity components. If ``interpolate_after`` is True, ``u_rot``
+        is defined on u-points and ``v_rot`` on v-points; otherwise, both are
+        defined at rho-points.
+
+    Notes
+    -----
+    The rotation follows the standard ROMS convention:
+    - ``u_rot =  u * cos(angle) + v * sin(angle)``
+    - ``v_rot =  v * cos(angle) - u * sin(angle)``
+    This function is versatile and can be used for both directions of rotation:
+    - Lat-lon → model grid: provide the grid angle.
+    - Model grid → lat-lon: provide the negative of the grid angle.
     """
-    if method == "additive":
-        field_interpolated = 0.5 * (field + field.shift(eta_rho=1)).isel(
-            eta_rho=slice(1, None)
-        )
-    elif method == "multiplicative":
-        field_interpolated = (field * field.shift(eta_rho=1)).isel(
-            eta_rho=slice(1, None)
-        )
-    else:
-        raise NotImplementedError(f"Unsupported method '{method}' specified.")
+    # Interpolate to rho-points
+    if interpolate_before:
+        u = interpolate_from_u_to_rho(u)
+        v = interpolate_from_v_to_rho(v)
 
-    vars_to_drop = ["lat_rho", "lon_rho", "eta_rho", "xi_rho"]
-    for var in vars_to_drop:
-        if var in field_interpolated.coords:
-            field_interpolated = field_interpolated.drop_vars(var)
+    # Rotate velocities to grid orientation
+    u_rot = u * np.cos(angle) + v * np.sin(angle)
+    v_rot = v * np.cos(angle) - u * np.sin(angle)
 
-    field_interpolated = field_interpolated.swap_dims({"eta_rho": "eta_v"})
+    # Interpolate to u- and v-points
+    if interpolate_after:
+        u_rot = interpolate_from_rho_to_u(u_rot)
+        v_rot = interpolate_from_rho_to_v(v_rot)
 
-    return field_interpolated
+    return u_rot, v_rot
 
 
 def transpose_dimensions(da: xr.DataArray) -> xr.DataArray:
@@ -1036,13 +1139,13 @@ def get_pkg_error_msg(purpose: str, package_name: str, option_name: str) -> str:
 
 def wrap_longitudes(ds: xr.Dataset, straddle: bool) -> xr.Dataset:
     """
-    Safely adjust longitude variables for datasets that may or may not cross
-    the dateline. Only modifies longitude-like variables that are present.
+    Safely adjust longitude coordinates for datasets that may or may not cross
+    the dateline. Only modifies longitude-like coordinates that are present.
 
     Parameters
     ----------
     ds : xr.Dataset
-        Dataset containing longitude variables (e.g., lon_rho, lon_u, lon_v).
+        Dataset containing longitude coordinates (e.g., lon_rho, lon_u, lon_v).
     straddle : bool
         - True: force longitudes into [-180, 180]
         - False: force longitudes into [0, 360]
@@ -1050,15 +1153,15 @@ def wrap_longitudes(ds: xr.Dataset, straddle: bool) -> xr.Dataset:
     Returns
     -------
     xr.Dataset
-        A new dataset with adjusted longitude values.
+        A new dataset with adjusted longitude coordinates.
     """
-    lon_vars = ["lon_rho", "lon_u", "lon_v"]
+    lon_coords = ["lon_rho", "lon_u", "lon_v"]
 
-    for lon_dim in lon_vars:
-        if lon_dim not in ds:
+    for lon_name in lon_coords:
+        if lon_name not in ds.coords:
             continue  # skip missing coordinate
 
-        lon = ds[lon_dim]
+        lon = ds.coords[lon_name]
 
         if straddle:
             # wrap into [-180, 180]
@@ -1068,8 +1171,10 @@ def wrap_longitudes(ds: xr.Dataset, straddle: bool) -> xr.Dataset:
             lon_wrapped = xr.where(lon < 0, lon + 360, lon)
 
         # preserve attributes
-        lon_wrapped.attrs.update(lon.attrs)
-        ds[lon_dim] = lon_wrapped
+        lon_wrapped.attrs = lon.attrs.copy()
+
+        # reassign explicitly as a coordinate
+        ds = ds.assign_coords({lon_name: lon_wrapped})
 
     return ds
 

{roms_tools-3.4.0.dist-info → roms_tools-3.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: roms-tools
-Version: 3.4.0
+Version: 3.5.0
 Summary: Tools for running and analysing UCLA-ROMS simulations
 Author-email: Nora Loose <nora.loose@gmail.com>, Thomas Nicholas <tom@cworthy.org>, Scott Eilerman <scott.eilerman@cworthy.org>
 License: Apache-2
@@ -64,7 +64,7 @@ Dynamic: license-file
 
 ## Overview
 
-A suite of Python tools for setting up and analyzing a [UCLA-ROMS](https://github.com/CESR-lab/ucla-roms) simulation with or without [MARBL biogeochemistry](https://marbl-ecosys.github.io/versions/latest_release/index.html).
+A suite of Python tools for setting up and analyzing a [UCLA-ROMS](https://github.com/CWorthy-ocean/ucla-roms) simulation with or without [MARBL biogeochemistry](https://marbl-ecosys.github.io/versions/latest_release/index.html).
 
 ## Installation
 
@@ -152,5 +152,4 @@ We also accept contributions in the form of Pull Requests.
 
 ## See also
 
-- [ROMS source code](https://github.com/CESR-lab/ucla-roms)
 - [C-Star](https://github.com/CWorthy-ocean/C-Star)