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

roms_tools/setup/cdr_release.py

@@ -5,6 +5,8 @@ from datetime import datetime
 from enum import StrEnum, auto
 from typing import Annotated, Literal
 
+import numpy as np
+import pandas as pd
 from annotated_types import Ge, Le
 from pydantic import (
     BaseModel,
@@ -16,10 +18,17 @@ from pydantic import (
 )
 from pydantic_core.core_schema import ValidationInfo
 
-from roms_tools.setup.utils import get_tracer_defaults, get_tracer_metadata_dict
+from roms_tools.setup.utils import (
+    convert_to_relative_days,
+    get_tracer_defaults,
+    get_tracer_metadata_dict,
+)
 
 NonNegativeFloat = Annotated[float, Ge(0)]
 
+# Show all columns when printing a DataFrame
+pd.set_option("display.max_columns", None)
+
 
 @dataclass
 class ValueArray(ABC):
@@ -272,10 +281,68 @@ class Release(BaseModel):
             if self.times[-1] < end_time:
                 self.times.append(end_time)
 
-    @staticmethod
-    def get_tracer_metadata():
+    @classmethod
+    def get_tracer_metadata(cls):
         return {}
 
+    @classmethod
+    def get_metadata(cls):
+        return pd.DataFrame(cls.get_tracer_metadata())
+
+    def _compute_integrated_tracers(
+        self,
+        roms_time_stamps: np.ndarray,
+        model_reference_date: datetime,
+        tracer_series_dict: dict[str, np.ndarray],
+    ) -> dict[str, float]:
+        """
+        Compute time-integrated tracer quantities over ROMS time steps using a left-hold rule.
+
+        This method performs a left-hold (stepwise constant) integration of tracer fluxes
+        over the intervals defined by the ROMS time stamps. It first interpolates the
+        tracer time series from the release schedule onto the ROMS time stamps, then
+        multiplies the value at the start of each interval by the duration of that interval.
+
+        Parameters
+        ----------
+        roms_time_stamps : np.ndarray
+            1D array of ROMS model time stamps in seconds since `model_reference_date`.
+            Must be strictly increasing and contain at least two entries.
+        model_reference_date : datetime
+            Reference datetime of the ROMS model calendar, used to compute relative times
+            for interpolation.
+        tracer_series_dict : dict[str, np.ndarray]
+            Dictionary mapping tracer names to 1D arrays of tracer flux values at the
+            release schedule times (`self.times`). Each array must have the same length
+            as `self.times`.
+
+        Returns
+        -------
+        dict[str, float]
+            Dictionary mapping each tracer name to its integrated quantity over the
+            ROMS time period. Integration is performed using the left-hold rule,
+            ignoring the last release point because it defines the end of the final interval.
+
+        Raises
+        ------
+        ValueError
+            If `roms_time_stamps` has fewer than two entries, since at least one interval
+            is required for integration.
+        """
+        if len(roms_time_stamps) < 2:
+            raise ValueError("Need at least two ROMS time stamps to define intervals.")
+
+        dt = np.diff(roms_time_stamps)
+        results = {}
+        for tracer, series in tracer_series_dict.items():
+            interp_values = np.interp(
+                roms_time_stamps,
+                convert_to_relative_days(self.times, model_reference_date) * 3600 * 24,
+                series,
+            )
+            results[tracer] = np.sum(interp_values[:-1] * dt)
+        return results
+
 
 class VolumeRelease(Release):
     """Represents a CDR release with volume flux and tracer concentrations.
@@ -389,9 +456,11 @@ class VolumeRelease(Release):
         num_times = len(self.times)
 
         for tracer_concentrations in self.tracer_concentrations.values():
-            tracer_concentrations.check_length(num_times)
+            if isinstance(tracer_concentrations, Concentration):
+                tracer_concentrations.check_length(num_times)
 
-        self.volume_fluxes.check_length(num_times)
+        if isinstance(self.volume_fluxes, Flux):
+            self.volume_fluxes.check_length(num_times)
 
         return self
 
@@ -410,7 +479,52 @@ class VolumeRelease(Release):
     @staticmethod
     def get_tracer_metadata():
         """Returns long names and expected units for the tracer concentrations."""
-        return get_tracer_metadata_dict(include_bgc=True, with_flux_units=False)
+        return get_tracer_metadata_dict(include_bgc=True, unit_type="concentration")
+
+    def _do_accounting(
+        self,
+        roms_time_stamps: np.ndarray,
+        model_reference_date: datetime,
+    ) -> dict[str, float]:
+        """
+        Compute time-integrated tracer quantities over ROMS time steps.
+
+        This method interpolates tracer flux time series from the CDR schedule
+        onto the provided ROMS time stamps (in seconds since model reference date),
+        then applies a "left-hold" rule: the interpolated value at t₀ is applied
+        across the full interval [t₀, t₁).
+
+        Parameters
+        ----------
+        roms_time_stamps : np.ndarray
+            1D array of ROMS time stamps (seconds since `model_reference_date`).
+            Must be strictly increasing.
+        model_reference_date : datetime
+            Reference date of the ROMS model calendar.
+
+        Returns
+        -------
+        dict[str, float]
+            Dictionary mapping tracer names to the total integrated quantity over
+            the entire ROMS time period. Each value is the sum of the interpolated
+            tracer fluxes multiplied by the corresponding ROMS time step durations.
+        """
+        tracer_series_dict = {}
+        volume_array = (
+            np.asarray(self.volume_fluxes.values)
+            if isinstance(self.volume_fluxes, Flux)
+            else np.asarray(self.volume_fluxes)
+        )
+        for tracer, conc in self.tracer_concentrations.items():
+            tracer_array = (
+                np.asarray(conc.values)
+                if isinstance(conc, Concentration)
+                else np.asarray(conc)
+            )
+            tracer_series_dict[tracer] = volume_array * tracer_array
+        return self._compute_integrated_tracers(
+            roms_time_stamps, model_reference_date, tracer_series_dict
+        )
 
     @model_serializer(mode="wrap")
     def _simplified_dump(self, pydantic_serializer) -> dict:
@@ -503,7 +617,8 @@ class TracerPerturbation(Release):
     def _check_tracer_flux_lengths(self):
         num_times = len(self.times)
         for flux in self.tracer_fluxes.values():
-            flux.check_length(num_times)
+            if isinstance(flux, Flux):
+                flux.check_length(num_times)
         return self
 
     def _extend_to_endpoints(self, start_time, end_time):
@@ -520,7 +635,45 @@ class TracerPerturbation(Release):
     @staticmethod
     def get_tracer_metadata():
         """Returns long names and expected units for the tracer fluxes."""
-        return get_tracer_metadata_dict(include_bgc=True, with_flux_units=True)
+        return get_tracer_metadata_dict(include_bgc=True, unit_type="flux")
+
+    def _do_accounting(
+        self,
+        roms_time_stamps: np.ndarray,
+        model_reference_date: datetime,
+    ) -> dict[str, float]:
+        """
+        Compute time-integrated tracer quantities over ROMS time steps.
+
+        This method interpolates tracer flux time series from the CDR schedule
+        onto the provided ROMS time stamps (in days since model reference date),
+        then applies a "left-hold" rule: the interpolated value at t₀ is applied
+        across the full interval [t₀, t₁).
+
+        Parameters
+        ----------
+        roms_time_stamps : np.ndarray
+            1D array of ROMS time stamps (days since `model_reference_date`).
+            Must be strictly increasing.
+        model_reference_date : datetime
+            Reference date of the ROMS model calendar.
+
+        Returns
+        -------
+        dict[str, float]
+            Dictionary mapping tracer names to the total integrated quantity over
+            the entire ROMS time period. Each value is the sum of the interpolated
+            tracer fluxes multiplied by the corresponding ROMS time step durations.
+        """
+        tracer_series_dict = {
+            tracer: np.asarray(flux.values)
+            if isinstance(flux, Flux)
+            else np.asarray(flux)
+            for tracer, flux in self.tracer_fluxes.items()
+        }
+        return self._compute_integrated_tracers(
+            roms_time_stamps, model_reference_date, tracer_series_dict
+        )
 
     @model_serializer(mode="wrap")
     def _simplified_dump(self, pydantic_serializer) -> dict:

roms_tools/setup/grid.py

@@ -1,9 +1,9 @@
 import importlib.metadata
 import logging
 import re
-import time
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 import xarray as xr
@@ -12,9 +12,10 @@ from matplotlib.axes import Axes
 
 from roms_tools.constants import MAXIMUM_GRID_SIZE, R_EARTH
 from roms_tools.plot import plot
-from roms_tools.setup.mask import _add_mask, _add_velocity_masks
-from roms_tools.setup.topography import _add_topography
+from roms_tools.setup.mask import add_mask, add_velocity_masks
+from roms_tools.setup.topography import add_topography
 from roms_tools.setup.utils import (
+    Timed,
     extract_single_value,
     gc_dist,
     get_target_coords,
@@ -64,6 +65,8 @@ class Grid:
         - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file. Can be a string or a Path object.
 
         The default is "ETOPO5", which does not require a path.
+    mask_shapefile: str | Path | None, optional
+        Path to a custom shapefile to use to determine the land mask; if None, use NaturalEarth 10m.
     hmin : float, optional
        The minimum ocean depth (in meters). The default is 5.0.
     N : int, optional
@@ -106,8 +109,10 @@ class Grid:
     """The bottom control parameter."""
     hc: float = 300.0
     """The critical depth (in meters)."""
-    topography_source: dict[str, str | Path | list[str | Path]] = None
+    topography_source: dict[str, str | Path | list[str | Path]] | None = None
     """Dictionary specifying the source of the topography data."""
+    mask_shapefile: str | Path | None = None
+    """Path to a custom shapefile to use to determine the landmask; if None, use NaturalEarth 10m."""
     hmin: float = 5.0
     """The minimum ocean depth (in meters)."""
     verbose: bool = False
@@ -129,7 +134,7 @@ class Grid:
         self._straddle()
 
         # Mask
-        self._create_mask(verbose=self.verbose)
+        self.update_mask(mask_shapefile=self.mask_shapefile, verbose=self.verbose)
 
         # Coarsen the dataset if needed
         self._coarsen()
@@ -165,28 +170,50 @@ class Grid:
                     "`topography_source` must include a 'path' key when the 'name' is not 'ETOPO5'."
                 )
 
-    def _create_mask(self, verbose=False) -> None:
-        if verbose:
-            start_time = time.time()
-            logging.info("=== Creating the mask ===")
-        ds = _add_mask(self.ds)
+    def update_mask(
+        self, mask_shapefile: str | Path | None = None, verbose: bool = False
+    ) -> None:
+        """
+        Update the land mask of the current grid dataset.
 
-        if verbose:
-            logging.info(f"Total time: {time.time() - start_time:.3f} seconds")
-            logging.info(
-                "========================================================================================================"
-            )
+        This method generates a land mask based on the provided coastline
+        shapefile, fills enclosed basins with lands and updates the dataset
+        stored in `self.ds`. If no shapefile is provided, a default dataset (Natural
+        Earth 10m) is used. The operation is optionally timed and logged.
 
-        self.ds = ds
+        Parameters
+        ----------
+        mask_shapefile : str or Path, optional
+            Path to a coastal shapefile to derive the land mask. If `None`,
+            the default Natural Earth 10m coastline dataset is used.
+        verbose : bool, default False
+            If True, prints timing and progress information.
+
+        Returns
+        -------
+        None
+            Updates the `self.ds` attribute in place with the new mask.
+
+        """
+        with Timed("=== Deriving the mask from coastlines ===", verbose=verbose):
+            ds = add_mask(self.ds, shapefile=mask_shapefile)
+            self.ds = ds
+            self.mask_shapefile = mask_shapefile
 
     def update_topography(
-        self, topography_source=None, hmin=None, verbose=False
+        self,
+        topography_source: dict | None = None,
+        hmin: float | None = None,
+        verbose: bool = False,
     ) -> None:
         """Update the grid dataset with processed topography.
 
-        This method performs several key operations, including regridding the topography, smoothing the
-        topography over the entire domain and locally.
-        The processed topography is then added to the grid's dataset.
+        This method performs the following operations:
+
+        1. Regrids the topography from the specified source.
+        2. Applies domain-wide and local smoothing.
+        3. Enforces the minimum depth constraint ``hmin``.
+        4. Updates the internal dataset (``self.ds``) with the processed bathymetry.
 
         Parameters
         ----------
@@ -215,36 +242,27 @@ class Grid:
         topography_source = topography_source or self.topography_source
         hmin = hmin or self.hmin
 
+        name = topography_source["name"]  # type: ignore[index]
+
         # Extract target coordinates for processing
         target_coords = get_target_coords(self)
 
-        # If verbose is enabled, start the timer and print the start message
-        if verbose:
-            start_time = time.time()
-            logging.info(
-                f"=== Generating the topography using {topography_source['name']} data and hmin = {hmin} meters ==="
-            )
-
-        # Add topography to the dataset
-        ds = _add_topography(
-            ds=self.ds,
-            target_coords=target_coords,
-            topography_source=topography_source,
-            hmin=hmin,
+        with Timed(
+            f"=== Generating the topography using {name} data and hmin = {hmin} meters ===",
             verbose=verbose,
-        )
-
-        # If verbose is enabled, print elapsed time and a separator
-        if verbose:
-            logging.info(f"Total time: {time.time() - start_time:.3f} seconds")
-            logging.info(
-                "========================================================================================================"
+        ):
+            ds = add_topography(
+                ds=self.ds,
+                target_coords=target_coords,
+                topography_source=topography_source,
+                hmin=hmin,
+                verbose=verbose,
             )
 
-        # Update the grid's dataset and related attributes
-        self.ds = ds
-        self.topography_source = topography_source
-        self.hmin = hmin
+            # Update the grid's dataset and related attributes
+            self.ds = ds
+            self.topography_source = topography_source
+            self.hmin = hmin
 
     def update_vertical_coordinate(
         self, N=None, theta_s=None, theta_b=None, hc=None, verbose=False
@@ -281,69 +299,61 @@ class Grid:
         theta_b = theta_b or self.theta_b
         hc = hc or self.hc
 
-        if verbose:
-            start_time = time.time()
-            logging.info(
-                f"=== Preparing the vertical coordinate system using N = {N}, theta_s = {theta_s}, theta_b = {theta_b}, hc = {hc} ==="
-            )
+        with Timed(
+            f"=== Preparing the vertical coordinate system using N = {N}, theta_s = {theta_s}, theta_b = {theta_b}, hc = {hc} ===",
+            verbose=verbose,
+        ):
+            ds = self.ds
+            # need to drop vertical coordinates because they could cause conflict if N changed
+            vars_to_drop = [
+                "layer_depth_rho",
+                "layer_depth_u",
+                "layer_depth_v",
+                "interface_depth_rho",
+                "interface_depth_u",
+                "interface_depth_v",
+                "sigma_r",
+                "sigma_w",
+                "Cs_w",
+                "Cs_r",
+            ]
 
-        ds = self.ds
-        # need to drop vertical coordinates because they could cause conflict if N changed
-        vars_to_drop = [
-            "layer_depth_rho",
-            "layer_depth_u",
-            "layer_depth_v",
-            "interface_depth_rho",
-            "interface_depth_u",
-            "interface_depth_v",
-            "sigma_r",
-            "sigma_w",
-            "Cs_w",
-            "Cs_r",
-        ]
-
-        for var in vars_to_drop:
-            if var in ds.variables:
-                ds = ds.drop_vars(var)
-
-        cs_r, sigma_r = sigma_stretch(theta_s, theta_b, N, "r")
-        cs_w, sigma_w = sigma_stretch(theta_s, theta_b, N, "w")
-
-        ds["sigma_r"] = sigma_r.astype(np.float32)
-        ds["sigma_r"].attrs["long_name"] = (
-            "Fractional vertical stretching coordinate at rho-points"
-        )
-        ds["sigma_r"].attrs["units"] = "nondimensional"
+            for var in vars_to_drop:
+                if var in ds.variables:
+                    ds = ds.drop_vars(var)
 
-        ds["Cs_r"] = cs_r.astype(np.float32)
-        ds["Cs_r"].attrs["long_name"] = "Vertical stretching function at rho-points"
-        ds["Cs_r"].attrs["units"] = "nondimensional"
+            cs_r, sigma_r = sigma_stretch(theta_s, theta_b, N, "r")
+            cs_w, sigma_w = sigma_stretch(theta_s, theta_b, N, "w")
 
-        ds["sigma_w"] = sigma_w.astype(np.float32)
-        ds["sigma_w"].attrs["long_name"] = (
-            "Fractional vertical stretching coordinate at w-points"
-        )
-        ds["sigma_w"].attrs["units"] = "nondimensional"
-
-        ds["Cs_w"] = cs_w.astype(np.float32)
-        ds["Cs_w"].attrs["long_name"] = "Vertical stretching function at w-points"
-        ds["Cs_w"].attrs["units"] = "nondimensional"
+            ds["sigma_r"] = sigma_r.astype(np.float32)
+            ds["sigma_r"].attrs["long_name"] = (
+                "Fractional vertical stretching coordinate at rho-points"
+            )
+            ds["sigma_r"].attrs["units"] = "nondimensional"
 
-        ds.attrs["theta_s"] = np.float32(theta_s)
-        ds.attrs["theta_b"] = np.float32(theta_b)
-        ds.attrs["hc"] = np.float32(hc)
+            ds["Cs_r"] = cs_r.astype(np.float32)
+            ds["Cs_r"].attrs["long_name"] = "Vertical stretching function at rho-points"
+            ds["Cs_r"].attrs["units"] = "nondimensional"
 
-        if verbose:
-            logging.info(f"Total time: {time.time() - start_time:.3f} seconds")
-            logging.info(
-                "========================================================================================================"
+            ds["sigma_w"] = sigma_w.astype(np.float32)
+            ds["sigma_w"].attrs["long_name"] = (
+                "Fractional vertical stretching coordinate at w-points"
             )
+            ds["sigma_w"].attrs["units"] = "nondimensional"
 
-        self.ds = ds
-        self.theta_s = theta_s
-        self.theta_b = theta_b
-        self.hc = hc
-        self.N = N
+            ds["Cs_w"] = cs_w.astype(np.float32)
+            ds["Cs_w"].attrs["long_name"] = "Vertical stretching function at w-points"
+            ds["Cs_w"].attrs["units"] = "nondimensional"
+
+            ds.attrs["theta_s"] = np.float32(theta_s)
+            ds.attrs["theta_b"] = np.float32(theta_b)
+            ds.attrs["hc"] = np.float32(hc)
+
+            self.ds = ds
+            self.theta_s = theta_s
+            self.theta_b = theta_b
+            self.hc = hc
+            self.N = N
 
     def _straddle(self) -> None:
         """Check if the Greenwich meridian goes through the domain.
@@ -603,7 +613,7 @@ class Grid:
         ds = xr.open_dataset(filepath)
 
         if not all(mask in ds for mask in ["mask_u", "mask_v"]):
-            ds = _add_velocity_masks(ds)
+            ds = add_velocity_masks(ds)
 
         # Create a new Grid instance without calling __init__ and __post_init__
         grid = cls.__new__(cls)
@@ -758,24 +768,30 @@ class Grid:
             "hmin",
         ]:
             if attr in ds.attrs:
-                a = float(ds.attrs[attr])
+                value = float(ds.attrs[attr])
             else:
-                a = None
+                value = None
 
-            object.__setattr__(grid, attr, a)
+            object.__setattr__(grid, attr, value)
 
         if "topography_source_name" in ds.attrs:
             if "topography_source_path" in ds.attrs:
-                a = {
+                topo_source = {
                     "name": ds.attrs["topography_source_name"],
                     "path": ds.attrs["topography_source_path"],
                 }
             else:
-                a = {"name": ds.attrs["topography_source_name"]}
+                topo_source = {"name": ds.attrs["topography_source_name"]}
         else:
-            a = None
+            topo_source = None
+        grid.topography_source = topo_source
 
-        object.__setattr__(grid, "topography_source", a)
+        if "mask_shapefile" in ds.attrs:
+            mask_shapefile = ds.attrs["mask_shapefile"]
+        else:
+            mask_shapefile = None
+
+        grid.mask_shapefile = mask_shapefile
 
         return grid
 
@@ -796,10 +812,7 @@ class Grid:
 
     @classmethod
     def from_yaml(
-        cls,
-        filepath: str | Path,
-        section_name: str = "Grid",
-        verbose: bool = False,
+        cls, filepath: str | Path, verbose: bool = False, **kwargs: Any
     ) -> "Grid":
         """Create an instance of the class from a YAML file.
 
@@ -807,10 +820,13 @@ class Grid:
         ----------
         filepath : Union[str, Path]
             The path to the YAML file from which the parameters will be read.
-        section_name : str, optional
-            The name of the YAML section containing the grid configuration. Defaults to "Grid".
         verbose : bool, optional
             Indicates whether to print grid generation steps with timing. Defaults to False.
+        **kwargs : Any
+            Additional keyword arguments:
+
+            - section_name : str, optional (default: "Grid")
+              The name of the YAML section containing the grid configuration.
 
         Returns
         -------
@@ -828,6 +844,8 @@ class Grid:
         Issues a warning if the ROMS-Tools version in the YAML header does not match the
         currently installed version.
         """
+        section_name: str = kwargs.pop("section_name", None) or "Grid"
+
         filepath = Path(filepath)
         # Read the entire file content
         with filepath.open("r") as file:
@@ -881,7 +899,7 @@ class Grid:
         attr_str = ", ".join(f"{k}={v!r}" for k, v in attr_dict.items())
         return f"{cls_name}({attr_str})"
 
-    def _create_horizontal_grid(self) -> xr.Dataset():
+    def _create_horizontal_grid(self) -> xr.Dataset:
         """Create the horizontal grid based on a Mercator projection and store it in the
         'ds' attribute.
 
@@ -899,41 +917,32 @@ class Grid:
         - Longitude values are adjusted to fall within the range [0, 360].
         - Grid rotation and translation are applied based on the specified parameters.
         """
-        if self.verbose:
-            start_time = time.time()
-            logging.info("=== Creating the horizontal grid ===")
-
-        self._raise_if_domain_size_too_large()
+        with Timed("=== Creating the horizontal grid ===", verbose=self.verbose):
+            self._raise_if_domain_size_too_large()
 
-        coords = self._make_initial_lon_lat_ds()
+            coords = self._make_initial_lon_lat_ds()
 
-        # rotate coordinate system
-        coords = _rotate(coords, self.rot)
+            # rotate coordinate system
+            coords = _rotate(coords, self.rot)
 
-        # translate coordinate system
-        coords = _translate(coords, self.center_lat, self.center_lon)
+            # translate coordinate system
+            coords = _translate(coords, self.center_lat, self.center_lon)
 
-        # compute 1/dx and 1/dy
-        coords["pm"], coords["pn"] = _compute_coordinate_metrics(coords)
+            # compute 1/dx and 1/dy
+            coords["pm"], coords["pn"] = _compute_coordinate_metrics(coords)
 
-        # compute angle of local grid positive x-axis relative to east
-        coords["angle"] = _compute_angle(coords)
+            # compute angle of local grid positive x-axis relative to east
+            coords["angle"] = _compute_angle(coords)
 
-        # make sure lons are in [0, 360] range
-        for lon in ["lon", "lonu", "lonv", "lonq"]:
-            coords[lon][coords[lon] < 0] = coords[lon][coords[lon] < 0] + 2 * np.pi
+            # make sure lons are in [0, 360] range
+            for lon in ["lon", "lonu", "lonv", "lonq"]:
+                coords[lon][coords[lon] < 0] = coords[lon][coords[lon] < 0] + 2 * np.pi
 
-        ds = self._create_grid_ds(coords)
+            ds = self._create_grid_ds(coords)
 
-        ds = self._add_global_metadata(ds)
+            ds = self._add_global_metadata(ds)
 
-        if self.verbose:
-            logging.info(f"Total time: {time.time() - start_time:.3f} seconds")
-            logging.info(
-                "========================================================================================================"
-            )
-
-        self.ds = ds
+            self.ds = ds
 
     def _add_global_metadata(self, ds):
         """Add global metadata and attributes to the dataset.