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

roms_tools/utils.py

@@ -1,340 +1,8 @@
-from numbers import Integral
-
-import numpy as np
 import xarray as xr
-from typing import Union
 from pathlib import Path
 import re
 import glob
-
-
-def partition(
-    ds: xr.Dataset, np_eta: int = 1, np_xi: int = 1
-) -> tuple[list[int], list[xr.Dataset]]:
-    """Partition a ROMS (Regional Ocean Modeling System) dataset into smaller spatial
-    tiles.
-
-    This function divides the input dataset into `np_eta` by `np_xi` tiles, where each tile
-    represents a subdomain of the original dataset. The partitioning is performed along
-    the spatial dimensions `eta_rho`, `xi_rho`, `eta_v`, `xi_u`, `eta_psi`, `xi_psi`, `eta_coarse`, and `xi_coarse`,
-    depending on which dimensions are present in the dataset.
-
-    Parameters
-    ----------
-    ds : xr.Dataset
-        The input ROMS dataset that is to be partitioned.
-
-    np_eta : int, optional
-        The number of partitions along the `eta` direction. Must be a positive integer. Default is 1.
-
-    np_xi : int, optional
-        The number of partitions along the `xi` direction. Must be a positive integer. Default is 1.
-
-    Returns
-    -------
-    tuple[list[int], list[xr.Dataset]]
-        A tuple containing two elements:
-
-        - A list of integers representing the file numbers associated with each partition.
-        - A list of `xarray.Dataset` objects, each representing a partitioned subdomain of the original dataset.
-
-    Raises
-    ------
-    ValueError
-        If `np_eta` or `np_xi` is not a positive integer, or if the dataset cannot be evenly partitioned
-        into the specified number of tiles.
-
-
-    Example
-    -------
-    >>> partitioned_file_numbers, partitioned_datasets = partition(
-    ...     ds, np_eta=2, np_xi=3
-    ... )
-    >>> print(partitioned_file_numbers)
-    [0, 1, 2, 3, 4, 5]
-    >>> print([ds.sizes for ds in partitioned_datasets])
-    [{'eta_rho': 50, 'xi_rho': 50}, {'eta_rho': 50, 'xi_rho': 50}, ...]
-
-    This example partitions the dataset into 2 tiles along the `eta` direction and 3 tiles
-    along the `xi` direction, resulting in a total of 6 partitions.
-    """
-
-    if (
-        not isinstance(np_eta, Integral)
-        or np_eta < 1
-        or not isinstance(np_xi, Integral)
-        or np_xi < 1
-    ):
-        raise ValueError("np_eta and np_xi must be positive integers")
-
-    partitionable_dims_maybe_present = [
-        "eta_rho",
-        "xi_rho",
-        "eta_v",
-        "xi_u",
-        "eta_psi",
-        "xi_psi",
-        "eta_coarse",
-        "xi_coarse",
-    ]
-    dims_to_partition = [d for d in partitionable_dims_maybe_present if d in ds.dims]
-
-    # if eta is periodic there are no ghost cells along those dimensions
-    if "eta_v" in ds.sizes and ds.sizes["eta_rho"] == ds.sizes["eta_v"]:
-        # TODO how are we supposed to know if eta is periodic if eta_v doesn't appear? partit.F doesn't say...
-        n_eta_ghost_cells = 0
-    else:
-        n_eta_ghost_cells = 1
-
-    # if xi is periodic there are no ghost cells along those dimensions
-    if "xi_u" in ds.sizes and ds.sizes["xi_rho"] == ds.sizes["xi_u"]:
-        n_xi_ghost_cells = 0
-    else:
-        n_xi_ghost_cells = 1
-
-    def integer_division_or_raise(a: int, b: int, dimension: str) -> int:
-        """Perform integer division and ensure that the division is exact.
-
-        Parameters
-        ----------
-        a : int
-            The numerator for the division.
-        b : int
-            The denominator for the division.
-        dimension : str
-            The name of the dimension being partitioned, used for error reporting.
-
-        Returns
-        -------
-        int
-            The result of the integer division.
-
-        Raises
-        ------
-        ValueError
-            If the division is not exact, indicating that the domain cannot be evenly divided
-            along the specified dimension.
-        """
-        remainder = a % b
-        if remainder == 0:
-            return a // b
-        else:
-            raise ValueError(
-                f"Dimension '{dimension}' of size {a} cannot be evenly divided into {b} partitions."
-            )
-
-    if "eta_rho" in dims_to_partition:
-        eta_rho_domain_size = integer_division_or_raise(
-            ds.sizes["eta_rho"] - 2 * n_eta_ghost_cells, np_eta, "eta_rho"
-        )
-
-    if "xi_rho" in dims_to_partition:
-        xi_rho_domain_size = integer_division_or_raise(
-            ds.sizes["xi_rho"] - 2 * n_xi_ghost_cells, np_xi, "xi_rho"
-        )
-
-    if "eta_v" in dims_to_partition:
-        eta_v_domain_size = integer_division_or_raise(
-            ds.sizes["eta_v"] - 1 * n_eta_ghost_cells, np_eta, "eta_v"
-        )
-
-    if "xi_u" in dims_to_partition:
-        xi_u_domain_size = integer_division_or_raise(
-            ds.sizes["xi_u"] - 1 * n_xi_ghost_cells, np_xi, "xi_u"
-        )
-
-    if "eta_psi" in dims_to_partition:
-        eta_psi_domain_size = integer_division_or_raise(
-            ds.sizes["eta_psi"] - 3 * n_eta_ghost_cells, np_eta, "eta_psi"
-        )
-
-    if "xi_psi" in dims_to_partition:
-        xi_psi_domain_size = integer_division_or_raise(
-            ds.sizes["xi_psi"] - 3 * n_xi_ghost_cells, np_xi, "xi_psi"
-        )
-
-    if "eta_coarse" in dims_to_partition:
-        eta_coarse_domain_size = integer_division_or_raise(
-            ds.sizes["eta_coarse"] - 2 * n_eta_ghost_cells, np_eta, "eta_coarse"
-        )
-    if "xi_coarse" in dims_to_partition:
-        xi_coarse_domain_size = integer_division_or_raise(
-            ds.sizes["xi_coarse"] - 2 * n_xi_ghost_cells, np_xi, "xi_coarse"
-        )
-
-    # unpartitioned dimensions should have sizes unchanged
-    partitioned_sizes = {
-        dim: [size] for dim, size in ds.sizes.items() if dim in dims_to_partition
-    }
-
-    # TODO refactor to use two functions for odd- and even-length dimensions
-    if "eta_v" in dims_to_partition:
-        partitioned_sizes["eta_v"] = [eta_v_domain_size] * (np_eta - 1) + [
-            eta_v_domain_size + n_eta_ghost_cells
-        ]
-    if "xi_u" in dims_to_partition:
-        partitioned_sizes["xi_u"] = [xi_u_domain_size] * (np_xi - 1) + [
-            xi_u_domain_size + n_xi_ghost_cells
-        ]
-
-    if np_eta > 1:
-        if "eta_rho" in dims_to_partition:
-            partitioned_sizes["eta_rho"] = (
-                [eta_rho_domain_size + n_eta_ghost_cells]
-                + [eta_rho_domain_size] * (np_eta - 2)
-                + [eta_rho_domain_size + n_eta_ghost_cells]
-            )
-        if "eta_psi" in dims_to_partition:
-            partitioned_sizes["eta_psi"] = (
-                [n_eta_ghost_cells + eta_psi_domain_size]
-                + [eta_psi_domain_size] * (np_eta - 2)
-                + [eta_psi_domain_size + 2 * n_eta_ghost_cells]
-            )
-        if "eta_coarse" in dims_to_partition:
-            partitioned_sizes["eta_coarse"] = (
-                [eta_coarse_domain_size + n_eta_ghost_cells]
-                + [eta_coarse_domain_size] * (np_eta - 2)
-                + [eta_coarse_domain_size + n_eta_ghost_cells]
-            )
-
-    if np_xi > 1:
-        if "xi_rho" in dims_to_partition:
-            partitioned_sizes["xi_rho"] = (
-                [xi_rho_domain_size + n_xi_ghost_cells]
-                + [xi_rho_domain_size] * (np_xi - 2)
-                + [xi_rho_domain_size + n_xi_ghost_cells]
-            )
-        if "xi_psi" in dims_to_partition:
-            partitioned_sizes["xi_psi"] = (
-                [n_xi_ghost_cells + xi_psi_domain_size]
-                + [xi_psi_domain_size] * (np_xi - 2)
-                + [xi_psi_domain_size + 2 * n_xi_ghost_cells]
-            )
-        if "xi_coarse" in dims_to_partition:
-            partitioned_sizes["xi_coarse"] = (
-                [xi_coarse_domain_size + n_xi_ghost_cells]
-                + [xi_coarse_domain_size] * (np_xi - 2)
-                + [xi_coarse_domain_size + n_xi_ghost_cells]
-            )
-
-    def cumsum(pmf):
-        """Implementation of cumsum which ensures the result starts with zero."""
-        cdf = np.empty(len(pmf) + 1, dtype=int)
-        cdf[0] = 0
-        np.cumsum(pmf, out=cdf[1:])
-        return cdf
-
-    file_numbers = []
-    partitioned_datasets = []
-    for i in range(np_eta):
-        for j in range(np_xi):
-            file_number = j + (i * np_xi)
-            file_numbers.append(file_number)
-
-            indexers = {}
-
-            if "eta_rho" in dims_to_partition:
-                eta_rho_partition_indices = cumsum(partitioned_sizes["eta_rho"])
-                indexers["eta_rho"] = slice(
-                    int(eta_rho_partition_indices[i]),
-                    int(eta_rho_partition_indices[i + 1]),
-                )
-            if "xi_rho" in dims_to_partition:
-                xi_rho_partition_indices = cumsum(partitioned_sizes["xi_rho"])
-                indexers["xi_rho"] = slice(
-                    int(xi_rho_partition_indices[j]),
-                    int(xi_rho_partition_indices[j + 1]),
-                )
-
-            if "eta_v" in dims_to_partition:
-                eta_v_partition_indices = cumsum(partitioned_sizes["eta_v"])
-                indexers["eta_v"] = slice(
-                    int(eta_v_partition_indices[i]),
-                    int(eta_v_partition_indices[i + 1]),
-                )
-            if "xi_u" in dims_to_partition:
-                xi_u_partition_indices = cumsum(partitioned_sizes["xi_u"])
-                indexers["xi_u"] = slice(
-                    int(xi_u_partition_indices[j]), int(xi_u_partition_indices[j + 1])
-                )
-            if "eta_psi" in dims_to_partition:
-                eta_psi_partition_indices = cumsum(partitioned_sizes["eta_psi"])
-                indexers["eta_psi"] = slice(
-                    int(eta_psi_partition_indices[i]),
-                    int(eta_psi_partition_indices[i + 1]),
-                )
-            if "xi_psi" in dims_to_partition:
-                xi_psi_partition_indices = cumsum(partitioned_sizes["xi_psi"])
-                indexers["xi_psi"] = slice(
-                    int(xi_psi_partition_indices[j]),
-                    int(xi_psi_partition_indices[j + 1]),
-                )
-
-            if "eta_coarse" in dims_to_partition:
-                eta_coarse_partition_indices = cumsum(partitioned_sizes["eta_coarse"])
-                indexers["eta_coarse"] = slice(
-                    int(eta_coarse_partition_indices[i]),
-                    int(eta_coarse_partition_indices[i + 1]),
-                )
-
-            if "xi_coarse" in dims_to_partition:
-                xi_coarse_partition_indices = cumsum(partitioned_sizes["xi_coarse"])
-                indexers["xi_coarse"] = slice(
-                    int(xi_coarse_partition_indices[j]),
-                    int(xi_coarse_partition_indices[j + 1]),
-                )
-
-            partitioned_ds = ds.isel(**indexers)
-
-            partitioned_datasets.append(partitioned_ds)
-
-    return file_numbers, partitioned_datasets
-
-
-def partition_netcdf(
-    filepath: Union[str, Path], np_eta: int = 1, np_xi: int = 1
-) -> None:
-    """Partition a ROMS NetCDF file into smaller spatial tiles and save them to disk.
-
-    This function divides the dataset in the specified NetCDF file into `np_eta` by `np_xi` tiles.
-    Each tile is saved as a separate NetCDF file.
-
-    Parameters
-    ----------
-    filepath : Union[str, Path]
-        The path to the input NetCDF file.
-
-    np_eta : int, optional
-        The number of partitions along the `eta` direction. Must be a positive integer. Default is 1.
-
-    np_xi : int, optional
-        The number of partitions along the `xi` direction. Must be a positive integer. Default is 1.
-
-    Returns
-    -------
-    List[Path]
-        A list of Path objects for the filenames that were saved.
-    """
-
-    # Ensure filepath is a Path object
-    filepath = Path(filepath)
-
-    # Open the dataset
-    ds = xr.open_dataset(filepath.with_suffix(".nc"))
-
-    # Partition the dataset
-    file_numbers, partitioned_datasets = partition(ds, np_eta=np_eta, np_xi=np_xi)
-
-    # Generate paths to the partitioned files
-    base_filepath = filepath.with_suffix("")
-    paths_to_partitioned_files = [
-        Path(f"{base_filepath}.{file_number}.nc") for file_number in file_numbers
-    ]
-
-    # Save the partitioned datasets to files
-    xr.save_mfdataset(partitioned_datasets, paths_to_partitioned_files)
-
-    return paths_to_partitioned_files
+import logging
 
 
 def _load_data(
@@ -634,3 +302,68 @@ def transpose_dimensions(da: xr.DataArray) -> xr.DataArray:
     transposed_da = da.transpose(*new_order)
 
     return transposed_da
+
+
+def save_datasets(dataset_list, output_filenames, use_dask=False, verbose=True):
+    """Save the list of datasets to netCDF4 files.
+
+    Parameters
+    ----------
+    dataset_list : list
+        List of datasets to be saved.
+    output_filenames : list
+        List of filenames for the output files.
+    use_dask : bool, optional
+        Whether to use Dask diagnostics (e.g., progress bars) when saving the datasets, by default False.
+    verbose : bool, optional
+        Whether to log information about the files being written. If True, logs the output filenames.
+        Defaults to True.
+
+    Returns
+    -------
+    List[Path]
+        A list of Path objects for the filenames that were saved.
+    """
+
+    saved_filenames = []
+
+    output_filenames = [f"{filename}.nc" for filename in output_filenames]
+    if verbose:
+        logging.info(
+            "Writing the following NetCDF files:\n%s", "\n".join(output_filenames)
+        )
+
+    if use_dask:
+        from dask.diagnostics import ProgressBar
+
+        with ProgressBar():
+            xr.save_mfdataset(dataset_list, output_filenames)
+    else:
+        xr.save_mfdataset(dataset_list, output_filenames)
+
+    saved_filenames.extend(Path(f) for f in output_filenames)
+
+    return saved_filenames
+
+
+def get_dask_chunks(location, chunk_size):
+    """Returns the appropriate Dask chunking dictionary based on grid location.
+
+    Parameters
+    ----------
+    location : str
+        The grid location, one of "rho", "u", or "v".
+    chunk_size : int
+        The chunk size to apply.
+
+    Returns
+    -------
+    dict
+        Dictionary specifying the chunking strategy.
+    """
+    chunk_mapping = {
+        "rho": {"eta_rho": chunk_size, "xi_rho": chunk_size},
+        "u": {"eta_rho": chunk_size, "xi_u": chunk_size},
+        "v": {"eta_v": chunk_size, "xi_rho": chunk_size},
+    }
+    return chunk_mapping.get(location, {})

roms_tools/vertical_coordinate.py

@@ -113,191 +113,112 @@ def compute_depth(zeta, h, hc, cs, sigma):
     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,
+    zeta: xr.DataArray | float = 0,
+    depth_type: str = "layer",
+    location: str = "rho",
     eta: int = None,
     xi: int = None,
 ) -> "xr.DataArray":
