roms-tools 2.0.0__py3-none-any.whl → 2.2.0__py3-none-any.whl

roms_tools/setup/surface_forcing.py

@@ -66,6 +66,10 @@ class SurfaceForcing:
         Reference date for the model. Default is January 1, 2000.
     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.
 
     Examples
     --------
@@ -88,6 +92,7 @@ class SurfaceForcing:
     use_coarse_grid: bool = False
     model_reference_date: datetime = datetime(2000, 1, 1)
     use_dask: bool = False
+    bypass_validation: bool = False
 
     ds: xr.Dataset = field(init=False, repr=False)
 
@@ -117,7 +122,7 @@ class SurfaceForcing:
                     data.ds[data.var_names[var_name]]
                 )
 
-        # rotation of velocities and interpolation to u/v points
+        # rotation of velocities
         if "uwnd" in self.variable_info and "vwnd" in self.variable_info:
             processed_fields["uwnd"], processed_fields["vwnd"] = rotate_velocities(
                 processed_fields["uwnd"],
@@ -134,7 +139,8 @@ class SurfaceForcing:
 
         ds = self._write_into_dataset(processed_fields, data, d_meta)
 
-        self._validate(ds)
+        if not self.bypass_validation:
+            self._validate(ds)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
         for var_name in ds.data_vars:
@@ -262,10 +268,8 @@ class SurfaceForcing:
         correction_data = self._get_correction_data()
         # choose same subdomain as forcing data so that we can use same mask
         coords_correction = {
-            correction_data.dim_names["latitude"]: data.ds[data.dim_names["latitude"]],
-            correction_data.dim_names["longitude"]: data.ds[
-                data.dim_names["longitude"]
-            ],
+            "lat": data.ds[data.dim_names["latitude"]],
+            "lon": data.ds[data.dim_names["longitude"]],
         }
         correction_data.choose_subdomain(
             coords_correction, straddle=self.target_coords["straddle"]
@@ -287,6 +291,8 @@ class SurfaceForcing:
 
         processed_fields["swrad"] = processed_fields["swrad"] * corr_factor
 
+        del corr_factor
+
         return processed_fields
 
     def _write_into_dataset(self, processed_fields, data, d_meta):
@@ -294,8 +300,9 @@ class SurfaceForcing:
         # save in new dataset
         ds = xr.Dataset()
 
-        for var_name in processed_fields.keys():
+        for var_name in list(processed_fields.keys()):
             ds[var_name] = processed_fields[var_name].astype(np.float32)
+            del processed_fields[var_name]
             ds[var_name].attrs["long_name"] = d_meta[var_name]["long_name"]
             ds[var_name].attrs["units"] = d_meta[var_name]["units"]
 
@@ -349,9 +356,14 @@ class SurfaceForcing:
         """
 
         for var_name in ds.data_vars:
-            # Only validate variables based on "validate" flag if use_dask is False
-            if not self.use_dask or self.variable_info[var_name]["validate"]:
-                nan_check(ds[var_name].isel(time=0), self.target_coords["mask"])
+            if self.variable_info[var_name]["validate"]:
+                if self.variable_info[var_name]["location"] == "rho":
+                    mask = self.target_coords["mask"]
+                elif self.variable_info[var_name]["location"] == "u":
+                    mask = self.target_coords["mask_u"]
+                elif self.variable_info[var_name]["location"] == "v":
+                    mask = self.target_coords["mask_v"]
+                nan_check(ds[var_name].isel(time=0), mask)
 
     def _add_global_metadata(self, ds=None):
 
@@ -541,7 +553,10 @@ class SurfaceForcing:
 
     @classmethod
     def from_yaml(
-        cls, filepath: Union[str, Path], use_dask: bool = False
+        cls,
+        filepath: Union[str, Path],
+        use_dask: bool = False,
+        bypass_validation: bool = False,
     ) -> "SurfaceForcing":
         """Create an instance of the SurfaceForcing class from a YAML file.
 
@@ -551,6 +566,10 @@ class SurfaceForcing:
             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
         -------
@@ -562,4 +581,6 @@ class SurfaceForcing:
         grid = Grid.from_yaml(filepath)
         params = _from_yaml(cls, filepath)
 
-        return cls(grid=grid, **params, use_dask=use_dask)
+        return cls(
+            grid=grid, **params, use_dask=use_dask, bypass_validation=bypass_validation
+        )

roms_tools/setup/tides.py

@@ -51,6 +51,10 @@ class TidalForcing:
         The reference date for the ROMS simulation. Default is datetime(2000, 1, 1).
     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.
 
     Examples
     --------
@@ -65,6 +69,7 @@ class TidalForcing:
     allan_factor: float = 2.0
     model_reference_date: datetime = datetime(2000, 1, 1)
     use_dask: bool = False
+    bypass_validation: bool = False
 
     ds: xr.Dataset = field(init=False, repr=False)
 
@@ -129,7 +134,8 @@ class TidalForcing:
 
         ds = self._add_global_metadata(ds)
 
-        self._validate(ds)
+        if not self.bypass_validation:
+            self._validate(ds)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
         for var_name in ds.data_vars:
@@ -264,9 +270,16 @@ class TidalForcing:
         The method utilizes `self.grid.ds.mask_rho` to determine the wet points in the domain.
         """
         for var_name in ds.data_vars:
-            # only validate variables based on "validate" flag if use_dask is false
-            if not self.use_dask or self.variable_info[var_name]["validate"]:
-                nan_check(ds[var_name].isel(ntides=0), self.grid.ds.mask_rho)
+            if self.variable_info[var_name]["validate"]:
+                if self.variable_info[var_name]["location"] == "rho":
+                    mask = self.grid.ds.mask_rho
+                elif self.variable_info[var_name]["location"] == "u":
+                    mask = self.grid.ds.mask_u
+                elif self.variable_info[var_name]["location"] == "v":
+                    mask = self.grid.ds.mask_v
+
+                da = ds[var_name].isel(ntides=0)
+                nan_check(da, mask)
 
     def plot(self, var_name, ntides=0) -> None:
         """Plot the specified tidal forcing variable for a given tidal constituent.
@@ -420,7 +433,10 @@ class TidalForcing:
 
     @classmethod
     def from_yaml(
-        cls, filepath: Union[str, Path], use_dask: bool = False
+        cls,
+        filepath: Union[str, Path],
+        use_dask: bool = False,
+        bypass_validation: bool = False,
     ) -> "TidalForcing":
         """Create an instance of the TidalForcing class from a YAML file.
 
@@ -430,6 +446,10 @@ class TidalForcing:
             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
         -------
@@ -440,7 +460,12 @@ class TidalForcing:
 
         grid = Grid.from_yaml(filepath)
         tidal_forcing_params = _from_yaml(cls, filepath)
-        return cls(grid=grid, **tidal_forcing_params, use_dask=use_dask)
+        return cls(
+            grid=grid,
+            **tidal_forcing_params,
+            use_dask=use_dask,
+            bypass_validation=bypass_validation
+        )
 
     def _correct_tides(self, data):
         """Apply tidal corrections to the dataset. This method corrects the dataset for

roms_tools/setup/topography.py

@@ -94,7 +94,18 @@ def _add_topography(
 
 
 def _get_topography_data(source):
+    """Load topography data based on the specified source.
 
+    Parameters
+    ----------
+    source : dict
+        A dictionary containing the source details (e.g., "name" and "path").
+
+    Returns
+    -------
+    data : object
+        The loaded topography dataset (ETOPO5 or SRTM15).
+    """
     kwargs = {"use_dask": False}
 
     if source["name"] == "ETOPO5":
@@ -115,7 +126,24 @@ def _get_topography_data(source):
 def _make_raw_topography(
     data, target_coords, method="linear", verbose=False
 ) -> xr.DataArray:
+    """Regrid topography data to match target coordinates.
+
+    Parameters
+    ----------
+    data : object
+        The dataset object containing the topography data.
+    target_coords : object
+        The target coordinates to which the data will be regridded.
+    method : str, optional
+        The regridding method to use, by default "linear".
+    verbose : bool, optional
+        If True, logs the time taken for regridding, by default False.
 
+    Returns
+    -------
+    xr.DataArray
+        The regridded topography data with the sign flipped (bathymetry positive).
+    """
     data.choose_subdomain(target_coords, buffer_points=3, verbose=verbose)
 
     if verbose:
@@ -134,9 +162,22 @@ def _make_raw_topography(
 
 
 def _smooth_topography_globally(hraw, factor) -> xr.DataArray:
+    """Apply global smoothing to the topography using a Gaussian filter.
+
+    Parameters
+    ----------
+    hraw : xr.DataArray
+        The raw topography data to be smoothed.
+    factor : float
+        The smoothing factor (controls the width of the Gaussian filter).
+
+    Returns
+    -------
+    xr.DataArray
+        The smoothed topography data.
+    """
     # since GCM-Filters assumes periodic domain, we extend the domain by one grid cell in each dimension
     # and set that margin to land
-
     mask = xr.ones_like(hraw)
     margin_mask = xr.concat([mask, 0 * mask.isel(eta_rho=-1)], dim="eta_rho")
     margin_mask = xr.concat(
@@ -164,7 +205,27 @@ def _smooth_topography_globally(hraw, factor) -> xr.DataArray:
 
 
 def _smooth_topography_locally(h, hmin=5, rmax=0.2):
-    """Smoothes topography locally to satisfy r < rmax."""
+    """Smooths topography locally to ensure the slope (r-factor) is below the specified
+    threshold.
+
+    This function applies a logarithmic transformation to the topography and iteratively smooths
+    it in four directions (eta, xi, and two diagonals) until the maximum slope parameter (r) is
+    below `rmax`. A threshold `hmin` is applied to prevent values from going below a minimum height.
+
+    Parameters
+    ----------
+    h : xarray.DataArray
+        The topography data to be smoothed.
+    hmin : float, optional
+        The minimum height threshold. Default is 5.
+    rmax : float, optional
+        The maximum allowable slope parameter (r-factor). Default is 0.2.
+
+    Returns
+    -------
+    xarray.DataArray
+        The smoothed topography data.
+    """
     # Compute rmax_log
     if rmax > 0.0:
         rmax_log = np.log((1.0 + rmax * 0.9) / (1.0 - rmax * 0.9))
@@ -174,65 +235,90 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
     # Apply hmin threshold
     h = xr.where(h < hmin, hmin, h)
 
-    # We will smooth logarithmically
+    # Perform logarithmic transformation of the height field
     h_log = np.log(h / hmin)
 
-    cf1 = 1.0 / 6
-    cf2 = 0.25
+    # Constants for smoothing
+    smoothing_factor_1 = 1.0 / 6
+    smoothing_factor_2 = 0.25
 
+    # Iterate until convergence
     for iter in count():
-        # Compute gradients in domain interior
+        # Compute gradients and smoothing for eta, xi, and diagonal directions
 
-        # in eta-direction
-        cff = h_log.diff("eta_rho").isel(xi_rho=slice(1, -1))
-        cr = np.abs(cff)
+        # Gradient in eta-direction
+        delta_eta = h_log.diff("eta_rho").isel(xi_rho=slice(1, -1))
+        abs_eta_gradient = np.abs(delta_eta)
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")  # Ignore division by zero warning
-            Op1 = xr.where(cr < rmax_log, 0, 1.0 * cff * (1 - rmax_log / cr))
-
-        # in xi-direction
-        cff = h_log.diff("xi_rho").isel(eta_rho=slice(1, -1))
-        cr = np.abs(cff)
+            eta_correction = xr.where(
+                abs_eta_gradient < rmax_log,
+                0,
+                delta_eta * (1 - rmax_log / abs_eta_gradient),
+            )
+
+        # Gradient in xi-direction
+        delta_xi = h_log.diff("xi_rho").isel(eta_rho=slice(1, -1))
+        abs_xi_gradient = np.abs(delta_xi)
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")  # Ignore division by zero warning
-            Op2 = xr.where(cr < rmax_log, 0, 1.0 * cff * (1 - rmax_log / cr))
-
-        # in diagonal direction
-        cff = (h_log - h_log.shift(eta_rho=1, xi_rho=1)).isel(
+            xi_correction = xr.where(
+                abs_xi_gradient < rmax_log,
+                0,
+                delta_xi * (1 - rmax_log / abs_xi_gradient),
+            )
+
+        # Gradient in first diagonal direction
+        delta_diag_1 = (h_log - h_log.shift(eta_rho=1, xi_rho=1)).isel(
             eta_rho=slice(1, None), xi_rho=slice(1, None)
         )
-        cr = np.abs(cff)
+        abs_diag_1_gradient = np.abs(delta_diag_1)
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")  # Ignore division by zero warning
-            Op3 = xr.where(cr < rmax_log, 0, 1.0 * cff * (1 - rmax_log / cr))
-
-        # in the other diagonal direction
-        cff = (h_log.shift(eta_rho=1) - h_log.shift(xi_rho=1)).isel(
+            diag_1_correction = xr.where(
+                abs_diag_1_gradient < rmax_log,
+                0,
+                delta_diag_1 * (1 - rmax_log / abs_diag_1_gradient),
+            )
+
+        # Gradient in second diagonal direction
+        delta_diag_2 = (h_log.shift(eta_rho=1) - h_log.shift(xi_rho=1)).isel(
             eta_rho=slice(1, None), xi_rho=slice(1, None)
         )
-        cr = np.abs(cff)
+        abs_diag_2_gradient = np.abs(delta_diag_2)
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")  # Ignore division by zero warning
-            Op4 = xr.where(cr < rmax_log, 0, 1.0 * cff * (1 - rmax_log / cr))
+            diag_2_correction = xr.where(
+                abs_diag_2_gradient < rmax_log,
+                0,
+                delta_diag_2 * (1 - rmax_log / abs_diag_2_gradient),
+            )
 
         # Update h_log in domain interior
-        h_log[1:-1, 1:-1] += cf1 * (
-            Op1[1:, :]
-            - Op1[:-1, :]
-            + Op2[:, 1:]
-            - Op2[:, :-1]
-            + cf2 * (Op3[1:, 1:] - Op3[:-1, :-1] + Op4[:-1, 1:] - Op4[1:, :-1])
+        h_log[1:-1, 1:-1] += smoothing_factor_1 * (
+            eta_correction[1:, :]
+            - eta_correction[:-1, :]
+            + xi_correction[:, 1:]
+            - xi_correction[:, :-1]
+            + smoothing_factor_2
+            * (
+                diag_1_correction[1:, 1:]
+                - diag_1_correction[:-1, :-1]
+                + diag_2_correction[:-1, 1:]
+                - diag_2_correction[1:, :-1]
+            )
         )
 
         # No gradient at the domain boundaries
         h_log = handle_boundaries(h_log)
 
-        # Update h
+        # Recompute the topography after smoothing
         h = hmin * np.exp(h_log)
+
         # Apply hmin threshold again
         h = xr.where(h < hmin, hmin, h)
 
-        # compute maximum slope parameter r
+        # Compute maximum slope parameter r
         r_eta, r_xi = _compute_rfactor(h)
         rmax0 = np.max([r_eta.max(), r_xi.max()])
         if rmax0 < rmax:
@@ -242,8 +328,23 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
 
 
 def _compute_rfactor(h):
-    """Computes slope parameter (or r-factor) r = |Delta h| / 2h in both horizontal grid
-    directions."""
+    """Computes the slope parameter (r-factor) in both horizontal directions.
+
+    The r-factor is calculated as |Δh| / (2h) in the eta and xi directions:
+        - r_eta = |h_i - h_{i-1}| / (h_i + h_{i+1})
+        - r_xi = |h_i - h_{i-1}| / (h_i + h_{i+1})
+
+    Parameters
+    ----------
+    h : xarray.DataArray
+        The topography data.
+
+    Returns
+    -------
+    tuple of xarray.DataArray
+        r_eta : r-factor in the eta direction.
+        r_xi : r-factor in the xi direction.
+    """
     # compute r_{i-1/2} = |h_i - h_{i-1}| / (h_i + h_{i+1})
     r_eta = np.abs(h.diff("eta_rho")) / (h + h.shift(eta_rho=1)).isel(
         eta_rho=slice(1, None)
@@ -256,6 +357,26 @@ def _compute_rfactor(h):
 
 
 def _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax):
+    """Adds topography metadata to the dataset.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        Dataset to update.
+    topography_source : dict
+        Dictionary with topography source information (requires 'name' key).
+    smooth_factor : float
+        Smoothing factor (unused in this function).
+    hmin : float
+        Minimum height threshold for smoothing.
+    rmax : float
+        Maximum slope parameter (unused in this function).
+
+    Returns
+    -------
+    xarray.Dataset
+        Updated dataset with added metadata.
+    """
     ds.attrs["topography_source"] = topography_source["name"]
     ds.attrs["hmin"] = hmin
 
@@ -263,6 +384,18 @@ def _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax):
 
 
 def nan_check(hraw):
+    """Checks for NaN values in the topography data.
+
+    Parameters
+    ----------
+    hraw : xarray.DataArray
+        Input topography data to check for NaN values.
+
+    Raises
+    ------
+    ValueError
+        If NaN values are found in the data, raises an error with a descriptive message.
+    """
     error_message = (
         "NaN values found in regridded topography. This likely occurs because the ROMS grid, including "
         "a small safety margin for interpolation, is not fully contained within the topography dataset's longitude/latitude range. Please ensure that the "