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

roms_tools/setup/initial_conditions.py

@@ -1,12 +1,12 @@
 import xarray as xr
 import numpy as np
-import yaml
 import importlib.metadata
-from dataclasses import dataclass, field, asdict
+from dataclasses import dataclass, field
 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,
@@ -16,6 +16,10 @@ 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,
 )
 from roms_tools.setup.regrid import LateralRegrid, VerticalRegrid
 from roms_tools.setup.plot import _plot, _section_plot, _profile_plot, _line_plot
@@ -91,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:
@@ -132,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
@@ -152,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]
@@ -172,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":
@@ -190,7 +206,7 @@ class InitialConditions:
                 processed_fields[var_name]
             )
 
-        return processed_fields, variable_info
+        return processed_fields
 
     def _input_checks(self):
 
@@ -344,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):
 
@@ -387,24 +480,29 @@ class InitialConditions:
         ds["Cs_w"] = self.grid.ds["Cs_w"]
 
         # Preserve absolute time coordinate for readability
-        ds = ds.assign_coords({"abs_time": ds["time"]})
+        abs_time = ds["time"]
+        attrs = [key for key in abs_time.attrs]
+        for attr in attrs:
+            del abs_time.attrs[attr]
+        abs_time.attrs["long_name"] = "absolute time"
+        ds = ds.assign_coords({"abs_time": abs_time})
 
         # Translate the time coordinate to days since the model reference date
         model_reference_date = np.datetime64(self.model_reference_date)
 
-        # Convert the time coordinate to the format expected by ROMS (days since model reference date)
+        # Convert the time coordinate to the format expected by ROMS (seconds since model reference date)
         ocean_time = (ds["time"] - model_reference_date).astype("float64") * 1e-9
         ds = ds.assign_coords(ocean_time=("time", ocean_time.data.astype("float64")))
         ds["ocean_time"].attrs[
             "long_name"
-        ] = f"seconds since {str(self.model_reference_date)}"
+        ] = f"relative time: seconds since {str(self.model_reference_date)}"
         ds["ocean_time"].attrs["units"] = "seconds"
         ds = ds.swap_dims({"time": "ocean_time"})
         ds = ds.drop_vars("time")
 
         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.
 
@@ -412,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
         ------
@@ -425,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
@@ -564,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(
@@ -574,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(
@@ -582,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(
@@ -606,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:
@@ -625,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:
@@ -753,48 +877,8 @@ class InitialConditions:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-        filepath = Path(filepath)
-
-        # Serialize Grid data
-        grid_data = asdict(self.grid)
-        grid_data.pop("ds", None)  # Exclude non-serializable fields
-        grid_data.pop("straddle", None)
-
-        # Include the version of roms-tools
-        try:
-            roms_tools_version = importlib.metadata.version("roms-tools")
-        except importlib.metadata.PackageNotFoundError:
-            roms_tools_version = "unknown"
-
-        # Create header
-        header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
-
-        grid_yaml_data = {"Grid": grid_data}
 
-        initial_conditions_data = {
-            "InitialConditions": {
-                "ini_time": self.ini_time.isoformat(),
-                "source": self.source,
-            }
-        }
-        # Include bgc_source if it's not None
-        if self.bgc_source is not None:
-            initial_conditions_data["InitialConditions"]["bgc_source"] = self.bgc_source
-
-        initial_conditions_data["InitialConditions"][
-            "model_reference_date"
-        ] = self.model_reference_date.isoformat()
-
-        yaml_data = {
-            **grid_yaml_data,
-            **initial_conditions_data,
-        }
-
-        with filepath.open("w") as file:
-            # Write header
-            file.write(header)
-            # Write YAML data
-            yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
+        _to_yaml(self, filepath)
 
     @classmethod
     def from_yaml(
@@ -815,35 +899,7 @@ class InitialConditions:
             An instance of the InitialConditions class.
         """
         filepath = Path(filepath)
-        # Read the entire file content
-        with filepath.open("r") as file:
-            file_content = file.read()
-
-        # Split the content into YAML documents
-        documents = list(yaml.safe_load_all(file_content))
-
-        initial_conditions_data = None
-
-        # Process the YAML documents
-        for doc in documents:
-            if doc is None:
-                continue
-            if "InitialConditions" in doc:
-                initial_conditions_data = doc["InitialConditions"]
-                break
-
-        if initial_conditions_data is None:
-            raise ValueError(
-                "No InitialConditions configuration found in the YAML file."
-            )
-
-        # Convert from string to datetime
-        for date_string in ["model_reference_date", "ini_time"]:
-            initial_conditions_data[date_string] = datetime.fromisoformat(
-                initial_conditions_data[date_string]
-            )
 
         grid = Grid.from_yaml(filepath)
-
-        # Create and return an instance of InitialConditions
-        return cls(grid=grid, **initial_conditions_data, use_dask=use_dask)
+        initial_conditions_params = _from_yaml(cls, filepath)
+        return cls(grid=grid, **initial_conditions_params, use_dask=use_dask)

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

@@ -58,16 +58,71 @@ def _plot(
     if straddle:
         lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
 
-    # Define projections
-    proj = ccrs.PlateCarree()
-
-    trans = ccrs.NearsidePerspective(
-        central_longitude=lon_deg.mean().values, central_latitude=lat_deg.mean().values
-    )
+    trans = _get_projection(lon_deg, lat_deg)
 
     lon_deg = lon_deg.values
     lat_deg = lat_deg.values
 
+    fig, ax = plt.subplots(1, 1, figsize=(13, 7), subplot_kw={"projection": trans})
+
+    _add_plot_to_ax(
+        ax, lon_deg, lat_deg, trans, field, depth_contours, c, title, kwargs=kwargs
+    )
+
+
+def _add_plot_to_ax(
+    ax,
+    lon_deg,
+    lat_deg,
+    trans,
+    field=None,
+    depth_contours=False,
+    c="red",
+    title="",
+    add_colorbar=True,
+    kwargs=None,
+):
+    """Plots a grid or field on a map with optional depth contours.
+
+    Parameters
+    ----------
+    ax : matplotlib.axes._axes.Axes
+        The axes on which to plot the data (Cartopy axis with projection).
+
+    lon_deg : np.ndarray
+        Longitude values in degrees.
+
+    lat_deg : np.ndarray
+        Latitude values in degrees.
+
+    trans : cartopy.crs.Projection
+        The projection for transforming coordinates.
+
+    field : xarray.DataArray, optional
+        Field data to plot (e.g., temperature, salinity). If None, only the grid is plotted.
+
+    depth_contours : bool, optional
+        If True, adds depth contours to the plot.
+
+    c : str, optional
+        Color of the grid boundary (default is 'red').
+
+    title : str, optional
+        Title of the plot.
+
+    add_colorbar : bool, optional
+        If True, add colobar.
+
+    kwargs : dict, optional
+        Additional keyword arguments passed to `pcolormesh` (e.g., colormap, limits).
+
+    Notes
+    -----
+    - If `field` is provided, a colorbar is added.
+    - If `depth_contours` is True, the field’s `layer_depth` is used to add contours.
+    """
+    proj = ccrs.PlateCarree()
+
     # find corners
     corners = [
         (lon_deg[0, 0], lat_deg[0, 0]),
@@ -80,14 +135,13 @@ def _plot(
     transformed_corners = [trans.transform_point(lo, la, proj) for lo, la in corners]
     transformed_lons, transformed_lats = zip(*transformed_corners)
 
-    fig, ax = plt.subplots(1, 1, figsize=(13, 7), subplot_kw={"projection": trans})
-
-    ax.plot(
-        list(transformed_lons) + [transformed_lons[0]],
-        list(transformed_lats) + [transformed_lats[0]],
-        "o-",
-        c=c,
-    )
+    if c is not None:
+        ax.plot(
+            list(transformed_lons) + [transformed_lons[0]],
+            list(transformed_lats) + [transformed_lats[0]],
+            "o-",
+            c=c,
+        )
 
     ax.coastlines(
         resolution="50m", linewidth=0.5, color="black"
@@ -104,13 +158,21 @@ def _plot(
             label = f"{field.Long_name} [{field.units}]"
         else:
             label = ""
-        plt.colorbar(p, label=label)
+        if add_colorbar:
+            plt.colorbar(p, label=label)
 
     if depth_contours:
         cs = ax.contour(lon_deg, lat_deg, field.layer_depth, transform=proj, colors="k")
         ax.clabel(cs, inline=True, fontsize=10)
 
 
+def _get_projection(lon, lat):
+
+    return ccrs.NearsidePerspective(
+        central_longitude=lon.mean().values, central_latitude=lat.mean().values
+    )
+
+
 def _section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
 
     if ax is None:
@@ -176,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