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

roms_tools/__init__.py

@@ -1,5 +1,5 @@
 from importlib.metadata import version as _version
-
+import logging  # noqa: F811
 
 try:
     __version__ = _version("roms_tools")
@@ -13,3 +13,6 @@ from roms_tools.setup.tides import TidalForcing  # noqa: F401
 from roms_tools.setup.surface_forcing import SurfaceForcing  # noqa: F401
 from roms_tools.setup.initial_conditions import InitialConditions  # noqa: F401
 from roms_tools.setup.boundary_forcing import BoundaryForcing  # noqa: F401
+
+# Configure logging when the package is imported
+logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")

roms_tools/_version.py

@@ -1,2 +1,2 @@
 # Do not change! Do not track in version control!
-__version__ = "1.6.1"
+__version__ = "1.6.2"

roms_tools/setup/boundary_forcing.py

@@ -1,9 +1,10 @@
 import xarray as xr
 import numpy as np
 import pandas as pd
+from scipy.ndimage import label
+import logging
 import yaml
 import importlib.metadata
-import warnings
 from typing import Dict, Union, List
 from dataclasses import dataclass, field, asdict
 from roms_tools.setup.grid import Grid
@@ -197,6 +198,13 @@ class BoundaryForcing:
                             ].isel(**bdry_coords[location][direction])
 
                 if not self.apply_2d_horizontal_fill:
+                    self._validate_1d_fill(
+                        processed_fields,
+                        variable_info,
+                        bdry_coords,
+                        direction,
+                        bdry_data.dim_names["depth"],
+                    )
                     processed_fields = apply_1d_horizontal_fill(processed_fields)
 
                 # vertical regridding
@@ -329,38 +337,46 @@ class BoundaryForcing:
                     "is_vector": False,
                     "vector_pair": None,
                     "is_3d": False,
+                    "validate": True,
                 },
-                "temp": default_info,
-                "salt": default_info,
+                "temp": {**default_info, "validate": True},
+                "salt": {**default_info, "validate": False},
                 "u": {
                     "location": "u",
                     "is_vector": True,
                     "vector_pair": "v",
                     "is_3d": True,
+                    "validate": True,
                 },
                 "v": {
                     "location": "v",
                     "is_vector": True,
                     "vector_pair": "u",
                     "is_3d": True,
+                    "validate": True,
                 },
                 "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,
                 },
             }
         elif self.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
 
@@ -506,9 +522,78 @@ class BoundaryForcing:
 
         return ds
 
+    def _validate_1d_fill(
+        self, processed_fields, variable_info, bdry_coords, direction, depth_dim
+    ):
+        """Check if any boundary is divided by land and issue a warning if so,
+        suggesting the use of 2D horizontal fill for safer regridding.
+
+        Parameters
+        ----------
+        processed_fields : dict
+            A dictionary where keys are variable names and values are `xarray.DataArray`
+            objects representing the processed data for each variable.
+
+        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.
+
+        bdry_coords : dict
+            A dictionary containing boundary coordinates for different directions (north, south,
+            east, west), used to slice the boundary-specific data for each variable.
+
+        direction : str
+            The boundary direction being processed (e.g., "north", "south", "east", or "west").
+
+        depth_dim : str
+            The dimension representing depth (e.g., 'z', 'depth', etc.), used when slicing 3D
+            data for a specific depth level.
+
+        Returns
+        -------
+        None
+            If a boundary is divided by land, a warning is issued. No return value is provided.
+        """
+
+        for var_name in processed_fields.keys():
+            # Only validate variables based on "validate" flag if use_dask is False
+            if not self.use_dask or variable_info[var_name]["validate"]:
+                location = variable_info[var_name]["location"]
+
+                # Select the appropriate mask based on variable location
+                if location == "rho":
+                    mask = self.grid.ds.mask_rho
+                elif location == "u":
+                    mask = self.grid.ds.mask_u
+                elif location == "v":
+                    mask = self.grid.ds.mask_v
+
+                mask = mask.isel(**bdry_coords[location][direction])
+
+                if variable_info[var_name]["is_3d"]:
+                    da = processed_fields[var_name].isel({depth_dim: 0, "time": 0})
+                else:
+                    da = processed_fields[var_name].isel({"time": 0})
+
+                wet_nans = xr.where(da.where(mask).isnull(), 1, 0)
+                # Apply label to find connected components of wet NaNs
+                labeled_array, num_features = label(wet_nans)
+                left_margin = labeled_array[0]
+                right_margin = labeled_array[-1]
+                if left_margin != 0:
+                    num_features = num_features - 1
+                if right_margin != 0:
+                    num_features = num_features - 1
+                if num_features > 0:
+                    logging.warning(
+                        f"For {var_name}, the {direction}ern boundary is divided by land. It would be safer (but slower) to use `apply_2d_horizontal_fill = True`."
+                    )
+
     def _validate(self, ds, variable_info, bdry_coords):
-        """Validate the dataset for NaN values at the first time step based on the fill
-        method used.
+        """Validate the dataset for NaN values at the first time step (bry_time=0) for
+        specified variables. If NaN values are found at wet points, this function raises
+        an error.
 
         Parameters
         ----------
@@ -530,14 +615,11 @@ class BoundaryForcing:
         Notes
         -----
         Validation is performed on the initial boundary time step (`bry_time=0`) for each
