roms-tools 1.7.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

roms_tools/setup/initial_conditions.py

@@ -6,6 +6,7 @@ from typing import Dict, Union, List, Optional
 from roms_tools.setup.grid import Grid
 from datetime import datetime
 from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
+from roms_tools.setup.vertical_coordinate import compute_depth
 from roms_tools.setup.utils import (
     nan_check,
     substitute_nans_by_fillvalue,
@@ -15,6 +16,8 @@ from roms_tools.setup.utils import (
     rotate_velocities,
     compute_barotropic_velocity,
     transpose_dimensions,
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
     _to_yaml,
     _from_yaml,
 )
@@ -63,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
     --------
@@ -84,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)
 
@@ -92,23 +100,18 @@ class InitialConditions:
         self._input_checks()
 
         processed_fields = {}
-        processed_fields, variable_info = self._process_data(
-            processed_fields, type="physics"
-        )
+        processed_fields = self._process_data(processed_fields, type="physics")
 
         if self.bgc_source is not None:
-            processed_fields, bgc_variable_info = self._process_data(
-                processed_fields, type="bgc"
-            )
+            processed_fields = self._process_data(processed_fields, type="bgc")
 
         d_meta = get_variable_metadata()
         ds = self._write_into_dataset(processed_fields, d_meta)
 
         ds = self._add_global_metadata(ds)
 
-        if self.bgc_source is not None:
-            variable_info = {**variable_info, **bgc_variable_info}
-        self._validate(ds, variable_info)
+        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:
@@ -133,7 +136,9 @@ class InitialConditions:
         data.extrapolate_deepest_to_bottom()
         data.apply_lateral_fill()
 
-        variable_info = self._set_variable_info(data, type=type)
+        self._set_variable_info(data, type=type)
+        attr_name = f"variable_info_{type}"
+        variable_info = getattr(self, attr_name)
         var_names = variable_info.keys()
 
         # lateral regridding
@@ -153,19 +158,31 @@ class InitialConditions:
                 interpolate=True,
             )
 
-        # vertical regridding
+        var_names_dict = {}
         for location in ["rho", "u", "v"]:
-            var_names = [
+            var_names_dict[location] = [
                 name
                 for name, info in variable_info.items()
                 if info["location"] == location and info["is_3d"]
             ]
-            if len(var_names) > 0:
+
+        # compute layer depth coordinates
+        if len(var_names_dict["u"]) > 0 or len(var_names_dict["v"]) > 0:
+            self._get_vertical_coordinates(
+                type="layer",
+                additional_locations=["u", "v"],
+            )
+        else:
+            if len(var_names_dict["rho"]) > 0:
+                self._get_vertical_coordinates(type="layer", additional_locations=[])
+        # vertical regridding
+        for location in ["rho", "u", "v"]:
+            if len(var_names_dict[location]) > 0:
                 vertical_regrid = VerticalRegrid(
                     self.grid.ds[f"layer_depth_{location}"],
                     data.ds[data.dim_names["depth"]],
                 )
-                for var_name in var_names:
+                for var_name in var_names_dict[location]:
                     if var_name in processed_fields:
                         processed_fields[var_name] = vertical_regrid.apply(
                             processed_fields[var_name]
@@ -173,10 +190,14 @@ class InitialConditions:
 
         # compute barotropic velocities
         if "u" in variable_info and "v" in variable_info:
-            for var_name in ["u", "v"]:
-                processed_fields[f"{var_name}bar"] = compute_barotropic_velocity(
-                    processed_fields[var_name],
-                    self.grid.ds[f"interface_depth_{var_name}"],
+            self._get_vertical_coordinates(
+                type="interface",
+                additional_locations=["u", "v"],
+            )
+            for location in ["u", "v"]:
+                processed_fields[f"{location}bar"] = compute_barotropic_velocity(
+                    processed_fields[location],
+                    self.grid.ds[f"interface_depth_{location}"],
                 )
 
         if type == "bgc":
@@ -191,7 +212,7 @@ class InitialConditions:
                 processed_fields[var_name]
             )
 
-        return processed_fields, variable_info
+        return processed_fields
 
     def _input_checks(self):
 
@@ -345,7 +366,84 @@ class InitialConditions:
                 else:
                     variable_info[var_name] = {**default_info, "validate": False}
 
-        return variable_info
+        object.__setattr__(self, f"variable_info_{type}", variable_info)
+
+    def _get_vertical_coordinates(self, type, additional_locations=["u", "v"]):
+        """Retrieve layer and interface depth coordinates.
+
+        This method computes and updates the layer and interface depth coordinates. It handles depth calculations for rho points and
+        additional specified locations (u and v).
+
+        Parameters
+        ----------
+        type : str
+            The type of depth coordinate to retrieve. Valid options are:
+            - "layer": Retrieves layer depth coordinates.
+            - "interface": Retrieves interface depth coordinates.
+
+        additional_locations : list of str, optional
+            Specifies additional locations to compute depth coordinates for. Default is ["u", "v"].
+            Valid options include:
+            - "u": Computes depth coordinates for u points.
+            - "v": Computes depth coordinates for v points.
+
+        Updates
+        -------
+        self.grid.ds : xarray.Dataset
+            The dataset is updated with the following vertical depth coordinates:
+            - f"{type}_depth_rho": Depth coordinates at rho points.
+            - f"{type}_depth_u": Depth coordinates at u points (if applicable).
+            - f"{type}_depth_v": Depth coordinates at v points (if applicable).
+        """
+
+        layer_vars = []
+        for location in ["rho"] + additional_locations:
+            layer_vars.append(f"{type}_depth_{location}")
+
+        if all(layer_var in self.grid.ds for layer_var in layer_vars):
+            # Vertical coordinate data already exists
+            pass
+
+        elif f"{type}_depth_rho" in self.grid.ds:
+            depth = self.grid.ds[f"{type}_depth_rho"]
+
+            if "u" in additional_locations or "v" in additional_locations:
+                # interpolation
+                if "u" in additional_locations:
+                    depth_u = interpolate_from_rho_to_u(depth)
+                    depth_u.attrs["long_name"] = f"{type} depth at u-points"
+                    depth_u.attrs["units"] = "m"
+                    self.grid.ds[f"{type}_depth_u"] = depth_u
+                if "v" in additional_locations:
+                    depth_v = interpolate_from_rho_to_v(depth)
+                    depth_v.attrs["long_name"] = f"{type} depth at v-points"
+                    depth_v.attrs["units"] = "m"
+                    self.grid.ds[f"{type}_depth_v"] = depth_v
+        else:
+            h = self.grid.ds["h"]
+            if type == "layer":
+                depth = compute_depth(
+                    0, h, self.grid.hc, self.grid.ds.Cs_r, self.grid.ds.sigma_r
+                )
+            else:
+                depth = compute_depth(
+                    0, h, self.grid.hc, self.grid.ds.Cs_w, self.grid.ds.sigma_w
+                )
+
+            depth.attrs["long_name"] = f"{type} depth at rho-points"
+            depth.attrs["units"] = "m"
+            self.grid.ds[f"{type}_depth_rho"] = depth
+
+            if "u" in additional_locations or "v" in additional_locations:
+                # interpolation
+                depth_u = interpolate_from_rho_to_u(depth)
+                depth_u.attrs["long_name"] = f"{type} depth at u-points"
+                depth_u.attrs["units"] = "m"
+                depth_v = interpolate_from_rho_to_v(depth)
+                depth_v.attrs["long_name"] = f"{type} depth at v-points"
+                depth_v.attrs["units"] = "m"
+                self.grid.ds[f"{type}_depth_u"] = depth_u
+                self.grid.ds[f"{type}_depth_v"] = depth_v
 
     def _write_into_dataset(self, processed_fields, d_meta):
 
@@ -410,7 +508,7 @@ class InitialConditions:
 
         return ds
 
-    def _validate(self, ds, variable_info):
+    def _validate(self, ds):
         """Validates the dataset by checking for NaN values in SSH at wet points, which
         would indicate missing raw data coverage over the target domain.
 
@@ -418,8 +516,6 @@ class InitialConditions:
         ----------
         ds : xarray.Dataset
             The dataset to validate.
-        variable_info : dict
-            A dictionary containing metadata about the variables, including whether to validate them.
 
         Raises
         ------
@@ -431,6 +527,10 @@ class InitialConditions:
         -----
         This check is only applied to the 2D variable SSH to improve performance.
         """
+        if self.bgc_source is not None:
+            variable_info = {**self.variable_info_physics, **self.variable_info_bgc}
+        else:
+            variable_info = self.variable_info_physics
 
         for var_name in variable_info:
             # Only validate variables based on "validate" flag if use_dask is False
@@ -570,9 +670,17 @@ class InitialConditions:
                 self.ds[var_name].load()
 
         field = self.ds[var_name].squeeze()
+        if s is not None:
+            layer_contours = False
 
         if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-            interface_depth = self.grid.ds.interface_depth_rho
+            if layer_contours:
+                if "interface_depth_rho" in self.grid.ds:
+                    interface_depth = self.grid.ds.interface_depth_rho
+                else:
+                    self.get_vertical_coordinates(
+                        type="interface", additional_locations=[]
+                    )
             layer_depth = self.grid.ds.layer_depth_rho
             mask = self.grid.ds.mask_rho
             field = field.assign_coords(
@@ -580,7 +688,13 @@ class InitialConditions:
             )
 
         elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-            interface_depth = self.grid.ds.interface_depth_u
+            if layer_contours:
+                if "interface_depth_u" in self.grid.ds:
+                    interface_depth = self.grid.ds.interface_depth_u
+                else:
+                    self.get_vertical_coordinates(
+                        type="interface", additional_locations=["u", "v"]
+                    )
             layer_depth = self.grid.ds.layer_depth_u
             mask = self.grid.ds.mask_u
             field = field.assign_coords(
@@ -588,7 +702,13 @@ class InitialConditions:
             )
 
         elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-            interface_depth = self.grid.ds.interface_depth_v
+            if layer_contours:
+                if "interface_depth_v" in self.grid.ds:
+                    interface_depth = self.grid.ds.interface_depth_v
+                else:
+                    self.get_vertical_coordinates(
+                        type="interface", additional_locations=["u", "v"]
+                    )
             layer_depth = self.grid.ds.layer_depth_v
             mask = self.grid.ds.mask_v
             field = field.assign_coords(
@@ -612,14 +732,16 @@ class InitialConditions:
                 title = title + f", eta_rho = {field.eta_rho[eta].item()}"
                 field = field.isel(eta_rho=eta)
                 layer_depth = layer_depth.isel(eta_rho=eta)
-                interface_depth = interface_depth.isel(eta_rho=eta)
+                if layer_contours:
+                    interface_depth = interface_depth.isel(eta_rho=eta)
                 if "s_rho" in field.dims:
                     field = field.assign_coords({"layer_depth": layer_depth})
             elif "eta_v" in field.dims:
                 title = title + f", eta_v = {field.eta_v[eta].item()}"
                 field = field.isel(eta_v=eta)
                 layer_depth = layer_depth.isel(eta_v=eta)
-                interface_depth = interface_depth.isel(eta_v=eta)
+                if layer_contours:
+                    interface_depth = interface_depth.isel(eta_v=eta)
                 if "s_rho" in field.dims:
                     field = field.assign_coords({"layer_depth": layer_depth})
             else:
@@ -631,14 +753,16 @@ class InitialConditions:
                 title = title + f", xi_rho = {field.xi_rho[xi].item()}"
                 field = field.isel(xi_rho=xi)
                 layer_depth = layer_depth.isel(xi_rho=xi)
-                interface_depth = interface_depth.isel(xi_rho=xi)
+                if layer_contours:
+                    interface_depth = interface_depth.isel(xi_rho=xi)
                 if "s_rho" in field.dims:
                     field = field.assign_coords({"layer_depth": layer_depth})
             elif "xi_u" in field.dims:
                 title = title + f", xi_u = {field.xi_u[xi].item()}"
                 field = field.isel(xi_u=xi)
                 layer_depth = layer_depth.isel(xi_u=xi)
-                interface_depth = interface_depth.isel(xi_u=xi)
+                if layer_contours:
+                    interface_depth = interface_depth.isel(xi_u=xi)
                 if "s_rho" in field.dims:
                     field = field.assign_coords({"layer_depth": layer_depth})
             else:
@@ -764,7 +888,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.
 
@@ -774,6 +901,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
         -------
@@ -784,4 +915,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

@@ -0,0 +1,115 @@
+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,
+    interpolate_from_rho_to_v,
+    handle_boundaries,
+)
+
+
+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
+
+    # 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
+    mask = _fill_enclosed_basins(mask.values)
+    # adjust mask boundaries by copying values from adjacent cells
+    mask = handle_boundaries(mask)
+
+    ds["mask_rho"] = xr.DataArray(mask.astype(np.int32), dims=("eta_rho", "xi_rho"))
+    ds["mask_rho"].attrs = {
+        "long_name": "Mask at rho-points",
+        "units": "land/water (0/1)",
+    }
+    ds = _add_velocity_masks(ds)
+
+    return ds
+
+
+def _fill_enclosed_basins(mask) -> np.ndarray:
+    """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)
+    # Find the largest region
+    lint = 0
+    lreg = 0
+    for ireg in range(nreg):
+        int_ = np.sum(reg == ireg)
+        if int_ > lint and mask[reg == ireg].sum() > 0:
+            lreg = ireg
+            lint = int_
+
+    # Remove regions other than the largest one
+    for ireg in range(nreg):
+        if ireg != lreg:
+            mask[reg == ireg] = 0
+
+    return mask
+
+
+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"
+    ).astype(np.int32)
+    ds["mask_v"] = interpolate_from_rho_to_v(
+        ds["mask_rho"], method="multiplicative"
+    ).astype(np.int32)
+
+    ds["mask_u"].attrs = {"long_name": "Mask at u-points", "units": "land/water (0/1)"}
+    ds["mask_v"].attrs = {"long_name": "Mask at v-points", "units": "land/water (0/1)"}
+
+    return ds