roms-tools 1.7.0__py3-none-any.whl → 2.0.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,
 )
@@ -92,23 +95,17 @@ 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)
+        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 +130,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 +152,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 +184,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 +206,7 @@ class InitialConditions:
                 processed_fields[var_name]
             )
 
-        return processed_fields, variable_info
+        return processed_fields
 
     def _input_checks(self):
 
@@ -345,7 +360,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 +502,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 +510,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 +521,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 +664,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 +682,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 +696,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 +726,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 +747,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:

roms_tools/setup/mask.py

@@ -0,0 +1,69 @@
+import xarray as xr
+import numpy as np
+import regionmask
+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):
+
+    land = regionmask.defined_regions.natural_earth_v5_0_0.land_10
+
+    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 in enclosed basins with land."""
+
+    # 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):
+
+    # 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

roms_tools/setup/plot.py

@@ -238,18 +238,14 @@ def _profile_plot(field, title="", ax=None):
     """
 
     depths_to_check = [
-        "layer_depth_rho",
-        "layer_depth_u",
-        "layer_depth_v",
-        "interface_depth_rho",
-        "interface_depth_u",
-        "interface_depth_v",
+        "layer_depth",
+        "interface_depth",
     ]
     try:
         depth_label = next(
             depth_label
-            for depth_label in depths_to_check
-            if depth_label in field.coords
+            for depth_label in field.coords
+            if any(depth_label.startswith(prefix) for prefix in depths_to_check)
         )
     except StopIteration:
         raise ValueError(

roms_tools/setup/regrid.py

@@ -28,20 +28,22 @@ class LateralRegrid:
             source_dim_names["longitude"]: target_coords["lon"],
         }
 
-    def apply(self, da):
+    def apply(self, da, method="linear"):
         """Fills missing values and regrids the variable.
 
         Parameters
         ----------
         da : xarray.DataArray
             Input data to fill and regrid.
+        method : str
+            Interpolation method to use.
 
         Returns
         -------
         xarray.DataArray
             Regridded data with filled values.
         """
