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

ci/environment.yml

@@ -10,5 +10,5 @@ dependencies:
   - pytest-xdist
   - flake8
   - black
-  - pre-commit
+  - pre-commit==3.8.0
   - coverage

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,7 @@ 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
+from roms_tools.setup.river_forcing import RiverForcing  # 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.7.0"

roms_tools/setup/boundary_forcing.py

@@ -1,11 +1,10 @@
 import xarray as xr
 import numpy as np
-import pandas as pd
-import yaml
+from scipy.ndimage import label
+import logging
 import importlib.metadata
-import warnings
 from typing import Dict, Union, List
-from dataclasses import dataclass, field, asdict
+from dataclasses import dataclass, field
 from roms_tools.setup.grid import Grid
 from roms_tools.setup.regrid import LateralRegrid, VerticalRegrid
 from datetime import datetime
@@ -21,6 +20,9 @@ from roms_tools.setup.utils import (
     one_dim_fill,
     nan_check,
     substitute_nans_by_fillvalue,
+    convert_to_roms_time,
+    _to_yaml,
+    _from_yaml,
 )
 from roms_tools.setup.plot import _section_plot, _line_plot
 import matplotlib.pyplot as plt
@@ -58,12 +60,12 @@ class BoundaryForcing:
           - "physics": for physical atmospheric forcing.
           - "bgc": for biogeochemical forcing.
 
-    model_reference_date : datetime, optional
-        Reference date for the model. Default is January 1, 2000.
     apply_2d_horizontal_fill: bool, optional
         Indicates whether to perform a two-dimensional horizontal fill on the source data prior to regridding to boundaries.
         If `False`, a one-dimensional horizontal fill is performed separately on each of the four regridded boundaries.
         Defaults to `False`.
+    model_reference_date : datetime, optional
+        Reference date for the model. Default is January 1, 2000.
     use_dask: bool, optional
         Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
 
@@ -92,8 +94,8 @@ class BoundaryForcing:
     )
     source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
     type: str = "physics"
-    model_reference_date: datetime = datetime(2000, 1, 1)
     apply_2d_horizontal_fill: bool = False
+    model_reference_date: datetime = datetime(2000, 1, 1)
     use_dask: bool = False
 
     ds: xr.Dataset = field(init=False, repr=False)
@@ -197,6 +199,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 +338,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
 
@@ -458,57 +475,88 @@ class BoundaryForcing:
         ds.attrs["hc"] = self.grid.ds.attrs["hc"]
 
         # Convert the time coordinate to the format expected by ROMS
-        if data.climatology:
-            ds.attrs["climatology"] = str(True)
-            # Preserve absolute time coordinate for readability
-            ds = ds.assign_coords(
-                {"abs_time": np.datetime64(self.model_reference_date) + ds["time"]}
-            )
-            # Convert to pandas TimedeltaIndex
-            timedelta_index = pd.to_timedelta(ds["time"].values)
+        ds, bry_time = convert_to_roms_time(
+            ds, self.model_reference_date, data.climatology
+        )
 
-            # Determine the start of the year for the base_datetime
-            start_of_year = datetime(self.model_reference_date.year, 1, 1)
+        ds = ds.assign_coords({"bry_time": bry_time})
+        ds = ds.swap_dims({"time": "bry_time"})
+        ds = ds.drop_vars("time")
 
-            # Calculate the offset from midnight of the new year
-            offset = self.model_reference_date - start_of_year
+        return ds
 
-            # Convert the timedelta to nanoseconds first, then to days
-            bry_time = xr.DataArray(
-                (timedelta_index - offset).view("int64") / 3600 / 24 * 1e-9,
-                dims="time",
-            )
+    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.
 
-        else:
-            # Preserve absolute time coordinate for readability
-            ds = ds.assign_coords({"abs_time": ds["time"]})
-            bry_time = (
-                (ds["time"] - np.datetime64(self.model_reference_date)).astype(
-                    "float64"
-                )
-                / 3600
-                / 24
-                * 1e-9
-            )
+        Parameters
+        ----------
+        processed_fields : dict
+            A dictionary where keys are variable names and values are `xarray.DataArray`
+            objects representing the processed data for each variable.
 
-        ds = ds.assign_coords({"bry_time": bry_time})
-        ds["bry_time"].attrs[
-            "long_name"
-        ] = f"days since {str(self.model_reference_date)}"
-        ds["bry_time"].encoding["units"] = "days"
-        ds["bry_time"].attrs["units"] = "days"
-        ds = ds.swap_dims({"time": "bry_time"})
-        ds = ds.drop_vars("time")
-        ds.encoding["unlimited_dims"] = "bry_time"
+        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.
 
-        if data.climatology:
-            ds["bry_time"].attrs["cycle_length"] = 365.25
+        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.
 
-        return ds
+        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 +578,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 +592,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 +670,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 +738,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,
@@ -780,45 +822,8 @@ class BoundaryForcing:
         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}
-
-        boundary_forcing_data = {
-            "BoundaryForcing": {
-                "start_time": self.start_time.isoformat(),
-                "end_time": self.end_time.isoformat(),
-                "boundaries": self.boundaries,
-                "source": self.source,
-                "type": self.type,
-                "model_reference_date": self.model_reference_date.isoformat(),
-            }
-        }
-
-        yaml_data = {
-            **grid_yaml_data,
-            **boundary_forcing_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(
@@ -839,36 +844,12 @@ class BoundaryForcing:
             An instance of the BoundaryForcing 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))
-
-        boundary_forcing_data = None
-
-        # Process the YAML documents
-        for doc in documents:
-            if doc is None:
-                continue
-            if "BoundaryForcing" in doc:
-                boundary_forcing_data = doc["BoundaryForcing"]
-                break
-
-        if boundary_forcing_data is None:
-            raise ValueError("No BoundaryForcing configuration found in the YAML file.")
-
-        # Convert from string to datetime
-        for date_string in ["model_reference_date", "start_time", "end_time"]:
-            boundary_forcing_data[date_string] = datetime.fromisoformat(
-                boundary_forcing_data[date_string]
-            )
 
         grid = Grid.from_yaml(filepath)
+        params = _from_yaml(cls, filepath)
 
         # Create and return an instance of InitialConditions
-        return cls(grid=grid, **boundary_forcing_data, use_dask=use_dask)
+        return cls(grid=grid, **params, use_dask=use_dask)
 
 
 def get_boundary_info():
@@ -952,6 +933,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"