roms-tools 3.1.1__py3-none-any.whl → 3.2.0__py3-none-any.whl

roms_tools/plot.py

@@ -1,6 +1,7 @@
 from typing import Any, Literal
 
 import cartopy.crs as ccrs
+import matplotlib.dates as mdates
 import matplotlib.pyplot as plt
 import numpy as np
 import xarray as xr
@@ -9,18 +10,16 @@ from matplotlib.figure import Figure
 
 from roms_tools.regrid import LateralRegridFromROMS, VerticalRegridFromROMS
 from roms_tools.utils import (
-    _generate_coordinate_range,
-    _remove_edge_nans,
+    generate_coordinate_range,
     infer_nominal_horizontal_resolution,
     normalize_longitude,
+    remove_edge_nans,
 )
 from roms_tools.vertical_coordinate import compute_depth_coordinates
 
 LABEL_COLOR = "k"
 LABEL_SZ = 10
 FONT_SZ = 10
-EDGE_POS_START = "start"
-EDGE_POS_END = "end"
 
 
 def _add_gridlines(ax: Axes) -> None:
@@ -212,7 +211,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 +230,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 +259,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 +292,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 +326,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 +344,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 +366,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 +386,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 +396,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 +413,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 +435,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"]
@@ -458,16 +496,16 @@ def line_plot(field, title="", ax=None):
 
 
 def _get_edge(
-    arr: xr.DataArray, dim_name: str, pos: Literal[EDGE_POS_START, EDGE_POS_END]
+    arr: xr.DataArray, dim_name: str, pos: Literal["start", "end"]
 ) -> xr.DataArray:
     """Extract the first ("start") or last ("end") slice along the given dimension."""
-    if pos == EDGE_POS_START:
+    if pos == "start":
         return arr.isel({dim_name: 0})
 
-    if pos == EDGE_POS_END:
+    if pos == "end":
         return arr.isel({dim_name: -1})
 
-    raise ValueError(f"pos must be {EDGE_POS_START} or {EDGE_POS_END}")
+    raise ValueError("pos must be `start` or `end`")
 
 
 def _add_boundary_to_ax(
@@ -499,23 +537,23 @@ def _add_boundary_to_ax(
 
     edges = [
         (
-            _get_edge(lon_deg, xi_dim, EDGE_POS_START),
-            _get_edge(lat_deg, xi_dim, EDGE_POS_START),
+            _get_edge(lon_deg, xi_dim, "start"),
+            _get_edge(lat_deg, xi_dim, "start"),
             r"$\eta$",
         ),  # left
         (
-            _get_edge(lon_deg, xi_dim, EDGE_POS_END),
-            _get_edge(lat_deg, xi_dim, EDGE_POS_END),
+            _get_edge(lon_deg, xi_dim, "end"),
+            _get_edge(lat_deg, xi_dim, "end"),
             r"$\eta$",
         ),  # right
         (
-            _get_edge(lon_deg, eta_dim, EDGE_POS_START),
-            _get_edge(lat_deg, eta_dim, EDGE_POS_START),
+            _get_edge(lon_deg, eta_dim, "start"),
+            _get_edge(lat_deg, eta_dim, "start"),
             r"$\xi$",
         ),  # bottom
         (
-            _get_edge(lon_deg, eta_dim, EDGE_POS_END),
-            _get_edge(lat_deg, eta_dim, EDGE_POS_END),
+            _get_edge(lon_deg, eta_dim, "end"),
+            _get_edge(lat_deg, eta_dim, "end"),
             r"$\xi$",
         ),  # top
     ]
@@ -775,11 +813,12 @@ 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,
     save_path: str | None = None,
-    cmap_name: str | None = "YlOrRd",
+    cmap_name: str = "YlOrRd",
     add_colorbar: bool = True,
 ) -> None:
     """Generate a plot of a 2D or 3D ROMS field for a horizontal or vertical slice.
@@ -838,6 +877,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.
 
@@ -1006,7 +1051,7 @@ def plot(
             title = title + f", lat = {lat}°N"
         else:
             resolution = infer_nominal_horizontal_resolution(grid_ds)
-            lats = _generate_coordinate_range(
+            lats = generate_coordinate_range(
                 field.lat.min().values, field.lat.max().values, resolution
             )
         lats = xr.DataArray(lats, dims=["lat"], attrs={"units": "°N"})
@@ -1016,7 +1061,7 @@ def plot(
             title = title + f", lon = {lon}°E"
         else:
             resolution = infer_nominal_horizontal_resolution(grid_ds, lat)
-            lons = _generate_coordinate_range(
+            lons = generate_coordinate_range(
                 field.lon.min().values, field.lon.max().values, resolution
             )
         lons = xr.DataArray(lons, dims=["lon"], attrs={"units": "°E"})
@@ -1033,11 +1078,11 @@ def plot(
         field = field.assign_coords({"layer_depth": layer_depth})
 
     if lat is not None:
-        field, layer_depth = _remove_edge_nans(
+        field, layer_depth = remove_edge_nans(
             field, "lon", layer_depth if "layer_depth" in locals() else None
         )
     if lon is not None:
-        field, layer_depth = _remove_edge_nans(
+        field, layer_depth = remove_edge_nans(
             field, "lat", layer_depth if "layer_depth" in locals() else None
         )
 
@@ -1086,14 +1131,15 @@ 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")
@@ -1205,3 +1251,58 @@ def plot_location(
 
     if include_legend:
         ax.legend(loc="center left", bbox_to_anchor=(1.1, 0.5))
+
+
+def plot_uptake_efficiency(ds: xr.Dataset) -> None:
+    """
+    Plot Carbon Dioxide Removal (CDR) uptake efficiency over time.
+
+    This function plots two estimates of uptake efficiency stored in the dataset:
+    1. `cdr_efficiency`, computed from CO2 flux differences.
+    2. `cdr_efficiency_from_delta_diff`, computed from DIC differences.
+
+    The x-axis shows absolute time, formatted as YYYY-MM-DD, and the y-axis shows
+    the uptake efficiency values. The plot includes a legend and grid for clarity.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        Dataset containing the following variables:
+        - "abs_time": array of timestamps (datetime-like)
+        - "cdr_efficiency": uptake efficiency from flux differences
+        - "cdr_efficiency_from_delta_diff": uptake efficiency from DIC differences
+
+    Raises
+    ------
+    ValueError
+        If required variables are missing or empty.
+
+    Returns
+    -------
+    None
+    """
+    required_vars = ["abs_time", "cdr_efficiency", "cdr_efficiency_from_delta_diff"]
+    for var in required_vars:
+        if var not in ds or ds[var].size == 0:
+            raise ValueError(f"Dataset must contain non-empty variable '{var}'.")
+
+    times = ds["abs_time"]
+
+    # Check for monotonically increasing times
+    if not np.all(times[1:] >= times[:-1]):
+        raise ValueError("abs_time must be strictly increasing.")
+
+    fig, ax = plt.subplots(figsize=(10, 4))
+
+    ax.plot(times, ds["cdr_efficiency"], label="from CO2 flux differences", lw=2)
+    ax.plot(
+        times, ds["cdr_efficiency_from_delta_diff"], label="from DIC differences", lw=2
+    )
+    ax.grid()
+    ax.set_title("CDR uptake efficiency")
+    ax.legend()
+
+    # Format x-axis as YYYY-MM-DD
+    ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m-%d"))
+    fig.autofmt_xdate()
+    plt.show()

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,
+    GLORYSDataset,
+    GLORYSDefaultDataset,
+    RawDataSource,
+    UnifiedBGCDataset,
+)
 from roms_tools.setup.utils import (
     add_time_info_to_ds,
     compute_barotropic_velocity,
@@ -56,7 +63,7 @@ class BoundaryForcing:
         If no time filtering is desired, set it to None. Default is None.
     boundaries : Dict[str, bool], optional
         Dictionary specifying which boundaries are forced (south, east, north, west). Default is all True.
-    source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool]
+    source : RawDataSource
         Dictionary specifying the source of the boundary forcing data. Keys include:
 
           - "name" (str): Name of the data source (e.g., "GLORYS").
@@ -64,7 +71,9 @@ class BoundaryForcing:
 
             - A single string (with or without wildcards).
             - A single Path object.
-            - A list of strings or Path objects containing multiple files.
+            - A list of strings or Path objects.
+            If omitted, the data will be streamed via the Copernicus Marine Toolkit.
+            Note: streaming is currently not recommended due to performance limitations.
           - "climatology" (bool): Indicates if the data is climatology data. Defaults to False.
 
     type : str
@@ -117,7 +126,7 @@ class BoundaryForcing:
         }
     )
     """Dictionary specifying which boundaries are forced (south, east, north, west)."""
-    source: dict[str, str | Path | list[str | Path]]
+    source: RawDataSource
     """Dictionary specifying the source of the boundary forcing data."""
     type: str = "physics"
     """Specifies the type of forcing data ("physics", "bgc")."""
@@ -150,7 +159,6 @@ class BoundaryForcing:
         if self.apply_2d_horizontal_fill:
             data.choose_subdomain(
                 target_coords,
-                buffer_points=20,  # lateral fill needs good buffer from data margin
             )
             # Enforce double precision to ensure reproducibility
             data.convert_to_float64()
@@ -181,8 +189,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]
@@ -290,14 +298,12 @@ class BoundaryForcing:
                     zeta_v = zeta_v.isel(**self.bdry_coords["v"][direction])
 
                 if not self.apply_2d_horizontal_fill and bdry_data.needs_lateral_fill:
-                    logging.info(
-                        f"Applying 1D horizontal fill to {direction}ern boundary."
-                    )
-                    self._validate_1d_fill(
-                        processed_fields,
-                        direction,
-                        bdry_data.dim_names["depth"],
-                    )
+                    if not self.bypass_validation:
+                        self._validate_1d_fill(
+                            processed_fields,
+                            direction,
+                            bdry_data.dim_names["depth"],
+                        )
                     for var_name in processed_fields:
                         processed_fields[var_name] = apply_1d_horizontal_fill(
                             processed_fields[var_name]
@@ -403,7 +409,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 +434,68 @@ 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,
+    ) -> GLORYSDataset | GLORYSDefaultDataset | CESMBGCDataset | UnifiedBGCDataset:
+        """Determine the correct `Dataset` type and return an instance.
+
+        Returns
+        -------
+        Dataset
+            The `Dataset` instance
+
+        """
+        dataset_map: dict[
+            str,
+            dict[
+                str,
+                dict[
+                    str,
+                    type[
+                        GLORYSDataset
+                        | GLORYSDefaultDataset
+                        | CESMBGCDataset
+                        | UnifiedBGCDataset
+                    ],
+                ],
+            ],
+        ] = {
+            "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"
 
-        return data
+        data_type = dataset_map[self.type][source_name][variant]
+
+        if isinstance(self.source["path"], bool):
+            raise ValueError('source["path"] cannot be a boolean here')
+
+        return data_type(
+            filename=self.source["path"],
+            start_time=self.start_time,
+            end_time=self.end_time,
+            climatology=self.source["climatology"],  # type: ignore[arg-type]
+            use_dask=self.use_dask,
+        )
 
     def _set_variable_info(self, data):
         """Sets up a dictionary with metadata for variables based on the type of data
@@ -731,6 +774,9 @@ class BoundaryForcing:
         None
             If a boundary is divided by land, a warning is issued. No return value is provided.
         """
+        if not hasattr(self, "_warned_directions"):
+            self._warned_directions = set()
+
         for var_name in processed_fields.keys():
             if self.variable_info[var_name]["validate"]:
                 location = self.variable_info[var_name]["location"]
@@ -753,16 +799,20 @@ class BoundaryForcing:
                 wet_nans = xr.where(da.where(mask).isnull(), 1, 0)
                 # Apply label to find connected components of wet NaNs
                 labeled_array, num_features = label(wet_nans)
+
                 left_margin = labeled_array[0]
                 right_margin = labeled_array[-1]
                 if left_margin != 0:
                     num_features = num_features - 1
                 if right_margin != 0:
                     num_features = num_features - 1
-                if num_features > 0:
+
+                if num_features > 0 and direction not in self._warned_directions:
                     logging.warning(
-                        f"For {var_name}, the {direction}ern boundary is divided by land. It would be safer (but slower) to use `apply_2d_horizontal_fill = True`."
+                        f"The {direction}ern boundary is divided by land. "
+                        "It would be safer (but slower and more memory-intensive) to use `apply_2d_horizontal_fill = True`."
                     )
+                    self._warned_directions.add(direction)
 
     def _validate(self, ds):
         """Validate the dataset for NaN values at the first time step (bry_time=0) for
@@ -797,8 +847,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