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

roms_tools/setup/datasets.py

@@ -35,10 +35,11 @@ class Dataset:
         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.
     start_time : Optional[datetime], optional
-        The start time for selecting relevant data. If not provided, the data is not filtered by start time.
+        Start time for selecting relevant data. If not provided, no time-based filtering is applied.
     end_time : Optional[datetime], optional
-        The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided,
-        or no filtering is applied if start_time is not provided.
+        End time for selecting relevant data. If not provided, the dataset selects the time entry
+        closest to `start_time` within the range `[start_time, start_time + 24 hours]`.
+        If `start_time` is also not provided, no time-based filtering is applied.
     dim_names: Dict[str, str], optional
         Dictionary specifying the names of dimensions in the dataset.
     var_names: Dict[str, str]
@@ -355,6 +356,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.
 
@@ -469,8 +511,32 @@ class Dataset:
         """
         pass
 
+    def convert_to_float64(self) -> None:
+        """Convert all data variables in the dataset to float64.
+
+        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. The dataset is modified in place.
+
+        Parameters
+        ----------
+        None
+
+        Returns
+        -------
+        None
+            This method modifies the dataset in place and does not return anything.
+        """
+        ds = self.ds.astype({var: "float64" for var in self.ds.data_vars})
+        object.__setattr__(self, "ds", ds)
+
     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 +554,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 +661,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:
@@ -1787,15 +1863,10 @@ class RiverDataset:
 
         Returns
         -------
-        indices : dict
+        indices : dict[str, list[tuple]]
             A dictionary containing the indices of the rivers that are within the threshold distance from
-            the target coordinates. The dictionary keys are:
-            - "station" : numpy.ndarray
-                The indices of the rivers that satisfy the distance threshold.
-            - "eta_rho" : numpy.ndarray
-                The indices of the `eta_rho` dimension corresponding to the selected stations.
-            - "xi_rho" : numpy.ndarray
-                The indices of the `xi_rho` dimension corresponding to the selected stations.
+            the target coordinates. The dictionary structure consists of river names as keys, and each value is a list of tuples. Each tuple represents
+            a pair of indices corresponding to the `eta_rho` and `xi_rho` grid coordinates of the river.
         """
 
         # Retrieve longitude and latitude of river mouths
@@ -1822,30 +1893,75 @@ class RiverDataset:
 
             # Find the indices of the closest grid cell to the river mouth
             indices = np.where(dist == dist_min)
+            stations = indices[0]
+            eta_rho_values = indices[1]
+            xi_rho_values = indices[2]
             names = (
-                self.ds[self.var_names["name"]]
-                .isel({self.dim_names["station"]: indices[0]})
+                ds[self.var_names["name"]]
+                .isel({self.dim_names["station"]: stations})
                 .values
             )
-            # Return the indices in a dictionary format
-            indices = {
-                "station": indices[0],
-                "eta_rho": indices[1],
-                "xi_rho": indices[2],
-                "name": names,
-            }
+            river_indices = {}
+            for i in range(len(stations)):
+                river_name = names[i]
+                river_indices[river_name] = [
+                    (int(eta_rho_values[i]), int(xi_rho_values[i]))
+                ]  # list of tuples
         else:
             ds = xr.Dataset()
-            indices = {
-                "station": [],
-                "eta_rho": [],
-                "xi_rho": [],
-                "name": [],
-            }
+            river_indices = {}
 
         object.__setattr__(self, "ds", ds)
 
-        return indices
+        return river_indices
+
+    def extract_named_rivers(self, indices):
+        """Extracts a subset of the dataset based on the provided river names in the
+        indices dictionary.
+
+        This method filters the dataset to include only the rivers specified in the `indices` dictionary.
+        The resulting subset is stored in the `ds` attribute of the class.
+
+        Parameters
+        ----------
+        indices : dict
+            A dictionary where the keys are river names (strings) and the values are dictionaries
+            containing river-related data (e.g., river indices, coordinates).
+
+        Returns
+        -------
+        None
+            The method modifies the `self.ds` attribute in place, setting it to the filtered dataset
+            containing only the data related to the specified rivers.
+
+        Raises
+        ------
+        ValueError
+            - If `indices` is not a dictionary.
+            - If any of the requested river names are not found in the dataset.
+        """
+
+        if not isinstance(indices, dict):
+            raise ValueError("`indices` must be a dictionary.")
+
+        river_names = list(indices.keys())
+
+        # Ensure the dataset is filtered based on the provided river names
+        ds_filtered = self.ds.where(
+            self.ds[self.var_names["name"]].isin(river_names), drop=True
+        )
+
+        # Check that all requested rivers exist in the dataset
+        filtered_river_names = set(ds_filtered[self.var_names["name"]].values)
+        missing_rivers = set(river_names) - filtered_river_names
+
+        if missing_rivers:
+            raise ValueError(
+                f"The following rivers were not found in the dataset: {missing_rivers}"
+            )
+
+        # Set the filtered dataset as the new `ds`
+        object.__setattr__(self, "ds", ds_filtered)
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -1994,7 +2110,7 @@ def _check_dataset(
 
 
 def _select_relevant_times(
-    ds, time_dim, start_time=None, end_time=None, climatology=False
+    ds, time_dim, start_time, end_time=None, climatology=False
 ) -> xr.Dataset:
     """Select a subset of the dataset based on the specified time range.
 
@@ -2011,11 +2127,10 @@ def _select_relevant_times(
         The input dataset to be filtered. Must contain a time dimension.
     time_dim: str
         Name of time dimension.
-    start_time : Optional[datetime], optional
-        The start time for selecting relevant data. If not provided, the data is not filtered by start time.
+    start_time : datetime
+        The start time for selecting relevant data.
     end_time : Optional[datetime], optional
-        The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided,
-        or no filtering is applied if start_time is not provided.
+        The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided.
     climatology : bool
         Indicates whether the dataset is climatological. Defaults to False.
 
@@ -2058,6 +2173,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,21 +8,19 @@ 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,
     interpolate_from_rho_to_v,
     get_target_coords,
     gc_dist,
+    _pop_grid_data,
 )
-from roms_tools.setup.utils import extract_single_value, save_datasets
+from roms_tools.setup.utils import extract_single_value
 from pathlib import Path
 
 
@@ -115,7 +113,7 @@ class Grid:
         # Coarsen the dataset if needed
         self._coarsen()
 
-        # Topography and mask
+        # Topography
         self.update_topography(
             topography_source=self.topography_source,
             hmin=self.hmin,
@@ -208,7 +206,7 @@ class Grid:
                 f"=== Generating the topography using {topography_source['name']} data and hmin = {hmin} meters ==="
             )
 
-        # Add topography and mask to the dataset
+        # Add topography to the dataset
         ds = _add_topography(
             ds=self.ds,
             target_coords=target_coords,
@@ -396,14 +394,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 +444,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 +474,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 +527,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 +551,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
 
@@ -754,9 +689,7 @@ class Grid:
         filepath = Path(filepath)
 
         data = asdict(self)
-        data.pop("ds", None)
-        data.pop("straddle", None)
-        data.pop("verbose", None)
+        data = _pop_grid_data(data)
 
         # Include the version of roms-tools
         try: