roms-tools 2.3.0__py3-none-any.whl → 2.5.0__py3-none-any.whl

roms_tools/setup/nesting.py

@@ -2,57 +2,59 @@ import numpy as np
 import xarray as xr
 from scipy.interpolate import griddata
 from dataclasses import dataclass, field
-from typing import Dict, Union
+from typing import Dict, Union, Any, Tuple
 from pathlib import Path
 import logging
 from scipy.interpolate import interp1d
 from roms_tools import Grid
 from roms_tools.plot import _plot_nesting
+from roms_tools.utils import save_datasets
+from roms_tools.setup.topography import _clip_depth
 from roms_tools.setup.utils import (
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
     get_boundary_coords,
     wrap_longitudes,
-    save_datasets,
     _to_yaml,
     _from_yaml,
 )
 
 
 @dataclass(frozen=True, kw_only=True)
-class Nesting:
-    """Represents the relationship between a parent and a child grid in nested ROMS
-    simulations. This class facilitates mapping the boundaries of the child grid onto
-    the parent grid indices and modifying the child grid topography such that it matches
-    the parent topography at the boundaries.
+class ChildGrid(Grid):
+    """Represents a ROMS child grid that is compatible with the provided parent grid.
+
+    This class establishes the relationship between a parent grid and a child grid in ROMS simulations.
+    It generates two datasets:
+
+    1. `ds`: Contains child grid variables, ensuring compatibility with the parent grid.
+       The child grid’s topography and mask are adjusted to match the parent grid at the boundaries.
+
+    2. `ds_nesting`: Contains boundary mappings, linking the child grid’s boundaries
+       to the corresponding parent grid indices.
 
     Parameters
     ----------
     parent_grid : Grid
-        The parent grid object, containing information about the larger domain.
-    child_grid : Grid
-        The child grid object, containing information about the nested domain.
-    boundaries : Dict[str, bool], optional
-        Dictionary specifying which boundaries of the child grid are used
-        in the nesting process. Keys are "south", "east", "north", and "west",
-        with boolean values indicating inclusion. Defaults to all boundaries (True).
-    child_prefix : str, optional
-        Prefix added to variable names in the generated dataset to distinguish
-        child grid information. Defaults to "child".
-    period : float, optional
-        The temporal resolution or output period for boundary variables in the child grid.
-        Defaults to hourly.
+        The parent grid object, providing the reference for the child grid's topography and mask.
+    boundaries : Dict[str, bool]
+        Specifies which child grid boundaries (south, east, north, west) should be adjusted for topography/mask
+        and included in `ds_nesting`. Defaults to all `True`.
+    metadata : Dict[str, Any]
+        Dictionary configuring the boundary nesting process, including:
+
+        - `"prefix"` (str): Prefix for variable names in `ds_nesting`. Defaults to `"child"`.
+        - `"period"` (float): Temporal resolution for boundary outputs in seconds. Defaults to 3600 (hourly).
 
     Attributes
     ----------
     ds : xr.Dataset
-        An xarray Dataset containing the index mappings between the child and parent grids
-        for each specified boundary. Includes metadata about grid points, angles,
-        and boundary variable outputs.
+        Dataset containing child grid variables aligned with the parent grid’s topography and mask at the boundaries.
+    ds_nesting : xr.Dataset
+        Dataset containing boundary mappings, where child grid boundaries are mapped onto parent grid indices.
     """
 
     parent_grid: Grid
-    child_grid: Grid
     boundaries: Dict[str, bool] = field(
         default_factory=lambda: {
             "south": True,
@@ -61,40 +63,96 @@ class Nesting:
             "west": True,
         }
     )
-    child_prefix: str = "child"
-    period: float = 3600.0
+    metadata: Dict[str, Any] = field(
+        default_factory=lambda: {"prefix": "child", "period": 3600.0}
+    )
 
     def __post_init__(self):
 
-        parent_grid_ds = self.parent_grid.ds
-        child_grid_ds = self.child_grid.ds
+        super().__post_init__()
+        self._map_child_boundaries_onto_parent_grid_indices()
+        self._modify_child_topography_and_mask()
 
-        # Adjust longitude for dateline crossing to prevent interpolation artifacts
-        for grid_ds in [parent_grid_ds, child_grid_ds]:
-            grid_ds = wrap_longitudes(grid_ds, straddle=self.parent_grid.straddle)
+    def _map_child_boundaries_onto_parent_grid_indices(self):
+        """Maps child grid boundary points onto absolute indices of the parent grid."""
+        # Prepare parent and child grid datasets by adjusting longitudes for dateline crossing
+        parent_grid_ds, child_grid_ds = self._prepare_grid_datasets()
 
         # Map child boundaries onto parent grid indices
-        ds = map_child_boundaries_onto_parent_grid_indices(
+        ds_nesting = map_child_boundaries_onto_parent_grid_indices(
             parent_grid_ds,
             child_grid_ds,
             self.boundaries,
-            self.child_prefix,
-            self.period,
+            self.metadata["prefix"],
+            self.metadata["period"],
         )
-        object.__setattr__(self, "ds", ds)
 
-        # Modify child topography and mask to match the parent grid
+        object.__setattr__(self, "ds_nesting", ds_nesting)
+
+    def _modify_child_topography_and_mask(self):
+        """Adjust the child grid's topography and mask to align with the parent grid.
+
+        Uses a weighted sum based on boundary distance and clips depth values to a
+        minimum.
+        """
+        # Prepare parent and child grid datasets by adjusting longitudes for dateline crossing
+        parent_grid_ds, child_grid_ds = self._prepare_grid_datasets()
+
         child_grid_ds = modify_child_topography_and_mask(
-            parent_grid_ds, child_grid_ds, self.boundaries
+            parent_grid_ds, child_grid_ds, self.boundaries, self.hmin
         )
 
-        # Convert longitudes back to [0, 360] range
-        for grid_ds in [parent_grid_ds, child_grid_ds]:
-            grid_ds = wrap_longitudes(grid_ds, straddle=False)
-        object.__setattr__(self.parent_grid, "ds", parent_grid_ds)
-        object.__setattr__(self.child_grid, "ds", child_grid_ds)
+        # Finalize grid datasets by adjusting longitudes back to [0, 360] range
+        parent_grid_ds, child_grid_ds = self._finalize_grid_datasets(
+            parent_grid_ds, child_grid_ds
+        )
+
+        object.__setattr__(self, "ds", child_grid_ds)
+
+    def update_topography(
+        self, topography_source=None, hmin=None, verbose=False
+    ) -> None:
+
+        """Update the child grid topography via the following steps:
+
+        - Regrids the topography based on the specified source.
+        - Applies global and local smoothing.
+        - Adjusts the child grid topography and mask to match the parent grid.
+        - Ensures the minimum depth constraint (`hmin`) is enforced.
+        - Updates the internal dataset (`self.ds`) with the processed topography.
+
+        Parameters
+        ----------
+        topography_source : dict, optional
+            A dictionary specifying the source of the topography data. Expected keys:
+            - `"name"` (str): Name of the topography dataset (e.g., `"SRTM15"`).
+            - `"path"` (str or Path): File path to the topography dataset.
+
+            If not provided, the existing topography source remains unchanged.
+
+        hmin : float, optional
+            The minimum allowable ocean depth (in meters). If not provided, the existing
+            value remains unchanged.
 
-    def plot(self, with_dim_names=False) -> None:
+        verbose : bool, optional
+            If `True`, prints detailed information about each processing step, including
+            timing and modifications. Defaults to `False`.
+
+        Returns
+        -------
+        None
+            This method updates the internal dataset (`self.ds`) in place, modifying the
+            topography variable. It does not return a value.
+        """
+
+        super().update_topography(
+            topography_source=topography_source, hmin=hmin, verbose=verbose
+        )
+
+        # Modify child topography and mask to match the parent grid
+        self._modify_child_topography_and_mask()
+
+    def plot_nesting(self, with_dim_names=False) -> None:
         """Plot the parent and child grids in a single figure.
 
         Returns
@@ -105,45 +163,21 @@ class Nesting:
 
         _plot_nesting(
             self.parent_grid.ds,
-            self.child_grid.ds,
+            self.ds,
             self.parent_grid.straddle,
             with_dim_names,
         )
 
-    def save(
+    def save_nesting(
         self,
         filepath: Union[str, Path],
-        filepath_child_grid: Union[str, Path],
-        np_eta: int = None,
-        np_xi: int = None,
     ) -> None:
-        """Save the nesting and child grid file to netCDF4 files. The child grid file is
-        required because the topography and mask of the child grid has been modified.
-
-        This method allows saving the nesting and child grid data either each as a single file or each partitioned into multiple files, based on the provided options. The dataset can be saved in two modes:
-
-        1. **Single File Mode (default)**:
-            - If both `np_eta` and `np_xi` are `None`, the entire dataset is saved as a single netCDF4 file.
-            - The file is named based on the provided `filepath`, with `.nc` automatically appended to the filename.
-
-        2. **Partitioned Mode**:
-            - If either `np_eta` or `np_xi` is specified, the dataset is partitioned spatially along the `eta` and `xi` axes into tiles.
-            - Each tile is saved as a separate netCDF4 file. Filenames will be modified with an index to represent each partition, e.g., `"filepath_YYYYMM.0.nc"`, `"filepath_YYYYMM.1.nc"`, etc.
+        """Save the nesting information to netCDF4 files.
 
         Parameters
         ----------
         filepath : Union[str, Path]
             The base path and filename for the output files. The filenames will include the specified path and the `.nc` extension.
-            If partitioning is used, additional indices will be appended to the filenames, e.g., `"filepath.0.nc"`, `"filepath.1.nc"`, etc.
-
-        filepath_child_grid : Union[str, Path]
-            The base path and filename for saving the childe grid file.
-
-        np_eta : int, optional
-            The number of partitions along the `eta` direction. If `None`, no spatial partitioning is performed along the `eta` axis.
-
-        np_xi : int, optional
-            The number of partitions along the `xi` direction. If `None`, no spatial partitioning is performed along the `xi` axis.
 
         Returns
         -------
@@ -153,20 +187,15 @@ class Nesting:
 
         # Ensure filepath is a Path object
         filepath = Path(filepath)
-        filepath_child_grid = Path(filepath_child_grid)
 
         # Remove ".nc" suffix if present
         if filepath.suffix == ".nc":
             filepath = filepath.with_suffix("")
-        if filepath_child_grid.suffix == ".nc":
-            filepath_child_grid = filepath_child_grid.with_suffix("")
 
-        dataset_list = [self.ds, self.child_grid.ds]
-        output_filenames = [str(filepath), str(filepath_child_grid)]
+        dataset_list = [self.ds_nesting]
+        output_filenames = [str(filepath)]
 
-        saved_filenames = save_datasets(
-            dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
-        )
+        saved_filenames = save_datasets(dataset_list, output_filenames)
 
         return saved_filenames
 
@@ -183,8 +212,8 @@ class Nesting:
         _to_yaml(self, filepath)
 
     @classmethod
-    def from_yaml(cls, filepath: Union[str, Path]) -> "Nesting":
-        """Create an instance of the Nesting class from a YAML file.
+    def from_yaml(cls, filepath: Union[str, Path]) -> "ChildGrid":
+        """Create an instance of the ChildGrid class from a YAML file.
 
         Parameters
         ----------
@@ -194,25 +223,115 @@ class Nesting:
         Returns
         -------
         Nesting
-            An instance of the Nesting class.
+            An instance of the ChildGrid class.
         """
         filepath = Path(filepath)
 
         parent_grid = Grid.from_yaml(filepath, "ParentGrid")
-        child_grid = Grid.from_yaml(filepath, "ChildGrid")
         params = _from_yaml(cls, filepath)
 
-        return cls(parent_grid=parent_grid, child_grid=child_grid, **params)
+        return cls(parent_grid=parent_grid, **params)
+
+    def _prepare_grid_datasets(self) -> Tuple[xr.Dataset, xr.Dataset]:
+        """Prepare parent and child grid datasets by adjusting longitudes for dateline
+        crossing.
+
+        This method ensures that longitudes are properly wrapped to avoid interpolation artifacts
+        and returns the updated parent and child grid datasets.
+
+        Returns
+        -------
+        Tuple[xr.Dataset, xr.Dataset]
+            A tuple containing:
+            - The modified parent grid dataset.
+            - The modified child grid dataset.
+        """
+        parent_grid_ds = wrap_longitudes(
+            self.parent_grid.ds, straddle=self.parent_grid.straddle
+        )
+        child_grid_ds = wrap_longitudes(self.ds, straddle=self.parent_grid.straddle)
+        return parent_grid_ds, child_grid_ds
+
+    def _finalize_grid_datasets(
+        self, parent_grid_ds: xr.Dataset, child_grid_ds: xr.Dataset
+    ) -> None:
+        """Finalize the grid datasets by converting longitudes back to the [0, 360]
+        range.
+
+        Parameters
+        ----------
+        parent_grid_ds : xr.Dataset
+            The parent grid dataset after modifications.
+
+        child_grid_ds : xr.Dataset
+            The child grid dataset after modifications.
+        """
+        parent_grid_ds = wrap_longitudes(parent_grid_ds, straddle=False)
+        child_grid_ds = wrap_longitudes(child_grid_ds, straddle=False)
+        return parent_grid_ds, child_grid_ds
+
+    @classmethod
+    def from_file(
+        cls, filepath: Union[str, Path], verbose: bool = False
+    ) -> "ChildGrid":
+        """This method is disabled in this subclass.
+
+        .. noindex::
+        """
+        raise NotImplementedError(
+            "The 'from_file' method is disabled in this subclass."
+        )
 
 
 def map_child_boundaries_onto_parent_grid_indices(
     parent_grid_ds,
     child_grid_ds,
     boundaries={"south": True, "east": True, "north": True, "west": True},
-    child_prefix="child",
+    prefix="child",
     period=3600.0,
     update_land_indices=True,
 ):
+    """Maps child grid boundary points onto absolute indices of the parent grid.
+
+    This function interpolates the spatial indices of the child grid boundaries onto
+    the parent grid, ensuring alignment between the two grids. It supports all four
+    boundaries (south, east, north, and west) and considers different grid locations
+    (`rho`, `u`, and `v`). Additionally, it updates land indices if they fall onto
+    land points in the parent grid.
+
+    Parameters
+    ----------
+    parent_grid_ds : xarray.Dataset
+        The parent grid dataset containing longitude, latitude, and mask variables.
+
+    child_grid_ds : xarray.Dataset
+        The child grid dataset containing longitude, latitude, mask, and angle variables.
+
+    boundaries : dict, optional
+        A dictionary specifying which child boundaries should be mapped onto the parent grid.
+        Keys should be `"south"`, `"east"`, `"north"`, and `"west"`, with boolean values
+        indicating whether to process each boundary. Defaults to mapping all boundaries.
+
+    prefix : str, optional
+        A string prefix for naming the output variables in the resulting dataset.
+        Defaults to `"child"`.
+
+    period : float, optional
+        The output period (in seconds) to be assigned to the mapped boundary indices.
+        Defaults to `3600.0` (1 hour).
+
+    update_land_indices : bool, optional
+        If `True`, updates indices that fall on land in the parent grid to nearby ocean points.
+        Defaults to `True`.
+
+    Returns
+    -------
+    xarray.Dataset
+        A dataset containing the mapped boundary indices for `rho`, `u`, and `v` grid points.
+        - For `rho` points: Contains mapped `xi` and `eta` indices.
+        - For `u` and `v` points: Contains mapped `xi`, `eta`, and angle values.
+        - Attributes include long names, output variable names, units, and output period.
+    """
 
     bdry_coords_dict = get_boundary_coords()
 
@@ -255,7 +374,7 @@ def map_child_boundaries_onto_parent_grid_indices(
                         i_eta, i_xi, grid_location, parent_grid_ds
                     )
 
-                var_name = f"{child_prefix}_{direction}_{suffix}"
+                var_name = f"{prefix}_{direction}_{suffix}"
                 if grid_location == "rho":
                     ds[var_name] = xr.concat([i_xi, i_eta], dim="two")
                     ds[var_name].attrs[
@@ -286,7 +405,7 @@ def map_child_boundaries_onto_parent_grid_indices(
 
     # Rename dimensions
     dims_to_rename = {
-        dim: f"{child_prefix}_{dim}" for dim in ds.dims if dim not in ["two", "three"]
+        dim: f"{prefix}_{dim}" for dim in ds.dims if dim not in ["two", "three"]
     }
     ds = ds.rename(dims_to_rename)
 
@@ -466,7 +585,39 @@ def modify_child_topography_and_mask(
     parent_grid_ds,
     child_grid_ds,
     boundaries={"south": True, "east": True, "north": True, "west": True},
+    hmin=5.0,
 ):
+    """Adjust the child grid topography and mask to align with the parent grid.
+
+    The child grid's topography is adjusted using a weighted sum based on the boundary distance,
+    and the depth values are clipped to enforce a minimum depth constraint.
+
+    Parameters
+    ----------
+    parent_grid_ds : xarray.Dataset
+        The parent grid dataset containing `h` (topography) and `mask_rho` (land-sea mask).
+
+    child_grid_ds : xarray.Dataset
+        The child grid dataset whose `h` and `mask_rho` will be modified.
+
+    boundaries : dict, optional
+        A dictionary specifying which boundaries should be modified. Expected keys:
+        - `"south"` (bool): Whether to adjust the southern boundary.
+        - `"east"` (bool): Whether to adjust the eastern boundary.
+        - `"north"` (bool): Whether to adjust the northern boundary.
+        - `"west"` (bool): Whether to adjust the western boundary.
+        Defaults to modifying all boundaries.
+
+    hmin : float, optional
+        The minimum allowable ocean depth (in meters). Depth values in the modified
+        child topography will be clipped to ensure they meet or exceed this value.
+        Defaults to `5.0` meters.
+
+    Returns
+    -------
+    xarray.Dataset
+        The updated child grid dataset with modified topography (`h`) and mask (`mask_rho`).
+    """
 
     # regrid parent topography and mask onto child grid
     points = np.column_stack(
@@ -492,6 +643,8 @@ def modify_child_topography_and_mask(
     child_grid_ds["h"] = (
         alpha * child_grid_ds["h"] + (1 - alpha) * h_parent_interpolated
     )
+    # Clip depth on modified child topography
+    child_grid_ds["h"] = _clip_depth(child_grid_ds["h"], hmin)
 
     child_mask = (
         alpha * child_grid_ds["mask_rho"] + (1 - alpha) * mask_parent_interpolated