roms-tools 2.6.0__py3-none-any.whl → 2.6.1__py3-none-any.whl

roms_tools/analysis/roms_output.py

@@ -9,11 +9,13 @@ from typing import Union, Optional
 from pathlib import Path
 import re
 import logging
+import warnings
 from datetime import datetime, timedelta
 from roms_tools import Grid
 from roms_tools.vertical_coordinate import (
     compute_depth_coordinates,
 )
+from roms_tools.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
 from roms_tools.analysis.utils import _validate_plot_inputs, _generate_coordinate_range
 
 
@@ -415,6 +417,145 @@ class ROMSOutput:
         if save_path:
             plt.savefig(save_path, dpi=300, bbox_inches="tight")
 
+    def regrid(self, var_names=None, horizontal_resolution=None, depth_levels=None):
+        """Regrid the dataset both horizontally and vertically.
+
+        This method selects the specified variables, interpolates them onto a lat-lon-z horizontal grid. The horizontal target resolution and vertical target depth levels are either specified or inferred dynamically.
+
+        Parameters
+        ----------
+        var_names : list of str, optional
+            List of variable names to be regridded. If None, all variables in the dataset
+            are used.
+        horizontal_resolution : float, optional
+            Target horizontal resolution in degrees. If None, the nominal horizontal resolution is inferred from the grid.
+        depth_levels : xarray.DataArray, numpy.ndarray, list, optional
+            Target depth levels. If None, depth levels are determined dynamically.
+            If provided as a list or numpy array, it is safely converted to an `xarray.DataArray`.
+
+        Returns
+        -------
+        xarray.Dataset
+            The regridded dataset.
+        """
+
+        if var_names is None:
+            var_names = list(self.ds.data_vars)
+
+        # Check that all var_names exist in self.ds
+        missing_vars = [var for var in var_names if var not in self.ds.data_vars]
+        if missing_vars:
+            raise ValueError(
+                f"The following variables are not found in the dataset: {', '.join(missing_vars)}"
+            )
+
+        # Retain only the variables in var_names and drop others
+        ds = self.ds[var_names]
+
+        # Prepare lateral regrid
+        lat_deg = self.grid.ds["lat_rho"]
+        lon_deg = self.grid.ds["lon_rho"]
+        if self.grid.straddle:
+            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
+
+        if horizontal_resolution is None:
+            horizontal_resolution = self._infer_nominal_horizontal_resolution()
+        lons = _generate_coordinate_range(
+            lon_deg.min().values, lon_deg.max().values, horizontal_resolution
+        )
+        lons = xr.DataArray(lons, dims=["lon"], attrs={"units": "°E"})
+        lats = _generate_coordinate_range(
+            lat_deg.min().values, lat_deg.max().values, horizontal_resolution
+        )
+        lats = xr.DataArray(lats, dims=["lat"], attrs={"units": "°N"})
+        target_coords = {"lat": lats, "lon": lons}
+
+        # Prepare vertical regrid
+        if depth_levels is None:
+            depth_levels, _ = self._compute_exponential_depth_levels()
+
+        # Ensure depth_levels is an xarray.DataArray
+        if not isinstance(depth_levels, xr.DataArray):
+            depth_levels = xr.DataArray(
+                np.asarray(depth_levels),
+                dims=["depth"],
+                attrs={"long_name": "Depth", "units": "m"},
+            )
+
+        depth_levels = depth_levels.astype(np.float32)
+
+        # Initialize list to hold regridded datasets
+        regridded_datasets = []
+
+        for loc, dims in [
+            ("rho", ("eta_rho", "xi_rho")),
+            ("u", ("eta_rho", "xi_u")),
+            ("v", ("eta_v", "xi_rho")),
+        ]:
+            var_names_loc = [
+                var_name
+                for var_name in var_names
+                if all(dim in ds[var_name].dims for dim in dims)
+            ]
+            if var_names_loc:
+                ds_loc = (
+                    ds[var_names_loc]
+                    .rename({f"lat_{loc}": "lat", f"lon_{loc}": "lon"})
+                    .where(self.grid.ds[f"mask_{loc}"])
+                )
+                self._get_depth_coordinates(depth_type="layer", locations=[loc])
+                layer_depth_loc = self.ds_depth_coords[f"layer_depth_{loc}"]
+                h_loc = self.grid.ds.h
+                if loc == "u":
+                    h_loc = interpolate_from_rho_to_u(h_loc)
+                elif loc == "v":
+                    h_loc = interpolate_from_rho_to_v(h_loc)
+
+                # Exclude the horizontal boundary cells since diagnostic variables may contain zeros there
+                ds_loc = ds_loc.isel({dims[0]: slice(1, -1), dims[1]: slice(1, -1)})
+                layer_depth_loc = layer_depth_loc.isel(
+                    {dims[0]: slice(1, -1), dims[1]: slice(1, -1)}
+                )
+                h_loc = h_loc.isel({dims[0]: slice(1, -1), dims[1]: slice(1, -1)})
+
+                # Lateral regridding
+                lateral_regrid = LateralRegridFromROMS(ds_loc, target_coords)
+                ds_loc = lateral_regrid.apply(ds_loc)
+                layer_depth_loc = lateral_regrid.apply(layer_depth_loc)
+                h_loc = lateral_regrid.apply(h_loc)
+                # Vertical regridding
+                vertical_regrid = VerticalRegridFromROMS(ds_loc)
+                for var_name in var_names_loc:
+                    if "s_rho" in ds_loc[var_name].dims:
+                        attrs = ds_loc[var_name].attrs
+                        regridded = vertical_regrid.apply(
+                            ds_loc[var_name],
+                            layer_depth_loc,
+                            depth_levels,
+                            mask_edges=False,
+                        )
+                        regridded = regridded.where(regridded.depth < h_loc)
+                        ds_loc[var_name] = regridded
+                        ds_loc[var_name].attrs = attrs
+
+                ds_loc = ds_loc.assign_coords({"depth": depth_levels})
+
+                # Collect regridded dataset for merging
+                regridded_datasets.append(ds_loc)
+
+        # Merge all regridded datasets
+        if regridded_datasets:
+            ds = xr.merge(regridded_datasets)
+
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=UserWarning)
+                ds = ds.rename({"abs_time": "time"}).set_index(time="time")
+            ds["time"].attrs = {"long_name": "Time"}
+            ds["lon"].attrs = {"long_name": "Longitude", "units": "Degrees East"}
+            ds["lat"].attrs = {"long_name": "Latitude", "units": "Degrees North"}
+
+            return ds
+
     def _get_depth_coordinates(self, depth_type="layer", locations=["rho"]):
         """Ensure depth coordinates are stored for a given location and depth type.
 
@@ -460,9 +601,11 @@ class ROMSOutput:
             zeta = 0
 
         for location in locations:
-            self.ds_depth_coords[
-                f"{depth_type}_depth_{location}"
-            ] = compute_depth_coordinates(self.grid.ds, zeta, depth_type, location)
+            var_name = f"{depth_type}_depth_{location}"
+            if var_name not in self.ds_depth_coords:
+                self.ds_depth_coords[var_name] = compute_depth_coordinates(
+                    self.grid.ds, zeta, depth_type, location
+                )
 
     def _load_model_output(self) -> xr.Dataset:
         """Load the model output."""
@@ -630,24 +773,40 @@ class ROMSOutput:
         return ds
 
     def _add_lat_lon_coords(self, ds: xr.Dataset) -> xr.Dataset:
-        """Add latitude and longitude coordinates to the dataset.
+        """Add latitude and longitude coordinates to the dataset based on the grid.
 
-        Adds "lat_rho" and "lon_rho" from the grid object to the dataset.
+        This method assigns latitude and longitude coordinates from the grid to the dataset.
+        It always adds the "lat_rho" and "lon_rho" coordinates. If the dataset contains the
+        "xi_u" or "eta_v" dimensions, it also adds the corresponding "lat_u", "lon_u",
+        "lat_v", and "lon_v" coordinates.
 
         Parameters
         ----------
         ds : xarray.Dataset
-            Dataset to update.
+            Input dataset to which latitude and longitude coordinates will be added.
 
         Returns
         -------
         xarray.Dataset
-            Dataset with "lat_rho" and "lon_rho" coordinates added.
+            Updated dataset with the appropriate latitude and longitude coordinates
+            assigned to "rho", "u", and "v" points if applicable.
         """
