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

roms_tools/setup/download.py

@@ -24,6 +24,16 @@ correction_data = pooch.create(
     },
 )
 
+# Create a Pooch object to manage the global topography data
+river_data = pooch.create(
+    # Use the default cache folder for the operating system
+    path=pooch.os_cache("roms-tools"),
+    base_url="https://github.com/CWorthy-ocean/roms-tools-data/raw/main/",
+    # The registry specifies the files that can be fetched
+    registry={
+        "dai_trenberth_may2019.nc": "sha256:793849e6aa60d1f6bdb480c345515fb2453d903c0a30599241b3d752f53715ab",
+    },
+)
 # Create a Pooch object to manage the test data
 pup_test_data = pooch.create(
     # Use the default cache folder for the operating system
@@ -77,6 +87,26 @@ def fetch_topo(topography_source: str) -> xr.Dataset:
     return ds
 
 
+def download_river_data(filename: str) -> xr.Dataset:
+    """Download river data file.
+
+    Parameters
+    ----------
+    filename : str
+        The name of the test data file to be downloaded. Available options:
+        - "dai_trenberth_may2019.nc"
+    Returns
+    -------
+    str
+        The path to the downloaded test data file.
+    """
+
+    # Fetch the file using Pooch, downloading if necessary
+    fname = river_data.fetch(filename)
+
+    return fname
+
+
 def download_correction_data(filename: str) -> str:
     """Download the correction data file.
 

roms_tools/setup/grid.py

@@ -13,7 +13,7 @@ from roms_tools.setup.plot import _plot, _section_plot, _profile_plot, _line_plo
 from roms_tools.setup.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
 from roms_tools.setup.vertical_coordinate import sigma_stretch, compute_depth
 from roms_tools.setup.utils import extract_single_value, save_datasets
-import warnings
+import logging
 from pathlib import Path
 
 RADIUS_OF_EARTH = 6371315.0  # in m
@@ -350,11 +350,7 @@ class Grid:
             _plot(self.ds, straddle=self.straddle)
 
     def plot_vertical_coordinate(
-        self,
-        varname="layer_depth_rho",
-        s=None,
-        eta=None,
-        xi=None,
+        self, varname="layer_depth_rho", s=None, eta=None, xi=None, ax=None
     ) -> None:
         """Plot the vertical coordinate system for a given eta-, xi-, or s-slice.
 
@@ -376,6 +372,8 @@ class Grid:
             The eta-index to plot. Default is None.
         xi : int, optional
             The xi-index to plot. Default is None.
+        ax : matplotlib.axes.Axes, optional
+            The axes to plot on. If None, a new figure is created. Note that this argument does not work for horizontal plots that display the eta- and xi-dimensions at the same time.
 
         Returns
         -------
@@ -477,12 +475,13 @@ class Grid:
                     interface_depth=interface_depth,
                     title=title,
                     kwargs=kwargs,
+                    ax=ax,
                 )
             else:
                 if "s_rho" in field.dims or "s_w" in field.dims:
-                    _profile_plot(field, title=title)
+                    _profile_plot(field, title=title, ax=ax)
                 else:
-                    _line_plot(field, title=title)
+                    _line_plot(field, title=title, ax=ax)
 
     def save(
         self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
@@ -728,9 +727,8 @@ class Grid:
                 roms_tools_version_current = "unknown"
 
             if roms_tools_version_header != roms_tools_version_current:
-                warnings.warn(
+                logging.warning(
                     f"Current roms-tools version ({roms_tools_version_current}) does not match the version in the YAML header ({roms_tools_version_header}).",
-                    UserWarning,
                 )
 
         if grid_data is None:

roms_tools/setup/initial_conditions.py

@@ -1,8 +1,7 @@
 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
@@ -16,6 +15,8 @@ from roms_tools.setup.utils import (
     rotate_velocities,
     compute_barotropic_velocity,
     transpose_dimensions,
+    _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,14 +92,13 @@ class InitialConditions:
         self._input_checks()
 
         processed_fields = {}
-        processed_fields = self._process_data(processed_fields, type="physics")
+        processed_fields, variable_info = self._process_data(
+            processed_fields, type="physics"
+        )
 
         if self.bgc_source is not None:
-            processed_fields = self._process_data(processed_fields, type="bgc")
-
-        for var_name in processed_fields.keys():
-            processed_fields[var_name] = transpose_dimensions(
-                processed_fields[var_name]
+            processed_fields, bgc_variable_info = self._process_data(
+                processed_fields, type="bgc"
             )
 
         d_meta = get_variable_metadata()
@@ -106,7 +106,9 @@ class InitialConditions:
 
         ds = self._add_global_metadata(ds)
 
-        self._validate(ds)
+        if self.bgc_source is not None:
+            variable_info = {**variable_info, **bgc_variable_info}
+        self._validate(ds, variable_info)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
         for var_name in ds.data_vars:
@@ -184,7 +186,12 @@ class InitialConditions:
                     {"time": processed_fields["temp"]["time"]}
                 )
 
-        return processed_fields
+        for var_name in processed_fields.keys():
+            processed_fields[var_name] = transpose_dimensions(
+                processed_fields[var_name]
+            )
+
+        return processed_fields, variable_info
 
     def _input_checks(self):
 
@@ -259,11 +266,21 @@ class InitialConditions:
         - `vector_pair`: For vector variables, this indicates the associated variable that forms the vector (e.g., 'u' and 'v').
         - `is_3d`: Indicates whether the variable is 3D (True for variables like 'temp' and 'salt') or 2D (False for 'zeta').
 
+        Parameters
+        ----------
+        data : object
+            The data object which contains variable names for the "bgc" type variables.
+
+        type : str, optional, default="physics"
+            The type of variable metadata to return. Can be one of:
+            - "physics": for physical variables such as temperature, salinity, and velocity components.
+            - "bgc": for biogeochemical variables (like ALK).
+
         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'.
+            about each variable, including 'location', 'is_vector', 'vector_pair', 'is_3d', and 'validate'.
         """
         default_info = {
             "location": "rho",
@@ -272,7 +289,6 @@ class InitialConditions:
             "is_3d": True,
         }
 
-        # Define a dictionary for variable names and their associated information
         if type == "physics":
             variable_info = {
                 "zeta": {
@@ -280,38 +296,54 @@ class InitialConditions:
                     "is_vector": False,
                     "vector_pair": None,
                     "is_3d": False,
+                    "validate": True,
                 },
-                "temp": default_info,
-                "salt": default_info,
+                "temp": {**default_info, "validate": False},
+                "salt": {**default_info, "validate": False},
                 "u": {
                     "location": "u",
                     "is_vector": True,
                     "vector_pair": "v",
                     "is_3d": True,
+                    "validate": False,
                 },
                 "v": {
                     "location": "v",
                     "is_vector": True,
                     "vector_pair": "u",
                     "is_3d": True,
+                    "validate": False,
                 },
                 "ubar": {
                     "location": "u",
                     "is_vector": True,
                     "vector_pair": "vbar",
                     "is_3d": False,
+                    "validate": False,
                 },
                 "vbar": {
                     "location": "v",
                     "is_vector": True,
                     "vector_pair": "ubar",
                     "is_3d": False,
+                    "validate": False,
+                },
+                "w": {
+                    "location": "rho",
+                    "is_vector": False,
+                    "vector_pair": None,
+                    "is_3d": True,
+                    "validate": False,
                 },
             }
-        elif type == "bgc":
+
+        if type == "bgc":
             variable_info = {}
             for var_name in data.var_names.keys():
-                variable_info[var_name] = default_info
+                if var_name == "ALK":
+                    variable_info[var_name] = {**default_info, "validate": True}
+                else:
+                    variable_info[var_name] = {**default_info, "validate": False}
 
         return variable_info
 
@@ -356,24 +388,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):
+    def _validate(self, ds, variable_info):
         """Validates the dataset by checking for NaN values in SSH at wet points, which
         would indicate missing raw data coverage over the target domain.
 
@@ -381,6 +418,8 @@ 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
         ------
@@ -393,8 +432,11 @@ class InitialConditions:
         This check is only applied to the 2D variable SSH to improve performance.
         """
 
-        ds["zeta"].load()
-        nan_check(ds["zeta"].squeeze(), self.grid.ds.mask_rho)
+        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"]:
+                ds[var_name].load()
+                nan_check(ds[var_name].squeeze(), self.grid.ds.mask_rho)
 
     def _add_global_metadata(self, ds):
 
@@ -425,6 +467,7 @@ class InitialConditions:
         xi=None,
         depth_contours=False,
         layer_contours=False,
+        ax=None,
     ) -> None:
         """Plot the initial conditions field for a given eta-, xi-, or s_rho- slice.
 
@@ -492,6 +535,8 @@ class InitialConditions:
             be added to the plot. This is particularly useful in vertical sections to
             visualize the layering of the water column. For clarity, the number of layer
             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 does not work for horizontal plots that display the eta- and xi-dimensions at the same time.
 
         Returns
         -------
@@ -639,13 +684,17 @@ class InitialConditions:
 
             if len(field.dims) == 2:
                 _section_plot(
-                    field, interface_depth=interface_depth, title=title, kwargs=kwargs
+                    field,
+                    interface_depth=interface_depth,
+                    title=title,
+                    kwargs=kwargs,
+                    ax=ax,
                 )
             else:
                 if "s_rho" in field.dims:
-                    _profile_plot(field, title=title)
+                    _profile_plot(field, title=title, ax=ax)
                 else:
-                    _line_plot(field, title=title)
+                    _line_plot(field, title=title, ax=ax)
 
     def save(
         self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
@@ -710,48 +759,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(
@@ -772,35 +781,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/plot.py

@@ -12,6 +12,32 @@ def _plot(
     title="",
     kwargs={},
 ):
+    """Plots a grid or field on a map with optional depth contours.
+
+    This function plots a map using Cartopy projections. It supports plotting a grid, a field, and adding depth contours if desired.
+    The projection can be customized, and the grid can be adjusted for domains straddling the 180° meridian.
+
+    Parameters
+    ----------
+    grid_ds : xarray.Dataset
+        The grid dataset containing coordinates (`lon_rho`, `lat_rho`).
+    field : xarray.DataArray, optional
+        The field to plot. If None, only the grid is plotted.
+    depth_contours : bool, optional
+        If True, adds depth contours to the plot.
+    straddle : bool, optional
+        If True, adjusts longitude values to straddle across the 180° meridian.
+    c : str, optional
+        Color for the boundary plot (default is 'red').
+    title : str, optional
+        Title of the plot.
+    kwargs : dict, optional
+        Additional keyword arguments to pass to `pcolormesh` (e.g., colormap or color limits).
+
+    Notes
+    -----
+    The function raises a `NotImplementedError` if the domain contains the North or South Pole.
+    """
 
     if field is None:
         lon_deg = grid_ds["lon_rho"]
@@ -32,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]),
@@ -54,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"
@@ -78,18 +158,25 @@ 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)
 
-    return fig
 
+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={}):
 
-    fig, ax = plt.subplots(1, 1, figsize=(9, 5))
+def _section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
+
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(9, 5))
 
     dims_to_check = ["eta_rho", "eta_u", "eta_v", "xi_rho", "xi_u", "xi_v"]
     try:
@@ -132,7 +219,23 @@ def _section_plot(field, interface_depth=None, title="", kwargs={}):
     ax.set_title(title)
 
 
-def _profile_plot(field, title=""):
+def _profile_plot(field, title="", ax=None):
+    """Plots a profile of the given field against depth.
+
+    Parameters
+    ----------
+    field : xarray.DataArray
+        Data to plot.
+    title : str, optional
+        Title of the plot.
+    ax : matplotlib.axes.Axes, optional
+        Axes to plot on. If None, a new figure is created.
+
+    Raises
+    ------
+    ValueError
+        If no expected depth coordinate is found in the field.
+    """
 
     depths_to_check = [
         "layer_depth_rho",
@@ -153,16 +256,33 @@ def _profile_plot(field, title=""):
             "None of the expected coordinates (layer_depth_rho, layer_depth_u, layer_depth_v, interface_depth_rho, interface_depth_u, interface_depth_v) found in field.coords"
         )
 
-    fig, ax = plt.subplots(1, 1, figsize=(4, 7))
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(4, 7))
     kwargs = {"y": depth_label, "yincrease": False}
     field.plot(**kwargs)
     ax.set_title(title)
     ax.grid()
 
 
-def _line_plot(field, title=""):
-
-    fig, ax = plt.subplots(1, 1, figsize=(7, 4))
-    field.plot()
+def _line_plot(field, title="", ax=None):
+    """Plots a line graph of the given field.
+
+    Parameters
+    ----------
+    field : xarray.DataArray
+        Data to plot.
+    title : str, optional
+        Title of the plot.
+    ax : matplotlib.axes.Axes, optional
+        Axes to plot on. If None, a new figure is created.
+
+    Returns
+    -------
+    None
+        Modifies the plot in-place.
+    """
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(7, 4))
+    field.plot(ax=ax)
     ax.set_title(title)
     ax.grid()