-        regridded = da.interp(self.coords, method="linear").drop_vars(
+        regridded = da.interp(self.coords, method=method).drop_vars(
             list(self.coords.keys())
         )
         return regridded

roms_tools/setup/surface_forcing.py

@@ -12,13 +12,13 @@ from roms_tools.setup.datasets import (
     CESMBGCSurfaceForcingDataset,
 )
 from roms_tools.setup.utils import (
+    get_target_coords,
     nan_check,
     substitute_nans_by_fillvalue,
     interpolate_from_climatology,
     get_variable_metadata,
     group_dataset,
     save_datasets,
-    get_target_coords,
     rotate_velocities,
     convert_to_roms_time,
     _to_yaml,
@@ -105,8 +105,8 @@ class SurfaceForcing:
 
         data.apply_lateral_fill()
 
-        variable_info = self._set_variable_info(data)
-        var_names = variable_info.keys()
+        self._set_variable_info(data)
+        var_names = self.variable_info.keys()
 
         processed_fields = {}
         # lateral regridding
@@ -118,7 +118,7 @@ class SurfaceForcing:
                 )
 
         # rotation of velocities and interpolation to u/v points
-        if "uwnd" in variable_info and "vwnd" in variable_info:
+        if "uwnd" in self.variable_info and "vwnd" in self.variable_info:
             processed_fields["uwnd"], processed_fields["vwnd"] = rotate_velocities(
                 processed_fields["uwnd"],
                 processed_fields["vwnd"],
@@ -134,7 +134,7 @@ class SurfaceForcing:
 
         ds = self._write_into_dataset(processed_fields, data, d_meta)
 
-        self._validate(ds, target_coords["mask"], variable_info)
+        self._validate(ds)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
         for var_name in ds.data_vars:
@@ -213,9 +213,8 @@ class SurfaceForcing:
 
         Returns
         -------
-        dict
-            A dictionary where the keys are variable names and the values are dictionaries of metadata
-            about each variable, including 'location', 'is_vector', 'vector_pair', and 'is_3d'.
+        None
+            This method updates the instance attribute `variable_info` with the metadata dictionary for the variables.
         """
         default_info = {
             "location": "rho",
@@ -256,7 +255,7 @@ class SurfaceForcing:
                 else:
                     variable_info[var_name] = {**default_info, "validate": False}
 
-        return variable_info
+        object.__setattr__(self, "variable_info", variable_info)
 
     def _apply_correction(self, processed_fields, data):
 
@@ -329,7 +328,7 @@ class SurfaceForcing:
 
         return ds
 
-    def _validate(self, ds, mask, variable_info):
+    def _validate(self, ds):
         """Validates the dataset by checking for NaN values at wet points, which would
         indicate missing raw data coverage over the target domain.
 
@@ -337,12 +336,6 @@ class SurfaceForcing:
         ----------
         ds : xarray.Dataset
             The dataset to validate.
-        mask : xarray.DataArray
-            Land mask (1=ocean, 0=land) to determine wet points in the domain.
-        variable_info : dict
-            A dictionary containing metadata about each variable (e.g., location,
-            whether it's a 3D variable, etc.). Used to retrieve information for
-            validating each variable.
 
         Raises
         ------
@@ -357,8 +350,8 @@ 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 variable_info[var_name]["validate"]:
-                nan_check(ds[var_name].isel(time=0), mask)
+            if not self.use_dask or self.variable_info[var_name]["validate"]:
+                nan_check(ds[var_name].isel(time=0), self.target_coords["mask"])
 
     def _add_global_metadata(self, ds=None):
 

roms_tools/setup/tides.py

@@ -85,8 +85,8 @@ class TidalForcing:
 
         data.apply_lateral_fill()
 
-        variable_info = self._set_variable_info()
-        var_names = variable_info.keys()
+        self._set_variable_info()
+        var_names = self.variable_info.keys()
 
         processed_fields = {}
         # lateral regridding
@@ -98,7 +98,7 @@ class TidalForcing:
                 )
 
         # rotation of velocities and interpolation to u/v points
-        vector_pairs = get_vector_pairs(variable_info)
+        vector_pairs = get_vector_pairs(self.variable_info)
         for pair in vector_pairs:
             u_component = pair[0]
             v_component = pair[1]
@@ -129,7 +129,7 @@ class TidalForcing:
 
         ds = self._add_global_metadata(ds)
 
-        self._validate(ds, variable_info)
+        self._validate(ds)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
         for var_name in ds.data_vars:
@@ -163,9 +163,8 @@ class TidalForcing:
 
         Returns
         -------
-        dict
-            A dictionary where the keys are variable names and the values are dictionaries of metadata
-            about each variable, including 'location', 'is_vector', 'vector_pair', and 'is_3d'.
+        None
+            This method updates the instance attribute `variable_info` with the metadata dictionary for the variables.
         """
         default_info = {
             "location": "rho",
@@ -210,7 +209,7 @@ class TidalForcing:
             },
         }
 
-        return variable_info
+        object.__setattr__(self, "variable_info", variable_info)
 
     def _write_into_dataset(self, processed_fields, d_meta):
 
@@ -243,7 +242,7 @@ class TidalForcing:
 
         return ds
 
-    def _validate(self, ds, variable_info):
+    def _validate(self, ds):
         """Validates the dataset by checking for NaN values at wet points for specified
         variables, which would indicate missing raw data coverage over the target
         domain.
@@ -252,8 +251,6 @@ class TidalForcing:
         ----------
         ds : xarray.Dataset
             The dataset to validate, containing tidal variables and a mask for wet points.
-        variable_info : dict
-            A dictionary containing metadata about the variables, including whether to validate them.
 
         Raises
         ------
@@ -268,7 +265,7 @@ class TidalForcing:
         """
         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 variable_info[var_name]["validate"]:
+            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)
 
     def plot(self, var_name, ntides=0) -> None: