roms-tools 3.1.2__py3-none-any.whl → 3.3.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:
@@ -97,12 +96,6 @@ def plot_2d_horizontal_field(
     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."
-        )
-
     trans = get_projection(lon_deg, lat_deg)
 
     if ax is None:
@@ -187,25 +180,8 @@ def plot_nesting(parent_grid_ds, child_grid_ds, parent_straddle, with_dim_names=
         with_dim_names=with_dim_names,
     )
 
-    vmax = 3
-    vmin = 0
-    cmap = plt.colormaps.get_cmap("Blues")
-    cmap.set_bad(color="gray")
-    kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
-
-    field = parent_grid_ds.mask_rho
-    field = field.where(field)
-
-    _add_field_to_ax(
-        ax,
-        parent_lon_deg,
-        parent_lat_deg,
-        field,
-        add_colorbar=False,
-        kwargs=kwargs,
-    )
-
     _add_gridlines(ax)
+    ax.coastlines()
 
     ax.legend(loc="best")
 
@@ -497,16 +473,16 @@ def line_plot(
 
 
 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(
@@ -538,23 +514,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
     ]
@@ -660,7 +636,7 @@ def _add_field_to_ax(
     proj = ccrs.PlateCarree()
 
     p = ax.pcolormesh(lon_deg, lat_deg, field, transform=proj, **kwargs)
-    if hasattr(field, "long_name"):
+    if hasattr(field, "long_name") and hasattr(field, "units"):
         label = f"{field.long_name} [{field.units}]"
     elif hasattr(field, "Long_name"):
         # this is the case for matlab generated grids
@@ -675,10 +651,36 @@ def _add_field_to_ax(
         ax.clabel(cs, inline=True, fontsize=FONT_SZ)
 
 
-def get_projection(lon, lat):
-    return ccrs.NearsidePerspective(
-        central_longitude=lon.mean().values, central_latitude=lat.mean().values
-    )
+def get_projection(lon: xr.DataArray, lat: xr.DataArray):
+    """
+    Return a Cartopy projection appropriate for the given lon/lat domain.
+
+    - Raises an error if the domain includes either pole.
+    - Uses PlateCarree if the longitudinal span > 90°.
+    - Otherwise uses NearsidePerspective centered on the domain.
+
+    Parameters
+    ----------
+    lon, lat : xr.DataArray
+        Longitude and latitude values in degrees.
+
+    Returns
+    -------
+    cartopy.crs.Projection
+        The chosen Cartopy projection.
+    """
+    # check if North or South pole are in domain
+    if lat.max() > 89 or lat.min() < -89:
+        raise NotImplementedError(
+            "Plotting is not implemented for the case that the domain contains the North or South pole."
+        )
+
+    if lon.max() - lon.min() > 90:
+        return ccrs.PlateCarree(central_longitude=lon.mean().values)
+    else:
+        return ccrs.NearsidePerspective(
+            central_longitude=lon.mean().values, central_latitude=lat.mean().values
+        )
 
 
 def _validate_plot_inputs(
@@ -817,9 +819,10 @@ def plot(
     yincrease: bool | None = None,
     use_coarse_grid: bool = False,
     with_dim_names: bool = False,
+    apply_mask: bool = True,
     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.
@@ -890,6 +893,9 @@ def plot(
     with_dim_names : bool, optional
         Add grid dimension names (`xi`, `eta`) to the outer plot edges. Only for 2D plots. Default is False.
 
+    apply_mask: bool, optional
+        Whether to apply the land mask to the field. Default is True.
+
     ax : Axes, optional
         Matplotlib axes object. If None, a new figure is created. Default is None.
 
@@ -972,7 +978,8 @@ def plot(
     field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
 
     # Mask the field
-    field = field.where(mask)
+    if apply_mask:
+        field = field.where(mask)
 
     # Assign eta and xi as coordinates
     coords_to_assign = {dim: field[dim] for dim in horizontal_dims.values()}
@@ -1052,7 +1059,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"})
@@ -1062,7 +1069,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"})
@@ -1079,11 +1086,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
         )
 
@@ -1252,3 +1259,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/setup/boundary_forcing.py

@@ -13,14 +13,14 @@ 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 (
+from roms_tools.setup.lat_lon_datasets import (
     CESMBGCDataset,
-    Dataset,
     GLORYSDataset,
     GLORYSDefaultDataset,
     UnifiedBGCDataset,
 )
 from roms_tools.setup.utils import (
+    RawDataSource,
     add_time_info_to_ds,
     compute_barotropic_velocity,
     compute_missing_bgc_variables,
@@ -63,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").
@@ -71,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
@@ -124,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")."""
@@ -157,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()
@@ -297,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]
@@ -435,7 +434,9 @@ class BoundaryForcing:
                 "Sea surface height will NOT be used to adjust depth coordinates."
             )
 
-    def _get_data(self) -> Dataset:
+    def _get_data(
+        self,
+    ) -> GLORYSDataset | GLORYSDefaultDataset | CESMBGCDataset | UnifiedBGCDataset:
         """Determine the correct `Dataset` type and return an instance.
 
         Returns
@@ -444,7 +445,21 @@ class BoundaryForcing:
             The `Dataset` instance
 
         """
-        dataset_map: dict[str, dict[str, dict[str, type[Dataset]]]] = {
+        dataset_map: dict[
+            str,
+            dict[
+                str,
+                dict[
+                    str,
+                    type[
+                        GLORYSDataset
+                        | GLORYSDefaultDataset
+                        | CESMBGCDataset
+                        | UnifiedBGCDataset
+                    ],
+                ],
+            ],
+        ] = {
             "physics": {
                 "GLORYS": {
                     "external": GLORYSDataset,
@@ -471,13 +486,16 @@ class BoundaryForcing:
 
         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"],
+            climatology=self.source["climatology"],  # type: ignore[arg-type]
             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
@@ -756,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"]
@@ -778,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

roms_tools/setup/cdr_forcing.py

@@ -2,19 +2,19 @@ import itertools
 import logging
 from collections import Counter
 from collections.abc import Iterator
-from datetime import datetime
+from datetime import datetime, timedelta
 from pathlib import Path
 from typing import Annotated
 
 import matplotlib.gridspec as gridspec
 import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
 import xarray as xr
 from pydantic import (
     BaseModel,
     Field,
     RootModel,
-    conlist,
     model_serializer,
     model_validator,
 )
@@ -40,6 +40,7 @@ from roms_tools.setup.utils import (
     from_yaml,
     gc_dist,
     get_target_coords,
+    get_tracer_metadata_dict,
     to_dict,
     validate_names,
     write_to_yaml,
@@ -103,14 +104,16 @@ class ReleaseSimulationManager(BaseModel):
 class ReleaseCollector(RootModel):
     """Collects and validates multiple releases against each other."""
 
-    root: conlist(
-        Annotated[
-            VolumeRelease | TracerPerturbation, Field(discriminator="release_type")
+    root: Annotated[
+        list[
+            Annotated[
+                VolumeRelease | TracerPerturbation, Field(discriminator="release_type")
+            ]
         ],
-        min_length=1,
-    ) = Field(alias="releases")
+        Field(alias="releases", min_length=1),
+    ]
 
-    _release_type: ReleaseType = None
+    _release_type: ReleaseType | None = None
 
     def __iter__(self) -> Iterator[Release]:
         return iter(self.root)
@@ -126,6 +129,9 @@ class ReleaseCollector(RootModel):
         else:
             raise TypeError(f"Invalid key type: {type(item)}. Must be int or str.")
 
+    def __len__(self):
+        return len(self.root)
+
     @model_validator(mode="before")
     @classmethod
     def unpack_dict(cls, data):
@@ -774,6 +780,62 @@ class CDRForcing(BaseModel):
         fig.subplots_adjust(hspace=0.45)
         fig.suptitle(f"Release distribution for: {release_name}")
 
+    def compute_total_cdr_source(self, dt: float) -> pd.DataFrame:
+        """
+        Compute integrated tracer quantities for all releases and return a DataFrame.
+
+        Parameters
+        ----------
+        dt : float
+            Time step in seconds for reconstructing ROMS time stamps.
+
+        Returns
+        -------
+        pd.DataFrame
+            DataFrame with one row per release and one row of units at the top.
+            Columns 'temp' and 'salt' are excluded from integrated totals.
+        """
+        # Reconstruct ROMS time stamps
+        _, rel_seconds = _reconstruct_roms_time_stamps(
+            self.start_time, self.end_time, dt, self.model_reference_date
+        )
+
+        # Collect accounting results for all releases
+        records = []
+        release_names = []
+        for release in self.releases:
+            result = release._do_accounting(rel_seconds, self.model_reference_date)
+            records.append(result)
+            release_names.append(getattr(release, "name", f"release_{len(records)}"))
+
+        # Build DataFrame: rows = releases, columns = tracer names
+        df = pd.DataFrame(records, index=release_names)
+
+        # Exclude temp and salt from units row and integrated totals
+        integrated_tracers = [col for col in df.columns if col not in ("temp", "salt")]
+
+        # Add a row of units only for integrated tracers
+        tracer_meta = get_tracer_metadata_dict(include_bgc=True, unit_type="integrated")
+        units_row = {
+            col: tracer_meta.get(col, {}).get("units", "") for col in integrated_tracers
+        }
+
+        df_units = pd.DataFrame([units_row], index=["units"])
+
+        # Keep only integrated_tracers columns in df, drop temp and salt
+        df_integrated = df[integrated_tracers]
+
+        # Concatenate units row on top
+        df_final = pd.concat([df_units, df_integrated])
+
+        # Store dt as metadata
+        df_final.attrs["time_step"] = dt
+        df_final.attrs["start_time"] = self.start_time
+        df_final.attrs["end_time"] = self.end_time
+        df_final.attrs["title"] = "Integrated tracer releases"
+
+        return df_final
+
     def save(
         self,
         filepath: str | Path,
@@ -1051,3 +1113,55 @@ def _map_3d_gaussian(
         distribution_3d /= distribution_3d.sum()
 
     return distribution_3d
+
+
+def _reconstruct_roms_time_stamps(
+    start_time: datetime,
+    end_time: datetime,
+    dt: float,
+    model_reference_date: datetime,
+) -> tuple[list[datetime], np.ndarray]:
+    """
+    Reconstruct ROMS time stamps between `start_time` and `end_time` with step `dt`.
+
+    Parameters
+    ----------
+    start_time : datetime
+        Beginning of the time series.
+    end_time : datetime
+        End of the time series (inclusive if it falls exactly on a step).
+    dt : float
+        Time step in seconds (can be fractional if needed).
+    model_reference_date : datetime
+        The reference date for ROMS time (elapsed time will be relative to this).
+
+    Returns
+    -------
+    times : list of datetime
+        Sequence of datetimes from `start_time` to `end_time`.
+    rel_days : np.ndarray
+        Array of elapsed times in **seconds** relative to `model_reference_date`.
+
+    Raises
+    ------
+    ValueError
+        If `end_time` is not after `start_time` or if `dt` is not positive.
+    """
+    if end_time <= start_time:
+        raise ValueError("end_time must be after start_time")
+    if dt <= 0:
+        raise ValueError("dt must be positive")
+
+    # Generate absolute times
+    delta = timedelta(seconds=dt)
+    times: list[datetime] = []
+    t = start_time
+    while t <= end_time:
+        times.append(t)
+        t += delta
+
+    # Convert to relative ROMS time (days since model_reference_date)
+    rel_days = convert_to_relative_days(times, model_reference_date)
+    rel_seconds = rel_days * 3600 * 24
+
+    return times, rel_seconds