roms-tools 3.1.0__py3-none-any.whl → 3.1.2__py3-none-any.whl

roms_tools/__init__.py

@@ -20,5 +20,9 @@ from roms_tools.setup.surface_forcing import SurfaceForcing  # noqa: F401
 from roms_tools.setup.tides import TidalForcing  # noqa: F401
 from roms_tools.tiling.partition import partition_netcdf  # noqa: F401
 
+
 # Configure logging when the package is imported
-logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
+LOG_FORMAT = "%(asctime)s - %(levelname)s - %(message)s"
+DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+logging.basicConfig(level=logging.INFO, format=LOG_FORMAT, datefmt=DATE_FORMAT)

roms_tools/constants.py

@@ -3,3 +3,4 @@ MAXIMUM_GRID_SIZE = 25000  # in km
 UPPER_BOUND_THETA_S = 10  # upper bound for surface vertical stretching parameter
 UPPER_BOUND_THETA_B = 10  # upper bound for bottom vertical stretching parameter
 NUM_TRACERS = 34  # Number of tracers (temperature, salinity, BGC tracers)
+MAX_DISTINCT_COLORS = 20  # Based on tab20 colormap

roms_tools/plot.py

@@ -212,7 +212,14 @@ def plot_nesting(parent_grid_ds, child_grid_ds, parent_straddle, with_dim_names=
     return fig
 
 
-def section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
+def section_plot(
+    field: xr.DataArray,
+    interface_depth: xr.DataArray | None = None,
+    title: str = "",
+    yincrease: bool | None = False,
+    kwargs: dict = {},
+    ax: Axes | None = None,
+):
     """Plots a vertical section of a field with optional interface depths.
 
     Parameters
@@ -224,6 +231,11 @@ def section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
         Defaults to None.
     title : str, optional
         Title of the plot. Defaults to an empty string.
+    yincrease : bool or None, optional
+        Whether to orient the y-axis with increasing values upward.
+        If True, y-values increase upward (standard).
+        If False, y-values decrease upward (inverted).
+        If None (default), behavior is equivalent to False (inverted axis).
     kwargs : dict, optional
         Additional keyword arguments to pass to `xarray.plot`. Defaults to an empty dictionary.
     ax : matplotlib.axes.Axes, optional
@@ -248,6 +260,8 @@ def section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
     """
     if ax is None:
         fig, ax = plt.subplots(1, 1, figsize=(9, 5))
+    if yincrease is None:
+        yincrease = False
 
     dims_to_check = ["eta_rho", "eta_v", "xi_rho", "xi_u", "lat", "lon"]
     try:
@@ -279,7 +293,7 @@ def section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
     # Handle NaNs on either horizontal end
     field = field.where(~field[depth_label].isnull(), drop=True)
 
-    more_kwargs = {"x": xdim, "y": depth_label, "yincrease": False}
+    more_kwargs = {"x": xdim, "y": depth_label, "yincrease": yincrease}
 
     field.plot(**kwargs, **more_kwargs, ax=ax)
 
@@ -313,7 +327,12 @@ def section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
         return fig
 
 
-def profile_plot(field, title="", ax=None):
+def profile_plot(
+    field: xr.DataArray,
+    title: str = "",
+    yincrease: bool | None = False,
+    ax: Axes | None = None,
+):
     """Plots a vertical profile of the given field against depth.
 
     This function generates a profile plot by plotting the field values against
@@ -326,6 +345,11 @@ def profile_plot(field, title="", ax=None):
         The field to plot, typically representing vertical profile data.
     title : str, optional
         Title of the plot. Defaults to an empty string.
+    yincrease : bool or None, optional
+        Whether to orient the y-axis with increasing values upward.
+        If True, y-values increase upward (standard).
+        If False, y-values decrease upward (inverted).
+        If None (default), behavior is equivalent to False (inverted axis).
     ax : matplotlib.axes.Axes, optional
         Pre-existing axes to draw the plot on. If None, a new figure and axes are created.
 
@@ -343,6 +367,9 @@ def profile_plot(field, title="", ax=None):
     -----
     - The y-axis is inverted to ensure that depth increases downward.
     """
+    if yincrease is None:
+        yincrease = False
+
     depths_to_check = [
         "layer_depth",
         "interface_depth",
@@ -360,8 +387,8 @@ def profile_plot(field, title="", ax=None):
 
     if ax is None:
         fig, ax = plt.subplots(1, 1, figsize=(4, 7))
-    kwargs = {"y": depth_label, "yincrease": False}
-    field.plot(**kwargs, linewidth=2)
+    kwargs = {"y": depth_label, "yincrease": yincrease}
+    field.plot(ax=ax, linewidth=2, **kwargs)
     ax.set_title(title)
     ax.set_ylabel("Depth [m]")
     ax.grid()
@@ -370,7 +397,12 @@ def profile_plot(field, title="", ax=None):
         return fig
 
 
-def line_plot(field, title="", ax=None):
+def line_plot(
+    field: xr.DataArray,
+    title: str = "",
+    ax: Axes | None = None,
+    yincrease: bool | None = False,
+):
     """Plots a line graph of the given field with grey vertical bars indicating NaN
     regions.
 
@@ -382,6 +414,11 @@ def line_plot(field, title="", ax=None):
         Title of the plot. Defaults to an empty string.
     ax : matplotlib.axes.Axes, optional
         Pre-existing axes to draw the plot on. If None, a new figure and axes are created.
+    yincrease : bool, optional
+        Whether to orient the y-axis with increasing values upward.
+        If True, y-values increase upward (standard).
+        If False, y-values decrease upward (inverted).
+        If None (default), behavior is equivalent to True (standard axis).
 
     Returns
     -------
@@ -399,10 +436,12 @@ def line_plot(field, title="", ax=None):
     -----
     - NaN regions are identified and marked using `axvspan` with a grey shade.
     """
+    if yincrease is None:
+        yincrease = True
     if ax is None:
         fig, ax = plt.subplots(1, 1, figsize=(7, 4))
 
-    field.plot(ax=ax, linewidth=2)
+    field.plot(ax=ax, linewidth=2, yincrease=yincrease)
 
     # Loop through the NaNs in the field and add grey vertical bars
     dims_to_check = ["eta_rho", "eta_v", "xi_rho", "xi_u", "lat", "lon"]
@@ -775,6 +814,7 @@ def plot(
     depth_contours: bool = False,
     layer_contours: bool = False,
     max_nr_layer_contours: int | None = 10,
+    yincrease: bool | None = None,
     use_coarse_grid: bool = False,
     with_dim_names: bool = False,
     ax: Axes | None = None,
@@ -838,6 +878,12 @@ def plot(
     max_nr_layer_contours : int, optional
         Maximum number of vertical layer contours to draw. Default is 10.
 
+    yincrease: bool, optional
+        If True, the y-axis values increase upward (standard orientation).
+        If False, the y-axis values decrease upward (inverted axis).
+        If None (default), the orientation is determined by the default behavior
+        of the underlying plotting function.
+
     use_coarse_grid : bool, optional
         Use precomputed coarse-resolution grid. Default is False.
 
@@ -1086,14 +1132,123 @@ def plot(
                 field,
                 interface_depth=interface_depth,
                 title=title,
+                yincrease=yincrease,
                 kwargs={**kwargs, "add_colorbar": add_colorbar},
                 ax=ax,
             )
         else:
             if "s_rho" in field.dims:
-                fig = profile_plot(field, title=title, ax=ax)
+                fig = profile_plot(field, title=title, yincrease=yincrease, ax=ax)
             else:
-                fig = line_plot(field, title=title, ax=ax)
+                fig = line_plot(field, title=title, ax=ax, yincrease=yincrease)
 
     if save_path:
         plt.savefig(save_path, dpi=300, bbox_inches="tight")
+
+
+def assign_category_colors(names: list[str]) -> dict[str, tuple]:
+    """
+    Assign a distinct color to each name using a Matplotlib categorical colormap.
+
+    Parameters
+    ----------
+    names : list[str]
+        List of category names (e.g., releases, rivers, etc.) to assign colors to.
+
+    Returns
+    -------
+    dict[str, tuple]
+        Dictionary mapping each name to a unique RGBA color.
+
+    Raises
+    ------
+    ValueError
+        If the number of names exceeds the selected colormap's capacity.
+
+    Notes
+    -----
+    Colormap selection is based on the number of items:
+    - <= 10: 'tab10'
+    - <= 20: 'tab20'
+    - > 20 : 'tab20b'
+    """
+    n = len(names)
+
+    if n <= 10:
+        cmap = plt.get_cmap("tab10")
+    elif n <= 20:
+        cmap = plt.get_cmap("tab20")
+
+    if n > cmap.N:
+        raise ValueError(
+            f"Too many categories ({n}) for selected colormap ({cmap.name}) "
+            f"which supports only {cmap.N} distinct entries."
+        )
+
+    return {name: cmap(i) for i, name in enumerate(names)}
+
+
+def plot_location(
+    grid_ds: xr.Dataset,
+    points: dict[str, dict],
+    ax: Axes,
+    include_legend: bool = True,
+) -> None:
+    """Plot named geographic points on a top-down map view.
+
+    Each point is represented as a marker on the map, optionally colored.
+    This function is generic and can be used for releases, rivers, etc.
+
+    Parameters
+    ----------
+    grid_ds : xr.Dataset
+        The grid dataset containing 'lon_rho' and 'lat_rho', and a 'straddle' attribute.
+
+    points : dict[str, dict]
+        Dictionary of points to plot. Keys are point names. Each value is a dict with:
+        - "lat": float, latitude in degrees
+        - "lon": float, longitude in degrees
+        - Optional "color": tuple or str, matplotlib color
+
+    ax : matplotlib.axes.Axes
+        The axis object to plot on.
+
+    include_legend : bool, default True
+        Whether to include a legend showing point names.
+
+    Returns
+    -------
+    None
+    """
+    lon_deg = grid_ds.lon_rho
+    lat_deg = grid_ds.lat_rho
+
+    if "straddle" not in grid_ds.attrs:
+        raise AttributeError("Grid dataset must have a 'straddle' attribute.")
+
+    straddle = grid_ds.attrs["straddle"] == "True"
+    if straddle:
+        lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
+
+    trans = get_projection(lon_deg, lat_deg)
+    proj = ccrs.PlateCarree()
+
+    for name, info in points.items():
+        lon = info["lon"]
+        lat = info["lat"]
+        color = info.get("color", "k")  # Default to black if no color specified
+
+        x, y = trans.transform_point(lon, lat, proj)
+
+        ax.plot(
+            x,
+            y,
+            marker="x",
+            markersize=8,
+            markeredgewidth=2,
+            label=name,
+            color=color,
+        )
+
+    if include_legend:
+        ax.legend(loc="center left", bbox_to_anchor=(1.1, 0.5))

roms_tools/regrid.py

@@ -251,7 +251,12 @@ class VerticalRegridFromROMS:
         ds : xarray.Dataset
             The dataset containing the ROMS output data, which must include the vertical coordinate `s_rho`.
         """
-        self.grid = xgcm.Grid(ds, coords={"s_rho": {"center": "s_rho"}}, periodic=False)
+        self.grid = xgcm.Grid(
+            ds,
+            coords={"s_rho": {"center": "s_rho"}},
+            periodic=False,
+            autoparse_metadata=False,
+        )
 
     def apply(self, da, depth_coords, target_depth_levels, mask_edges=True):
         """Applies vertical regridding from ROMS to the specified target depth levels.

roms_tools/setup/boundary_forcing.py

@@ -1,5 +1,6 @@
 import importlib.metadata
 import logging
+from collections import defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
@@ -12,7 +13,13 @@ from scipy.ndimage import label
 from roms_tools import Grid
 from roms_tools.plot import line_plot, section_plot
 from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
-from roms_tools.setup.datasets import CESMBGCDataset, GLORYSDataset, UnifiedBGCDataset
+from roms_tools.setup.datasets import (
+    CESMBGCDataset,
+    Dataset,
+    GLORYSDataset,
+    GLORYSDefaultDataset,
+    UnifiedBGCDataset,
+)
 from roms_tools.setup.utils import (
     add_time_info_to_ds,
     compute_barotropic_velocity,
@@ -181,8 +188,8 @@ class BoundaryForcing:
             }
         )
 
-        for direction in ["south", "east", "north", "west"]:
-            if self.boundaries[direction]:
+        for direction, is_enabled in self.boundaries.items():
+            if is_enabled:
                 bdry_target_coords = {
                     "lat": target_coords["lat"].isel(
                         **self.bdry_coords["vector"][direction]
@@ -403,7 +410,10 @@ class BoundaryForcing:
         if "name" not in self.source:
             raise ValueError("`source` must include a 'name'.")
         if "path" not in self.source:
-            raise ValueError("`source` must include a 'path'.")
+            if self.source["name"] != "GLORYS":
+                raise ValueError("`source` must include a 'path'.")
+
+            self.source["path"] = GLORYSDefaultDataset.dataset_name
 
         # Set 'climatology' to False if not provided in 'source'
         self.source = {
@@ -425,34 +435,49 @@ class BoundaryForcing:
                 "Sea surface height will NOT be used to adjust depth coordinates."
             )
 
-    def _get_data(self):
-        data_dict = {
-            "filename": self.source["path"],
-            "start_time": self.start_time,
-            "end_time": self.end_time,
-            "climatology": self.source["climatology"],
-            "use_dask": self.use_dask,
+    def _get_data(self) -> Dataset:
+        """Determine the correct `Dataset` type and return an instance.
+
+        Returns
+        -------
+        Dataset
+            The `Dataset` instance
+
+        """
+        dataset_map: dict[str, dict[str, dict[str, type[Dataset]]]] = {
+            "physics": {
+                "GLORYS": {
+                    "external": GLORYSDataset,
+                    "default": GLORYSDefaultDataset,
+                },
+            },
+            "bgc": {
+                "CESM_REGRIDDED": defaultdict(lambda: CESMBGCDataset),
+                "UNIFIED": defaultdict(lambda: UnifiedBGCDataset),
+            },
         }
 
-        if self.type == "physics":
-            if self.source["name"] == "GLORYS":
-                data = GLORYSDataset(**data_dict)
-            else:
-                raise ValueError(
-                    'Only "GLORYS" is a valid option for source["name"] when type is "physics".'
-                )
+        source_name = str(self.source["name"])
+        if source_name not in dataset_map[self.type]:
+            tpl = 'Valid options for source["name"] for type {} include: {}'
+            msg = tpl.format(self.type, " and ".join(dataset_map[self.type].keys()))
+            raise ValueError(msg)
 
-        elif self.type == "bgc":
-            if self.source["name"] == "CESM_REGRIDDED":
-                data = CESMBGCDataset(**data_dict)
-            elif self.source["name"] == "UNIFIED":
-                data = UnifiedBGCDataset(**data_dict)
-            else:
-                raise ValueError(
-                    'Only "CESM_REGRIDDED" and "UNIFIED" are valid options for source["name"] when type is "bgc".'
-                )
+        has_no_path = "path" not in self.source
+        has_default_path = self.source.get("path") == GLORYSDefaultDataset.dataset_name
+        use_default = has_no_path or has_default_path
+
+        variant = "default" if use_default else "external"
+
+        data_type = dataset_map[self.type][source_name][variant]
 
-        return data
+        return data_type(
+            filename=self.source["path"],
+            start_time=self.start_time,
+            end_time=self.end_time,
+            climatology=self.source["climatology"],
+            use_dask=self.use_dask,
+        )  # type: ignore
 
     def _set_variable_info(self, data):
         """Sets up a dictionary with metadata for variables based on the type of data
@@ -797,8 +822,8 @@ class BoundaryForcing:
                 elif location == "v":
                     mask = self.grid.ds.mask_v
 
-                for direction in ["south", "east", "north", "west"]:
-                    if self.boundaries[direction]:
+                for direction, is_enabled in self.boundaries.items():
+                    if is_enabled:
                         bdry_var_name = f"{var_name}_{direction}"
 
                         # Check for NaN values at the first time step using the nan_check function