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

roms_tools/utils.py

@@ -10,6 +10,7 @@ from pathlib import Path
 from typing import TypeAlias
 
 import numpy as np
+import pandas as pd
 import xarray as xr
 
 from roms_tools.constants import R_EARTH
@@ -476,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})
 
-    field_interpolated = field_interpolated.swap_dims({"xi_rho": "xi_u"})
+    return interp
 
-    return field_interpolated
+
+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, method="additive"):
-    """Interpolates the given field from rho points to v points.
+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"],
+    )
 
-    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 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,
+    )
+
+
+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 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:
@@ -1031,3 +1135,197 @@ def get_pkg_error_msg(purpose: str, package_name: str, option_name: str) -> str:
           • `pip install {package_name}` or
           • `conda install {package_name}`
         Alternatively, install `roms-tools` with conda to include all dependencies.""")
+
+
+def wrap_longitudes(ds: xr.Dataset, straddle: bool) -> xr.Dataset:
+    """
+    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 coordinates (e.g., lon_rho, lon_u, lon_v).
+    straddle : bool
+        - True: force longitudes into [-180, 180]
+        - False: force longitudes into [0, 360]
+
+    Returns
+    -------
+    xr.Dataset
+        A new dataset with adjusted longitude coordinates.
+    """
+    lon_coords = ["lon_rho", "lon_u", "lon_v"]
+
+    for lon_name in lon_coords:
+        if lon_name not in ds.coords:
+            continue  # skip missing coordinate
+
+        lon = ds.coords[lon_name]
+
+        if straddle:
+            # wrap into [-180, 180]
+            lon_wrapped = xr.where(lon > 180, lon - 360, lon)
+        else:
+            # wrap into [0, 360]
+            lon_wrapped = xr.where(lon < 0, lon + 360, lon)
+
+        # preserve attributes
+        lon_wrapped.attrs = lon.attrs.copy()
+
+        # reassign explicitly as a coordinate
+        ds = ds.assign_coords({lon_name: lon_wrapped})
+
+    return ds
+
+
+def interpolate_from_climatology(
+    field: xr.DataArray | xr.Dataset,
+    time_dim: str,
+    time_coord: str,
+    time: xr.DataArray | pd.DatetimeIndex,
+) -> xr.DataArray | xr.Dataset:
+    """Interpolates a climatological field to specified time points.
+
+    This function interpolates the input `field` based on `day_of_year` values
+    extracted from the provided `time` points. If `field` is an `xarray.Dataset`,
+    interpolation is applied to all its data variables individually.
+
+    Parameters
+    ----------
+    field : xarray.DataArray or xarray.Dataset
+        The input field to be interpolated.
+        - If `field` is an `xarray.DataArray`, it must have a time dimension identified by `time_dim_name`.
+        - If `field` is an `xarray.Dataset`, all variables within the dataset are interpolated along `time_dim_name`.
+        The time dimension is assumed to represent `day_of_year` for climatological purposes.
+    time_dim : str
+        The name of the time dimension in `field`.
+    time_coord : str
+        The name of the time coordinate in `field`.
+    time : xarray.DataArray or pandas.DatetimeIndex
+        The target time points for interpolation. These are internally converted to `day_of_year`
+        before performing interpolation.
+
+    Returns
+    -------
+    xarray.DataArray or xarray.Dataset
+        The interpolated field, maintaining the same type (`xarray.DataArray` or `xarray.Dataset`)
+        but aligned to the specified `time` values.
+
+    Notes
+    -----
+    - This function assumes that `field` represents a climatological dataset, where time is expressed as `day_of_year` (1-365).
+    - The `time` input is automatically converted to `day_of_year`, so manual conversion is not required before calling this function.
+    """
+
+    def np_times_to_fractional_days(
+        np_times: np.ndarray | pd.DatetimeIndex | np.datetime64 | pd.Timestamp,
+    ) -> np.ndarray:
+        """Convert datetime(s) to fractional day-of-year values."""
+        pd_times = pd.to_datetime(np_times)
+
+        # scalar input -> make it a 1-element array
+        if np.isscalar(pd_times):
+            pd_times = np.array([pd_times])
+
+        fractional_days = pd_times.dayofyear + (
+            pd_times.hour / 24 + pd_times.minute / 1440 + pd_times.second / 86400
+        )
+        return (
+            fractional_days.values
+            if hasattr(fractional_days, "values")
+            else np.array(fractional_days)
+        )
+
+    def interpolate_single_field(data_array: xr.DataArray) -> xr.DataArray:
+        if isinstance(time, xr.DataArray):
+            day_of_year = time.dt.dayofyear
+        else:
+            day_of_year = np_times_to_fractional_days(time)
+
+        data_array_interpolated = interpolate_cyclic_time(
+            data_array, time_dim, time_coord, day_of_year
+        )
+
+        # expand dims if single element
+        if day_of_year.size == 1:
+            data_array_interpolated = data_array_interpolated.expand_dims({time_dim: 1})
+        return data_array_interpolated
+
+    if isinstance(field, xr.DataArray):
+        return interpolate_single_field(field)
+    elif isinstance(field, xr.Dataset):
+        interpolated_data_vars = {
+            var: interpolate_single_field(data_array)
+            for var, data_array in field.data_vars.items()
+        }
+        return xr.Dataset(interpolated_data_vars, attrs=field.attrs)
+
+    else:
+        raise TypeError("Input 'field' must be an xarray.DataArray or xarray.Dataset.")
+
+
+def interpolate_cyclic_time(
+    data_array: xr.DataArray,
+    time_dim: str,
+    time_coord: str,
+    day_of_year: int | float | np.ndarray | xr.DataArray | Sequence[int | float],
+) -> xr.DataArray:
+    """Interpolates a DataArray cyclically across the start and end of the year.
+
+    This function extends the data cyclically by appending the last time step
+    (shifted back by one year) at the beginning and the first time step
+    (shifted forward by one year) at the end. It then performs linear interpolation
+    to match the specified `day_of_year` values.
+
+    Parameters
+    ----------
+    data_array : xr.DataArray
+        The input data array containing a time-like dimension.
+    time_dim : str
+        The name of the time dimension in the dataset.
+    time_coord : str
+        The name of the time coordinate in the dataset.
+    day_of_year : Union[int, float, np.ndarray, xr.DataArray, Sequence[Union[int, float]]]
+        The target day(s) of the year for interpolation. This can be:
+        - A single integer or float representing the day of the year.
+        - A NumPy array or xarray DataArray containing multiple days.
+        - A list or tuple of integers or floats for multiple target days.
+
+    Returns
+    -------
+    xr.DataArray
+        The interpolated DataArray, ensuring cyclic continuity across year boundaries.
+
+    Notes
+    -----
+    - This function is useful for interpolating climatological data, where the time axis
+      represents a repeating annual cycle.
+    - The `day_of_year` values should be within the range [1, 365] or [1, 366] for leap years.
+    """
+    # Concatenate across the beginning and end of the year
+    time_concat = xr.concat(
+        [
+            data_array[time_coord][-1] - 365.25,  # Shift last time backward
+            data_array[time_coord],
+            data_array[time_coord][0] + 365.25,  # Shift first time forward
+        ],
+        dim=time_dim,
+    )
+
+    data_array_concat = xr.concat(
+        [
+            data_array.isel(**{time_dim: -1}),  # Append last value at the beginning
+            data_array,
+            data_array.isel(**{time_dim: 0}),  # Append first value at the end
+        ],
+        dim=time_dim,
+    )
+    data_array_concat[time_dim] = time_concat
+
+    # Interpolate to specified times
+    data_array_interpolated = data_array_concat.interp(
+        **{time_dim: day_of_year}, method="linear"
+    )
+
+    return data_array_interpolated

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: roms-tools
-Version: 3.3.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)