-        ds = ds.assign_coords(
-            {"lat_rho": self.grid.ds["lat_rho"], "lon_rho": self.grid.ds["lon_rho"]}
-        )
+        coords_to_add = {
+            "lat_rho": self.grid.ds["lat_rho"],
+            "lon_rho": self.grid.ds["lon_rho"],
+        }
 
+        if "xi_u" in ds.dims:
+            coords_to_add.update(
+                {"lat_u": self.grid.ds["lat_u"], "lon_u": self.grid.ds["lon_u"]}
+            )
+        if "eta_v" in ds.dims:
+            coords_to_add.update(
+                {"lat_v": self.grid.ds["lat_v"], "lon_v": self.grid.ds["lon_v"]}
+            )
+
+        # Add all necessary coordinates in one go
+        ds = ds.assign_coords(coords_to_add)
         return ds
 
     def _infer_nominal_horizontal_resolution(self, lat=None):
@@ -693,3 +852,52 @@ class ROMSOutput:
         resolution_in_degrees = resolution_in_m / (meters_per_degree * np.cos(lat_rad))
 
         return resolution_in_degrees
+
+    def _compute_exponential_depth_levels(self, Nz=None, depth=None, h=None):
+        """Compute vertical grid center and face depths using an exponential profile.
+
+        Parameters
+        ----------
+        Nz : int, optional
+            Number of vertical levels. Defaults to `len(self.ds.s_rho)`.
+
+        depth : float, optional
+            Total depth of the domain. Defaults to `grid.ds.h.max().values`.
+
+        h : float, optional
+            Scaling parameter for the exponential profile. Defaults to `Nz / 4.5`.
+
+        Returns
+        -------
+        tuple of numpy.ndarray
+            Depth values at the vertical grid centers (`z_centers`) and grid faces (`z_faces`),
+            both rounded to two decimal places.
+        """
+        if Nz is None:
+            Nz = len(self.ds.s_rho)
+        if depth is None:
+            depth = self.grid.ds.h.max().values
+        if h is None:
+            h = Nz / 4.5
+
+        k = np.arange(1, Nz + 2)
+
+        # Define the exponential profile function
+        def exponential_profile(k, Nz, h):
+            return np.exp(k / h)
+
+        z_faces = np.vectorize(exponential_profile)(k, Nz, h)
+
+        # Normalize
+        z_faces -= z_faces[0]
+        z_faces *= depth / z_faces[-1]
+        z_faces[0] = 0.0
+
+        # Calculate center depths (average between adjacent face depths)
+        z_centers = (z_faces[:-1] + z_faces[1:]) / 2
+
+        # Round both z_faces and z_centers to two decimal places
+        z_faces = np.round(z_faces, 2)
+        z_centers = np.round(z_centers, 2)
+
+        return z_centers, z_faces

