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

roms_tools/setup/initial_conditions.py

@@ -66,6 +66,10 @@ class InitialConditions:
         The reference date for the model. Defaults to 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
     --------
@@ -87,6 +91,7 @@ class InitialConditions:
     bgc_source: Optional[Dict[str, Union[str, Path, List[Union[str, Path]]]]] = None
     model_reference_date: datetime = datetime(2000, 1, 1)
     use_dask: bool = False
+    bypass_validation: bool = False
 
     ds: xr.Dataset = field(init=False, repr=False)
 
@@ -105,7 +110,8 @@ class InitialConditions:
 
         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:
@@ -527,10 +533,15 @@ class InitialConditions:
             variable_info = self.variable_info_physics
 
         for var_name in variable_info:
-            # Only validate variables based on "validate" flag if use_dask is False
-            if not self.use_dask or variable_info[var_name]["validate"]:
+            if variable_info[var_name]["validate"]:
+                if variable_info[var_name]["location"] == "rho":
+                    mask = self.grid.ds.mask_rho
+                elif variable_info[var_name]["location"] == "u":
+                    mask = self.grid.ds.mask_u
+                elif variable_info[var_name]["location"] == "v":
+                    mask = self.grid.ds.mask_v
                 ds[var_name].load()
-                nan_check(ds[var_name].squeeze(), self.grid.ds.mask_rho)
+                nan_check(ds[var_name].squeeze(), mask)
 
     def _add_global_metadata(self, ds):
 
@@ -882,7 +893,10 @@ class InitialConditions:
 
     @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,
     ) -> "InitialConditions":
         """Create an instance of the InitialConditions class from a YAML file.
 
@@ -892,6 +906,10 @@ 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
         -------
@@ -902,4 +920,9 @@ class InitialConditions:
 
         grid = Grid.from_yaml(filepath)
         initial_conditions_params = _from_yaml(cls, filepath)
-        return cls(grid=grid, **initial_conditions_params, use_dask=use_dask)
+        return cls(
+            grid=grid,
+            **initial_conditions_params,
+            use_dask=use_dask,
+            bypass_validation=bypass_validation,
+        )

roms_tools/setup/mask.py

@@ -1,6 +1,7 @@
 import xarray as xr
 import numpy as np
 import regionmask
+import warnings
 from scipy.ndimage import label
 from roms_tools.setup.utils import (
     interpolate_from_rho_to_u,
@@ -10,10 +11,26 @@ from roms_tools.setup.utils import (
 
 
 def _add_mask(ds):
-
+    """Adds a land/water mask to the dataset at rho-points.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        Input dataset containing latitude and longitude coordinates at rho-points.
+
+    Returns
+    -------
+    xarray.Dataset
+        The original dataset with an added 'mask_rho' variable, representing land/water mask.
+    """
     land = regionmask.defined_regions.natural_earth_v5_0_0.land_10
 
-    land_mask = land.mask(ds["lon_rho"], ds["lat_rho"])
+    # Suppress specific warning
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", message="No gridpoint belongs to any region.*"
+        )
+        land_mask = land.mask(ds["lon_rho"], ds["lat_rho"])
     mask = land_mask.isnull()
 
     # fill enclosed basins with land
@@ -32,7 +49,21 @@ def _add_mask(ds):
 
 
 def _fill_enclosed_basins(mask) -> np.ndarray:
-    """Fills in enclosed basins with land."""
+    """Fills enclosed basins in the mask with land (value = 1).
+
+    This function identifies the largest connected region in the mask, which is assumed to represent
+    the land, and sets all other regions to water (value = 0).
+
+    Parameters
+    ----------
+    mask : np.ndarray
+        A binary array representing the land/water mask (land = 1, water = 0).
+
+    Returns
+    -------
+    np.ndarray
+        The modified mask with enclosed basins filled with land (1).
+    """
 
     # Label connected regions in the mask
     reg, nreg = label(mask)
@@ -54,7 +85,22 @@ def _fill_enclosed_basins(mask) -> np.ndarray:
 
 
 def _add_velocity_masks(ds):
-
+    """Adds velocity masks for u- and v-points based on the rho-point mask.
+
+    This function generates masks for u- and v-points by interpolating the rho-point land/water mask.
+    The interpolation method used is "multiplicative", which scales the rho-point mask to the respective
+    u- and v-points.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        The dataset containing the land/water mask at rho-points (`mask_rho`).
+
+    Returns
+    -------
+    xarray.Dataset
+        The input dataset with added velocity masks (`mask_u` and `mask_v`) for u- and v-points.
+    """
     # add u- and v-masks
     ds["mask_u"] = interpolate_from_rho_to_u(
         ds["mask_rho"], method="multiplicative"