-        variable in the dataset. If the `apply_2d_horizontal_fill` attribute is set to False,
-        a warning is issued instead of a strict NaN check, as the data may not be reliably validated.
-        Conversely, if `apply_2d_horizontal_fill` is True, a strict NaN check is performed, raising
-        a ValueError if any NaN values are detected.
+        variable in the dataset.
         """
-        if self.apply_2d_horizontal_fill:
-            # Strict NaN check with ValueError makes sense to be applied
-            for var_name in variable_info:
+        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"]:
                 location = variable_info[var_name]["location"]
 
                 # Select the appropriate mask based on variable location
@@ -547,38 +629,29 @@ class BoundaryForcing:
                     mask = self.grid.ds.mask_u
                 elif location == "v":
                     mask = self.grid.ds.mask_v
-                else:
-                    continue  # Skip if location is not recognized
 
                 for direction in ["south", "east", "north", "west"]:
                     if self.boundaries[direction]:
                         bdry_var_name = f"{var_name}_{direction}"
 
                         # Check for NaN values at the first time step using the nan_check function
+                        if self.apply_2d_horizontal_fill:
+                            error_message = None
+                        else:
+                            error_message = (
+                                f"{bdry_var_name} consists entirely of NaNs after regridding. "
+                                f"This may be due to the {direction}ern boundary being on land in the "
+                                f"{self.source['name']} data, which could have a coarser resolution than the ROMS domain. "
+                                f"Try setting `apply_2d_horizontal_fill = True` to resolve this issue."
+                            )
+
                         nan_check(
                             ds[bdry_var_name].isel(bry_time=0),
                             mask.isel(**bdry_coords[location][direction]),
+                            error_message=error_message,
                         )
-        else:
-            # Can't apply strict NaN check because land values haven't been filled before regridding step; instead warn user
-            for direction in ["south", "east", "north", "west"]:
-                if self.boundaries[direction]:
-                    for var_name in variable_info:
-                        bdry_var_name = f"{var_name}_{direction}"
-                        if ds[bdry_var_name].isel(bry_time=0).isnull().any().values:
-                            warnings.warn(
-                                f"NaN values detected in regridded variables along the {direction}ern boundary. This may indicate that the entire boundary is on land in the source data, or that the source data does not cover this boundary.",
-                                UserWarning,
-                            )
-                            # Break after the first warning for this direction to avoid duplicates
-                            break
 
-    def plot(
-        self,
-        var_name,
-        time=0,
-        layer_contours=False,
-    ) -> None:
+    def plot(self, var_name, time=0, layer_contours=False, ax=None) -> None:
         """Plot the boundary forcing field for a given time-slice.
 
         Parameters
@@ -634,6 +707,8 @@ class BoundaryForcing:
             If True, contour lines representing the boundaries between vertical layers will
             be added to the plot. 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.
 
         Returns
         -------
@@ -700,10 +775,14 @@ class BoundaryForcing:
                 interface_depth = None
 
             _section_plot(
-                field, interface_depth=interface_depth, title=title, kwargs=kwargs
+                field,
+                interface_depth=interface_depth,
+                title=title,
+                kwargs=kwargs,
+                ax=ax,
             )
         else:
-            _line_plot(field, title=title)
+            _line_plot(field, title=title, ax=ax)
 
     def save(
         self,
@@ -805,6 +884,7 @@ class BoundaryForcing:
                 "boundaries": self.boundaries,
                 "source": self.source,
                 "type": self.type,
+                "apply_2d_horizontal_fill": self.apply_2d_horizontal_fill,
                 "model_reference_date": self.model_reference_date.isoformat(),
             }
         }
@@ -952,6 +1032,7 @@ def apply_1d_horizontal_fill(processed_fields: dict) -> dict:
             raise ValueError(
                 f"No valid horizontal dimension found for variable '{var_name}'."
             )
+
         # Forward and backward fill in the horizontal direction
         filled = one_dim_fill(
             processed_fields[var_name], selected_horizontal_dim, direction="forward"

roms_tools/setup/datasets.py

@@ -6,7 +6,7 @@ from datetime import datetime, timedelta
 import numpy as np
 from typing import Dict, Optional, Union, List
 from pathlib import Path
-import warnings
+import logging
 from roms_tools.setup.utils import (
     assign_dates_to_climatology,
     interpolate_from_climatology,
@@ -410,7 +410,7 @@ class Dataset:
                             ds[time_dim].where(before_start, drop=True).max()
                         )
                     else:
-                        warnings.warn("No records found at or before the start_time.")
+                        logging.warning("No records found at or before the start_time.")
                         closest_before_start = ds[time_dim].min()
 
                     # Identify records after or at end_time
@@ -420,7 +420,7 @@ class Dataset:
                             ds[time_dim].where(after_end, drop=True).min()
                         )
                     else:
-                        warnings.warn("No records found at or after the end_time.")
+                        logging.warning("No records found at or after the end_time.")
                         closest_after_end = ds[time_dim].max()
 
                     # Select records within the time range and add the closest before/after
@@ -451,11 +451,11 @@ class Dataset:
                     if ds.sizes[time_dim] > 1:
                         # Pick the time closest to self.start_time
                         ds = ds.isel({time_dim: 0})
-                    print(
+                    logging.info(
                         f"Selected time entry closest to the specified start_time ({self.start_time}) within the range [{self.start_time}, {self.start_time + timedelta(hours=24)}]: {ds[time_dim].values}"
                     )
         else:
-            warnings.warn(
+            logging.warning(
                 "Dataset does not contain any time information. Please check if the time dimension "
                 "is correctly named or if the dataset includes time data."
             )

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

@@ -91,14 +91,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 +105,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 +185,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 +265,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 +288,6 @@ class InitialConditions:
             "is_3d": True,
         }
 
-        # Define a dictionary for variable names and their associated information
         if type == "physics":
             variable_info = {
                 "zeta": {
@@ -280,38 +295,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
 
@@ -373,7 +404,7 @@ class InitialConditions:
 
         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 +412,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 +426,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 +461,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 +529,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 +678,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