-    """Compute vertical depth coordinates efficiently for a specified grid location and
-    optional indices.
+    """Compute vertical depth coordinates for a given ROMS grid location.
 
-    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.
+    This function calculates depth coordinates (layer or interface) at a specified grid
+    location (`rho`, `u`, or `v`). It optimizes computations by slicing the bathymetry (`h`)
+    and free-surface elevation (`zeta`) before performing vertical coordinate calculations,
+    reducing memory usage and improving efficiency.
 
     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`).
+        ROMS grid dataset containing bathymetry (`h`), stretching curves (`Cs`),
+        and sigma-layer parameters (`sigma`). The dataset's attributes should include
+        the critical depth parameter (`hc`).
 
-    depth_type : str
+    zeta : xr.DataArray or float, optional
+        Free-surface elevation. If set to `0` (default), the static sea level is assumed.
+
+    depth_type : str, optional
         Type of depth coordinates to compute:
-        - `"layer"`: Depth at the center of layers.
+        - `"layer"` (default): Depth at the center of vertical 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.
+    location : str, optional
+        Grid location for depth computation:
+        - `"rho"` (default): Depth at cell centers (rho points).
+        - `"u"`: Depth at eastward velocity points (u points).
+        - `"v"`: Depth at northward velocity points (v points).
 
     eta : int, optional
-        Meridional (north-south) index to extract a slice. If not provided, all meridional indices
-        are included.
+        Meridional (north-south) index to extract a zonal 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.
+        Zonal (east-west) index to extract a meridional 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).
+        Computed depth coordinates. The shape of the output depends on the given indices:
+        - Full 3D (or 4D if `zeta` includes time) depth field if no indices are specified.
+        - 2D (or 3D if `zeta` includes time) slice if either `eta` or `xi` are specified.
+        - 1D (or 2D if `zeta` includes time) slice if both `eta` and `xi` are specified.
 
     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.
+    - The function first interpolates `h` and `zeta` to the specified grid location (`rho`, `u`, or `v`).
+    - Spatial slicing (`eta`, `xi`) is performed before depth computation to optimize efficiency.
+    - Depth calculations rely on the ROMS vertical stretching curves (`Cs`) and sigma-layers.
     """
 
+    # Select the appropriate depth computation parameters
+    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"]
+    else:
+        raise ValueError(
+            f"Invalid depth_type: {depth_type}. Choose 'layer' or 'interface'."
+        )
+
     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:
+        if isinstance(zeta, xr.DataArray):
             zeta = interpolate_from_rho_to_u(zeta)
     elif location == "v":
         h = interpolate_from_rho_to_v(h)
-        if zeta is not None:
+        if isinstance(zeta, xr.DataArray):
             zeta = interpolate_from_rho_to_v(zeta)
 
-    # Slice spatially based on the location's specific dimensions
+    # Slice spatially based on indices
     if eta is not None:
         if location == "v":
             h = h.isel(eta_v=eta)
-            if zeta is not None:
+            if isinstance(zeta, xr.DataArray):
                 zeta = zeta.isel(eta_v=eta)
-        else:  # Default to "rho" or "u"
+        else:  # Applies to "rho" or "u"
             h = h.isel(eta_rho=eta)
-            if zeta is not None:
+            if isinstance(zeta, xr.DataArray):
                 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:
+            if isinstance(zeta, xr.DataArray):
                 zeta = zeta.isel(xi_u=xi)
-        else:  # Default to "rho" or "v"
+        else:  # Applies to "rho" or "v"
             h = h.isel(xi_rho=xi)
-            if zeta is not None:
+            if isinstance(zeta, xr.DataArray):
                 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"}

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

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