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

roms_tools/{setup/plot.py → plot.py}

@@ -1,37 +1,35 @@
 import cartopy.crs as ccrs
 import matplotlib.pyplot as plt
 import xarray as xr
+import numpy as np
 
 
 def _plot(
-    grid_ds,
-    field=None,
+    field,
     depth_contours=False,
-    straddle=False,
     c="red",
     title="",
     with_dim_names=False,
+    plot_data=True,
     kwargs={},
 ):
     """Plots a grid or field on a map with optional depth contours.
 
     This function plots a map using Cartopy projections. It supports plotting a grid, a field, and adding depth contours if desired.
-    The projection can be customized, and the grid can be adjusted for domains straddling the 180° meridian.
 
     Parameters
     ----------
-    grid_ds : xarray.Dataset
-        The grid dataset containing coordinates (`lon_rho`, `lat_rho`).
     field : xarray.DataArray, optional
         The field to plot. If None, only the grid is plotted.
     depth_contours : bool, optional
         If True, adds depth contours to the plot.
-    straddle : bool, optional
-        If True, adjusts longitude values to straddle across the 180° meridian.
     c : str, optional
         Color for the boundary plot (default is 'red').
     title : str, optional
         Title of the plot.
+    plot_data : bool, optional
+        If True, plots the provided field data on the map. If False, only the grid
+        boundaries and optional depth contours are plotted. Default is True.
     kwargs : dict, optional
         Additional keyword arguments to pass to `pcolormesh` (e.g., colormap or color limits).
 
@@ -40,24 +38,15 @@ def _plot(
     The function raises a `NotImplementedError` if the domain contains the North or South Pole.
     """
 
-    if field is None:
-        lon_deg = grid_ds["lon_rho"]
-        lat_deg = grid_ds["lat_rho"]
-
-    else:
-
-        field = field.squeeze()
-        lon_deg = field.lon
-        lat_deg = field.lat
+    field = field.squeeze()
+    lon_deg = field.lon
+    lat_deg = field.lat
 
-        # check if North or South pole are in domain
-        if lat_deg.max().values > 89 or lat_deg.min().values < -89:
-            raise NotImplementedError(
-                "Plotting is not implemented for the case that the domain contains the North or South pole."
-            )
-
-    if straddle:
-        lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
+    # check if North or South pole are in domain
+    if lat_deg.max().values > 89 or lat_deg.min().values < -89:
+        raise NotImplementedError(
+            "Plotting is not implemented for the case that the domain contains the North or South pole."
+        )
 
     trans = _get_projection(lon_deg, lat_deg)
 
@@ -71,7 +60,7 @@ def _plot(
             ax, lon_deg, lat_deg, trans, c, with_dim_names=with_dim_names
         )
 
