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

roms_tools/setup/datasets.py

@@ -1,13 +1,12 @@
 import time
-import re
 import xarray as xr
 from dataclasses import dataclass, field
-import glob
 from datetime import datetime, timedelta
 import numpy as np
 from typing import Dict, Optional, Union, List
 from pathlib import Path
 import logging
+from roms_tools.utils import _load_data
 from roms_tools.setup.utils import (
     assign_dates_to_climatology,
     interpolate_from_climatology,
@@ -16,7 +15,7 @@ from roms_tools.setup.utils import (
     one_dim_fill,
     gc_dist,
 )
-from roms_tools.setup.download import (
+from roms_tools.download import (
     download_correction_data,
     download_topo,
     download_river_data,
@@ -356,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.
 
@@ -471,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
@@ -489,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.
@@ -593,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:
@@ -1945,138 +2000,6 @@ class DaiRiverDataset(RiverDataset):
 # shared functions
 
 
-def _load_data(filename, dim_names, use_dask, decode_times=True):
-    """Load dataset from the specified file.
-
-    Parameters
-    ----------
-    filename : Union[str, Path, List[Union[str, Path]]]
-        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.
-    dim_names: Dict[str, str], optional
-        Dictionary specifying the names of dimensions in the dataset.
-    use_dask: bool
-        Indicates whether to use dask for chunking. If True, data is loaded with dask; if False, data is loaded eagerly. Defaults to False.
-    decode_times: bool, optional
-        If True, decode times encoded in the standard NetCDF datetime format into datetime objects. Otherwise, leave them encoded as numbers.
-        Defaults to True.
-
-    Returns
-    -------
-    ds : xr.Dataset
-        The loaded xarray Dataset containing the forcing data.
-
-    Raises
-    ------
-    FileNotFoundError
-        If the specified file does not exist.
-    ValueError
-        If a list of files is provided but dim_names["time"] is not available or use_dask=False.
-    """
-
-    # Precompile the regex for matching wildcard characters
-    wildcard_regex = re.compile(r"[\*\?\[\]]")
-
-    # Convert Path objects to strings
-    if isinstance(filename, (str, Path)):
-        filename_str = str(filename)
-    elif isinstance(filename, list):
-        filename_str = [str(f) for f in filename]
-    else:
-        raise ValueError("filename must be a string, Path, or a list of strings/Paths.")
-    # Handle the case when filename is a string
-    contains_wildcard = False
-    if isinstance(filename_str, str):
-        contains_wildcard = bool(wildcard_regex.search(filename_str))
-        if contains_wildcard:
-            matching_files = glob.glob(filename_str)
-            if not matching_files:
-                raise FileNotFoundError(
-                    f"No files found matching the pattern '{filename_str}'."
-                )
-        else:
-            matching_files = [filename_str]
-
-    # Handle the case when filename is a list
-    elif isinstance(filename_str, list):
-        contains_wildcard = any(wildcard_regex.search(f) for f in filename_str)
-        if contains_wildcard:
-            matching_files = []
-            for f in filename_str:
-                files = glob.glob(f)
-                if not files:
-                    raise FileNotFoundError(
-                        f"No files found matching the pattern '{f}'."
-                    )
-                matching_files.extend(files)
-        else:
-            matching_files = filename_str
-
-    # Check if time dimension is available when multiple files are provided
-    if isinstance(filename_str, list) and "time" not in dim_names:
-        raise ValueError(
-            "A list of files is provided, but time dimension is not available. "
-            "A time dimension must be available to concatenate the files."
-        )
-
-    # Determine the kwargs for combining datasets
-    if contains_wildcard or len(matching_files) == 1:
-        # If there is a wildcard or just one file, use by_coords
-        kwargs = {"combine": "by_coords"}
-    else:
-        # Otherwise, use nested combine based on time
-        kwargs = {"combine": "nested", "concat_dim": dim_names["time"]}
-
-    # Base kwargs used for dataset combination
-    combine_kwargs = {
-        "coords": "minimal",
-        "compat": "override",
-        "combine_attrs": "override",
-    }
-
-    if use_dask:
-
-        chunks = {
-            dim_names["latitude"]: -1,
-            dim_names["longitude"]: -1,
-        }
-        if "depth" in dim_names:
-            chunks[dim_names["depth"]] = -1
-        if "time" in dim_names:
-            chunks[dim_names["time"]] = 1
-
-        ds = xr.open_mfdataset(
-            matching_files,
-            decode_times=decode_times,
-            chunks=chunks,
-            **combine_kwargs,
-            **kwargs,
-        )
-
-        # Rechunk the dataset along the tidal constituent dimension ("ntides") after loading
-        # because the original dataset does not have a chunk size of 1 along this dimension.
-        if "ntides" in dim_names:
-            ds = ds.chunk({dim_names["ntides"]: 1})
-
-    else:
-        ds_list = []
-        for file in matching_files:
-            ds = xr.open_dataset(file, decode_times=decode_times, chunks=None)
-            ds_list.append(ds)
-
-        if kwargs["combine"] == "by_coords":
-            ds = xr.combine_by_coords(ds_list, **combine_kwargs)
-        elif kwargs["combine"] == "nested":
-            ds = xr.combine_nested(
-                ds_list, concat_dim=kwargs["concat_dim"], **combine_kwargs
-            )
-
-    if "time" in dim_names and dim_names["time"] not in ds.dims:
-        ds = ds.expand_dims(dim_names["time"])
-
-    return ds
-
-
 def _check_dataset(
     ds: xr.Dataset,
     dim_names: Dict[str, str],
@@ -2191,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,17 +8,18 @@ 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.setup.plot import _plot, _section_plot
+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,
 )
-from roms_tools.setup.vertical_coordinate import sigma_stretch, compute_depth
-from roms_tools.setup.utils import extract_single_value, save_datasets
+from roms_tools.setup.utils import extract_single_value
 from pathlib import Path
 
 
@@ -392,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
@@ -412,13 +413,16 @@ class Grid:
             This method does not return any value. It generates and displays a plot.
         """
 
+        field = self.ds.h.where(self.ds.mask_rho)
+        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)
+        field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
+
         if bathymetry:
             if title is None:
                 title = "ROMS grid and bathymetry"
-            field = self.ds.h.where(self.ds.mask_rho)
-            field = field.assign_coords(
-                {"lon": self.ds.lon_rho, "lat": self.ds.lat_rho}
-            )
 
             vmax = field.max().values
             vmin = field.min().values
@@ -427,9 +431,7 @@ class Grid:
             kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
             _plot(
-                self.ds,
                 field=field,
-                straddle=self.straddle,
                 title=title,
                 with_dim_names=with_dim_names,
                 kwargs=kwargs,
@@ -438,10 +440,7 @@ class Grid:
             if title is None:
                 title = "ROMS grid"
             _plot(
-                self.ds,
-                straddle=self.straddle,
-                title=title,
-                with_dim_names=with_dim_names,
+                field=field, title=title, with_dim_names=with_dim_names, plot_data=False
             )
 
     def plot_vertical_coordinate(
@@ -474,54 +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"]
-        h = h.assign_coords({"lon": self.ds.lon_rho, "lat": self.ds.lat_rho})
+        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)
-
-            layer_depth.attrs["long_name"] = "Layer depth"
-            layer_depth.attrs["units"] = "m"
+            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)
 
-            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(
-                self.ds,
-                field=layer_depth.where(self.ds.mask_rho),
-                straddle=self.straddle,
+                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}
@@ -533,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
         -------
@@ -575,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