roms-tools 2.4.0__py3-none-any.whl → 2.6.0__py3-none-any.whl

roms_tools/regrid.py

@@ -1,7 +1,9 @@
+import xgcm
 import xarray as xr
+import warnings
 
 
-class LateralRegrid:
+class LateralRegridToROMS:
     """Handles lateral regridding of data onto a new spatial grid."""
 
     def __init__(self, target_coords, source_dim_names):
@@ -49,7 +51,92 @@ class LateralRegrid:
         return regridded
 
 
-class VerticalRegrid:
+class LateralRegridFromROMS:
+    """Regrids data from a curvilinear ROMS grid onto latitude-longitude coordinates
+    using xESMF.
+
+    It requires the `xesmf` library, which can be installed by installing `roms-tools` via conda.
+
+    Parameters
+    ----------
+    source_grid_ds : xarray.Dataset
+        The source dataset containing the curvilinear ROMS grid with 'lat_rho' and 'lon_rho'.
+
+    target_coords : dict
+        A dictionary containing 'lat' and 'lon' arrays representing the target
+        latitude and longitude coordinates for regridding.
+
+    method : str, optional
+        The regridding method to use. Default is "bilinear". Other options include "nearest_s2d" and "conservative".
+
+    Raises
+    ------
+    ImportError
+        If xESMF is not installed.
+    """
+
+    def __init__(self, ds_in, target_coords, method="bilinear"):
+        """Initializes the regridder with the source and target grids.
+
+        Parameters
+        ----------
+        ds_in : xarray.Dataset or xarray.DataArray
+            The source dataset or dataarray containing the curvilinear ROMS grid with coordinates 'lat' and 'lon'.
+
+        target_coords : dict
+            A dictionary containing 'lat' and 'lon' arrays representing the target latitude
+            and longitude coordinates for regridding.
+
+        method : str, optional
+            The regridding method to use. Default is "bilinear". Other options include
+            "nearest_s2d" and "conservative".
+
+        Raises
+        ------
+        ImportError
+            If xESMF is not installed.
+        """
+
+        try:
+            import xesmf as xe
+
+        except ImportError:
+            raise ImportError(
+                "xesmf is required for this regridding task. Please install `roms-tools` via conda, which includes xesmf."
+            )
+
+        ds_out = xr.Dataset()
+        ds_out["lat"] = target_coords["lat"]
+        ds_out["lon"] = target_coords["lon"]
+
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=UserWarning, module="xesmf")
+            self.regridder = xe.Regridder(
+                ds_in, ds_out, method=method, unmapped_to_nan=True
+            )
+
+    def apply(self, da):
+        """Applies the regridding to the provided data array.
+
+        Parameters
+        ----------
+        da : xarray.DataArray
+            The data array to regrid. This should have the same dimension names as the
+            source grid (e.g., 'lat' and 'lon').
+
+        Returns
+        -------
+        xarray.DataArray
+            The regridded data array.
+        """
+
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=UserWarning, module="xesmf")
+            regridded = self.regridder(da, keep_attrs=True)
+        return regridded
+
+
+class VerticalRegridToROMS:
     """Interpolates data onto new vertical (depth) coordinates.
 
     Parameters
@@ -106,13 +193,13 @@ class VerticalRegrid:
             }
         )
 
-    def apply(self, var, fill_nans=True):
+    def apply(self, da, fill_nans=True):
         """Interpolates the variable onto the new depth grid using precomputed
         coefficients for linear interpolation between layers.
 
         Parameters
         ----------
-        var : xarray.DataArray
+        da : xarray.DataArray
             The input data to be regridded along the depth dimension. This should be
             an array with the same depth coordinates as the original grid.
         fill_nans : bool, optional
@@ -130,16 +217,74 @@ class VerticalRegrid:
 
         dims = {"dim": self.depth_dim}
 
-        var_below = var.where(self.coeff["is_below"]).sum(**dims)
-        var_above = var.where(self.coeff["is_above"]).sum(**dims)
+        da_below = da.where(self.coeff["is_below"]).sum(**dims)
+        da_above = da.where(self.coeff["is_above"]).sum(**dims)
 
-        result = var_below + (var_above - var_below) * self.coeff["factor"]
+        result = da_below + (da_above - da_below) * self.coeff["factor"]
         if fill_nans:
-            result = result.where(self.coeff["upper_mask"], var.isel({dims["dim"]: 0}))
-            result = result.where(self.coeff["lower_mask"], var.isel({dims["dim"]: -1}))
+            result = result.where(self.coeff["upper_mask"], da.isel({dims["dim"]: 0}))
+            result = result.where(self.coeff["lower_mask"], da.isel({dims["dim"]: -1}))
         else:
             result = result.where(self.coeff["upper_mask"]).where(
                 self.coeff["lower_mask"]
             )
 
         return result
+
+
+class VerticalRegridFromROMS:
+    """A class for regridding data from the ROMS vertical coordinate system to target
+    depth levels.
+
+    This class uses the `xgcm` package to perform the transformation from the ROMS depth coordinates to
+    a user-defined set of target depth levels. It assumes that the input dataset `ds` contains the necessary
+    vertical coordinate information (`s_rho`).
+
+    Attributes
+    ----------
+    grid : xgcm.Grid
+        The grid object used for regridding, initialized with the given dataset `ds`.
+    """
+
+    def __init__(self, ds):
+        """Initializes the `VerticalRegridFromROMS` object by creating an `xgcm.Grid`
+        instance.
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            The dataset containing the ROMS output data, which must include the vertical coordinate `s_rho`.
+        """
+        self.grid = xgcm.Grid(ds, coords={"s_rho": {"center": "s_rho"}}, periodic=False)
+
+    def apply(self, da, depth_coords, target_depth_levels):
+        """Applies vertical regridding from ROMS to the specified target depth levels.
+
+        This method transforms the input data array `da` from the ROMS vertical coordinate (`s_rho`)
+        to a set of target depth levels defined by `target_depth_levels`.
+
+        Parameters
+        ----------
+        da : xarray.DataArray
+            The data array containing the ROMS output field to be regridded. It must have a vertical
+            dimension corresponding to `s_rho`.
+
+        depth_coords : array-like
+            The depth coordinates of the input data array `da` (typically the `s_rho` coordinate in ROMS).
+
+        target_depth_levels : array-like
+            The target depth levels to which the input data `da` will be regridded.
+
+        Returns
+        -------
+        xarray.DataArray
+            A new `xarray.DataArray` containing the regridded data at the specified target depth levels.
+        """
+
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning, module="xgcm")
+            transformed = self.grid.transform(
+                da, "s_rho", target_depth_levels, target_data=depth_coords
+            )
+
+        return transformed

roms_tools/setup/boundary_forcing.py

@@ -3,13 +3,13 @@ import numpy as np
 from scipy.ndimage import label
 import logging
 import importlib.metadata
-from typing import Dict, Union, List
+from typing import Dict, Union, List, Optional
 from dataclasses import dataclass, field
 from datetime import datetime
 import matplotlib.pyplot as plt
 from pathlib import Path
 from roms_tools import Grid
-from roms_tools.regrid import LateralRegrid, VerticalRegrid
+from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
 from roms_tools.utils import save_datasets
 from roms_tools.vertical_coordinate import compute_depth
 from roms_tools.plot import _section_plot, _line_plot
@@ -35,7 +35,7 @@ from roms_tools.setup.utils import (
 )
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class BoundaryForcing:
     """Represents boundary forcing input data for ROMS.
 
@@ -43,10 +43,14 @@ class BoundaryForcing:
     ----------
     grid : Grid
         Object representing the grid information.
-    start_time : datetime
-        Start time of the desired boundary forcing data.
-    end_time : datetime
-        End time of the desired boundary forcing data.
+    start_time : datetime, optional
+        The start time of the desired surface forcing data. This time is used to filter the dataset
+        to include only records on or after this time, with a single record at or before this time.
+        If no time filtering is desired, set it to None. Default is None.
+    end_time : datetime, optional
+        The end time of the desired surface forcing data. This time is used to filter the dataset
+        to include only records on or before this time, with a single record at or after this time.
+        If no time filtering is desired, set it to None. Default is None.
     boundaries : Dict[str, bool], optional
         Dictionary specifying which boundaries are forced (south, east, north, west). Default is all True.
     source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool]
@@ -96,8 +100,8 @@ class BoundaryForcing:
     """
 
     grid: Grid
-    start_time: datetime
-    end_time: datetime
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
     boundaries: Dict[str, bool] = field(
         default_factory=lambda: {
             "south": True,
@@ -120,17 +124,20 @@ class BoundaryForcing:
 
         self._input_checks()
         # Dataset for depth coordinates
-        object.__setattr__(self, "ds_depth_coords", xr.Dataset())
+        self.ds_depth_coords = xr.Dataset()
 
         target_coords = get_target_coords(self.grid)
 
         data = self._get_data()
 
         if self.apply_2d_horizontal_fill:
+
             data.choose_subdomain(
                 target_coords,
                 buffer_points=20,  # lateral fill needs good buffer from data margin
             )
+            # Enforce double precision to ensure reproducibility
+            data.convert_to_float64()
             data.extrapolate_deepest_to_bottom()
             data.apply_lateral_fill()
 
@@ -158,6 +165,8 @@ class BoundaryForcing:
                 )
 
                 if not self.apply_2d_horizontal_fill:
+                    # Enforce double precision to ensure reproducibility
+                    bdry_data.convert_to_float64()
                     bdry_data.extrapolate_deepest_to_bottom()
 
                 processed_fields = {}
@@ -176,7 +185,7 @@ class BoundaryForcing:
                     lat = target_coords["lat"].isel(
                         **self.bdry_coords["vector"][direction]
                     )
-                    lateral_regrid = LateralRegrid(
+                    lateral_regrid = LateralRegridToROMS(
                         {"lat": lat, "lon": lon}, bdry_data.dim_names
                     )
                     for var_name in vector_var_names:
@@ -205,7 +214,7 @@ class BoundaryForcing:
                     lat = target_coords["lat"].isel(
                         **self.bdry_coords["rho"][direction]
                     )
-                    lateral_regrid = LateralRegrid(
+                    lateral_regrid = LateralRegridToROMS(
                         {"lat": lat, "lon": lon}, bdry_data.dim_names
                     )
                     for var_name in tracer_var_names:
@@ -283,7 +292,7 @@ class BoundaryForcing:
                 # vertical regridding
                 for location in ["rho", "u", "v"]:
                     if len(var_names_dict[location]) > 0:
-                        vertical_regrid = VerticalRegrid(
+                        vertical_regrid = VerticalRegridToROMS(
                             self.ds_depth_coords[f"layer_depth_{location}_{direction}"],
                             bdry_data.ds[bdry_data.dim_names["depth"]],
                         )
@@ -326,9 +335,21 @@ class BoundaryForcing:
         for var_name in ds.data_vars:
             ds[var_name] = substitute_nans_by_fillvalue(ds[var_name])
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def _input_checks(self):
+        # Check that start_time and end_time are both None or none of them is
+        if (self.start_time is None) != (self.end_time is None):
+            raise ValueError(
+                "Both `start_time` and `end_time` must be provided together as datetime objects or both should be None."
+            )
+
+        # Trigger a warning if both are None
+        if self.start_time is None and self.end_time is None:
+            logging.warning(
+                "Both `start_time` and `end_time` are None. No time filtering will be applied to the source data."
+            )
+
         # Validate the 'type' parameter
         if self.type not in ["physics", "bgc"]:
             raise ValueError("`type` must be either 'physics' or 'bgc'.")
@@ -340,11 +361,10 @@ class BoundaryForcing:
             raise ValueError("`source` must include a 'path'.")
 
         # Set 'climatology' to False if not provided in 'source'
-        object.__setattr__(
-            self,
-            "source",
-            {**self.source, "climatology": self.source.get("climatology", False)},
-        )
+        self.source = {
+            **self.source,
+            "climatology": self.source.get("climatology", False),
+        }
 
         # Ensure adjust_depth_for_sea_surface_height is only used with type="physics"
         if self.type == "bgc" and self.adjust_depth_for_sea_surface_height:
@@ -352,7 +372,7 @@ class BoundaryForcing:
                 "adjust_depth_for_sea_surface_height is not applicable for BGC fields. "
                 "Setting it to False."
             )
-            object.__setattr__(self, "adjust_depth_for_sea_surface_height", False)
+            self.adjust_depth_for_sea_surface_height = False
         elif self.adjust_depth_for_sea_surface_height:
             logging.info("Sea surface height will be used to adjust depth coordinates.")
         else:
@@ -474,7 +494,7 @@ class BoundaryForcing:
                 else:
                     variable_info[var_name] = {**default_info, "validate": False}
 
-        object.__setattr__(self, "variable_info", variable_info)
+        self.variable_info = variable_info
 
     def _write_into_dataset(self, direction, processed_fields, ds=None):
         if ds is None:
@@ -542,7 +562,7 @@ class BoundaryForcing:
 
         bdry_coords = get_boundary_coords()
 
-        object.__setattr__(self, "bdry_coords", bdry_coords)
+        self.bdry_coords = bdry_coords
 
     def _get_depth_coordinates(
         self,
@@ -840,12 +860,6 @@ class BoundaryForcing:
 
         field = self.ds[var_name].isel(bry_time=time)
 
-        if self.use_dask:
-            from dask.diagnostics import ProgressBar
-
-            with ProgressBar():
-                field = field.load()
-
         title = field.long_name
         var_name_wo_direction, direction = var_name.split("_")
         location = self.variable_info[var_name_wo_direction]["location"]
@@ -860,10 +874,17 @@ class BoundaryForcing:
 
         mask = mask.isel(**self.bdry_coords[location][direction])
 
+        # Load the data
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                field = field.load()
+
         if "s_rho" in field.dims:
             layer_depth = self.ds_depth_coords[f"layer_depth_{location}_{direction}"]
             if self.adjust_depth_for_sea_surface_height:
-                layer_depth = layer_depth.isel(time=time)
+                layer_depth = layer_depth.isel(time=time).load()
             field = field.assign_coords({"layer_depth": layer_depth})
         if var_name.startswith(("u", "v", "ubar", "vbar", "zeta")):
             vmax = max(field.max().values, -field.min().values)
@@ -967,7 +988,6 @@ class BoundaryForcing:
         cls,
         filepath: Union[str, Path],
         use_dask: bool = False,
-        bypass_validation: bool = False,
     ) -> "BoundaryForcing":
         """Create an instance of the BoundaryForcing class from a YAML file.
 
@@ -977,10 +997,6 @@ class BoundaryForcing:
             The path to the YAML file from which the parameters will be read.
         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.
-        bypass_validation: bool, optional
-            Indicates whether to skip validation checks in the processed data. When set to True,
-            the validation process that ensures no NaN values exist at wet points
-            in the processed dataset is bypassed. Defaults to False.
 
         Returns
         -------
@@ -993,9 +1009,7 @@ class BoundaryForcing:
         params = _from_yaml(cls, filepath)
 
         # Create and return an instance of InitialConditions
-        return cls(
-            grid=grid, **params, use_dask=use_dask, bypass_validation=bypass_validation
-        )
+        return cls(grid=grid, **params, use_dask=use_dask)
 
 
 def apply_1d_horizontal_fill(data_array: xr.DataArray) -> xr.DataArray: