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

roms_tools/__init__.py

@@ -10,6 +10,7 @@ except ImportError:  # pragma: no cover
 # grid must be imported first
 from roms_tools.setup.grid import Grid  # noqa: I001, F401
 from roms_tools.analysis.roms_output import ROMSOutput  # noqa: F401
+from roms_tools.analysis.cdr_ensemble import Ensemble  # noqa: F401
 from roms_tools.setup.boundary_forcing import BoundaryForcing  # noqa: F401
 from roms_tools.setup.cdr_forcing import CDRForcing  # noqa: F401
 from roms_tools.setup.cdr_release import TracerPerturbation, VolumeRelease  # noqa: F401
@@ -19,6 +20,8 @@ from roms_tools.setup.river_forcing import RiverForcing  # noqa: F401
 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
+from roms_tools.setup.datasets import get_glorys_bounds  # noqa: F401
+from roms_tools.tiling.join import open_partitions, join_netcdf  # noqa: F401
 
 
 # Configure logging when the package is imported

roms_tools/analysis/cdr_analysis.py

@@ -0,0 +1,203 @@
+import logging
+
+import numpy as np
+import xarray as xr
+
+
+def compute_cdr_metrics(ds: xr.Dataset, grid_ds: xr.Dataset) -> xr.Dataset:
+    """
+    Compute Carbon Dioxide Removal (CDR) metrics from model output.
+
+    Calculates CDR uptake efficiency using two methods:
+      1. Flux-based: from area-integrated CO2 flux differences.
+      2. DIC difference-based: from volume-integrated DIC differences.
+
+    Copies selected tracer and flux variables and computes grid cell areas
+    and averaging window durations.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Model output with required variables:
+        'avg_begin_time', 'avg_end_time', 'ALK_source', 'DIC_source',
+        'FG_CO2', 'FG_ALT_CO2', 'hDIC', 'hDIC_ALT_CO2'.
+
+    grid_ds : xr.Dataset
+        Grid dataset with 'pm', 'pn' (inverse grid spacing).
+
+    Returns
+    -------
+    ds_cdr : xr.Dataset
+        Dataset containing:
+        - 'area', 'window_length'
+        - copied flux/tracer variables
+        - 'cdr_efficiency' and 'cdr_efficiency_from_delta_diff' (dimensionless)
+
+    Raises
+    ------
+    KeyError
+        If required variables are missing from `ds` or `grid_ds`.
+    """
+    # Define required variables
+    ds_vars = [
+        "avg_begin_time",
+        "avg_end_time",
+        "ALK_source",
+        "FG_CO2",
+        "FG_ALT_CO2",
+        "hDIC",
+        "hDIC_ALT_CO2",
+    ]
+    grid_vars = ["pm", "pn"]
+
+    # Check that all required variables exist
+    missing_ds = [var for var in ds_vars if var not in ds]
+    missing_grid = [var for var in grid_vars if var not in grid_ds]
+
+    if missing_ds:
+        raise KeyError(f"Missing required variables in ds: {missing_ds}")
+    if missing_grid:
+        raise KeyError(f"Missing required variables in grid_ds: {missing_grid}")
+
+    ds_cdr = xr.Dataset()
+
+    # Copy relevant variables
+    vars_to_copy = ["FG_CO2", "FG_ALT_CO2", "hDIC", "hDIC_ALT_CO2"]
+    for var_name in vars_to_copy:
+        ds_cdr[var_name] = ds[var_name]
+
+    # Grid cell area
+    ds_cdr["area"] = 1 / (grid_ds["pm"] * grid_ds["pn"])
+    ds_cdr["area"].attrs.update(
+        long_name="Grid cell area",
+        units="m^2",
+    )
+
+    # Duration of each averaging window
+    ds_cdr["window_length"] = ds["avg_end_time"] - ds["avg_begin_time"]
+    ds_cdr["window_length"].attrs.update(
+        long_name="Duration of each averaging window",
+        units="s",
+    )
+
+    _validate_source(ds)
+
+    # Cumulative alkalinity source
+    source = (
+        (
+            (ds["ALK_source"] - ds["DIC_source"]).sum(
+                dim=["s_rho", "eta_rho", "xi_rho"]
+            )
+            * ds_cdr["window_length"]
+        )
+        .cumsum(dim="time")
+        .compute()
+    )
+
+    # Cumulative flux-based uptake (Method 1)
+    flux = (
+        (
+            ((ds["FG_CO2"] - ds["FG_ALT_CO2"]) * ds_cdr["area"]).sum(
+                dim=["eta_rho", "xi_rho"]
+            )
+            * ds_cdr["window_length"]
+        )
+        .cumsum(dim="time")
+        .compute()
+    )
+
+    # DIC difference-based uptake (Method 2)
+    diff_DIC = (
+        ((ds["hDIC"] - ds["hDIC_ALT_CO2"]) * ds_cdr["area"])
+        .sum(dim=["s_rho", "eta_rho", "xi_rho"])
+        .compute()
+    )
+
+    # Normalize by cumulative source with safe division (NaN where source=0)
+    with np.errstate(divide="ignore", invalid="ignore"):
+        uptake_efficiency_flux = (flux / source).where(np.isfinite(flux / source))
+        uptake_efficiency_diff = (diff_DIC / source).where(
+            np.isfinite(diff_DIC / source)
+        )
+
+    _validate_uptake_efficiency(uptake_efficiency_flux, uptake_efficiency_diff)
+
+    # Store results with metadata
+    ds_cdr["cdr_efficiency"] = uptake_efficiency_flux
+    ds_cdr["cdr_efficiency"].attrs.update(
+        long_name="CDR uptake efficiency (from flux differences)",
+        units="nondimensional",
+        description="Carbon Dioxide Removal efficiency computed using area-integrated CO2 flux differences",
+    )
+    ds_cdr["cdr_efficiency_from_delta_diff"] = uptake_efficiency_diff
+    ds_cdr["cdr_efficiency_from_delta_diff"].attrs.update(
+        long_name="CDR uptake efficiency (from DIC differences)",
+        units="nondimensional",
+        description="Carbon Dioxide Removal efficiency computed using volume-integrated DIC differences",
+    )
+
+    return ds_cdr
+
+
+def _validate_uptake_efficiency(
+    uptake_efficiency_flux: xr.DataArray,
+    uptake_efficiency_diff: xr.DataArray,
+) -> float:
+    """
+    Compute and log the maximum absolute difference between two uptake efficiency estimates.
+
+    Parameters
+    ----------
+    uptake_efficiency_flux : xr.DataArray
+        Uptake computed from fluxes.
+    uptake_efficiency_diff : xr.DataArray
+        Uptake computed from DIC differences.
+
+    Returns
+    -------
+    max_abs_diff : float
+        Maximum absolute difference between uptake_flux and uptake_diff.
+    """
+    abs_diff = np.abs(uptake_efficiency_flux - uptake_efficiency_diff)
+    max_abs_diff = float(abs_diff.max())
+
+    logging.info(
+        "Max absolute difference between flux-based and DIC-based uptake efficiency: %.3e",
+        max_abs_diff,
+    )
+
+    return max_abs_diff
+
+
+def _validate_source(ds: xr.Dataset):
+    """
+    Validate that ALK_source and DIC_source in a ROMS dataset respect release constraints.
+
+    - 'ALK_source' must be non-negative (≥ 0).
+    - 'DIC_source' must be non-positive (≤ 0).
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Dataset expected to contain 'ALK_source' and 'DIC_source'.
+
+    Raises
+    ------
+    KeyError
+        If 'ALK_source' or 'DIC_source' are missing from the dataset.
+    ValueError
+        If 'ALK_source' or 'DIC_source' violate the release constraints.
+    """
+    constraints = {
+        "ALK_source": lambda x: x >= 0,
+        "DIC_source": lambda x: x <= 0,
+    }
+
+    for var, check in constraints.items():
+        if var not in ds.data_vars:
+            raise KeyError(f"Dataset is missing required variable '{var}'.")
+        if not check(ds[var]).all():
+            sign = "negative" if var == "ALK_source" else "positive"
+            raise ValueError(
+                f"'{var}' contains {sign} values, which violates release constraints."
+            )

roms_tools/analysis/cdr_ensemble.py

@@ -0,0 +1,198 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import xarray as xr
+
+
+@dataclass
+class Ensemble:
+    """
+    Represents an ensemble of CDR (Carbon Dioxide Removal) experiments.
+
+    Loads, aligns, and analyzes efficiency metrics across multiple members.
+
+    Parameters
+    ----------
+    members : dict[str, str | xr.Dataset]
+        Dictionary mapping member names to either file paths (NetCDF) or
+        xarray.Dataset objects containing the CDR metrics.
+    """
+
+    members: dict[str, str | xr.Dataset]
+    """Dictionary mapping member names to CDR metrics."""
+    ds: xr.Dataset = field(init=False)
+    """xarray Dataset containing aligned efficiencies for all ensemble members."""
+
+    def __post_init__(self):
+        """
+        Loads datasets, extracts efficiencies, aligns times, stores in self.ds, and
+        plots ensemble curves.
+        """
+        datasets = self._load_members()
+        effs = {name: self._extract_efficiency(ds) for name, ds in datasets.items()}
+        aligned = self._align_times(effs)
+        self.ds = self._compute_statistics(aligned)
+
+    def _load_members(self) -> dict[str, xr.Dataset]:
+        """
+        Loads ensemble member datasets.
+
+        Converts any file paths in `self.members` to xarray Datasets.
+        Members that are already xarray Datasets are left unchanged.
+
+        Returns
+        -------
+        dict[str, xr.Dataset]
+            Dictionary mapping member names to xarray.Dataset objects.
+        """
+        return {
+            name: xr.open_dataset(path) if isinstance(path, str | Path) else path
+            for name, path in self.members.items()
+        }
+
+    def _extract_efficiency(self, ds: xr.Dataset) -> xr.DataArray:
+        """
+        Extracts the CDR efficiency metric and reindex to time since release start.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            Dataset containing a "cdr_efficiency" variable and "abs_time" coordinate.
+
+        Returns
+        -------
+        xr.DataArray
+            Efficiency reindexed to a relative time axis in days since release start.
+
+        Raises
+        ------
+        ValueError
+            If 'abs_time' coordinate is missing or there are no valid efficiency values.
+        """
+        eff = ds["cdr_efficiency"]
+
+        # Check that abs_time exists
+        if "abs_time" in eff.coords:
+            abs_time = eff.coords["abs_time"]
+        elif "abs_time" in ds.data_vars:
+            abs_time = ds["abs_time"]
+        else:
+            raise ValueError(
+                "Dataset must contain an 'abs_time' coordinate or data variable."
+            )
+
+        # Drop NaNs to find first valid time
+        valid_mask = ~np.isnan(eff.values)
+        if not valid_mask.any():
+            raise ValueError("No valid efficiency values found in dataset.")
+
+        release_start = abs_time.values[valid_mask][0]
+
+        # Compute relative time in days
+        time_rel = (abs_time - release_start).astype("timedelta64[D]")
+
+        # Assign new time coordinate and drop abs_time if it was a data variable
+        eff_rel = eff.assign_coords(time=time_rel)
+        eff_rel.time.attrs["long_name"] = "time since release start"
+
+        if hasattr(eff_rel, "coords") and "abs_time" in eff_rel.coords:
+            eff_rel = eff_rel.drop_vars("abs_time")
+        elif hasattr(eff_rel, "variables") and "abs_time" in eff_rel.variables:
+            eff_rel = eff_rel.drop_vars("abs_time")
+
+        return eff_rel
+
+    def _align_times(self, effs: dict[str, xr.DataArray]) -> xr.Dataset:
+        """
+        Align all ensemble members to a common time axis.
+
+        Each member is reindexed to the union of all time coordinates.
+        Times outside the original range of a mamber are filled with NaN.
+
+        Parameters
+        ----------
+        effs : dict[str, xr.DataArray]
+            Dictionary mapping member names to efficiency DataArrays
+            (reindexed relative to their release start).
+
+        Returns
+        -------
+        xr.Dataset
+            Dataset containing all members aligned to a shared time coordinate, with
+            NaNs for missing times.
+        """
+        all_times = np.unique(
+            np.concatenate([eff.time.values for eff in effs.values()])
+        )
+        aligned = {name: eff.reindex(time=all_times) for name, eff in effs.items()}
+        return xr.Dataset(aligned)
+
+    def _compute_statistics(self, ds: xr.Dataset) -> xr.Dataset:
+        """
+        Computes ensemble statistics: mean and standard deviation.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            Dataset containing aligned ensemble member efficiencies.
+
+        Returns
+        -------
+        xr.Dataset
+            Dataset with additional variables "mean" and "std" representing
+            the ensemble mean and standard deviation across members.
+        """
+        da = ds.to_dataarray("member")  # stack into (member, time)
+        ds["ensemble_mean"] = da.mean(dim="member")
+        ds["ensemble_std"] = da.std(dim="member")
+        return ds
+
+    def plot(
+        self,
+        save_path: str | None = None,
+    ) -> None:
+        """
+        Plots ensemble members with mean ± standard deviation shading.
+
+        Displays individual member efficiency time series along with the ensemble
+        mean and ±1 standard deviation as a shaded region.
+
+        Parameters
+        ----------
+        save_path : str, optional
+            Path to save the generated plot. If None, the plot is shown interactively.
+            Default is None.
+
+        Returns
+        -------
+        None
+            This method does not return any value. It generates and displays a plot.
+        """
+        fig, ax = plt.subplots(figsize=(8, 5))
+
+        time = self.ds.time.values / np.timedelta64(1, "D")  # converts to float days
+
+        # Individual ensemble members
+        for name in self.members.keys():
+            ax.plot(time, self.ds[name], lw=2, label=name)
+
+        # Mean ± std
+        ax.plot(time, self.ds.ensemble_mean, color="black", lw=2, label="ensemble mean")
+        ax.fill_between(
+            time,
+            self.ds.ensemble_mean - self.ds.ensemble_std,
+            self.ds.ensemble_mean + self.ds.ensemble_std,
+            color="gray",
+            alpha=0.3,
+        )
+
+        ax.set_xlabel("Time since release start [days]")
+        ax.set_ylabel("CDR Efficiency")
+        ax.set_title("Ensemble of interventions")
+        ax.legend()
+        ax.grid()
+
+        if save_path:
+            plt.savefig(save_path, dpi=300, bbox_inches="tight")

roms_tools/analysis/roms_output.py

@@ -10,14 +10,15 @@ import xarray as xr
 from matplotlib.axes import Axes
 
 from roms_tools import Grid
-from roms_tools.plot import plot
+from roms_tools.analysis.cdr_analysis import compute_cdr_metrics
+from roms_tools.plot import plot, plot_uptake_efficiency
 from roms_tools.regrid import LateralRegridFromROMS, VerticalRegridFromROMS
 from roms_tools.utils import (
-    _generate_coordinate_range,
-    _load_data,
+    generate_coordinate_range,
     infer_nominal_horizontal_resolution,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
+    load_data,
 )
 from roms_tools.vertical_coordinate import (
     compute_depth_coordinates,
@@ -71,6 +72,27 @@ class ROMSOutput:
         # Dataset for depth coordinates
         self.ds_depth_coords = xr.Dataset()
 
+    def cdr_metrics(self) -> None:
+        """
+        Compute and plot Carbon Dioxide Removal (CDR) metrics.
+
+        If the CDR metrics dataset (`self.ds_cdr`) does not already exist,
+        it computes the metrics using model output and grid information.
+        Afterwards, it generates a plot of the computed metrics.
+
+        Notes
+        -----
+        Metrics include:
+          - Grid cell area
+          - Selected tracer and flux variables
+          - Uptake efficiency computed from flux differences and DIC differences
+        """
+        if not hasattr(self, "ds_cdr"):
+            # Compute metrics and store
+            self.ds_cdr = compute_cdr_metrics(self.ds, self.grid.ds)
+
+        plot_uptake_efficiency(self.ds_cdr)
+
     def plot(
         self,
         var_name: str,
@@ -170,7 +192,7 @@ class ROMSOutput:
         """
         # Check if variable exists
         if var_name not in self.ds:
-            raise ValueError(f"Variable '{var_name}' is not found in the dataset.")
+            raise ValueError(f"Variable '{var_name}' is not found in self.ds.")
 
         # Pick the variable
         field = self.ds[var_name]
@@ -269,11 +291,11 @@ class ROMSOutput:
 
         if horizontal_resolution is None:
             horizontal_resolution = infer_nominal_horizontal_resolution(self.grid.ds)
-        lons = _generate_coordinate_range(
+        lons = generate_coordinate_range(
             lon_deg.min().values, lon_deg.max().values, horizontal_resolution
         )
         lons = xr.DataArray(lons, dims=["lon"], attrs={"units": "°E"})
-        lats = _generate_coordinate_range(
+        lats = generate_coordinate_range(
             lat_deg.min().values, lat_deg.max().values, horizontal_resolution
         )
         lats = xr.DataArray(lats, dims=["lat"], attrs={"units": "°N"})
@@ -333,21 +355,22 @@ class ROMSOutput:
                 layer_depth_loc = lateral_regrid.apply(layer_depth_loc)
                 h_loc = lateral_regrid.apply(h_loc)
                 # Vertical regridding
-                vertical_regrid = VerticalRegridFromROMS(ds_loc)
-                for var_name in var_names_loc:
-                    if "s_rho" in ds_loc[var_name].dims:
-                        attrs = ds_loc[var_name].attrs
-                        regridded = vertical_regrid.apply(
-                            ds_loc[var_name],
-                            layer_depth_loc,
-                            depth_levels,
-                            mask_edges=False,
-                        )
-                        regridded = regridded.where(regridded.depth < h_loc)
-                        ds_loc[var_name] = regridded
-                        ds_loc[var_name].attrs = attrs
-
-                ds_loc = ds_loc.assign_coords({"depth": depth_levels})
+                if "s_rho" in ds_loc.dims:
+                    vertical_regrid = VerticalRegridFromROMS(ds_loc)
+                    for var_name in var_names_loc:
+                        if "s_rho" in ds_loc[var_name].dims:
+                            attrs = ds_loc[var_name].attrs
+                            regridded = vertical_regrid.apply(
+                                ds_loc[var_name],
+                                layer_depth_loc,
+                                depth_levels,
+                                mask_edges=False,
+                            )
+                            regridded = regridded.where(regridded.depth < h_loc)
+                            ds_loc[var_name] = regridded
+                            ds_loc[var_name].attrs = attrs
+
+                    ds_loc = ds_loc.assign_coords({"depth": depth_levels})
 
                 # Collect regridded dataset for merging
                 regridded_datasets.append(ds_loc)
@@ -418,10 +441,12 @@ class ROMSOutput:
     def _load_model_output(self) -> xr.Dataset:
         """Load the model output."""
         # Load the dataset
-        ds = _load_data(
+        ds = load_data(
             self.path,
             dim_names={"time": "time"},
             use_dask=self.use_dask,
+            decode_times=False,
+            decode_timedelta=False,
             time_chunking=True,
             force_combine_nested=True,
         )
@@ -514,35 +539,39 @@ class ROMSOutput:
 
         Notes
         -----
+        - Missing attributes trigger a warning instead of an exception.
         - `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']})."
-            )
+        required_exact = ["theta_s", "theta_b", "hc"]
+        required_close = ["Cs_r", "Cs_w"]
 
-        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']})."
-            )
+        # Check exact equality
+        for param in required_exact:
+            value = ds.attrs.get(param, None)
+            if value is None:
+                logging.warning(
+                    f"Dataset is missing attribute '{param}'. Skipping this check."
+                )
+                continue
+            if not np.array_equal(getattr(self.grid, param), value):
+                raise ValueError(
+                    f"{param} from grid ({getattr(self.grid, param)}) does not match dataset ({value})."
+                )
 
-        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']})."
-            )
+        # Check numerical closeness
+        for param in required_close:
+            value = ds.attrs.get(param, None)
+            if value is None:
+                logging.warning(
+                    f"Dataset is missing attribute '{param}'. Skipping this check."
+                )
+                continue
+            grid_value = getattr(self.grid.ds, param)
+            if not np.allclose(grid_value, value):
+                raise ValueError(
+                    f"{param} from grid ({grid_value}) is not close to dataset ({value})."
+                )
 
     def _add_absolute_time(self, ds: xr.Dataset) -> xr.Dataset:
         """Add absolute time as a coordinate to the dataset.
@@ -560,6 +589,11 @@ class ROMSOutput:
         xarray.Dataset
             Dataset with "abs_time" added and "time" removed.
         """
+        if self.model_reference_date is None:
+            raise ValueError(
+                "`model_reference_date` must be set before computing absolute time."
+            )
+
         ocean_time_seconds = ds["ocean_time"].values
 
         abs_time = np.array(

roms_tools/download.py

@@ -86,6 +86,10 @@ pup_test_data = pooch.create(
         "eastpac25km_rst.19980106000000.nc": "8f56d72bd8daf72eb736cc6705f93f478f4ad0ae4a95e98c4c9393a38e032f4c",
         "eastpac25km_rst.19980126000000.nc": "20ad9007c980d211d1e108c50589183120c42a2d96811264cf570875107269e4",
         "epac25km_grd.nc": "ec26c69cda4c4e96abde5b7756c955a7e1074931ab5a0641f598b099778fb617",
+        "GSHHS_l_L1.dbf": "181236ffbf553a83d2afedc5fe5e1f2fea64190f56ea366fc7a8ff5aa6163663",
+        "GSHHS_l_L1.prj": "98aaf3d1c0ecadf1a424a4536de261c3daf4e373697cb86c40c43b989daf52eb",
+        "GSHHS_l_L1.shp": "bc76f101f9b8671f90e734b4026da91c20066fc627cc8b5889ba22d90cbf97e9",
+        "GSHHS_l_L1.shx": "72879354892d80d6c39c612f645661ec0edc75f3f9f8f74b19d9387ae0327377",
     },
 )