roms_tools/analysis/utils.py

@@ -107,12 +107,13 @@ def _validate_plot_inputs(field, s, eta, xi, depth, lat, lon, include_boundary):
 
 
 def _generate_coordinate_range(min, max, resolution):
-    """Generate an array of target latitude or longitude coordinates within a specified
-    range.
+    """Generate an array of target coordinates (e.g., latitude or longitude) within a
+    specified range, with a resolution that is rounded to the nearest value of the form
+    `1/n` (or integer).
 
     This method generates an array of target coordinates between the provided `min` and `max`
-    values, with a specified resolution (spacing) between each coordinate. Both `min` and `max`
-    are included in the generated range.
+    values, ensuring that both `min` and `max` are included in the resulting range. The resolution
+    is rounded to the nearest fraction of the form `1/n` or an integer, based on the input.
 
     Parameters
     ----------
@@ -123,15 +124,46 @@ def _generate_coordinate_range(min, max, resolution):
         The maximum value (in degrees) of the coordinate range (inclusive).
 
     resolution : float
-        The spacing (in degrees) between each coordinate in the array.
+        The spacing (in degrees) between each coordinate in the array. The resolution will
+        be rounded to the nearest value of the form `1/n` or an integer, depending on the size
+        of the resolution value.
 
     Returns
     -------
     numpy.ndarray
-        An array of target coordinates generated from the specified range and resolution.
+        An array of target coordinates generated from the specified range, with the resolution
+        rounded to a suitable fraction (e.g., `1/n`) or integer, depending on the input resolution.
     """
 
-    # Generate the array including both min and max values
-    target = np.arange(min, max + resolution, resolution)
+    # Find the closest fraction of the form 1/n or integer to match the resolution
+    resolution_rounded = None
+    min_diff = float("inf")  # Initialize the minimum difference as infinity
 
-    return target
+    # Search for the best fraction or integer approximation to the resolution
+    for n in range(1, 1000):  # Try fractions 1/n, where n ranges from 1 to 999
+        if resolution <= 1:
+            fraction = (
+                1 / n
+            )  # For small resolutions (<= 1), consider fractions of the form 1/n
+        else:
+            fraction = n  # For larger resolutions (>1), consider integers (n)
+
+        diff = abs(
+            fraction - resolution
+        )  # Calculate the difference between the fraction and the resolution
+
+        if diff < min_diff:  # If the current fraction is a better approximation
+            min_diff = diff
+            resolution_rounded = fraction  # Update the best fraction (or integer) found
+
+    # Adjust the start and end of the range to include integer values
+    start_int = np.floor(min)  # Round the minimum value down to the nearest integer
+    end_int = np.ceil(max)  # Round the maximum value up to the nearest integer
+
+    # Generate the array of target coordinates, including both the min and max values
+    target = np.arange(start_int, end_int + resolution_rounded, resolution_rounded)
+
+    # Truncate any values that exceed max (including small floating point errors)
+    target = target[target <= end_int + 1e-10]
+
+    return target.astype(np.float32)

roms_tools/regrid.py

@@ -257,7 +257,7 @@ class VerticalRegridFromROMS:
         """
         self.grid = xgcm.Grid(ds, coords={"s_rho": {"center": "s_rho"}}, periodic=False)
 
-    def apply(self, da, depth_coords, target_depth_levels):
+    def apply(self, da, depth_coords, target_depth_levels, mask_edges=True):
         """Applies vertical regridding from ROMS to the specified target depth levels.
 
         This method transforms the input data array `da` from the ROMS vertical coordinate (`s_rho`)
@@ -275,6 +275,9 @@ class VerticalRegridFromROMS:
         target_depth_levels : array-like
             The target depth levels to which the input data `da` will be regridded.
 
+        mask_edges: bool, optional
+            If activated, target values outside the range of depth_coords are masked with nan. Defaults to True.
+
         Returns
         -------
         xarray.DataArray
@@ -284,7 +287,11 @@ class VerticalRegridFromROMS:
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=FutureWarning, module="xgcm")
             transformed = self.grid.transform(
-                da, "s_rho", target_depth_levels, target_data=depth_coords
+                da,
+                "s_rho",
+                target_depth_levels,
+                target_data=depth_coords,
+                mask_edges=mask_edges,
             )
 
         return transformed

roms_tools/setup/river_forcing.py

@@ -441,21 +441,21 @@ class RiverForcing:
         return river_indices
 
     def _write_indices_into_dataset(self, ds):
-        """Adds river location indices to the dataset as the "river_location" variable.
+        """Adds river location indices to the dataset as the "river_flux" variable.
 
-        This method creates a new "river_location" variable
+        This method creates a new "river_flux" variable
         using river station indices from `self.indices` and assigns it to the dataset.
         The indices specify the river station locations in terms of eta_rho and xi_rho grid cell indices.
 
         Parameters
         ----------
         ds : xarray.Dataset
-            The dataset to which the "river_location" variable will be added.
+            The dataset to which the "river_flux" variable will be added.
 
         Returns
         -------
         xarray.Dataset
-            The modified dataset with the "river_location" variable added.
+            The modified dataset with the "river_flux" variable added.
         """
 
         river_locations = xr.zeros_like(self.grid.ds.h)
