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

roms_tools/setup/datasets.py

@@ -355,6 +355,47 @@ class Dataset:
         # Set the computed resolution as an attribute
         object.__setattr__(self, "resolution", resolution)
 
+    def compute_minimal_grid_spacing(self, ds: xr.Dataset):
+        """Compute the minimal grid spacing in a dataset based on latitude and longitude
+        spacing, considering Earth's radius.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            Dataset containing latitude and longitude dimensions.
+
+        Returns
+        -------
+        minimal_spacing : float
+            The smallest horizontal grid spacing derived from the latitude
+            and longitude differences, in meters.
+        """
+
+        r_earth = 6371315.0
+        lat_dim = self.dim_names["latitude"]
+        lon_dim = self.dim_names["longitude"]
+
+        # Get latitude and longitude values from the dataset
+        latitudes = ds[lat_dim].values
+        longitudes = ds[lon_dim].values
+
+        # Compute differences along latitude and longitude
+        lat_diff = np.abs(np.diff(latitudes)).min()  # Minimal latitude spacing
+        lon_diff = np.abs(np.diff(longitudes)).min()  # Minimal longitude spacing
+
+        # Latitude spacing is constant at all longitudes
+        min_lat_spacing = (2 * np.pi * r_earth * lat_diff) / 360
+
+        # Longitude spacing varies with latitude
+        min_lon_spacing = (
+            2 * np.pi * r_earth * lon_diff * np.cos(np.radians(latitudes.min()))
+        ) / 360
+
+        # The minimal spacing is the smaller of the two
+        minimal_spacing = min(min_lat_spacing, min_lon_spacing)
+
+        return minimal_spacing
+
     def check_if_global(self, ds) -> bool:
         """Checks if the dataset covers the entire globe in the longitude dimension.
 
@@ -470,7 +511,12 @@ class Dataset:
         pass
 
     def choose_subdomain(
-        self, target_coords, buffer_points=20, return_copy=False, verbose=False
+        self,
+        target_coords,
+        buffer_points=20,
+        return_copy=False,
+        return_coords_only=False,
+        verbose=False,
     ):
         """Selects a subdomain from the xarray Dataset based on specified target
         coordinates, extending the selection by a defined buffer. Adjusts longitude
@@ -488,6 +534,9 @@ class Dataset:
         return_subdomain : bool, optional
             If True, returns the subset of the original dataset representing the chosen
             subdomain. If False, assigns the subset to `self.ds`. Defaults to False.
+        return_coords_only : bool, optional
+            If True, returns a new xarray.Dataset containing only the latitude and longitude
+            of the subdomain. Defaults to False.
         verbose : bool, optional
             If True, print message if dataset is concatenated along longitude dimension.
             Defaults to False.
@@ -592,6 +641,13 @@ class Dataset:
         else:
             subdomain[self.dim_names["longitude"]] = xr.where(lon < 0, lon + 360, lon)
 
+        if return_coords_only:
+            # Create and return a dataset with only latitudes and longitudes
+            coords_ds = subdomain[
+                [self.dim_names["latitude"], self.dim_names["longitude"]]
+            ]
+            return coords_ds
+
         if return_copy:
             return Dataset.from_ds(self, subdomain)
         else:
@@ -2058,6 +2114,7 @@ def _select_relevant_times(
                 )
             if not end_time:
                 # Interpolate from climatology for initial conditions
+                ds["time"] = ds["time"].dt.days
                 ds = interpolate_from_climatology(ds, time_dim, start_time)
         else:
             time_type = get_time_type(ds[time_dim])

roms_tools/setup/grid.py

@@ -8,13 +8,10 @@ import matplotlib.pyplot as plt
 import yaml
 import importlib.metadata
 from typing import Dict, Union, List
+from roms_tools.utils import save_datasets
 from roms_tools.setup.topography import _add_topography
 from roms_tools.setup.mask import _add_mask, _add_velocity_masks
-from roms_tools.vertical_coordinate import (
-    sigma_stretch,
-    compute_depth,
-    add_depth_coordinates_to_dataset,
-)
+from roms_tools.vertical_coordinate import compute_depth_coordinates, sigma_stretch
 from roms_tools.plot import _plot, _section_plot
 from roms_tools.setup.utils import (
     interpolate_from_rho_to_u,
@@ -22,7 +19,7 @@ from roms_tools.setup.utils import (
     get_target_coords,
     gc_dist,
 )
-from roms_tools.setup.utils import extract_single_value, save_datasets
+from roms_tools.setup.utils import extract_single_value
 from pathlib import Path
 
 
@@ -396,14 +393,14 @@ class Grid:
         object.__setattr__(self, "ds", ds)
 
     def plot(
-        self, bathymetry: bool = False, title: str = None, with_dim_names: bool = False
+        self, bathymetry: bool = True, title: str = None, with_dim_names: bool = False
     ) -> None:
         """Plot the grid.
 
         Parameters
         ----------
         bathymetry : bool, optional
-            Whether or not to plot the bathymetry. Default is False.
+            Whether or not to plot the bathymetry. Default is True.
         title : str, optional
             The title of the plot. If not provided, it will be set to a default.
         with_dim_names : bool, optional
@@ -446,41 +443,6 @@ class Grid:
                 field=field, title=title, with_dim_names=with_dim_names, plot_data=False
             )
 
-    def compute_depth_coordinates(
-        self, depth_type: str, locations: list[str] = ["rho", "u", "v"]
-    ):
-        """Compute and update vertical depth coordinates.
-
-        Calculates vertical depth coordinates (layer or interface) for specified locations (e.g., rho, u, v points)
-        and updates them in the dataset (`self.ds`).
-
-        Parameters
-        ----------
-        depth_type : str
-            The type of depth coordinate to compute. Valid options:
-            - "layer": Compute layer depth coordinates.
-            - "interface": Compute interface depth coordinates.
-        locations : list[str], optional
-            Locations for which to compute depth coordinates. Default is ["rho", "u", "v"].
-            Valid options include:
-            - "rho": Depth coordinates at rho points.
-            - "u": Depth coordinates at u points.
-            - "v": Depth coordinates at v points.
-
-        Updates
-        -------
-        self.ds : xarray.Dataset
-            The dataset (`self.ds`) is updated with the following depth coordinate variables:
-            - f"{depth_type}_depth_rho": Depth coordinates at rho points.
-            - f"{depth_type}_depth_u": Depth coordinates at u points (if included in `locations`).
-            - f"{depth_type}_depth_v": Depth coordinates at v points (if included in `locations`).
-
-        Notes
-        -----
-        This method uses the `compute_and_update_depth_coordinates` function to perform calculations and updates.
-        """
-        add_depth_coordinates_to_dataset(self.ds, self.ds, depth_type, locations)
-
     def plot_vertical_coordinate(
         self,
         s=None,
@@ -511,56 +473,48 @@ class Grid:
             If not exactly one of s, eta, xi is specified.
         """
 
-        title = "Layer depth at rho-points"
-
         if sum(s is not None for s in [s, eta, xi]) != 1:
             raise ValueError("Exactly one of s, eta, or xi must be specified.")
 
-        h = self.ds["h"]
-        lat_deg = self.ds.lat_rho
-        lon_deg = self.ds.lon_rho
-        if self.straddle:
-            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
-        h = h.assign_coords({"lon": lon_deg, "lat": lat_deg})
+        depth = compute_depth_coordinates(
+            self.ds, zeta=0, depth_type="layer", location="rho", eta=eta, xi=xi
+        )
 
-        # slice the bathymetry as desired
+        title = "Layer depth at rho-points"
         if eta is not None:
-            title = title + f", eta_rho = {h.eta_rho[eta].item()}"
-            h = h.isel(eta_rho=eta)
+            title = title + f", eta_rho = {self.ds.eta_rho[eta].item()}"
         if xi is not None:
-            title = title + f", xi_rho = {h.xi_rho[xi].item()}"
-            h = h.isel(xi_rho=xi)
-
-        if eta is None and xi is None:
-            layer_depth = compute_depth(0, h, self.hc, self.ds.Cs_r, self.ds.sigma_r)
-            title = title + f", s_rho = {layer_depth.s_rho[s].item()}"
-            layer_depth = layer_depth.isel(s_rho=s)
+            title = title + f", xi_rho = {self.ds.xi_rho[xi].item()}"
+        if s is not None:
+            title = title + f", s_rho = {depth.s_rho[s].item()}"
+            depth = depth.isel(s_rho=s)
 
-            layer_depth.attrs["long_name"] = "Layer depth"
-            layer_depth.attrs["units"] = "m"
-
-            vmax = layer_depth.max().values
-            vmin = layer_depth.min().values
+            vmax = depth.max().values
+            vmin = depth.min().values
             cmap = plt.colormaps.get_cmap("YlGnBu")
             cmap.set_bad(color="gray")
             kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
+            lat_deg = self.ds.lat_rho
+            lon_deg = self.ds.lon_rho
+            if self.straddle:
+                lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
+            depth = depth.assign_coords({"lon": lon_deg, "lat": lat_deg})
+
             _plot(
-                field=layer_depth.where(self.ds.mask_rho),
+                field=depth.where(self.ds.mask_rho),
                 depth_contours=False,
                 title=title,
                 kwargs=kwargs,
             )
         else:
-            layer_depth = compute_depth(0, h, self.hc, self.ds.Cs_r, self.ds.sigma_r)
-            layer_depth.attrs["long_name"] = "Layer depth"
-            layer_depth.attrs["units"] = "m"
-            field = xr.zeros_like(layer_depth)
-            field = field.assign_coords({"layer_depth": layer_depth})
-
-            interface_depth = compute_depth(
-                0, h, self.hc, self.ds.Cs_w, self.ds.sigma_w
+            field = xr.zeros_like(depth)
+            field = field.assign_coords({"layer_depth": depth})
+
+            interface_depth = compute_depth_coordinates(
+                self.ds, zeta=0, depth_type="interface", eta=eta, xi=xi
             )
+
             cmap = plt.colormaps.get_cmap("YlGnBu")
             cmap.set_bad(color="gray")
             kwargs = {"vmax": 0.0, "vmin": 0.0, "cmap": cmap, "add_colorbar": False}
@@ -572,31 +526,13 @@ class Grid:
                 kwargs=kwargs,
             )
 
-    def save(
-        self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
-    ) -> None:
+    def save(self, filepath: Union[str, Path]) -> None:
         """Save the grid information to a netCDF4 file.
 
-        This method supports saving the dataset in two modes:
-
-          1. **Single File Mode (default)**:
-
-            If both `np_eta` and `np_xi` are `None`, the entire dataset is saved as a single netCDF4 file
-            with the base filename specified by `filepath.nc`.
-
-          2. **Partitioned Mode**:
-
-            - If either `np_eta` or `np_xi` is specified, the dataset is divided into spatial tiles along the eta-axis and xi-axis.
-            - Each spatial tile is saved as a separate netCDF4 file.
-
         Parameters
         ----------
         filepath : Union[str, Path]
             The base path or filename where the dataset should be saved.
-        np_eta : int, optional
-            The number of partitions along the `eta` direction. If `None`, no spatial partitioning is performed.
-        np_xi : int, optional
-            The number of partitions along the `xi` direction. If `None`, no spatial partitioning is performed.
 
         Returns
         -------
@@ -614,9 +550,7 @@ class Grid:
         dataset_list = [self.ds.load()]
         output_filenames = [str(filepath)]
 
-        saved_filenames = save_datasets(
-            dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
-        )
+        saved_filenames = save_datasets(dataset_list, output_filenames)
 
         return saved_filenames