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

roms_tools/regrid.py

@@ -295,3 +295,79 @@ class VerticalRegridFromROMS:
             )
 
         return transformed
+
+
+class VerticalRegrid:
+    """Regrid ROMS variables along the vertical, using spatially varying coordinates.
+
+    This class uses the `xgcm` package to transform data from a ROMS vertical coordinate
+    system (`s_rho`) to a user-defined set of target depth levels, where both the source
+    and target coordinates can vary spatially (i.e., 2D fields in horizontal space).
+
+    Attributes
+    ----------
+    grid : xgcm.Grid
+        The XGCM grid object used for vertical regridding, initialized with the input dataset `ds`.
+    """
+
+    def __init__(self, ds: "xr.Dataset"):
+        """Initialize the VerticalRegrid object with a ROMS dataset.
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            The ROMS dataset containing the vertical coordinate `s_rho` and the variable(s)
+            to be regridded.
+        """
+        self.grid = xgcm.Grid(
+            ds,
+            coords={"s_rho": {"center": "s_rho"}},
+            periodic=False,
+            autoparse_metadata=False,
+        )
+
+    def apply(
+        self,
+        da: "xr.DataArray",
+        source_depth_coords: "xr.DataArray",
+        target_depth_coords: "xr.DataArray",
+        mask_edges: bool = True,
+    ) -> "xr.DataArray":
+        """Regrid a ROMS variable from source vertical coordinates to target vertical coordinates.
+
+        This method supports spatially varying vertical coordinates for both source and target,
+        meaning that the depth levels can vary across the horizontal grid.
+
+        Parameters
+        ----------
+        da : xarray.DataArray
+            The data array to regrid. Must have a vertical dimension corresponding to `s_rho`.
+
+        source_depth_coords : array-like (1D or 2D)
+            Depth coordinates of the source data. Can be a 1D array (same for all horizontal points)
+            or a 2D array (varying in horizontal space).
+
+        target_depth_coords : array-like (1D or 2D)
+            Desired depth coordinates of the regridded data. Can also be 1D or 2D.
+
+        mask_edges : bool, optional
+            If True, target values outside the range of source depth coordinates are masked with NaN.
+            Defaults to True.
+
+        Returns
+        -------
+        xarray.DataArray
+            A new `DataArray` containing the regridded variable at the target depth coordinates.
+        """
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning, module="xgcm")
+            transformed = self.grid.transform(
+                da,
+                "s_rho",
+                target=target_depth_coords,
+                target_data=source_depth_coords,
+                target_dim="s_rho",
+                mask_edges=mask_edges,
+            )
+
+        return transformed

roms_tools/setup/boundary_forcing.py

@@ -11,17 +11,19 @@ import xarray as xr
 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.lat_lon_datasets import (
+from roms_tools.datasets.lat_lon_datasets import (
     CESMBGCDataset,
     GLORYSDataset,
     GLORYSDefaultDataset,
     UnifiedBGCDataset,
 )
+from roms_tools.fill import one_dim_fill
+from roms_tools.plot import line_plot, section_plot
+from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
 from roms_tools.setup.utils import (
     RawDataSource,
     add_time_info_to_ds,
+    check_and_set_boundaries,
     compute_barotropic_velocity,
     compute_missing_bgc_variables,
     from_yaml,
@@ -30,8 +32,6 @@ from roms_tools.setup.utils import (
     get_variable_metadata,
     group_dataset,
     nan_check,
-    one_dim_fill,
-    rotate_velocities,
     substitute_nans_by_fillvalue,
     to_dict,
     write_to_yaml,
@@ -39,6 +39,7 @@ from roms_tools.setup.utils import (
 from roms_tools.utils import (
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
+    rotate_velocities,
     save_datasets,
     transpose_dimensions,
 )
@@ -61,8 +62,9 @@ class BoundaryForcing:
         The end time of the desired surface forcing data. This time is used to filter the dataset
         to include only records on or before this time, with a single record at or after this time.
         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.
+    boundaries : dict[str, bool], optional
+        Specifies which grid boundaries ('south', 'east', 'north', 'west') are active and to be processed.
+        if not provided, valid (non-land) boundaries are enabled automatically.
     source : RawDataSource
         Dictionary specifying the source of the boundary forcing data. Keys include:
 
@@ -86,10 +88,6 @@ class BoundaryForcing:
         Indicates whether to perform a two-dimensional horizontal fill on the source data prior to regridding to boundaries.
         If `False`, a one-dimensional horizontal fill is performed separately on each of the four regridded boundaries.
         Defaults to `False`.
-    adjust_depth_for_sea_surface_height : bool, optional
-        Whether to account for sea surface height (`zeta`) variations when computing depth coordinates.
-        This adjustment is only applicable for `type="physics"`, as for biogeochemical fields usually `zeta` is not available.
-        Defaults to `False`.
     model_reference_date : datetime, optional
         Reference date for the model. Default is January 1, 2000.
     use_dask: bool, optional
@@ -117,14 +115,7 @@ class BoundaryForcing:
     """The start time of the desired surface forcing data."""
     end_time: datetime | None = None
     """The end time of the desired surface forcing data."""
-    boundaries: dict[str, bool] = field(
-        default_factory=lambda: {
-            "south": True,
-            "east": True,
-            "north": True,
-            "west": True,
-        }
-    )
+    boundaries: dict[str, bool] | None = None
     """Dictionary specifying which boundaries are forced (south, east, north, west)."""
     source: RawDataSource
     """Dictionary specifying the source of the boundary forcing data."""
@@ -133,9 +124,6 @@ class BoundaryForcing:
     apply_2d_horizontal_fill: bool = False
     """Whether to perform a two-dimensional horizontal fill on the source data prior to
     regridding to boundaries."""
-    adjust_depth_for_sea_surface_height: bool = False
-    """Whether to account for sea surface height (`zeta`) variations when computing
-    depth coordinates."""
     model_reference_date: datetime = datetime(2000, 1, 1)
     """Reference date for the model."""
     use_dask: bool = False
@@ -146,12 +134,18 @@ class BoundaryForcing:
     ds: xr.Dataset = field(init=False, repr=False)
     """An xarray Dataset containing post-processed variables ready for input into
     ROMS."""
+    adjust_depth_for_sea_surface_height: bool = field(init=False)
+    """Whether to account for sea surface height when computing depth coordinates."""
+    ds_depth_coords: xr.Dataset = field(init=False, repr=False)
+    """An xarray Dataset containing the depth coordinates."""
 
     def __post_init__(self):
-        self._input_checks()
-        # Dataset for depth coordinates
+        # Initialize depth coordinates
+        self.adjust_depth_for_sea_surface_height = False
         self.ds_depth_coords = xr.Dataset()
 
+        self._input_checks()
+
         target_coords = get_target_coords(self.grid)
 
         data = self._get_data()
@@ -279,7 +273,7 @@ class BoundaryForcing:
                         processed_fields["u"],
                         processed_fields["v"],
                         angle,
-                        interpolate=True,
+                        interpolate_after=True,
                     )
                     if self.adjust_depth_for_sea_surface_height:
                         zeta_u = interpolate_from_rho_to_u(zeta_vector)
@@ -388,51 +382,58 @@ class BoundaryForcing:
 
         self.ds = ds
 
-    def _input_checks(self):
-        # Check that start_time and end_time are both None or none of them is
+    def _input_checks(self) -> None:
+        """Validate and normalize user-provided input parameters."""
+        # -------------------------------------------------------
+        # Time range checks
+        # -------------------------------------------------------
         if (self.start_time is None) != (self.end_time is None):
             raise ValueError(
                 "Both `start_time` and `end_time` must be provided together as datetime objects or both should be None."
             )
 
-        # Trigger a warning if both are None
         if self.start_time is None and self.end_time is None:
             logging.warning(
                 "Both `start_time` and `end_time` are None. No time filtering will be applied to the source data."
             )
 
-        # Validate the 'type' parameter
-        if self.type not in ["physics", "bgc"]:
+        # -------------------------------------------------------
+        # Type check
+        # -------------------------------------------------------
+        if self.type not in {"physics", "bgc"}:
             raise ValueError("`type` must be either 'physics' or 'bgc'.")
 
-        # Ensure 'source' dictionary contains required keys
+        # -------------------------------------------------------
+        # Source configuration checks
+        # -------------------------------------------------------
         if "name" not in self.source:
             raise ValueError("`source` must include a 'name'.")
+
         if "path" not in self.source:
             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 = {
-            **self.source,
-            "climatology": self.source.get("climatology", False),
-        }
+        # Assign default value
+        self.source["climatology"] = self.source.get("climatology", False)
+
+        # -------------------------------------------------------
+        # Boundary selection defaults and validation
+        # -------------------------------------------------------
+
+        self.boundaries = check_and_set_boundaries(
+            self.boundaries, self.grid.ds.mask_rho
+        )
 
-        # Ensure adjust_depth_for_sea_surface_height is only used with type="physics"
+        # -------------------------------------------------------
+        # Depth adjustment checks
+        # -------------------------------------------------------
         if self.type == "bgc" and self.adjust_depth_for_sea_surface_height:
             logging.warning(
                 "adjust_depth_for_sea_surface_height is not applicable for BGC fields. "
                 "Setting it to False."
             )
             self.adjust_depth_for_sea_surface_height = False
-        elif self.adjust_depth_for_sea_surface_height:
-            logging.info("Sea surface height will be used to adjust depth coordinates.")
-        else:
-            logging.info(
-                "Sea surface height will NOT be used to adjust depth coordinates."
-            )
 
     def _get_data(
         self,
@@ -1061,7 +1062,14 @@ class BoundaryForcing:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-        forcing_dict = to_dict(self, exclude=["use_dask"])
+        forcing_dict = to_dict(
+            self,
+            exclude=[
+                "ds_depth_coords",
+                "adjust_depth_for_sea_surface_height",
+                "use_dask",
+            ],
+        )
         write_to_yaml(forcing_dict, filepath)
 
     @classmethod

roms_tools/setup/cdr_release.py

@@ -573,7 +573,7 @@ class TracerPerturbation(Release):
 
     tracer_fluxes : dict, optional
 
-        Dictionary of tracer names and their non-negative flux values. The flux values can be either
+        Dictionary of tracer names and their flux values. The flux values can be either
         a float (constant in time) or a list of float (time-varying).
 
         - Constant: applies uniformly across the entire simulation period.
@@ -589,9 +589,7 @@ class TracerPerturbation(Release):
     """
 
     times: list[datetime] = Field([])
-    tracer_fluxes: dict[str, Flux | NonNegativeFloat | list[NonNegativeFloat]] = Field(
-        {}
-    )
+    tracer_fluxes: dict[str, Flux | float | list[float]] = Field({})
     """Dictionary of tracer names and their non-negative flux values."""
 
     release_type: Literal[ReleaseType.tracer_perturbation] = (

roms_tools/setup/grid.py

@@ -1,3 +1,4 @@
+import copy
 import importlib.metadata
 import logging
 import re
@@ -19,12 +20,14 @@ from roms_tools.setup.utils import (
     extract_single_value,
     gc_dist,
     get_target_coords,
-    interpolate_from_rho_to_u,
-    interpolate_from_rho_to_v,
     pop_grid_data,
     write_to_yaml,
 )
-from roms_tools.utils import save_datasets
+from roms_tools.utils import (
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
+    save_datasets,
+)
 from roms_tools.vertical_coordinate import compute_depth_coordinates, sigma_stretch
 
 
@@ -206,7 +209,8 @@ class Grid:
         hmin: float | None = None,
         verbose: bool = False,
     ) -> None:
-        """Update the grid dataset with processed topography.
+        """
+        Update the grid dataset with processed topography.
 
         This method performs the following operations:
 
@@ -218,29 +222,45 @@ class Grid:
         Parameters
         ----------
         topography_source : dict, optional
-            A dictionary specifying the source of the topography data. The dictionary should
-            contain the following keys:
-            - "name" (str): The name of the topography data source (e.g., "SRTM15").
-            - "path" (Union[str, Path): The path to the raw data file.
+            Dictionary specifying the topography data source with keys:
+
+            - ``"name"`` (str): Name of the topography dataset (e.g. ``"SRTM15"``)
+            - ``"path"`` (str or pathlib.Path): Path to the raw data file
 
             If not provided, `topography_source` will remain unchanged (i.e., the existing value will not be overwritten).
 
         hmin : float, optional
-            The minimum ocean depth (in meters).
+            Minimum ocean depth in meters.
             If not provided, `hmin` will remain unchanged (i.e., the existing value will not be overwritten).
 
         verbose : bool, optional
-            If True, the method will print detailed information about the grid generation process,
-            including the timing of each step. Defaults to False.
+            If True, print detailed information about the processing steps and timing.
+            Defaults to False.
 
         Returns
         -------
         None
-            This method updates the internal dataset (`self.ds`) in place by adding or overwriting the
-            topography variable. It does not return any value.
+            Updates ``self.ds`` in place.
         """
-        topography_source = topography_source or self.topography_source
-        hmin = hmin or self.hmin
+        if topography_source is None:
+            topography_source = self.topography_source
+
+        if hmin is None:
+            hmin = self.hmin
+
+        # This can only happen for externally generated grids read via Grid.from_file()
+        if topography_source is None:
+            raise ValueError(
+                "Topography source information is not available. "
+                "Please provide `topography_source` explicitly when calling "
+                "`update_topography()`."
+            )
+        # This can only happen for externally generated grids read via Grid.from_file()
+        if hmin is None:
+            raise ValueError(
+                "Minimal ocean depth is not available. "
+                "Please provide `hmin` explicitly when calling `update_topography()`."
+            )
 
         name = topography_source["name"]  # type: ignore[index]
 
@@ -1249,6 +1269,17 @@ class Grid:
 
         return z_centers, z_faces
 
+    def copy_with_ds(self, ds: xr.Dataset) -> "Grid":
+        """
+        Return a copy of this Grid with the given Dataset.
+
+        Grid metadata is preserved; only the backing xarray Dataset
+        is replaced. The original Grid is not modified.
+        """
+        new = copy.copy(self)  # shallow copy of metadata
+        new.ds = ds
+        return new
+
 
 def _rotate(coords, rot):
     """Rotate grid counterclockwise relative to surface of Earth by rot degrees."""