roms-tools 2.4.0__py3-none-any.whl → 2.6.0__py3-none-any.whl

roms_tools/setup/initial_conditions.py

@@ -8,7 +8,7 @@ from pathlib import Path
 import logging
 from datetime import datetime
 from roms_tools import Grid
-from roms_tools.regrid import LateralRegrid, VerticalRegrid
+from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
 from roms_tools.plot import _plot, _section_plot, _profile_plot, _line_plot
 from roms_tools.utils import (
     transpose_dimensions,
@@ -34,7 +34,7 @@ from roms_tools.setup.utils import (
 )
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class InitialConditions:
     """Represents initial conditions for ROMS, including physical and biogeochemical
     data.
@@ -115,7 +115,7 @@ class InitialConditions:
 
         self._input_checks()
         # Dataset for depth coordinates
-        object.__setattr__(self, "ds_depth_coords", xr.Dataset())
+        self.ds_depth_coords = xr.Dataset()
 
         processed_fields = {}
         processed_fields = self._process_data(processed_fields, type="physics")
@@ -135,7 +135,7 @@ class InitialConditions:
         for var_name in ds.data_vars:
             ds[var_name] = substitute_nans_by_fillvalue(ds[var_name])
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def _process_data(self, processed_fields, type="physics"):
 
@@ -150,6 +150,8 @@ class InitialConditions:
             target_coords,
             buffer_points=20,  # lateral fill needs good buffer from data margin
         )
+        # Enforce double precision to ensure reproducibility
+        data.convert_to_float64()
         data.extrapolate_deepest_to_bottom()
         data.apply_lateral_fill()
 
@@ -159,7 +161,7 @@ class InitialConditions:
         var_names = variable_info.keys()
 
         # lateral regridding
-        lateral_regrid = LateralRegrid(target_coords, data.dim_names)
+        lateral_regrid = LateralRegridToROMS(target_coords, data.dim_names)
 
         for var_name in var_names:
             if var_name in data.var_names.keys():
@@ -204,7 +206,7 @@ class InitialConditions:
         # Vertical regridding
         for location in ["rho", "u", "v"]:
             if len(var_names_dict[location]) > 0:
-                vertical_regrid = VerticalRegrid(
+                vertical_regrid = VerticalRegridToROMS(
                     self.ds_depth_coords[f"layer_depth_{location}"],
                     data.ds[data.dim_names["depth"]],
                 )
@@ -234,20 +236,21 @@ class InitialConditions:
         return processed_fields
 
     def _input_checks(self):
+        # Check that ini_time is not None
+        if self.ini_time is None:
+            raise ValueError(
+                "`ini_time` must be a valid datetime object and cannot be None."
+            )
 
         if "name" not in self.source.keys():
             raise ValueError("`source` must include a 'name'.")
         if "path" not in self.source.keys():
             raise ValueError("`source` must include a 'path'.")
         # set self.source["climatology"] to False if not provided
-        object.__setattr__(
-            self,
-            "source",
-            {
-                **self.source,
-                "climatology": self.source.get("climatology", False),
-            },
-        )
+        self.source = {
+            **self.source,
+            "climatology": self.source.get("climatology", False),
+        }
         if self.bgc_source is not None:
             if "name" not in self.bgc_source.keys():
                 raise ValueError(
@@ -258,14 +261,10 @@ class InitialConditions:
                     "`bgc_source` must include a 'path' if it is provided."
                 )
             # set self.bgc_source["climatology"] to False if not provided
-            object.__setattr__(
-                self,
-                "bgc_source",
-                {
-                    **self.bgc_source,
-                    "climatology": self.bgc_source.get("climatology", False),
-                },
-            )
+            self.bgc_source = {
+                **self.bgc_source,
+                "climatology": self.bgc_source.get("climatology", False),
+            }
         if self.adjust_depth_for_sea_surface_height:
             logging.info("Sea surface height will be used to adjust depth coordinates.")
         else:
@@ -648,7 +647,7 @@ class InitialConditions:
             visualize the layering of the water column. For clarity, the number of layer
             contours displayed is limited to a maximum of 10. Default is False.
         ax : matplotlib.axes.Axes, optional
-            The axes to plot on. If None, a new figure is created. Note that this argument does not work for horizontal plots that display the eta- and xi-dimensions at the same time.
+            The axes to plot on. If None, a new figure is created. Note that this argument does not work for 2D horizontal plots. Default is None.
 
         Returns
         -------
@@ -680,13 +679,6 @@ class InitialConditions:
                 "Conflicting input: For 2D fields, specify only one dimension, either 'eta' or 'xi', not both."
             )
 
-        # Load the data
-        if self.use_dask:
-            from dask.diagnostics import ProgressBar
-
-            with ProgressBar():
-                self.ds[var_name].load()
-
         field = self.ds[var_name].squeeze()
 
         # Get correct mask and horizontal coordinates
@@ -708,11 +700,18 @@ class InitialConditions:
 
         field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
 
+        # Load the data
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                self.ds[var_name].load()
+
         # Retrieve depth coordinates
         if s is not None:
             layer_contours = False
         # Note that `layer_depth_{loc}` has already been computed during `__post_init__`.
-        layer_depth = self.ds_depth_coords[f"layer_depth_{loc}"].squeeze()
+        layer_depth = self.ds_depth_coords[f"layer_depth_{loc}"].squeeze().load()
 
         # Slice the field as desired
         def _slice_and_assign(
@@ -898,7 +897,6 @@ class InitialConditions:
         cls,
         filepath: Union[str, Path],
         use_dask: bool = False,
-        bypass_validation: bool = False,
     ) -> "InitialConditions":
         """Create an instance of the InitialConditions class from a YAML file.
 
@@ -908,10 +906,6 @@ class InitialConditions:
             The path to the YAML file from which the parameters will be read.
         use_dask: bool, optional
             Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
-        bypass_validation: bool, optional
-            Indicates whether to skip validation checks in the processed data. When set to True,
-            the validation process that ensures no NaN values exist at wet points
-            in the processed dataset is bypassed. Defaults to False.
 
         Returns
         -------
@@ -926,5 +920,4 @@ class InitialConditions:
             grid=grid,
             **initial_conditions_params,
             use_dask=use_dask,
-            bypass_validation=bypass_validation,
         )

roms_tools/setup/nesting.py

@@ -2,13 +2,14 @@ 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,
@@ -19,40 +20,41 @@ from roms_tools.setup.utils import (
 )
 
 
-@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.
+@dataclass(kw_only=True)
+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
+        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
+        )
+
+        # 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
+        )
+
+        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.
+
+        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
         )
 
-        # 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)
+        # Modify child topography and mask to match the parent grid
+        self._modify_child_topography_and_mask()
 
-    def plot(self, with_dim_names=False) -> None:
+    def plot_nesting(self, with_dim_names=False) -> None:
         """Plot the parent and child grids in a single figure.
 
         Returns
@@ -105,27 +163,22 @@ 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],
     ) -> 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.
+        """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.
 
-        filepath_child_grid : Union[str, Path]
-            The base path and filename for saving the childe grid file.
-
         Returns
         -------
         List[Path]
@@ -134,16 +187,13 @@ 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)
 
@@ -162,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
         ----------
@@ -173,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()
 
@@ -234,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[
@@ -265,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)
 
@@ -445,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(
@@ -471,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