roms-tools 2.2.1__py3-none-any.whl → 2.3.0__py3-none-any.whl

roms_tools/utils.py

@@ -4,6 +4,8 @@ import numpy as np
 import xarray as xr
 from typing import Union
 from pathlib import Path
+import re
+import glob
 
 
 def partition(
@@ -333,3 +335,302 @@ def partition_netcdf(
     xr.save_mfdataset(partitioned_datasets, paths_to_partitioned_files)
 
     return paths_to_partitioned_files
+
+
+def _load_data(
+    filename,
+    dim_names,
+    use_dask,
+    time_chunking=True,
+    decode_times=True,
+    force_combine_nested=False,
+):
+    """Load dataset from the specified file.
+
+    Parameters
+    ----------
+    filename : Union[str, Path, List[Union[str, Path]]]
+        The path to the data file(s). Can be a single string (with or without wildcards), a single Path object,
+        or a list of strings or Path objects containing multiple files.
+    dim_names : Dict[str, str], optional
+        Dictionary specifying the names of dimensions in the dataset.
+        Required only for lat-lon datasets to map dimension names like "latitude" and "longitude".
+        For ROMS datasets, this parameter can be omitted, as default ROMS dimensions ("eta_rho", "xi_rho", "s_rho") are assumed.
+    use_dask: bool
+        Indicates whether to use dask for chunking. If True, data is loaded with dask; if False, data is loaded eagerly. Defaults to False.
+    time_chunking : bool, optional
+        If True and `use_dask=True`, the data will be chunked along the time dimension with a chunk size of 1.
+        If False, the data will not be chunked explicitly along the time dimension, but will follow the default auto chunking scheme. This option is useful for ROMS restart files.
+        Defaults to True.
+    decode_times: bool, optional
+        If True, decode times encoded in the standard NetCDF datetime format into datetime objects. Otherwise, leave them encoded as numbers.
+        Defaults to True.
+    force_combine_nested: bool, optional
+        If True, forces the use of nested combination (`combine_nested`) regardless of whether wildcards are used.
+        Defaults to False.
+
+    Returns
+    -------
+    ds : xr.Dataset
+        The loaded xarray Dataset containing the forcing data.
+
+    Raises
+    ------
+    FileNotFoundError
+        If the specified file does not exist.
+    ValueError
+        If a list of files is provided but dim_names["time"] is not available or use_dask=False.
+    """
+    if dim_names is None:
+        dim_names = {}
+
+    # Precompile the regex for matching wildcard characters
+    wildcard_regex = re.compile(r"[\*\?\[\]]")
+
+    # Convert Path objects to strings
+    if isinstance(filename, (str, Path)):
+        filename_str = str(filename)
+    elif isinstance(filename, list):
+        filename_str = [str(f) for f in filename]
+    else:
+        raise ValueError("filename must be a string, Path, or a list of strings/Paths.")
+
+    # Handle the case when filename is a string
+    contains_wildcard = False
+    if isinstance(filename_str, str):
+        contains_wildcard = bool(wildcard_regex.search(filename_str))
+        if contains_wildcard:
+            matching_files = glob.glob(filename_str)
+            if not matching_files:
+                raise FileNotFoundError(
+                    f"No files found matching the pattern '{filename_str}'."
+                )
+        else:
+            matching_files = [filename_str]
+
+    # Handle the case when filename is a list
+    elif isinstance(filename_str, list):
+        contains_wildcard = any(wildcard_regex.search(f) for f in filename_str)
+        if contains_wildcard:
+            matching_files = []
+            for f in filename_str:
+                files = glob.glob(f)
+                if not files:
+                    raise FileNotFoundError(
+                        f"No files found matching the pattern '{f}'."
+                    )
+                matching_files.extend(files)
+        else:
+            matching_files = filename_str
+
+    # Sort the matching files
+    matching_files = sorted(matching_files)
+
+    # Check if time dimension is available when multiple files are provided
+    if isinstance(filename_str, list) and "time" not in dim_names:
+        raise ValueError(
+            "A list of files is provided, but time dimension is not available. "
+            "A time dimension must be available to concatenate the files."
+        )
+
+    # Determine the kwargs for combining datasets
+    if force_combine_nested:
+        kwargs = {"combine": "nested", "concat_dim": dim_names["time"]}
+    elif contains_wildcard or len(matching_files) == 1:
+        kwargs = {"combine": "by_coords"}
+    else:
+        kwargs = {"combine": "nested", "concat_dim": dim_names["time"]}
+
+    # Base kwargs used for dataset combination
+    combine_kwargs = {
+        "coords": "minimal",
+        "compat": "override",
+        "combine_attrs": "override",
+    }
+
+    if use_dask:
+
+        if "latitude" in dim_names and "longitude" in dim_names:
+            # for lat-lon datasets
+            chunks = {
+                dim_names["latitude"]: -1,
+                dim_names["longitude"]: -1,
+            }
+        else:
+            # For ROMS datasets
+            chunks = {
+                "eta_rho": -1,
+                "eta_v": -1,
+                "xi_rho": -1,
+                "xi_u": -1,
+                "s_rho": -1,
+            }
+
+        if "depth" in dim_names:
+            chunks[dim_names["depth"]] = -1
+        if "time" in dim_names and time_chunking:
+            chunks[dim_names["time"]] = 1
+
+        ds = xr.open_mfdataset(
+            matching_files,
+            decode_times=decode_times,
+            chunks=chunks,
+            **combine_kwargs,
+            **kwargs,
+        )
+
+        # Rechunk the dataset along the tidal constituent dimension ("ntides") after loading
+        # because the original dataset does not have a chunk size of 1 along this dimension.
+        if "ntides" in dim_names:
+            ds = ds.chunk({dim_names["ntides"]: 1})
+
+    else:
+        ds_list = []
+        for file in matching_files:
+            ds = xr.open_dataset(file, decode_times=decode_times, chunks=None)
+            ds_list.append(ds)
+
+        if kwargs["combine"] == "by_coords":
+            ds = xr.combine_by_coords(ds_list, **combine_kwargs)
+        elif kwargs["combine"] == "nested":
+            ds = xr.combine_nested(
+                ds_list, concat_dim=kwargs["concat_dim"], **combine_kwargs
+            )
+
+    if "time" in dim_names and dim_names["time"] not in ds.dims:
+        ds = ds.expand_dims(dim_names["time"])
+
+    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.
+
+    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.
+
+    Returns
+    -------
+    field_interpolated : xr.DataArray
+        The interpolated data array on the u grid with the dimension "xi_u".
+    """
+
+    if method == "additive":
+        field_interpolated = 0.5 * (field + field.shift(xi_rho=1)).isel(
+            xi_rho=slice(1, None)
+        )
+    elif method == "multiplicative":
+        field_interpolated = (field * field.shift(xi_rho=1)).isel(xi_rho=slice(1, None))
+    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)
+
+    field_interpolated = field_interpolated.swap_dims({"xi_rho": "xi_u"})
+
+    return field_interpolated
+
+
+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.
+
+    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.
+
+    Returns
+    -------
+    field_interpolated : xr.DataArray
+        The interpolated data array on the v grid with the dimension "eta_v".
+    """
+
+    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.")
+
+    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)
+
+    field_interpolated = field_interpolated.swap_dims({"eta_rho": "eta_v"})
+
+    return field_interpolated
+
+
+def transpose_dimensions(da: xr.DataArray) -> xr.DataArray:
+    """Transpose the dimensions of an xarray.DataArray to ensure that 'time', any
+    dimension starting with 's_', 'eta_', and 'xi_' are ordered first, followed by the
+    remaining dimensions in their original order.
+
+    Parameters
+    ----------
+    da : xarray.DataArray
+        The input DataArray whose dimensions are to be reordered.
+
+    Returns
+    -------
+    xarray.DataArray
+        The DataArray with dimensions reordered so that 'time', 's_*', 'eta_*',
+        and 'xi_*' are first, in that order, if they exist.
+    """
+
+    # List of preferred dimension patterns
+    preferred_order = ["time", "s_", "eta_", "xi_"]
+
+    # Get the existing dimensions in the DataArray
+    dims = list(da.dims)
+
+    # Collect dimensions that match any of the preferred patterns
+    matched_dims = []
+    for pattern in preferred_order:
+        # Find dimensions that start with the pattern
+        matched_dims += [dim for dim in dims if dim.startswith(pattern)]
+
+    # Create a new order: first the matched dimensions, then the rest
+    remaining_dims = [dim for dim in dims if dim not in matched_dims]
+    new_order = matched_dims + remaining_dims
+
+    # Transpose the DataArray to the new order
+    transposed_da = da.transpose(*new_order)
+
+    return transposed_da

roms_tools/vertical_coordinate.py

@@ -0,0 +1,306 @@
+import numpy as np
+import xarray as xr
+from roms_tools.utils import (
+    transpose_dimensions,
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
+)
+
+
+def compute_cs(sigma, theta_s, theta_b):
+    """Compute the S-coordinate stretching curves according to Shchepetkin and
+    McWilliams (2009).
+
+    Parameters
+    ----------
+    sigma : np.ndarray or float
+        The sigma-coordinate values.
+    theta_s : float
+        The surface control parameter.
+    theta_b : float
+        The bottom control parameter.
+
+    Returns
+    -------
+    C : np.ndarray or float
+        The stretching curve values.
+
+    Raises
+    ------
+    ValueError
+        If theta_s or theta_b are not within the valid range.
+    """
+    if not (0 < theta_s <= 10):
+        raise ValueError("theta_s must be between 0 and 10.")
+    if not (0 < theta_b <= 4):
+        raise ValueError("theta_b must be between 0 and 4.")
+
+    C = (1 - np.cosh(theta_s * sigma)) / (np.cosh(theta_s) - 1)
+    C = (np.exp(theta_b * C) - 1) / (1 - np.exp(-theta_b))
+
+    return C
+
+
+def sigma_stretch(theta_s, theta_b, N, type):
+    """Compute sigma and stretching curves based on the type and parameters.
+
+    Parameters
+    ----------
+    theta_s : float
+        The surface control parameter.
+    theta_b : float
+        The bottom control parameter.
+    N : int
+        The number of vertical levels.
+    type : str
+        The type of sigma ('w' for vertical velocity points, 'r' for rho-points).
+
+    Returns
+    -------
+    cs : xr.DataArray
+        The stretching curve values.
+    sigma : xr.DataArray
+        The sigma-coordinate values.
+
+    Raises
+    ------
+    ValueError
+        If the type is not 'w' or 'r'.
+    """
+    if type == "w":
+        k = xr.DataArray(np.arange(N + 1), dims="s_w")
+        sigma = (k - N) / N
+    elif type == "r":
+        k = xr.DataArray(np.arange(1, N + 1), dims="s_rho")
+        sigma = (k - N - 0.5) / N
+    else:
+        raise ValueError(
+            "Type must be either 'w' for vertical velocity points or 'r' for rho-points."
+        )
+
+    cs = compute_cs(sigma, theta_s, theta_b)
+
+    return cs, sigma
+
+
+def compute_depth(zeta, h, hc, cs, sigma):
+    """Compute the depth at different sigma levels.
+
+    Parameters
+    ----------
+    zeta : xr.DataArray or scalar
+        The sea surface height.
+    h : xr.DataArray
+        The depth of the sea bottom.
+    hc : float
+        The critical depth.
+    cs : xr.DataArray
+        The stretching curve values.
+    sigma : xr.DataArray
+        The sigma-coordinate values.
+
+    Returns
+    -------
+    z : xr.DataArray
+        The depth at different sigma levels.
+    """
+
+    z = (hc * sigma + h * cs) / (hc + h)
+    z = zeta + (zeta + h) * z
+
+    z = -transpose_dimensions(z)
+
+    return z
+
+
+def add_depth_coordinates_to_dataset(
+    ds: "xr.Dataset",
+    grid_ds: "xr.Dataset",
+    depth_type: str,
+    locations: list[str] = ["rho", "u", "v"],
+) -> None:
+    """Add computed vertical depth coordinates to a dataset for specified grid
+    locations.
+
+    This function computes vertical depth coordinates (layer or interface) and updates
+    the provided dataset with these coordinates for the specified grid locations. If
+    the dataset already contains depth coordinates for all specified locations, the function
+    does nothing.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Target dataset to which computed depth coordinates will be added.
+        If the `zeta` variable is not present, static vertical coordinates are used.
+
+    grid_ds : xr.Dataset
+        Grid dataset containing bathymetry, stretching curves, and parameters.
+
+    depth_type : str
+        Type of depth coordinates to compute. Options are:
+        - "layer": Layer depth coordinates.
+        - "interface": Interface depth coordinates.
+
+    locations : list of str, optional
+        List of locations for which to compute depth coordinates. Default is ["rho", "u", "v"].
+    """
+    required_vars = [f"{depth_type}_depth_{loc}" for loc in locations]
+
+    if all(var in ds for var in required_vars):
+        return  # Depth coordinates already exist
+
+    # Compute or interpolate depth coordinates
+    if f"{depth_type}_depth_rho" in ds:
+        depth_rho = ds[f"{depth_type}_depth_rho"]
+    else:
+        h = grid_ds["h"]
+        zeta = ds.get("zeta", 0)
+        if depth_type == "layer":
+            Cs = grid_ds["Cs_r"]
+            sigma = grid_ds["sigma_r"]
+        elif depth_type == "interface":
+            Cs = grid_ds["Cs_w"]
+            sigma = grid_ds["sigma_w"]
+        depth_rho = compute_depth(zeta, h, grid_ds.attrs["hc"], Cs, sigma)
+        depth_rho.attrs.update(
+            {"long_name": f"{depth_type} depth at rho-points", "units": "m"}
+        )
+        ds[f"{depth_type}_depth_rho"] = depth_rho
+
+    # Interpolate depth to other locations
+    for loc in locations:
+        if loc == "rho":
+            continue
+
+        interp_func = (
+            interpolate_from_rho_to_u if loc == "u" else interpolate_from_rho_to_v
+        )
+        depth_loc = interp_func(depth_rho)
+        depth_loc.attrs.update(
+            {"long_name": f"{depth_type} depth at {loc}-points", "units": "m"}
+        )
+        ds[f"{depth_type}_depth_{loc}"] = depth_loc
+
+
+def compute_depth_coordinates(
+    ds: "xr.Dataset",
+    grid_ds: "xr.Dataset",
+    depth_type: str,
+    location: str,
+    s: int = None,
+    eta: int = None,
+    xi: int = None,
+) -> "xr.DataArray":
+    """Compute vertical depth coordinates efficiently for a specified grid location and
+    optional indices.
+
+    This function calculates vertical depth coordinates (layer or interface) for a given grid
+    location (`rho`, `u`, or `v`). It performs spatial slicing (meridional or zonal) on the
+    bathymetry and free-surface elevation (`zeta`) before computing depth coordinates. This
+    approach minimizes computational overhead by reducing the dataset size before performing
+    vertical coordinate calculations.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Dataset containing optional `zeta` (free-surface elevation). If `zeta` is not present,
+        static vertical coordinates are computed.
+
+    grid_ds : xr.Dataset
+        Grid dataset containing bathymetry (`h`), stretching curves (`Cs`), and sigma-layer
+        parameters (`sigma`). The attributes of this dataset should include the critical depth (`hc`).
+
+    depth_type : str
+        Type of depth coordinates to compute:
+        - `"layer"`: Depth at the center of layers.
+        - `"interface"`: Depth at layer interfaces.
+
+    location : str
+        Grid location for the computation. Options are:
+        - `"rho"`: Depth at rho points (cell centers).
+        - `"u"`: Depth at u points (eastward velocity points).
+        - `"v"`: Depth at v points (northward velocity points).
+
+    s : int, optional
+        Vertical index to extract a single layer or interface slice. If not provided, all vertical
+        layers are included.
+
+    eta : int, optional
+        Meridional (north-south) index to extract a slice. If not provided, all meridional indices
+        are included.
+
+    xi : int, optional
+        Zonal (east-west) index to extract a slice. If not provided, all zonal indices are included.
+
+    Returns
+    -------
+    xr.DataArray
+        A DataArray containing the computed depth coordinates. If no indices are specified, the
+        array will have the full dimensionality of the depth coordinates. The dimensions of the
+        output depend on the provided indices:
+        - Full 3D (or 4D if `zeta` includes time) depth coordinates if no indices are provided.
+        - Reduced dimensionality for specified slices (e.g., 2D for a single vertical slice).
+
+    Notes
+    -----
+    - To ensure computational efficiency, spatial slicing (based on `eta` and `xi`) is performed
+      before computing depth coordinates. This reduces memory usage and processing time.
+    - Depth coordinates are interpolated to the specified grid location (`rho`, `u`, or `v`) if
+      necessary.
+    - If depth coordinates for the specified location and configuration already exist in `ds`,
+      they are not recomputed.
+    """
+
+    h = grid_ds["h"]
+    zeta = ds.get("zeta", None)
+
+    # Interpolate h and zeta to the specified location
+    if location == "u":
+        h = interpolate_from_rho_to_u(h)
+        if zeta is not None:
+            zeta = interpolate_from_rho_to_u(zeta)
+    elif location == "v":
+        h = interpolate_from_rho_to_v(h)
+        if zeta is not None:
+            zeta = interpolate_from_rho_to_v(zeta)
+
+    # Slice spatially based on the location's specific dimensions
+    if eta is not None:
+        if location == "v":
+            h = h.isel(eta_v=eta)
+            if zeta is not None:
+                zeta = zeta.isel(eta_v=eta)
+        else:  # Default to "rho" or "u"
+            h = h.isel(eta_rho=eta)
+            if zeta is not None:
+                zeta = zeta.isel(eta_rho=eta)
+    if xi is not None:
+        if location == "u":
+            h = h.isel(xi_u=xi)
+            if zeta is not None:
+                zeta = zeta.isel(xi_u=xi)
+        else:  # Default to "rho" or "v"
+            h = h.isel(xi_rho=xi)
+            if zeta is not None:
+                zeta = zeta.isel(xi_rho=xi)
+
+    # Compute depth
+    if depth_type == "layer":
+        Cs = grid_ds["Cs_r"]
+        sigma = grid_ds["sigma_r"]
+    elif depth_type == "interface":
+        Cs = grid_ds["Cs_w"]
+        sigma = grid_ds["sigma_w"]
+    depth = compute_depth(zeta, h, grid_ds.attrs["hc"], Cs, sigma)
+
+    # Slice vertically
+    if s is not None:
+        vertical_dim = "s_rho" if "s_rho" in depth.dims else "s_w"
+        depth = depth.isel({vertical_dim: s})
+
+    # Add metadata
+    depth.attrs.update(
+        {"long_name": f"{depth_type} depth at {location}-points", "units": "m"}
+    )
+
+    return depth

{roms_tools-2.2.1.dist-info → roms_tools-2.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: roms-tools
-Version: 2.2.1
+Version: 2.3.0
 Summary: Tools for running and analysing UCLA-ROMS simulations
 Author-email: Nora Loose <nora.loose@gmail.com>, Thomas Nicholas <tom@cworthy.org>
 License: Apache-2