@@ -474,7 +474,7 @@ class RiverForcing:
 
         river_locations.attrs["long_name"] = "River ID plus local volume fraction"
         river_locations.attrs["units"] = "none"
-        ds["river_location"] = river_locations
+        ds["river_flux"] = river_locations
 
         ds = ds.drop_vars(["lat_rho", "lon_rho"])
 

roms_tools/tests/test_analysis/test_roms_output.py

@@ -1,6 +1,7 @@
 import pytest
 from pathlib import Path
 import xarray as xr
+import numpy as np
 import os
 import logging
 from datetime import datetime
@@ -486,3 +487,57 @@ def test_figure_gets_saved(roms_output_from_restart_file, tmp_path):
 
     assert filename.exists()
     filename.unlink()
+
+
+@pytest.mark.skipif(xesmf is None, reason="xesmf required")
+def test_regrid_all_variables(roms_output_from_restart_file):
+    ds_regridded = roms_output_from_restart_file.regrid()
+    assert isinstance(ds_regridded, xr.Dataset)
+    assert set(ds_regridded.data_vars).issubset(
+        set(roms_output_from_restart_file.ds.data_vars)
+    )
+    assert "lon" in ds_regridded.coords
+    assert "lat" in ds_regridded.coords
+    assert "depth" in ds_regridded.coords
+    assert "time" in ds_regridded.coords
+
+
+@pytest.mark.skipif(xesmf is None, reason="xesmf required")
+def test_regrid_specific_variables(roms_output_from_restart_file):
+    var_names = ["temp", "salt"]
+    ds_regridded = roms_output_from_restart_file.regrid(var_names=var_names)
+    assert isinstance(ds_regridded, xr.Dataset)
+    assert set(ds_regridded.data_vars) == set(var_names)
+
+    ds = roms_output_from_restart_file.regrid(var_names=[])
+    assert ds is None
+
+
+@pytest.mark.skipif(xesmf is None, reason="xesmf required")
+def test_regrid_missing_variable_raises_error(roms_output_from_restart_file):
+    with pytest.raises(
+        ValueError, match="The following variables are not found in the dataset"
+    ):
+        roms_output_from_restart_file.regrid(var_names=["fake_variable"])
+
+
+@pytest.mark.skipif(xesmf is None, reason="xesmf required")
+def test_regrid_with_custom_horizontal_resolution(roms_output_from_restart_file):
+    ds_regridded = roms_output_from_restart_file.regrid(horizontal_resolution=0.1)
+    assert isinstance(ds_regridded, xr.Dataset)
+    assert "lon" in ds_regridded.coords
+    assert "lat" in ds_regridded.coords
+
+    assert np.allclose(ds_regridded.lon.diff(dim="lon"), 0.1, atol=1e-4)
+    assert np.allclose(ds_regridded.lat.diff(dim="lat"), 0.1, atol=1e-4)
+
+
+@pytest.mark.skipif(xesmf is None, reason="xesmf required")
+def test_regrid_with_custom_depth_levels(roms_output_from_restart_file):
+    depth_levels = xr.DataArray(
+        np.linspace(0, 500, 51), dims=["depth"], attrs={"units": "m"}
+    )
+    ds_regridded = roms_output_from_restart_file.regrid(depth_levels=depth_levels)
+    assert isinstance(ds_regridded, xr.Dataset)
+    assert "depth" in ds_regridded.coords
+    np.allclose(ds_regridded.depth, depth_levels, atol=0.0)

roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/.zmetadata

@@ -58,7 +58,7 @@
             ],
             "long_name": "River ID (1-based Fortran indexing)"
         },
