roms-tools 3.0.0__py3-none-any.whl → 3.1.0__py3-none-any.whl

roms_tools/setup/initial_conditions.py

@@ -3,27 +3,26 @@ import logging
 from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Dict, List, Optional, Union
 
-import matplotlib.pyplot as plt
 import numpy as np
 import xarray as xr
+from matplotlib.axes import Axes
 
 from roms_tools import Grid
-from roms_tools.plot import _line_plot, _plot, _profile_plot, _section_plot
+from roms_tools.plot import plot
 from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
 from roms_tools.setup.datasets import CESMBGCDataset, GLORYSDataset, UnifiedBGCDataset
 from roms_tools.setup.utils import (
-    _from_yaml,
-    _to_dict,
-    _write_to_yaml,
     compute_barotropic_velocity,
     compute_missing_bgc_variables,
+    from_yaml,
     get_target_coords,
     get_variable_metadata,
     nan_check,
     rotate_velocities,
     substitute_nans_by_fillvalue,
+    to_dict,
+    write_to_yaml,
 )
 from roms_tools.utils import (
     get_dask_chunks,
@@ -34,7 +33,6 @@ from roms_tools.utils import (
 )
 from roms_tools.vertical_coordinate import (
     compute_depth,
-    compute_depth_coordinates,
 )
 
 
@@ -104,19 +102,32 @@ class InitialConditions:
     """
 
     grid: Grid
+    """Object representing the grid information."""
     ini_time: datetime
-    source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
-    bgc_source: Optional[Dict[str, Union[str, Path, List[Union[str, Path]]]]] = None
+    """The date and time at which the initial conditions are set."""
+    source: dict[str, str | Path | list[str | Path]]
+    """Dictionary specifying the source of the physical initial condition data."""
+    bgc_source: dict[str, str | Path | list[str | Path]] | None = None
+    """Dictionary specifying the source of the biogeochemical (BGC) initial condition
+    data."""
     model_reference_date: datetime = datetime(2000, 1, 1)
+    """The reference date for the model."""
     adjust_depth_for_sea_surface_height: bool = False
+    """Whether to account for sea surface height variations when computing depth
+    coordinates."""
     use_dask: bool = False
+    """Whether to use dask for processing."""
     horizontal_chunk_size: int = 50
+    """The chunk size used for horizontal partitioning for the vertical regridding when
+    `use_dask = True`."""
     bypass_validation: bool = False
+    """Whether to skip validation checks in the processed data."""
 
     ds: xr.Dataset = field(init=False, repr=False)
+    """An xarray Dataset containing post-processed variables ready for input into
+    ROMS."""
 
     def __post_init__(self):
-
         self._input_checks()
         # Dataset for depth coordinates
         self.ds_depth_coords = xr.Dataset()
@@ -148,7 +159,6 @@ class InitialConditions:
         self.ds = ds
 
     def _process_data(self, processed_fields, type="physics"):
-
         target_coords = get_target_coords(self.grid)
 
         if type == "physics":
@@ -298,7 +308,6 @@ class InitialConditions:
             )
 
     def _get_data(self):
-
         if self.source["name"] == "GLORYS":
             data = GLORYSDataset(
                 filename=self.source["path"],
@@ -311,9 +320,7 @@ class InitialConditions:
         return data
 
     def _get_bgc_data(self):
-
         if self.bgc_source["name"] == "CESM_REGRIDDED":
-
             data = CESMBGCDataset(
                 filename=self.bgc_source["path"],
                 start_time=self.ini_time,
@@ -485,13 +492,11 @@ class InitialConditions:
             self.ds_depth_coords[key] = depth
 
     def _write_into_dataset(self, processed_fields, d_meta):
-
         # save in new dataset
         ds = xr.Dataset()
 
         for var_name in processed_fields:
             if var_name in d_meta:
-
                 # drop auxiliary variables
                 ds[var_name] = processed_fields[var_name].astype(np.float32)
                 ds[var_name].attrs["long_name"] = d_meta[var_name]["long_name"]
@@ -541,9 +546,9 @@ class InitialConditions:
         # 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"relative time: seconds since {str(self.model_reference_date)}"
+        ds["ocean_time"].attrs["long_name"] = (
+            f"relative time: seconds since {self.model_reference_date!s}"
+        )
         ds["ocean_time"].attrs["units"] = "seconds"
         ds = ds.swap_dims({"time": "ocean_time"})
         ds = ds.drop_vars("time")
@@ -586,7 +591,6 @@ class InitialConditions:
                 nan_check(ds[var_name].squeeze(), mask)
 
     def _add_global_metadata(self, ds):
-
         ds.attrs["title"] = "ROMS initial conditions file created by ROMS-Tools"
         # Include the version of roms-tools
         try:
@@ -611,13 +615,14 @@ class InitialConditions:
 
     def plot(
         self,
-        var_name,
-        s=None,
-        eta=None,
-        xi=None,
-        depth_contours=False,
-        layer_contours=False,
-        ax=None,
+        var_name: str,
+        s: int | None = None,
+        eta: int | None = None,
+        xi: int | None = None,
+        depth_contours: bool = False,
+        layer_contours: bool = False,
+        ax: Axes | None = None,
+        save_path: str | None = None,
     ) -> None:
         """Plot the initial conditions field for a given eta-, xi-, or s_rho- slice.
 
@@ -687,6 +692,9 @@ class InitialConditions:
             contours displayed is limited to a maximum of 10. Default is False.
         ax : matplotlib.axes.Axes, optional
             The axes to plot on. If None, a new figure is created. Note that this argument is ignored for 2D horizontal plots. Default is None.
+        save_path : str, optional
+            Path to save the generated plot. If None, the plot is shown interactively.
+            Default is None.
 
         Returns
         -------
@@ -700,44 +708,8 @@ class InitialConditions:
             If the field specified by `var_name` is 3D and none of `s`, `eta`, or `xi` are specified.
             If the field specified by `var_name` is 2D and both `eta` and `xi` are specified.
         """
-
-        field = self.ds[var_name].squeeze()
-
-        if len(field.dims) == 3:
-            if not any([s is not None, eta is not None, xi is not None]):
-                raise ValueError(
-                    "Invalid input: For 3D fields, you must specify at least one of the dimensions 's', 'eta', or 'xi'."
-                )
-            if all([s is not None, eta is not None, xi is not None]):
-                raise ValueError(
-                    "Ambiguous input: For 3D fields, specify at most two of 's', 'eta', or 'xi'. Specifying all three is not allowed."
-                )
-
-        if len(field.dims) == 2 and all([eta is not None, xi is not None]):
-            raise ValueError(
-                "Conflicting input: For 2D fields, specify only one dimension, either 'eta' or 'xi', not both."
-            )
-
-        field = self.ds[var_name].squeeze()
-
-        # Get correct mask and horizontal coordinates
-        if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-            loc = "rho"
-        elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-            loc = "u"
-        elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-            loc = "v"
-        else:
-            ValueError("provided field does not have two horizontal dimension")
-
-        mask = self.grid.ds[f"mask_{loc}"]
-        lat_deg = self.grid.ds[f"lat_{loc}"]
-        lon_deg = self.grid.ds[f"lon_{loc}"]
-
-        if self.grid.straddle:
-            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
-
-        field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
+        if var_name not in self.ds:
+            raise ValueError(f"Variable '{var_name}' is not found in the dataset.")
 
         # Load the data
         if self.use_dask:
@@ -746,150 +718,35 @@ class InitialConditions:
             with ProgressBar():
                 self.ds[var_name].load()
 
-        # Retrieve depth coordinates
-        if s is not None:
-            layer_contours = False
-        # Note that `layer_depth_{loc}` has already been computed during `__post_init__`.
-        layer_depth = self.ds_depth_coords[f"layer_depth_{loc}"].squeeze().load()
-
-        # Slice the field as desired
-        def _slice_and_assign(
-            field,
-            mask,
-            layer_depth,
-            title,
-            dim_name,
-            dim_values,
-            idx,
-        ):
-            if dim_name in field.dims:
-                title = title + f", {dim_name} = {dim_values[idx].item()}"
-                field = field.isel(**{dim_name: idx})
-                mask = mask.isel(**{dim_name: idx})
-                layer_depth = layer_depth.isel(**{dim_name: idx})
-                if "s_rho" in field.dims:
-                    field = field.assign_coords({"layer_depth": layer_depth})
-            else:
-                raise ValueError(
-                    f"None of the expected dimensions ({dim_name}) found in field."
-                )
-            return field, mask, layer_depth, title
-
-        title = field.long_name
-        if s is not None:
-            title = title + f", s_rho = {field.s_rho[s].item()}"
-            field = field.isel(s_rho=s)
-            layer_depth = layer_depth.isel(s_rho=s)
-            field = field.assign_coords({"layer_depth": layer_depth})
+        if self.adjust_depth_for_sea_surface_height:
+            zeta = self.ds.zeta.squeeze().load()
         else:
-            depth_contours = False
-
-        if eta is not None:
-            field, mask, layer_depth, title = _slice_and_assign(
-                field,
-                mask,
-                layer_depth,
-                title,
-                "eta_rho" if "eta_rho" in field.dims else "eta_v",
-                field.eta_rho if "eta_rho" in field.dims else field.eta_v,
-                eta,
-            )
+            zeta = 0
 
-        if xi is not None:
-            field, mask, layer_depth, title = _slice_and_assign(
-                field,
-                mask,
-                layer_depth,
-                title,
-                "xi_rho" if "xi_rho" in field.dims else "xi_u",
-                field.xi_rho if "xi_rho" in field.dims else field.xi_u,
-                xi,
-            )
+        field = self.ds[var_name].squeeze()
 
-        # Choose colorbar
         if var_name in ["u", "v", "w", "ubar", "vbar", "zeta"]:
-            vmax = max(field.max().values, -field.min().values)
-            vmin = -vmax
-            cmap = plt.colormaps.get_cmap("RdBu_r")
+            cmap_name = "RdBu_r"
+        elif var_name in ["temp", "salt"]:
+            cmap_name = "YlOrRd"
         else:
-            vmax = field.max().values
-            vmin = field.min().values
-            if var_name in ["temp", "salt"]:
-                cmap = plt.colormaps.get_cmap("YlOrRd")
-            else:
-                cmap = plt.colormaps.get_cmap("YlGn")
-        cmap.set_bad(color="gray")
-        kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
-
-        if eta is None and xi is None:
-            _plot(
-                field=field.where(mask),
-                depth_contours=depth_contours,
-                title=title,
-                kwargs=kwargs,
-                c="g",
-            )
-        else:
-            if len(field.dims) == 2:
-                if layer_contours:
-                    if loc == "rho":
-                        # interface_depth_rho has not been computed yet
-                        interface_depth = compute_depth_coordinates(
-                            self.grid.ds,
-                            self.ds.zeta,
-                            depth_type="interface",
-                            location=loc,
-                            eta=eta,
-                            xi=xi,
-                        )
-                    elif loc == "u":
-                        index_kwargs = {}
-                        if eta is not None:
-                            index_kwargs["eta_rho"] = eta
-                        if xi is not None:
-                            index_kwargs["xi_u"] = xi
-
-                        interface_depth = (
-                            self.ds_depth_coords[f"interface_depth_{loc}"]
-                            .isel(**index_kwargs)
-                            .squeeze()
-                        )
-                    elif loc == "v":
-                        index_kwargs = {}
-                        if eta is not None:
-                            index_kwargs["eta_v"] = eta
-                        if xi is not None:
-                            index_kwargs["xi_rho"] = xi
-
-                        interface_depth = (
-                            self.ds_depth_coords[f"interface_depth_{loc}"]
-                            .isel(**index_kwargs)
-                            .squeeze()
-                        )
-
-                    # restrict number of layer_contours to 10 for the sake of plot clearity
-                    nr_layers = len(interface_depth["s_w"])
-                    selected_layers = np.linspace(
-                        0, nr_layers - 1, min(nr_layers, 10), dtype=int
-                    )
-                    interface_depth = interface_depth.isel(s_w=selected_layers)
-                else:
-                    interface_depth = None
-
-                _section_plot(
-                    field,
-                    interface_depth=interface_depth,
-                    title=title,
-                    kwargs=kwargs,
-                    ax=ax,
-                )
-            else:
-                if "s_rho" in field.dims:
-                    _profile_plot(field.where(mask), title=title, ax=ax)
-                else:
-                    _line_plot(field.where(mask), title=title, ax=ax)
+            cmap_name = "YlGn"
+
+        plot(
+            field=field,
+            grid_ds=self.grid.ds,
+            zeta=zeta,
+            s=s,
+            eta=eta,
+            xi=xi,
+            depth_contours=depth_contours,
+            layer_contours=layer_contours,
+            ax=ax,
+            save_path=save_path,
+            cmap_name=cmap_name,
+        )
 
-    def save(self, filepath: Union[str, Path]) -> None:
+    def save(self, filepath: str | Path) -> None:
         """Save the initial conditions information to one netCDF4 file.
 
         Parameters
@@ -902,7 +759,6 @@ class InitialConditions:
         Path
             A `Path` object representing the location of the saved file.
         """
-
         # Ensure filepath is a Path object
         filepath = Path(filepath)
 
@@ -919,7 +775,7 @@ class InitialConditions:
 
         return saved_filenames
 
-    def to_yaml(self, filepath: Union[str, Path]) -> None:
+    def to_yaml(self, filepath: str | Path) -> None:
         """Export the parameters of the class to a YAML file, including the version of
         roms-tools.
 
@@ -928,14 +784,13 @@ class InitialConditions:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-
-        forcing_dict = _to_dict(self)
-        _write_to_yaml(forcing_dict, filepath)
+        forcing_dict = to_dict(self, exclude=["use_dask"])
+        write_to_yaml(forcing_dict, filepath)
 
     @classmethod
     def from_yaml(
         cls,
-        filepath: Union[str, Path],
+        filepath: str | Path,
         use_dask: bool = False,
     ) -> "InitialConditions":
         """Create an instance of the InitialConditions class from a YAML file.
@@ -955,7 +810,7 @@ class InitialConditions:
         filepath = Path(filepath)
 
         grid = Grid.from_yaml(filepath)
-        initial_conditions_params = _from_yaml(cls, filepath)
+        initial_conditions_params = from_yaml(cls, filepath)
         return cls(
             grid=grid,
             **initial_conditions_params,

roms_tools/setup/mask.py

@@ -66,7 +66,6 @@ def _fill_enclosed_basins(mask) -> np.ndarray:
     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