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

roms_tools/datasets/lat_lon_datasets.py

@@ -661,6 +661,18 @@ class LatLonDataset:
         if "depth" in self.dim_names:
             self.ds = extrapolate_deepest_to_bottom(self.ds, self.dim_names["depth"])
 
+    def rotate_velocities_to_east_and_north(self) -> None:
+        """
+        Rotate velocity components to east/north directions.
+
+        For lat-lon datasets, velocity components are already defined in
+        earth-relative east/north coordinates. Therefore, no rotation is
+        required and this method is a no-op.
+        This method is provided for API compatibility with ROMSDataset,
+        where an explicit rotation using the grid angle is necessary.
+        """
+        return None
+
     @classmethod
     def from_ds(cls, original_dataset: LatLonDataset, ds: xr.Dataset) -> LatLonDataset:
         """Substitute the internal dataset of a LatLonDataset object with a new xarray

roms_tools/datasets/roms_dataset.py

@@ -18,7 +18,12 @@ from roms_tools.datasets.utils import (
     validate_start_end_time,
 )
 from roms_tools.fill import LateralFill
-from roms_tools.utils import load_data, wrap_longitudes
+from roms_tools.utils import (
+    get_dask_chunks,
+    load_data,
+    rotate_velocities,
+    wrap_longitudes,
+)
 from roms_tools.vertical_coordinate import (
     compute_depth_coordinates,
 )
@@ -578,6 +583,12 @@ class ROMSDataset:
         )
         self.ds = subdomain
 
+        subdomain_grid_ds = choose_subdomain(
+            self.grid.ds, self.grid.ds, target_coords, buffer_points
+        )
+
+        self.grid = self.grid.copy_with_ds(subdomain_grid_ds)
+
     def convert_to_float64(self) -> None:
         """Convert all data variables in the dataset to float64.
 
@@ -664,6 +675,45 @@ class ROMSDataset:
                 if filler is not None:
                     self.ds[var_name] = filler.apply(var)
 
+    def rotate_velocities_to_east_and_north(
+        self,
+        velocity_pairs: tuple[tuple[str, str], ...] = (("u", "v"),),
+    ) -> None:
+        """
+        Rotate model-grid velocity components to earth-relative east/north directions.
+
+        Parameters
+        ----------
+        velocity_pairs : tuple of (str, str), optional
+            Pairs of velocity variable keys (as used in ``self.var_names``) to rotate,
+            e.g. ("u", "v") or ("ubar", "vbar"). By default, only the 3D velocities
+            ("u", "v") are rotated.
+        """
+        if self.var_names is None:
+            return
+
+        angle = -self.grid.ds["angle"]
+
+        for u_key, v_key in velocity_pairs:
+            if u_key not in self.var_names or v_key not in self.var_names:
+                continue
+
+            u_name = self.var_names[u_key]
+            v_name = self.var_names[v_key]
+
+            self.ds[u_name], self.ds[v_name] = rotate_velocities(
+                self.ds[u_name],
+                self.ds[v_name],
+                angle,
+                interpolate_before=True,
+            )
+
+        if self.use_dask:
+            chunks = get_dask_chunks(self.dim_names)
+            # Only keep chunks for dimensions that exist in the dataset
+            chunks = {dim: size for dim, size in chunks.items() if dim in self.ds.dims}
+            self.ds = self.ds.chunk(chunks)
+
 
 def choose_subdomain(
     ds: xr.Dataset,
@@ -699,69 +749,106 @@ def choose_subdomain(
     ValueError
         If the selected latitude or longitude range does not intersect with the dataset.
     """
-    # Adjust longitude range if needed to match the expected range
-    ds = wrap_longitudes(ds, target_coords["straddle"])
-
+    # Extract lat/lon min/max from target
     lat_min = target_coords["lat"].min().values
     lat_max = target_coords["lat"].max().values
     lon_min = target_coords["lon"].min().values
     lon_max = target_coords["lon"].max().values
 
-    # Extract grid spacing (in meters)
-    dx = 0.5 * ((1 / ds_grid.pm).mean() + (1 / ds_grid.pn).mean())  # meters
-    buffer = dx * buffer_points  # buffer distance in meters
+    # Compute buffer in degrees
+    dx = 0.5 * ((1 / ds_grid.pm).mean() + (1 / ds_grid.pn).mean())
+    buffer = dx * buffer_points
+    lat_center = np.deg2rad(0.5 * (lat_min + lat_max))
+    margin_lat = buffer / 111_320.0
+    margin_lon = buffer / (111_320.0 * np.cos(lat_center))
+
+    lon_min_buf = lon_min - margin_lon
+    lon_max_buf = lon_max + margin_lon
+
+    # Normalize buffered bounds to target convention
+    if target_coords["straddle"]:
+        # [-180, 180]
+        if lon_min_buf < -180:
+            lon_min_buf += 360
+        if lon_max_buf > 180:
+            lon_max_buf -= 360
+    else:
+        # [0, 360]
+        if lon_min_buf < 0:
+            lon_min_buf += 360
+        if lon_max_buf >= 360:
+            lon_max_buf -= 360
+
+    # Wrap dataset longitudes to target convention
+    ds = wrap_longitudes(ds, target_coords["straddle"])
 
-    lat = np.deg2rad(0.5 * (lat_min + lat_max))
+    # Rho points
+    location = "rho"
+    eta_dim, xi_dim = "eta_rho", "xi_rho"
+    lat_coord, lon_coord = f"lat_{location}", f"lon_{location}"
+    _check_latlon_coords(ds, eta_dim, xi_dim, location)
+    ds_lon = ds[lon_coord]
+
+    if lon_max_buf < lon_min_buf:  # crosses dateline
+        subset_mask_lon = (ds_lon >= lon_min_buf) | (ds_lon <= lon_max_buf)
+    else:
+        subset_mask_lon = (ds_lon >= lon_min_buf) & (ds_lon <= lon_max_buf)
+
+    # Full mask including latitude
+    subset_mask = (
+        (ds[lat_coord] >= lat_min - margin_lat)
+        & (ds[lat_coord] <= lat_max + margin_lat)
+        & subset_mask_lon
+    )
 
-    deg_per_meter_lat = 1 / 111_320.0
-    margin_lat = buffer * deg_per_meter_lat
+    eta_mask = subset_mask.any(dim=xi_dim)
+    xi_mask = subset_mask.any(dim=eta_dim)
+    eta_indices = np.where(eta_mask)[0]
+    xi_indices = np.where(xi_mask)[0]
+    first_eta, last_eta = eta_indices[0], eta_indices[-1]
+    first_xi, last_xi = xi_indices[0], xi_indices[-1]
+
+    # Subset rho points
+    ds = ds.isel(
+        **{
+            "eta_rho": slice(first_eta, last_eta + 1),
+            "xi_rho": slice(first_xi, last_xi + 1),
+        }
+    )
 
-    deg_per_meter_lon = 1 / (111_320.0 * np.cos(lat))
-    margin_lon = buffer * deg_per_meter_lon
+    # Subset u points only if these dimensions exist
+    if "xi_u" in ds.dims:
+        ds = ds.isel(
+            **{
+                "xi_u": slice(first_xi, last_xi),
+            }
+        )
 
-    mapping = {
-        "rho": ("eta_rho", "xi_rho"),
-        "u": ("eta_rho", "xi_u"),
-        "v": ("eta_v", "xi_rho"),
-    }
+    # Subset v points only if these dimensions exist
+    if "eta_v" in ds.dims:
+        ds = ds.isel(
+            **{
+                "eta_v": slice(first_eta, last_eta),
+            }
+        )
 
-    for location in ["rho", "u", "v"]:
-        eta_dim, xi_dim = mapping[location]
+    return ds
 
-        if eta_dim in ds.dims and xi_dim in ds.dims:
-            # Check that lat/lon coordinates exist
-            lat_coord = f"lat_{location}"
-            lon_coord = f"lon_{location}"
-            if lat_coord not in ds.coords or lon_coord not in ds.coords:
-                raise ValueError(
-                    f"Dataset is missing expected coordinates for location '{location}': "
-                    f"expected '{lat_coord}' and '{lon_coord}'"
-                )
 
-            # Build subset mask
-            subset_mask = (
-                (ds[lat_coord] > lat_min - margin_lat)
-                & (ds[lat_coord] < lat_max + margin_lat)
-                & (ds[lon_coord] > lon_min - margin_lon)
-                & (ds[lon_coord] < lon_max + margin_lon)
-            )
+def _check_latlon_coords(
+    ds: xr.Dataset, eta_dim: str, xi_dim: str, location: str
+) -> None:
+    """
+    Ensure latitude and longitude coordinates exist for a given grid location.
 
-            # Reduce along xi_dim
-            eta_mask = subset_mask.any(dim=xi_dim)
-            eta_indices = np.where(eta_mask)[0]
-            first_eta, last_eta = eta_indices[0], eta_indices[-1]
-
-            # Reduce along eta_dim
-            xi_mask = subset_mask.any(dim=eta_dim)
-            xi_indices = np.where(xi_mask)[0]
-            first_xi, last_xi = xi_indices[0], xi_indices[-1]
-
-            # Subset the dataset
-            ds = ds.isel(
-                **{
-                    eta_dim: slice(first_eta, last_eta + 1),
-                    xi_dim: slice(first_xi, last_xi + 1),
-                }
-            )
+    Raises ValueError if the expected coordinates are missing.
+    """
+    if eta_dim in ds.dims and xi_dim in ds.dims:
+        lat_coord = f"lat_{location}"
+        lon_coord = f"lon_{location}"
 
-    return ds
+        if lat_coord not in ds.coords or lon_coord not in ds.coords:
+            raise ValueError(
+                f"Dataset missing coordinates for location '{location}': "
+                f"expected '{lat_coord}' and '{lon_coord}'"
+            )

roms_tools/datasets/utils.py

@@ -36,9 +36,21 @@ def extrapolate_deepest_to_bottom(ds: xr.Dataset, depth_dim: str) -> xr.Dataset:
 
 
 def convert_to_float64(ds: xr.Dataset) -> xr.Dataset:
-    """Convert all non-mask data variables to float64.
+    """Convert all data variables in the dataset to float64.
 
-    Variables whose names start with ``"mask_"`` are left unchanged.
+    This method updates the dataset by converting all of its data variables to the
+    `float64` data type, ensuring consistency for numerical operations that require
+    high precision. Variables whose names start with ``"mask_"`` are left unchanged.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Input dataset
+
+    Returns
+    -------
+    xr.Dataset:
+        Input dataset with data variables converted to double precision.
     """
     dtype_map = {
         name: ("float64" if not name.startswith("mask_") else var.dtype)

roms_tools/regrid.py

@@ -295,3 +295,79 @@ class VerticalRegridFromROMS:
             )
 
         return transformed
+
+
+class VerticalRegrid:
+    """Regrid ROMS variables along the vertical, using spatially varying coordinates.
+
+    This class uses the `xgcm` package to transform data from a ROMS vertical coordinate
+    system (`s_rho`) to a user-defined set of target depth levels, where both the source
+    and target coordinates can vary spatially (i.e., 2D fields in horizontal space).
+
+    Attributes
+    ----------
+    grid : xgcm.Grid
+        The XGCM grid object used for vertical regridding, initialized with the input dataset `ds`.
+    """
+
+    def __init__(self, ds: "xr.Dataset"):
+        """Initialize the VerticalRegrid object with a ROMS dataset.
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            The ROMS dataset containing the vertical coordinate `s_rho` and the variable(s)
+            to be regridded.
+        """
+        self.grid = xgcm.Grid(
+            ds,
+            coords={"s_rho": {"center": "s_rho"}},
+            periodic=False,
+            autoparse_metadata=False,
+        )
+
+    def apply(
+        self,
+        da: "xr.DataArray",
+        source_depth_coords: "xr.DataArray",
+        target_depth_coords: "xr.DataArray",
+        mask_edges: bool = True,
+    ) -> "xr.DataArray":
+        """Regrid a ROMS variable from source vertical coordinates to target vertical coordinates.
+
+        This method supports spatially varying vertical coordinates for both source and target,
+        meaning that the depth levels can vary across the horizontal grid.
+
+        Parameters
+        ----------
+        da : xarray.DataArray
+            The data array to regrid. Must have a vertical dimension corresponding to `s_rho`.
+
+        source_depth_coords : array-like (1D or 2D)
+            Depth coordinates of the source data. Can be a 1D array (same for all horizontal points)
+            or a 2D array (varying in horizontal space).
+
+        target_depth_coords : array-like (1D or 2D)
+            Desired depth coordinates of the regridded data. Can also be 1D or 2D.
+
+        mask_edges : bool, optional
+            If True, target values outside the range of source depth coordinates are masked with NaN.
+            Defaults to True.
+
+        Returns
+        -------
+        xarray.DataArray
+            A new `DataArray` containing the regridded variable at the target depth coordinates.
+        """
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning, module="xgcm")
+            transformed = self.grid.transform(
+                da,
+                "s_rho",
+                target=target_depth_coords,
+                target_data=source_depth_coords,
+                target_dim="s_rho",
+                mask_edges=mask_edges,
+            )
+
+        return transformed

roms_tools/setup/boundary_forcing.py

@@ -32,7 +32,6 @@ from roms_tools.setup.utils import (
     get_variable_metadata,
     group_dataset,
     nan_check,
-    rotate_velocities,
     substitute_nans_by_fillvalue,
     to_dict,
     write_to_yaml,
@@ -40,6 +39,7 @@ from roms_tools.setup.utils import (
 from roms_tools.utils import (
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
+    rotate_velocities,
     save_datasets,
     transpose_dimensions,
 )
@@ -273,7 +273,7 @@ class BoundaryForcing:
                         processed_fields["u"],
                         processed_fields["v"],
                         angle,
-                        interpolate=True,
+                        interpolate_after=True,
                     )
                     if self.adjust_depth_for_sea_surface_height:
                         zeta_u = interpolate_from_rho_to_u(zeta_vector)

roms_tools/setup/grid.py

@@ -1,3 +1,4 @@
+import copy
 import importlib.metadata
 import logging
 import re
@@ -19,12 +20,14 @@ from roms_tools.setup.utils import (
     extract_single_value,
     gc_dist,
     get_target_coords,
-    interpolate_from_rho_to_u,
-    interpolate_from_rho_to_v,
     pop_grid_data,
     write_to_yaml,
 )
-from roms_tools.utils import save_datasets
+from roms_tools.utils import (
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
+    save_datasets,
+)
 from roms_tools.vertical_coordinate import compute_depth_coordinates, sigma_stretch
 
 
@@ -1266,6 +1269,17 @@ class Grid:
 
         return z_centers, z_faces
 
+    def copy_with_ds(self, ds: xr.Dataset) -> "Grid":
+        """
+        Return a copy of this Grid with the given Dataset.
+
+        Grid metadata is preserved; only the backing xarray Dataset
+        is replaced. The original Grid is not modified.
+        """
+        new = copy.copy(self)  # shallow copy of metadata
+        new.ds = ds
+        return new
+
 
 def _rotate(coords, rot):
     """Rotate grid counterclockwise relative to surface of Earth by rot degrees."""