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

roms_tools/__init__.py

@@ -15,6 +15,7 @@ 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.analysis.roms_output import ROMSOutput  # noqa: F401
 
 # Configure logging when the package is imported
 logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")

roms_tools/analysis/roms_output.py

@@ -0,0 +1,586 @@
+import xarray as xr
+import numpy as np
+import matplotlib.pyplot as plt
+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,
+)
+
+
+@dataclass(frozen=True, kw_only=True)
+class ROMSOutput:
+    """Represents ROMS model output.
+
+    Parameters
+    ----------
+    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.
+
+    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.
+    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
+    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)
+        self._check_vertical_coordinate(ds)
+        ds = self._add_absolute_time(ds)
+        ds = self._add_lat_lon_coords(ds)
+        object.__setattr__(self, "ds", ds)
+
+    def plot(
+        self,
+        var_name,
+        time=0,
+        s=None,
+        eta=None,
+        xi=None,
+        depth_contours=False,
+        layer_contours=False,
+        ax=None,
+    ) -> None:
+        """Plot a ROMS output field for a given vertical (s_rho) or horizontal (eta, xi)
+        slice.
+
+        Parameters
+        ----------
+        var_name : str
+            Name of the variable to plot. Supported options include:
+
+                - Oceanographic fields: "temp", "salt", "zeta", "u", "v", "w", etc.
+                - Biogeochemical tracers: "PO4", "NO3", "O2", "DIC", "ALK", etc.
+
+        time : int, optional
+            Index of the time dimension to plot. Default is 0.
+        s : int, optional
+            The index of the vertical layer (`s_rho`) to plot. If not specified, the plot
+            will represent a horizontal slice (eta- or xi- plane). Default is None.
+        eta : int, optional
+            The eta-index to plot. Used for vertical sections or horizontal slices.
+            Default is None.
+        xi : int, optional
+            The xi-index to plot. Used for vertical sections or horizontal slices.
+            Default is None.
+        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.
+            Default is False.
+        layer_contours : bool, optional
+            If True, contour lines representing the boundaries between vertical layers will
+            be added to the plot. This is particularly useful in vertical sections to
+            visualize the layering of the water column. For clarity, the number of layer
+            contours displayed is limited to a maximum of 10. Default is False.
+        ax : matplotlib.axes.Axes, optional
+            The axes to plot on. If None, a new figure is created. Note that this argument does not work for horizontal plots that display the eta- and xi-dimensions at the same time.
+
+        Returns
+        -------
+        None
+            This method does not return any value. It generates and displays a plot.
+
+        Raises
+        ------
+        ValueError
+            If the specified `var_name` is not one of the valid options.
+            If the field specified by `var_name` is 3D and none of `s`, `eta`, or `xi` are specified.
+            If the field specified by `var_name` is 2D and both `eta` and `xi` are specified.
+        """
+
+        # Input checks
+        if var_name not in self.ds:
+            raise ValueError(f"Variable '{var_name}' is not found in dataset.")
+
+        if "time" in self.ds[var_name].dims:
+            if time >= len(self.ds[var_name].time):
+                raise ValueError(
+                    f"Invalid time index: The specified time index ({time}) exceeds the maximum index "
+                    f"({len(self.ds[var_name].time) - 1}) for the 'time' dimension in variable '{var_name}'."
+                )
+            field = self.ds[var_name].isel(time=time)
+        else:
+            if time > 0:
+                raise ValueError(
+                    f"Invalid input: The variable '{var_name}' does not have a 'time' dimension, "
+                    f"but a time index ({time}) greater than 0 was provided."
+                )
+            field = self.ds[var_name]
+
+        if len(field.dims) == 3:
+            if not any([s is not None, eta is not None, xi is not None]):
+                raise ValueError(
+                    "Invalid input: For 3D fields, you must specify at least one of the dimensions 's', 'eta', or 'xi'."
+                )
+            if all([s is not None, eta is not None, xi is not None]):
+                raise ValueError(
+                    "Ambiguous input: For 3D fields, specify at most two of 's', 'eta', or 'xi'. Specifying all three is not allowed."
+                )
+
+        if len(field.dims) == 2 and all([eta is not None, xi is not None]):
+            raise ValueError(
+                "Conflicting input: For 2D fields, specify only one dimension, either 'eta' or 'xi', not both."
+            )
+
+        # Load the data
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                field.load()
+
+        # Get correct mask and spatial coordinates
+        if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
+            loc = "rho"
+        elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
+            loc = "u"
+        elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
+            loc = "v"
+        else:
+            ValueError("provided field does not have two horizontal dimension")
+
+        mask = self.grid.ds[f"mask_{loc}"]
+        lat_deg = self.grid.ds[f"lat_{loc}"]
+        lon_deg = self.grid.ds[f"lon_{loc}"]
+
+        if self.grid.straddle:
+            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
+
+        field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
+
+        # Retrieve depth coordinates
+        compute_layer_depth = (depth_contours or s is None) and len(field.dims) > 2
+        compute_interface_depth = layer_contours and s is None
+
+        if compute_layer_depth:
+            layer_depth = compute_depth_coordinates(
+                self.ds.isel(time=time),
+                self.grid.ds,
+                depth_type="layer",
+                location=loc,
+                s=s,
+                eta=eta,
+                xi=xi,
+            )
+        if compute_interface_depth:
+            interface_depth = compute_depth_coordinates(
+                self.ds.isel(time=time),
+                self.grid.ds,
+                depth_type="interface",
+                location=loc,
+                s=s,
+                eta=eta,
+                xi=xi,
+            )
+
+        # Slice the field as desired
+        title = field.long_name
+        if s is not None:
+            title = title + f", s_rho = {field.s_rho[s].item()}"
+            field = field.isel(s_rho=s)
+        else:
+            depth_contours = False
+
+        def _process_dimension(field, mask, dim_name, dim_values, idx, title):
+            if dim_name in field.dims:
+                title = title + f", {dim_name} = {dim_values[idx].item()}"
+                field = field.isel(**{dim_name: idx})
+                mask = mask.isel(**{dim_name: idx})
+            else:
+                raise ValueError(
+                    f"None of the expected dimensions ({dim_name}) found in field."
+                )
+            return field, mask, title
+
+        if eta is not None:
+            field, mask, title = _process_dimension(
+                field,
+                mask,
+                "eta_rho" if "eta_rho" in field.dims else "eta_v",
+                field.eta_rho if "eta_rho" in field.dims else field.eta_v,
+                eta,
+                title,
+            )
+
+        if xi is not None:
+            field, mask, title = _process_dimension(
+                field,
+                mask,
+                "xi_rho" if "xi_rho" in field.dims else "xi_u",
+                field.xi_rho if "xi_rho" in field.dims else field.xi_u,
+                xi,
+                title,
+            )
+
+        # Format to exclude seconds
+        formatted_time = np.datetime_as_string(field.abs_time.values, unit="m")
+        title = title + f", time: {formatted_time}"
+
+        if compute_layer_depth:
+            field = field.assign_coords({"layer_depth": layer_depth})
+
+        # Choose colorbar
+        if var_name in ["u", "v", "w", "ubar", "vbar", "zeta"]:
+            vmax = max(field.where(mask).max().values, -field.where(mask).min().values)
+            vmin = -vmax
+            cmap = plt.colormaps.get_cmap("RdBu_r")
+        else:
+            vmax = field.where(mask).max().values
+            vmin = field.where(mask).min().values
+            if var_name in ["temp", "salt"]:
+                cmap = plt.colormaps.get_cmap("YlOrRd")
+            else:
+                cmap = plt.colormaps.get_cmap("YlGn")
+        cmap.set_bad(color="gray")
+        kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
+
+        # Plotting
+        if eta is None and xi is None:
+            _plot(
+                field=field.where(mask),
+                depth_contours=depth_contours,
+                title=title,
+                kwargs=kwargs,
+                c="g",
+            )
+        else:
+            if len(field.dims) == 2:
+                if not layer_contours:
+                    interface_depth = None
+                else:
+                    # restrict number of layer_contours to 10 for the sake of plot clearity
+                    nr_layers = len(interface_depth["s_w"])
+                    selected_layers = np.linspace(
+                        0, nr_layers - 1, min(nr_layers, 10), dtype=int
+                    )
+                    interface_depth = interface_depth.isel(s_w=selected_layers)
+                _section_plot(
+                    field.where(mask),
+                    interface_depth=interface_depth,
+                    title=title,
+                    kwargs=kwargs,
+                    ax=ax,
+                )
+            else:
+                if "s_rho" in field.dims:
+                    _profile_plot(field.where(mask), title=title, ax=ax)
+                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.
+
+        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.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."
+                )
+        elif Path(self.path).is_file():
+            filetype = "file"
+            force_combine_nested = False
+        elif Path(self.path).is_dir():
+            filetype = "dir"
+            force_combine_nested = True
+        else:
+            raise FileNotFoundError(
+                f"The specified path '{self.path}' is neither a file, nor a list of files, nor a directory."
+            )
+
+        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}'.")
+
+        # Load the dataset
+        ds = _load_data(
+            filename,
+            dim_names={"time": "time"},
+            use_dask=self.use_dask,
+            time_chunking=time_chunking,
+            force_combine_nested=force_combine_nested,
+        )
+
+        return ds
+
+    def _infer_model_reference_date_from_metadata(self, ds: xr.Dataset) -> None:
+        """Infer and validate the model reference date from `ocean_time` metadata.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            Dataset with an `ocean_time` variable and a `long_name` attribute
+            in the format `Time since YYYY/MM/DD`.
+
+        Raises
+        ------
+        ValueError
+            If `self.model_reference_date` is not set and the reference date cannot
+            be inferred, or if the inferred date does not match `self.model_reference_date`.
+
+        Warns
+        -----
+        UserWarning
+            If `self.model_reference_date` is set but the reference date cannot be inferred.
+        """
+        # Check if 'long_name' exists in the attributes of 'ocean_time'
+        if "long_name" in ds.ocean_time.attrs:
+            input_string = ds.ocean_time.attrs["long_name"]
+            match = re.search(r"(\d{4})/(\d{2})/(\d{2})", input_string)
+
+            if match:
+                # If a match is found, extract year, month, day and create the inferred date
+                year, month, day = map(int, match.groups())
+                inferred_date = datetime(year, month, day)
+
+                if hasattr(self, "model_reference_date") and self.model_reference_date:
+                    # Check if the inferred date matches the provided model reference date
+                    if self.model_reference_date != inferred_date:
+                        raise ValueError(
+                            f"Mismatch between `self.model_reference_date` ({self.model_reference_date}) "
+                            f"and inferred reference date ({inferred_date})."
+                        )
+                else:
+                    # Set the model reference date if not already set
+                    object.__setattr__(self, "model_reference_date", inferred_date)
+            else:
+                # Handle case where no match is found
+                if hasattr(self, "model_reference_date") and self.model_reference_date:
+                    logging.warning(
+                        "Could not infer the model reference date from the metadata. "
+                        "`self.model_reference_date` will be used.",
+                    )
+                else:
+                    raise ValueError(
+                        "Model reference date could not be inferred from the metadata, "
+                        "and `self.model_reference_date` is not set."
+                    )
+        else:
+            # Handle case where 'long_name' attribute doesn't exist
+            if hasattr(self, "model_reference_date") and self.model_reference_date:
+                logging.warning(
+                    "`long_name` attribute not found in ocean_time. "
+                    "`self.model_reference_date` will be used instead.",
+                )
+            else:
+                raise ValueError(
+                    "Model reference date could not be inferred from the metadata, "
+                    "and `self.model_reference_date` is not set."
+                )
+
+    def _check_vertical_coordinate(self, ds: xr.Dataset) -> None:
+        """Check that the vertical coordinate parameters in the dataset are consistent
+        with the model grid.
+
+        This method compares the vertical coordinate parameters (`theta_s`, `theta_b`, `hc`, `Cs_r`, `Cs_w`) in
+        the provided dataset (`ds`) with those in the model grid (`self.grid`). The first three parameters are
+        checked for exact equality, while the last two are checked for numerical closeness.
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            The dataset containing vertical coordinate parameters in its attributes, such as `theta_s`, `theta_b`,
+            `hc`, `Cs_r`, and `Cs_w`.
+
+        Raises
+        ------
+        ValueError
+            If the vertical coordinate parameters do not match the expected values (based on exact or approximate equality).
+
+        Notes
+        -----
+        - `theta_s`, `theta_b`, and `hc` are checked for exact equality using `np.array_equal`.
+        - `Cs_r` and `Cs_w` are checked for numerical closeness using `np.allclose`.
+        """
+
+        # Check exact equality for theta_s, theta_b, and hc
+        if not np.array_equal(self.grid.theta_s, ds.attrs["theta_s"]):
+            raise ValueError(
+                f"theta_s from grid ({self.grid.theta_s}) does not match dataset ({ds.attrs['theta_s']})."
+            )
+
+        if not np.array_equal(self.grid.theta_b, ds.attrs["theta_b"]):
+            raise ValueError(
+                f"theta_b from grid ({self.grid.theta_b}) does not match dataset ({ds.attrs['theta_b']})."
+            )
+
+        if not np.array_equal(self.grid.hc, ds.attrs["hc"]):
+            raise ValueError(
+                f"hc from grid ({self.grid.hc}) does not match dataset ({ds.attrs['hc']})."
+            )
+
+        # Check numerical closeness for Cs_r and Cs_w
+        if not np.allclose(self.grid.ds.Cs_r, ds.attrs["Cs_r"]):
+            raise ValueError(
+                f"Cs_r from grid ({self.grid.ds.Cs_r}) is not close to dataset ({ds.attrs['Cs_r']})."
+            )
+
+        if not np.allclose(self.grid.ds.Cs_w, ds.attrs["Cs_w"]):
+            raise ValueError(
+                f"Cs_w from grid ({self.grid.ds.Cs_w}) is not close to dataset ({ds.attrs['Cs_w']})."
+            )
+
+    def _add_absolute_time(self, ds: xr.Dataset) -> xr.Dataset:
+        """Add absolute time as a coordinate to the dataset.
+
+        Computes "abs_time" based on "ocean_time" and a reference date,
+        and adds it as a coordinate.
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            Dataset containing "ocean_time" in seconds since the model reference date.
+
+        Returns
+        -------
+        xarray.Dataset
+            Dataset with "abs_time" added and "time" removed.
+        """
+        ocean_time_seconds = ds["ocean_time"].values
+
+        abs_time = np.array(
+            [
+                self.model_reference_date + timedelta(seconds=seconds)
+                for seconds in ocean_time_seconds
+            ]
+        )
+
+        abs_time = xr.DataArray(
+            abs_time, dims=["time"], coords={"time": ds["ocean_time"]}
+        )
+        abs_time.attrs["long_name"] = "absolute time"
+        ds = ds.assign_coords({"abs_time": abs_time})
+        ds = ds.drop_vars("time")
+
+        return ds
+
+    def _add_lat_lon_coords(self, ds: xr.Dataset) -> xr.Dataset:
+        """Add latitude and longitude coordinates to the dataset.
+
+        Adds "lat_rho" and "lon_rho" from the grid object to the dataset.
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            Dataset to update.
+
+        Returns
+        -------
+        xarray.Dataset
+            Dataset with "lat_rho" and "lon_rho" coordinates added.
+        """
+        ds = ds.assign_coords(
+            {"lat_rho": self.grid.ds["lat_rho"], "lon_rho": self.grid.ds["lon_rho"]}
+        )
+
+        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/{setup/download.py → download.py}

@@ -59,6 +59,9 @@ pup_test_data = pooch.create(
         "grid_created_with_matlab.nc": "fd537ef8159fabb18e38495ec8d44e2fa1b7fb615fcb1417dd4c0e1bb5f4e41d",
         "etopo5_coarsened_and_shifted.nc": "9a5cb4b38c779d22ddb0ad069b298b9722db34ca85a89273eccca691e89e6f96",
         "srtm15_coarsened.nc": "48bc8f4beecfdca9c192b13f4cbeef1455f49d8261a82563aaec5757e100dff9",
+        "eastpac25km_rst.19980106000000.nc": "8f56d72bd8daf72eb736cc6705f93f478f4ad0ae4a95e98c4c9393a38e032f4c",
+        "eastpac25km_rst.19980126000000.nc": "20ad9007c980d211d1e108c50589183120c42a2d96811264cf570875107269e4",
+        "epac25km_grd.nc": "ec26c69cda4c4e96abde5b7756c955a7e1074931ab5a0641f598b099778fb617",
     },
 )