-    if field is not None:
+    if plot_data:
         _add_field_to_ax(ax, lon_deg, lat_deg, field, depth_contours, kwargs=kwargs)
 
     ax.coastlines(
@@ -287,7 +276,7 @@ def _section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
     field.plot(**kwargs, **more_kwargs, ax=ax)
 
     if interface_depth is not None:
-        layer_key = "s_rho" if "s_rho" in interface_depth else "s_w"
+        layer_key = "s_rho" if "s_rho" in interface_depth.dims else "s_w"
 
         for i in range(len(interface_depth[layer_key])):
             ax.plot(
@@ -339,7 +328,8 @@ def _profile_plot(field, title="", ax=None):
 
 
 def _line_plot(field, title="", ax=None):
-    """Plots a line graph of the given field.
+    """Plots a line graph of the given field, with grey vertical bars where NaNs are
+    located.
 
     Parameters
     ----------
@@ -358,6 +348,22 @@ def _line_plot(field, title="", ax=None):
     if ax is None:
         fig, ax = plt.subplots(1, 1, figsize=(7, 4))
     field.plot(ax=ax)
+
+    # Loop through the NaNs in the field and add grey vertical bars
+    nan_mask = np.isnan(field.values)
+    nan_indices = np.where(nan_mask)[0]
+
+    if len(nan_indices) > 0:
+        # Add grey vertical bars for each NaN region
+        start_idx = nan_indices[0]
+        for idx in range(1, len(nan_indices)):
+            if nan_indices[idx] != nan_indices[idx - 1] + 1:
+                ax.axvspan(start_idx, nan_indices[idx - 1] + 1, color="gray", alpha=0.3)
+                start_idx = nan_indices[idx]
+        # Add the last region of NaNs
+        ax.axvspan(start_idx, nan_indices[-1] + 1, color="gray", alpha=0.3)
+
+    # Set plot title and grid
     ax.set_title(title)
     ax.grid()
 

roms_tools/setup/boundary_forcing.py

@@ -5,11 +5,19 @@ import logging
 import importlib.metadata
 from typing import Dict, Union, List
 from dataclasses import dataclass, field
-from roms_tools.setup.grid import Grid
-from roms_tools.setup.regrid import LateralRegrid, VerticalRegrid
 from datetime import datetime
+import matplotlib.pyplot as plt
+from pathlib import Path
+from roms_tools import Grid
+from roms_tools.regrid import LateralRegrid, VerticalRegrid
+from roms_tools.vertical_coordinate import compute_depth
+from roms_tools.plot import _section_plot, _line_plot
+from roms_tools.utils import (
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
+    transpose_dimensions,
+)
 from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
-from roms_tools.setup.vertical_coordinate import compute_depth
 from roms_tools.setup.utils import (
     get_variable_metadata,
     group_dataset,
@@ -17,20 +25,14 @@ from roms_tools.setup.utils import (
     get_target_coords,
     rotate_velocities,
     compute_barotropic_velocity,
-    transpose_dimensions,
     one_dim_fill,
     nan_check,
     substitute_nans_by_fillvalue,
-    interpolate_from_rho_to_u,
-    interpolate_from_rho_to_v,
     convert_to_roms_time,
     get_boundary_coords,
     _to_yaml,
     _from_yaml,
 )
-from roms_tools.setup.plot import _section_plot, _line_plot
-import matplotlib.pyplot as plt
-from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -827,11 +829,20 @@ class BoundaryForcing:
         var_name_wo_direction, direction = var_name.split("_")
         location = self.variable_info[var_name_wo_direction]["location"]
 
+        # Find correct mask
+        if location == "rho":
+            mask = self.grid.ds.mask_rho
+        elif location == "u":
+            mask = self.grid.ds.mask_u
+        elif location == "v":
+            mask = self.grid.ds.mask_v
+
+        mask = mask.isel(**self.bdry_coords[location][direction])
+
         if "s_rho" in field.dims:
             field = field.assign_coords(
                 {"layer_depth": self.grid.ds[f"layer_depth_{location}_{direction}"]}
             )
-        # chose colorbar
         if var_name.startswith(("u", "v", "ubar", "vbar", "zeta")):
             vmax = max(field.max().values, -field.min().values)
             vmin = -vmax
@@ -872,14 +883,14 @@ class BoundaryForcing:
                 interface_depth = None
 
             _section_plot(
-                field,
+                field.where(mask),
                 interface_depth=interface_depth,
                 title=title,
                 kwargs=kwargs,
                 ax=ax,
             )
         else:
-            _line_plot(field, title=title, ax=ax)
+            _line_plot(field.where(mask), title=title, ax=ax)
 
     def save(
         self,

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,
@@ -1945,138 +1944,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],

roms_tools/setup/grid.py

@@ -10,14 +10,18 @@ import importlib.metadata
 from typing import Dict, Union, List
 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 (
+    sigma_stretch,
+    compute_depth,
+    add_depth_coordinates_to_dataset,
+)
+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 pathlib import Path
 
@@ -412,13 +416,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 +434,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,12 +443,44 @@ 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 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,
@@ -480,7 +517,11 @@ class Grid:
             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})
+        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})
 
         # slice the bathymetry as desired
         if eta is not None:
@@ -505,9 +546,7 @@ class Grid:
             kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
             _plot(
-                self.ds,
                 field=layer_depth.where(self.ds.mask_rho),
-                straddle=self.straddle,
                 depth_contours=False,
                 title=title,
                 kwargs=kwargs,