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

ci/environment.yml

@@ -8,6 +8,7 @@ dependencies:
   - zarr
   - pytest
   - pytest-xdist
+  - h5py
   - flake8
   - black
   - pre-commit==3.8.0

roms_tools/__init__.py

@@ -14,7 +14,8 @@ from roms_tools.setup.surface_forcing import SurfaceForcing  # noqa: F401
 from roms_tools.setup.initial_conditions import InitialConditions  # noqa: F401
 from roms_tools.setup.boundary_forcing import BoundaryForcing  # noqa: F401
 from roms_tools.setup.river_forcing import RiverForcing  # noqa: F401
-from roms_tools.setup.nesting import Nesting  # noqa: F401
+from roms_tools.setup.nesting import ChildGrid  # noqa: F401
+from roms_tools.tiling.partition import partition_netcdf  # noqa: F401
 from roms_tools.analysis.roms_output import ROMSOutput  # noqa: F401
 
 # Configure logging when the package is imported

roms_tools/analysis/roms_output.py

@@ -5,14 +5,12 @@ from roms_tools.utils import _load_data
 from dataclasses import dataclass, field
 from typing import Union, Optional
 from pathlib import Path
-import os
 import re
 import logging
 from datetime import datetime, timedelta
 from roms_tools import Grid
 from roms_tools.plot import _plot, _section_plot, _profile_plot, _line_plot
 from roms_tools.vertical_coordinate import (
-    add_depth_coordinates_to_dataset,
     compute_depth_coordinates,
 )
 
@@ -26,34 +24,25 @@ class ROMSOutput:
     grid : Grid
         Object representing the grid information.
     path : Union[str, Path, List[Union[str, Path]]]
-        Directory, filename, or list of filenames with model output.
-    type : str
-        Specifies the type of model output. Options are:
-
-          - "restart": for restart files.
-          - "average": for time-averaged files.
-          - "snapshot": for snapshot files.
-
+        Filename, or list of filenames with model output.
     model_reference_date : datetime, optional
         If not specified, this is inferred from metadata of the model output
         If specified and does not coincide with metadata, a warning is raised.
+    adjust_depth_for_sea_surface_height : bool, optional
+        Whether to account for sea surface height variations when computing depth coordinates.
+        Defaults to `False`.
     use_dask: bool, optional
         Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
     """
 
     grid: Grid
     path: Union[str, Path]
-    type: Union[str, Path]
     use_dask: bool = False
     model_reference_date: Optional[datetime] = None
+    adjust_depth_for_sea_surface_height: Optional[bool] = False
     ds: xr.Dataset = field(init=False, repr=False)
 
     def __post_init__(self):
-        # Validate `type`
-        if self.type not in {"restart", "average", "snapshot"}:
-            raise ValueError(
-                f"Invalid type '{self.type}'. Must be one of 'restart', 'average', or 'snapshot'."
-            )
 
         ds = self._load_model_output()
         self._infer_model_reference_date_from_metadata(ds)
@@ -62,6 +51,9 @@ class ROMSOutput:
         ds = self._add_lat_lon_coords(ds)
         object.__setattr__(self, "ds", ds)
 
+        # Dataset for depth coordinates
+        object.__setattr__(self, "ds_depth_coords", xr.Dataset())
+
     def plot(
         self,
         var_name,
@@ -69,6 +61,7 @@ class ROMSOutput:
         s=None,
         eta=None,
         xi=None,
+        include_boundary=False,
         depth_contours=False,
         layer_contours=False,
         ax=None,
@@ -95,6 +88,11 @@ class ROMSOutput:
         xi : int, optional
             The xi-index to plot. Used for vertical sections or horizontal slices.
             Default is None.
+        include_boundary : bool, optional
+            Whether to include the outermost grid cells along the `eta`- and `xi`-boundaries in the plot.
+            In diagnostic ROMS output fields, these boundary cells are set to zero, so excluding them can improve visualization.
+            This option is only relevant for 2D horizontal plots (`eta=None`, `xi=None`).
+            Default is False.
         depth_contours : bool, optional
             If True, depth contours will be overlaid on the plot, showing lines of constant
             depth. This is typically used for plots that show a single vertical layer.
@@ -184,26 +182,36 @@ class ROMSOutput:
         compute_layer_depth = (depth_contours or s is None) and len(field.dims) > 2
         compute_interface_depth = layer_contours and s is None
 
+        # Compute depth coordinates directly instead of using .ds_depth_coords.
+        # Many cases below require only a 1D or 2D slice, making direct computation
+        # more efficient than triggering a full 3D depth computation just to extract a subset.
+        # This is especially beneficial when using Dask or if .ds_depth_coords is precomputed.
+        if self.adjust_depth_for_sea_surface_height:
+            zeta = self.ds.zeta.isel(time=time)
+        else:
+            zeta = 0
         if compute_layer_depth:
             layer_depth = compute_depth_coordinates(
-                self.ds.isel(time=time),
                 self.grid.ds,
+                zeta,
                 depth_type="layer",
                 location=loc,
-                s=s,
                 eta=eta,
                 xi=xi,
             )
+            if s is not None:
+                layer_depth = layer_depth.isel(s_rho=s)
         if compute_interface_depth:
             interface_depth = compute_depth_coordinates(
-                self.ds.isel(time=time),
                 self.grid.ds,
+                zeta,
                 depth_type="interface",
                 location=loc,
-                s=s,
                 eta=eta,
                 xi=xi,
             )
+            if s is not None:
+                interface_depth = interface_depth.isel(s_w=s)
 
         # Slice the field as desired
         title = field.long_name
@@ -251,6 +259,32 @@ class ROMSOutput:
         if compute_layer_depth:
             field = field.assign_coords({"layer_depth": layer_depth})
 
+        if not include_boundary:
+            slice_dict = None
+
+            if eta is None and xi is None:
+                slice_dict = {
+                    "rho": {"eta_rho": slice(1, -1), "xi_rho": slice(1, -1)},
+                    "u": {"eta_rho": slice(1, -1), "xi_u": slice(1, -1)},
+                    "v": {"eta_v": slice(1, -1), "xi_rho": slice(1, -1)},
+                }
+            elif eta is None:
+                slice_dict = {
+                    "rho": {"eta_rho": slice(1, -1)},
+                    "u": {"eta_rho": slice(1, -1)},
+                    "v": {"eta_v": slice(1, -1)},
+                }
+            elif xi is None:
+                slice_dict = {
+                    "rho": {"xi_rho": slice(1, -1)},
+                    "u": {"xi_u": slice(1, -1)},
+                    "v": {"xi_rho": slice(1, -1)},
+                }
+            if slice_dict is not None:
+                if loc in slice_dict:
+                    field = field.isel(**slice_dict[loc])
+                    mask = mask.isel(**slice_dict[loc])
+
         # Choose colorbar
         if var_name in ["u", "v", "w", "ubar", "vbar", "zeta"]:
             vmax = max(field.where(mask).max().values, -field.where(mask).min().values)
@@ -299,8 +333,8 @@ class ROMSOutput:
                 else:
                     _line_plot(field.where(mask), title=title, ax=ax)
 
-    def compute_depth_coordinates(self, depth_type="layer", locations=["rho"]):
-        """Compute and update vertical depth coordinates.
+    def _get_depth_coordinates(self, depth_type="layer", locations=["rho"]):
+        """Ensure depth coordinates are stored for a given location and depth type.
 
         Calculates vertical depth coordinates (layer or interface) for specified locations (e.g., rho, u, v points)
         and updates them in the dataset (`self.ds`).
@@ -320,58 +354,44 @@ class ROMSOutput:
 
         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`).
+        self.ds_depth_coords : xarray.Dataset
+
+        Raises
+        ------
+        ValueError
+            If `adjust_depth_for_sea_surface_height` is enabled but `zeta` is missing from `self.ds`.
 
         Notes
         -----
-        This method uses the `compute_and_update_depth_coordinates` function to perform calculations and updates.
+        - This method relies on the `compute_depth_coordinates` function to perform calculations.
+        - If `adjust_depth_for_sea_surface_height` is `True`, the method accounts for variations
+          in sea surface height (`zeta`).
         """
 