-        "river_location/.zarray": {
+        "river_flux/.zarray": {
             "chunks": [
                 20,
                 20
@@ -80,7 +80,7 @@
             ],
             "zarr_format": 2
         },
-        "river_location/.zattrs": {
+        "river_flux/.zattrs": {
             "_ARRAY_DIMENSIONS": [
                 "eta_rho",
                 "xi_rho"

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/.zmetadata

@@ -86,7 +86,7 @@
             ],
             "long_name": "River ID (1-based Fortran indexing)"
         },
-        "river_location/.zarray": {
+        "river_flux/.zarray": {
             "chunks": [
                 20,
                 20
@@ -108,7 +108,7 @@
             ],
             "zarr_format": 2
         },
-        "river_location/.zattrs": {
+        "river_flux/.zattrs": {
             "_ARRAY_DIMENSIONS": [
                 "eta_rho",
                 "xi_rho"

roms_tools/tests/test_setup/test_river_forcing.py

@@ -236,7 +236,7 @@ class TestRiverForcingGeneral:
         for name in indices.keys():
             for (eta_rho, xi_rho) in indices[name]:
                 assert coast[eta_rho, xi_rho]
-                assert river_forcing.ds["river_location"][eta_rho, xi_rho] > 0
+                assert river_forcing.ds["river_flux"][eta_rho, xi_rho] > 0
 
     def test_missing_source_name(self, iceland_test_grid):
         with pytest.raises(ValueError, match="`source` must include a 'name'."):
@@ -475,14 +475,14 @@ class TestRiverForcingWithPrescribedIndices:
 
         # check that all values are integers for single cell rivers
         non_zero_values = river_forcing_with_prescribed_single_cell_indices.ds[
-            "river_location"
+            "river_flux"
         ]
         is_integer = non_zero_values == np.floor(non_zero_values)
         assert (is_integer).all()
 
         # check that not all values are integers for multi cell rivers
         non_zero_values = river_forcing_with_prescribed_multi_cell_indices.ds[
-            "river_location"
+            "river_flux"
         ]
         is_integer = non_zero_values == np.floor(non_zero_values)
         assert not (is_integer).all()

roms_tools/tests/test_vertical_coordinate.py

@@ -0,0 +1,73 @@
+import pytest
+import numpy as np
+import xarray as xr
+
+from roms_tools.vertical_coordinate import (
+    compute_cs,
+    sigma_stretch,
+    compute_depth,
+    compute_depth_coordinates,
+)
+
+
+def test_compute_cs():
+    sigma = np.linspace(-1, 0, 10)
+    theta_s, theta_b = 5, 2
+    cs = compute_cs(sigma, theta_s, theta_b)
+    assert cs.shape == sigma.shape
+    assert np.all(cs <= 0) and np.all(cs >= -1)
+
+    with pytest.raises(ValueError, match="theta_s must be between 0 and 10"):
+        compute_cs(sigma, 15, 2)
+
+    with pytest.raises(ValueError, match="theta_b must be between 0 and 4"):
+        compute_cs(sigma, 5, 5)
+
+
+def test_sigma_stretch():
+    theta_s, theta_b, N = 5, 2, 10
+    cs, sigma = sigma_stretch(theta_s, theta_b, N, "r")
+    assert cs.shape == sigma.shape
+    assert isinstance(cs, xr.DataArray)
+    assert isinstance(sigma, xr.DataArray)
+
+    with pytest.raises(
+        ValueError,
+        match="Type must be either 'w' for vertical velocity points or 'r' for rho-points.",
+    ):
+        sigma_stretch(theta_s, theta_b, N, "invalid")
+
+
+def test_compute_depth():
+    zeta = xr.DataArray(0.5)
+    h = xr.DataArray(10.0)
+    hc = 5.0
+    cs = xr.DataArray(np.linspace(-1, 0, 10), dims="s_rho")
+    sigma = xr.DataArray(np.linspace(-1, 0, 10), dims="s_rho")
+
+    depth = compute_depth(zeta, h, hc, cs, sigma)
+    assert depth.shape == sigma.shape
+    assert isinstance(depth, xr.DataArray)
+
+
+def test_compute_depth_coordinates():
+    grid_ds = xr.Dataset(
+        {
+            "h": xr.DataArray([[10, 20], [30, 40]], dims=("eta_rho", "xi_rho")),
+            "Cs_r": xr.DataArray(np.linspace(-1, 0, 10), dims="s_rho"),
+            "sigma_r": xr.DataArray(np.linspace(-1, 0, 10), dims="s_rho"),
+            "Cs_w": xr.DataArray(np.linspace(-1, 0, 11), dims="s_w"),
+            "sigma_w": xr.DataArray(np.linspace(-1, 0, 11), dims="s_w"),
+        },
+        attrs={"hc": 5.0},
+    )
+
+    depth = compute_depth_coordinates(grid_ds, depth_type="layer", location="rho")
+    assert isinstance(depth, xr.DataArray)
+    assert "eta_rho" in depth.dims and "xi_rho" in depth.dims
+
+    with pytest.raises(ValueError, match="Invalid depth_type"):
+        compute_depth_coordinates(grid_ds, depth_type="invalid")
+
+    with pytest.raises(ValueError, match="Invalid location"):
+        compute_depth_coordinates(grid_ds, location="invalid")

roms_tools/vertical_coordinate.py

@@ -171,6 +171,12 @@ def compute_depth_coordinates(
     - 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.
     """
+    # Validate location
+    valid_locations = {"rho", "u", "v"}
+    if location not in valid_locations:
+        raise ValueError(
+            f"Invalid location: {location}. Must be one of {valid_locations}."
+        )
 
     # Select the appropriate depth computation parameters
     if depth_type == "layer":

{roms_tools-2.6.0.dist-info → roms_tools-2.6.1.dist-info}/METADATA

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

{roms_tools-2.6.0.dist-info → roms_tools-2.6.1.dist-info}/RECORD

@@ -3,11 +3,11 @@ ci/environment.yml,sha256=Ehxy6nYiVQXoS7EGlmNm2G0ZPHg6VFBGY1IflApIhIY,207
 roms_tools/__init__.py,sha256=jRghiteCoPjJvJjkFI36ocGyqzcTN5m-5eCa_DNQ9Dw,988
 roms_tools/download.py,sha256=W6S_DFVJDXat2w9MfyyHyushrswbpUI2hxegSuua1XE,6248
 roms_tools/plot.py,sha256=33ft1wN0kc_vIvyy_sIoY-nc0k4THXWLc_k7wEavNq8,17578
-roms_tools/regrid.py,sha256=BNtsJS5Oxmg0ivDQ-MLgVumeeLX15LZ7cp_RHAhUIOI,10222
+roms_tools/regrid.py,sha256=av9fROSNxlDeczOB544zUjFRUTiUKO42wbfZ90mpuD0,10476
 roms_tools/utils.py,sha256=eveBkWuDsXNJADFMWFgRMHdbXkZlTyVK9dN2YAnjYJ4,13323
-roms_tools/vertical_coordinate.py,sha256=txObO4Duf-8Ef2r4JgK3C1PsDmaF7KeRngDTKx7MPHg,7034
-roms_tools/analysis/roms_output.py,sha256=Wm91jUICcmKFOaayt7ADWUopowBcfuh2dyk6Id65ICo,27915
-roms_tools/analysis/utils.py,sha256=RwV7SDBjx-pzJOJ7fG3q0qwfbwXThkcI6wAgeIe1bzY,5431
+roms_tools/vertical_coordinate.py,sha256=uIxZl7rwY-fSCqXWhm6TXrOsLK4pMOMXOZB8VuP9xwg,7253
+roms_tools/analysis/roms_output.py,sha256=DUSmT2YRoqd1fowuPnoCxdukuO_NC3ONphMWsu43_nE,36560
+roms_tools/analysis/utils.py,sha256=K1Z1VyZUWKth1GMT8dumE1uUVcxrcO6rAm0Yfh8DtIg,7207
 roms_tools/setup/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 roms_tools/setup/boundary_forcing.py,sha256=LpNBurvcUhiemmubglI5aEgbOuPpjCDBj3C8XmRV60k,43609
 roms_tools/setup/datasets.py,sha256=WDEkMbLrzmdHxU5G9x_ECJASYVpmPa2i3oXWg9Rvxc4,84822
@@ -16,14 +16,15 @@ roms_tools/setup/grid.py,sha256=DWNPXoG2D76u6QMjOf_f5byvzGXT8gFOfiantUvnVjg,4940
 roms_tools/setup/initial_conditions.py,sha256=uVLkqpo1l1uw_LNjugypNVX2Tio5Hmi4txJoGC7P8G8,36012
 roms_tools/setup/mask.py,sha256=IK2SrVnMJwZjE4jNFtzMQhp1c5c8SUO45OIpCIrNSis,3413
 roms_tools/setup/nesting.py,sha256=j8l2zVCfhxNNtN4ZiSpipSlPXYmaoZWIPw6zImafsuk,26511
-roms_tools/setup/river_forcing.py,sha256=vH8sGGZ4fmcb8gDFfnwd1s9DflYksNI8e61WQyP1MxY,32824
+roms_tools/setup/river_forcing.py,sha256=G-WoajtMloqS9gzt3aERSC8NObvC5ODi8XD_Q0jhYFU,32804
 roms_tools/setup/surface_forcing.py,sha256=HwIiA0qS64gBcq2ktqZx8ZhTIFRjweRZZe3haMTjCOg,23980
 roms_tools/setup/tides.py,sha256=2eFzQMVnlcHsG83AEhTUfGgQbO0CqI_pDb8YeZ5APTU,27088
 roms_tools/setup/topography.py,sha256=s1dSF0ZWCNdrZ25yv-pUcCgufxRmGgy3yr_VhdRM7LU,14126
 roms_tools/setup/utils.py,sha256=S9_alsNEORLjEA74LVbgaxhWyVKa-PqFSPida9MQNxw,40911
 roms_tools/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 roms_tools/tests/test_regrid.py,sha256=-wzZ31BkUdSn81jq5NF1wnuaBRfd0aiYgQZEv5E_h9w,4682
-roms_tools/tests/test_analysis/test_roms_output.py,sha256=bckL5onorFl7tx9Mu5gZ1SyPmJya0j5IaiF4TcNjj8s,16556
+roms_tools/tests/test_vertical_coordinate.py,sha256=4D2jruuxBwUYk1fSeaJKICgSjzpixhjd03Y4lyivMkQ,2348
+roms_tools/tests/test_analysis/test_roms_output.py,sha256=EtxS9b5JHoREch99jbhs4QDg_3zmXDmVzXplbVrFocA,18811
 roms_tools/tests/test_setup/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 roms_tools/tests/test_setup/test_boundary_forcing.py,sha256=gET86PiGnMBejTBe69-inXpaEZL8YH4FMqyT_b8V-Zk,21183
 roms_tools/tests/test_setup/test_datasets.py,sha256=EInjmCLFNWJoO8nIQJhDtmFmMFmjcrr5L_5vJYm-qno,16435
@@ -31,7 +32,7 @@ roms_tools/tests/test_setup/test_fill.py,sha256=gDHuM58d3ECQE317ZntChYt4hWCfo4eV
 roms_tools/tests/test_setup/test_grid.py,sha256=e1S8TYt21TVczaEfyeYoAU-qxUhAgiugkHA3EDAGtIQ,14095
 roms_tools/tests/test_setup/test_initial_conditions.py,sha256=4LiULWmuktHoOty94asiFXgGNysag2okiosxG00om9M,15697
 roms_tools/tests/test_setup/test_nesting.py,sha256=WUhyP9mlMYwAJQgbgqjEU1zOyb8QD3XMgSZvHR_9LE8,18764
-roms_tools/tests/test_setup/test_river_forcing.py,sha256=gtqq67gdo4M0r5NgZT2C6gwYRikbjMuMwfE29VyLoIM,23839
+roms_tools/tests/test_setup/test_river_forcing.py,sha256=on3j7pSMESVlFHKXNVBL4H7qqMc_DmMupAtDcUmZ014,23827
 roms_tools/tests/test_setup/test_surface_forcing.py,sha256=uZAZLbX6zK_e5CYOVZvZLp_i28A1o372naeaOeF83vE,27626
 roms_tools/tests/test_setup/test_tides.py,sha256=ACFXytda3Am984QMKGxtML00KPX1LvLWlSL0FpZTqyc,8085
 roms_tools/tests/test_setup/test_topography.py,sha256=EAF-zCHfo6XnXQfBTrSZLuZDVzegWHHlrf9oBmvY0hM,4963
@@ -948,16 +949,16 @@ roms_tools/tests/test_setup/test_data/initial_conditions_with_bgc_from_climatolo
 roms_tools/tests/test_setup/test_data/initial_conditions_with_bgc_from_climatology.zarr/zooC/0.0.0.0,sha256=t8ItMb1eQSD9ZijnctYDpcZSvDigD2-5YYGHMEMmnHM,194
 roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/.zattrs,sha256=RBNvo1WzZ4oRRq0W9-hknpT7T8If536DEMBg9hyq_4o,2
 roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/.zgroup,sha256=I4N0bme0vMJ2Kz8QDwbD-i1fFJq1qOXaXTNSFGSgGVk,24
-roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/.zmetadata,sha256=f-y-IgLl25pRtOCGZI_F7Sol9wUwZiI7iXGohzo9whs,6141
+roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/.zmetadata,sha256=yUONMD7prhzTzs4Haer6nX06txMntOJilUf6iaBMEJU,6133
 roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/abs_time/.zarray,sha256=AVrYaEeCVeJWUi2FbvisrHwjFyifr4Vt3Yag7Q-n33g,312
 roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/abs_time/.zattrs,sha256=0hG3RGr8ZCkuViKEvgrkYvnehkAk3wndb4vesT1jYeI,177
 roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/abs_time/0,sha256=PvO30V_gimagXglEZjQSDgwcOrEIoWmt9JQbY4toiic,48
 roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/nriver/.zarray,sha256=zti2U0p_i_zbPjPqTErqhxNnmnUtoQPZ_4r1-9yDH7Y,312
 roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/nriver/.zattrs,sha256=Ng-_DbPNix2DMxTo-GGKyyBnH6mo3teb2SuqtQWmokE,109
 roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/nriver/0,sha256=irfTfL897AhvCrDL98Vo7ZX6zJZhKFiktdZNWOdDlSU,64
-roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_location/.zarray,sha256=jYVuViCHVLNS6VB-K_4O32zkhyPcFnoUdLM8pA1nGXk,339
-roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_location/.zattrs,sha256=Ulaw9zPnPcNer1cdhavwQtQ_tUPPTDzTrQGcHP8u9y4,149
-roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_location/0.0,sha256=_gdp3GhMsDlFGIuL7bmmlXcm1aD6AFF1fAmIjanpwWI,196
+roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_flux/.zarray,sha256=jYVuViCHVLNS6VB-K_4O32zkhyPcFnoUdLM8pA1nGXk,339
+roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_flux/.zattrs,sha256=Ulaw9zPnPcNer1cdhavwQtQ_tUPPTDzTrQGcHP8u9y4,149
+roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_flux/0.0,sha256=_gdp3GhMsDlFGIuL7bmmlXcm1aD6AFF1fAmIjanpwWI,196
 roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_name/.zarray,sha256=IiztAgS5VBT0bsRnDObags8I7F_tyQbMajEszHXfHb8,364
 roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_name/.zattrs,sha256=EPrut8rW6fC_MGEIOeqFacMyRlxf36MZsFCgP-nZ9d4,84
 roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_name/0,sha256=fuD1GoUrOX7tb_9WSj42rwDVUExgniEIJ9XrblUos0E,95
@@ -975,7 +976,7 @@ roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_n
 roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_name/0,sha256=m5RTy6Ka6ScsKdL_pma6qpevroqCf8LJzuYXCFl9Mds,48
 roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/.zattrs,sha256=jFoQJqWC46fw5uLFNg8B7WYcGSGnc9XJAWTk3ZKejOY,29
 roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/.zgroup,sha256=I4N0bme0vMJ2Kz8QDwbD-i1fFJq1qOXaXTNSFGSgGVk,24
-roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/.zmetadata,sha256=FZTi-ZnRzmPH_XNilKJIiDIbJ9clP29eBHwKFgMxiEY,6901
+roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/.zmetadata,sha256=zKm4eTj9Anuc0qm_Ip8PXeBofwHA5cnl9J1CQyVlT7Q,6893
 roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/abs_time/.zarray,sha256=bRSAng1vaRGTtFKNWhv1GcxwVKYhRsJnjbtncEbdWNA,314
 roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/abs_time/.zattrs,sha256=gDDHVjAFysUhH72DUxR7UwPLMOXFxalOA1fCWa6ImCQ,177
 roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/abs_time/0,sha256=SZIHK8deA_aE9wtgz1iLQ8ybjbVQK4Ym1FEsVryztu0,112
@@ -985,9 +986,9 @@ roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/month/0,sha256
 roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/nriver/.zarray,sha256=zti2U0p_i_zbPjPqTErqhxNnmnUtoQPZ_4r1-9yDH7Y,312
 roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/nriver/.zattrs,sha256=Ng-_DbPNix2DMxTo-GGKyyBnH6mo3teb2SuqtQWmokE,109
 roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/nriver/0,sha256=irfTfL897AhvCrDL98Vo7ZX6zJZhKFiktdZNWOdDlSU,64
-roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_location/.zarray,sha256=jYVuViCHVLNS6VB-K_4O32zkhyPcFnoUdLM8pA1nGXk,339
-roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_location/.zattrs,sha256=Ulaw9zPnPcNer1cdhavwQtQ_tUPPTDzTrQGcHP8u9y4,149
-roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_location/0.0,sha256=_gdp3GhMsDlFGIuL7bmmlXcm1aD6AFF1fAmIjanpwWI,196
+roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_flux/.zarray,sha256=jYVuViCHVLNS6VB-K_4O32zkhyPcFnoUdLM8pA1nGXk,339
+roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_flux/.zattrs,sha256=Ulaw9zPnPcNer1cdhavwQtQ_tUPPTDzTrQGcHP8u9y4,149
+roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_flux/0.0,sha256=_gdp3GhMsDlFGIuL7bmmlXcm1aD6AFF1fAmIjanpwWI,196
 roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_name/.zarray,sha256=IiztAgS5VBT0bsRnDObags8I7F_tyQbMajEszHXfHb8,364
 roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_name/.zattrs,sha256=EPrut8rW6fC_MGEIOeqFacMyRlxf36MZsFCgP-nZ9d4,84
 roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_name/0,sha256=fuD1GoUrOX7tb_9WSj42rwDVUExgniEIJ9XrblUos0E,95
@@ -1065,8 +1066,8 @@ roms_tools/tests/test_setup/test_data/tidal_forcing.zarr/v_Re/.zattrs,sha256=2z7
 roms_tools/tests/test_setup/test_data/tidal_forcing.zarr/v_Re/0.0.0,sha256=33Gl8otBmgqVarmAnZuEqTYS2_hVJUJh-iN1HzvaDuo,96
 roms_tools/tests/test_tiling/test_partition.py,sha256=b6EepZndVDv1B6Qt5_MbDfrFF2LtR0BF7i1t30xHEvA,7977
 roms_tools/tiling/partition.py,sha256=ZxDNGIKXZf_7eEzw9cxGP2XR_WBhZ4WCeIMl7_IdskA,12302
-roms_tools-2.6.0.dist-info/licenses/LICENSE,sha256=yiff76E4xRioW2bHhlPpyYpstmePQBx2bF8HhgQhSsg,11318
-roms_tools-2.6.0.dist-info/METADATA,sha256=QehIrzcMOtZk3NgvS1nFrp2ZvQfMrI9IgQUFiErw1b4,4698
-roms_tools-2.6.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-roms_tools-2.6.0.dist-info/top_level.txt,sha256=aAf4T4nYQSkay5iKJ9kmTjlDgd4ETdp9OSlB4sJdt8Y,19
-roms_tools-2.6.0.dist-info/RECORD,,
+roms_tools-2.6.1.dist-info/licenses/LICENSE,sha256=yiff76E4xRioW2bHhlPpyYpstmePQBx2bF8HhgQhSsg,11318
+roms_tools-2.6.1.dist-info/METADATA,sha256=-nmxtLvdMwgledQl2Fuf1ir4Se9NsYEi5wTWx8rMQGo,4698
+roms_tools-2.6.1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+roms_tools-2.6.1.dist-info/top_level.txt,sha256=aAf4T4nYQSkay5iKJ9kmTjlDgd4ETdp9OSlB4sJdt8Y,19
+roms_tools-2.6.1.dist-info/RECORD,,