-        add_depth_coordinates_to_dataset(self.ds, self.grid.ds, depth_type, locations)
-
-    def _load_model_output(self) -> xr.Dataset:
-        """Load the model output based on the type."""
-        if isinstance(self.path, list):
-            filetype = "list"
-            force_combine_nested = True
-            # Check if all items in the list are files
-            if not all(Path(item).is_file() for item in self.path):
-                raise FileNotFoundError(
-                    "All items in the provided list must be valid files."
+        if self.adjust_depth_for_sea_surface_height:
+            if "zeta" not in self.ds:
+                raise ValueError(
+                    "`zeta` is required in provided ROMS output when `adjust_depth_for_sea_surface_height` is enabled."
                 )
-        elif Path(self.path).is_file():
-            filetype = "file"
-            force_combine_nested = False
-        elif Path(self.path).is_dir():
-            filetype = "dir"
-            force_combine_nested = True
+            zeta = self.ds.zeta
         else:
-            raise FileNotFoundError(
-                f"The specified path '{self.path}' is neither a file, nor a list of files, nor a directory."
-            )
+            zeta = 0
 
-        time_chunking = True
-        if self.type == "restart":
-            time_chunking = False
-            filename = _validate_and_set_filenames(self.path, filetype, "rst")
-        elif self.type == "average":
-            filename = _validate_and_set_filenames(self.path, filetype, "avg")
-        elif self.type == "snapshot":
-            filename = _validate_and_set_filenames(self.path, filetype, "his")
-        else:
-            raise ValueError(f"Unsupported type '{self.type}'.")
+        for location in locations:
+            self.ds_depth_coords[
+                f"{depth_type}_depth_{location}"
+            ] = compute_depth_coordinates(self.grid.ds, zeta, depth_type, location)
+
+    def _load_model_output(self) -> xr.Dataset:
+        """Load the model output."""
 
         # Load the dataset
         ds = _load_data(
-            filename,
+            self.path,
             dim_names={"time": "time"},
             use_dask=self.use_dask,
-            time_chunking=time_chunking,
-            force_combine_nested=force_combine_nested,
+            time_chunking=True,
+            force_combine_nested=True,
         )
 
         return ds
@@ -547,40 +567,3 @@ class ROMSOutput:
         )
 
         return ds
-
-
-def _validate_and_set_filenames(
-    filenames: Union[str, list], filetype: str, string: str
-) -> Union[str, list]:
-    """Validates and adjusts the filename or list of filenames based on the specified
-    type and checks for the presence of a string in the filename.
-
-    Parameters
-    ----------
-    filenames : Union[str, list]
-        A single filename (str), a list of filenames, or a directory path.
-    filetype : str
-        The type of input: 'file' for a single file, 'list' for a list of files, or 'dir' for a directory.
-    string : str
-        The string that should be present in each filename.
-
-    Returns
-    -------
-    Union[str, list]
-        The validated filename(s). If a directory is provided, the function returns the adjusted file pattern.
-    """
-    if filetype == "file":
-        if string not in os.path.basename(filenames):
-            logging.warning(
-                f"The file '{filenames}' does not appear to contain '{string}' in the name."
-            )
-    elif filetype == "list":
-        for file in filenames:
-            if string not in os.path.basename(file):
-                logging.warning(
-                    f"The file '{file}' does not appear to contain '{string}' in the name."
-                )
-    elif filetype == "dir":
-        filenames = os.path.join(filenames, f"*{string}.*.nc")
-
-    return filenames

roms_tools/plot.py

@@ -19,8 +19,8 @@ def _plot(
 
     Parameters
     ----------
-    field : xarray.DataArray, optional
-        The field to plot. If None, only the grid is plotted.
+    field : xarray.DataArray
+        The field to plot.
     depth_contours : bool, optional
         If True, adds depth contours to the plot.
     c : str, optional
@@ -284,6 +284,7 @@ def _section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
             )
 
     ax.set_title(title)
+    ax.set_ylabel("Depth [m]")
 
 
 def _profile_plot(field, title="", ax=None):
@@ -324,6 +325,7 @@ def _profile_plot(field, title="", ax=None):
     kwargs = {"y": depth_label, "yincrease": False}
     field.plot(**kwargs)
     ax.set_title(title)
+    ax.set_ylabel("Depth [m]")
     ax.grid()