roms-tools 2.5.0__py3-none-any.whl → 2.6.1__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,81 @@ 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, mask_edges=True):
+        """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.
+
+        mask_edges: bool, optional
+            If activated, target values outside the range of depth_coords are masked with nan. Defaults to True.
+
+        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,
+                mask_edges=mask_edges,
+            )
+
+        return transformed

roms_tools/setup/boundary_forcing.py

@@ -9,7 +9,7 @@ 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.
 
@@ -124,7 +124,7 @@ 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)
 
@@ -185,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:
@@ -214,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:
@@ -292,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"]],
                         )
@@ -335,7 +335,7 @@ 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
@@ -361,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:
@@ -373,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:
@@ -495,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:
@@ -563,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,
@@ -861,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"]
@@ -881,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)

roms_tools/setup/datasets.py

@@ -25,7 +25,7 @@ from roms_tools.setup.fill import LateralFill
 # lat-lon datasets
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class Dataset:
     """Represents forcing data on original grid.
 
@@ -132,8 +132,8 @@ class Dataset:
         self.infer_horizontal_resolution(ds)
 
         # Check whether the data covers the entire globe
-        object.__setattr__(self, "is_global", self.check_if_global(ds))
-        object.__setattr__(self, "ds", ds)
+        self.is_global = self.check_if_global(ds)
+        self.ds = ds
 
         if self.apply_post_processing:
             self.post_process()
@@ -354,7 +354,7 @@ class Dataset:
         resolution = np.mean([lat_resolution, lon_resolution])
 
         # Set the computed resolution as an attribute
-        object.__setattr__(self, "resolution", resolution)
+        self.resolution = resolution
 
     def compute_minimal_grid_spacing(self, ds: xr.Dataset):
         """Compute the minimal grid spacing in a dataset based on latitude and longitude
@@ -528,7 +528,7 @@ class Dataset:
             This method modifies the dataset in place and does not return anything.
         """
         ds = self.ds.astype({var: "float64" for var in self.ds.data_vars})
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def choose_subdomain(
         self,
@@ -671,7 +671,7 @@ class Dataset:
         if return_copy:
             return Dataset.from_ds(self, subdomain)
         else:
-            object.__setattr__(self, "ds", subdomain)
+            self.ds = subdomain
 
     def apply_lateral_fill(self):
         """Apply lateral fill to variables using the dataset's mask and grid dimensions.
@@ -757,7 +757,7 @@ class Dataset:
         dataset = cls.__new__(cls)
 
         # Directly set the provided dataset as the 'ds' attribute
-        object.__setattr__(dataset, "ds", ds)
+        dataset.ds = ds
 
         # Copy all other attributes from the original data instance
         for attr in vars(original_dataset):
@@ -767,7 +767,7 @@ class Dataset:
         return dataset
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class TPXODataset(Dataset):
     """Represents tidal data on the original grid from the TPXO dataset.
 
@@ -858,15 +858,11 @@ class TPXODataset(Dataset):
             {"nx": "longitude", "ny": "latitude", self.dim_names["ntides"]: "ntides"}
         )
 
-        object.__setattr__(
-            self,
-            "dim_names",
-            {
-                "latitude": "latitude",
-                "longitude": "longitude",
-                "ntides": "ntides",
-            },
-        )
+        self.dim_names = {
+            "latitude": "latitude",
+            "longitude": "longitude",
+            "ntides": "ntides",
+        }
 
         return ds
 
@@ -909,15 +905,15 @@ class TPXODataset(Dataset):
             ds["mask"] = mask
             ds = ds.drop_vars(["depth"])
 
-            object.__setattr__(self, "ds", ds)
+            self.ds = ds
 
         # Remove "depth" from var_names
         updated_var_names = {**self.var_names}  # Create a copy of the dictionary
         updated_var_names.pop("depth", None)  # Remove "depth" if it exists
-        object.__setattr__(self, "var_names", updated_var_names)
+        self.var_names = updated_var_names
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class GLORYSDataset(Dataset):
     """Represents GLORYS data on original grid.
 
@@ -996,7 +992,7 @@ class GLORYSDataset(Dataset):
         self.ds["mask_vel"] = mask_vel
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class CESMDataset(Dataset):
     """Represents CESM data on original grid.
 
@@ -1077,12 +1073,12 @@ class CESMDataset(Dataset):
             # Update dimension names
             updated_dim_names = self.dim_names.copy()
             updated_dim_names["time"] = "time"
-            object.__setattr__(self, "dim_names", updated_dim_names)
+            self.dim_names = updated_dim_names
 
         return ds
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class CESMBGCDataset(CESMDataset):
     """Represents CESM BGC data on original grid.
 
@@ -1184,12 +1180,12 @@ class CESMBGCDataset(CESMDataset):
             if "z_t_150m" in ds.variables:
                 ds = ds.drop_vars("z_t_150m")
             # update dataset
-            object.__setattr__(self, "ds", ds)
+            self.ds = ds
 
             # Update dim_names with "depth": "depth" key-value pair
             updated_dim_names = self.dim_names.copy()
             updated_dim_names["depth"] = "depth"
-            object.__setattr__(self, "dim_names", updated_dim_names)
+            self.dim_names = updated_dim_names
 
         mask = xr.where(
             self.ds[self.var_names["PO4"]]
@@ -1202,7 +1198,7 @@ class CESMBGCDataset(CESMDataset):
         self.ds["mask"] = mask
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class CESMBGCSurfaceForcingDataset(CESMDataset):
     """Represents CESM BGC surface forcing data on original grid.
 
@@ -1258,7 +1254,7 @@ class CESMBGCSurfaceForcingDataset(CESMDataset):
 
         if "z_t" in self.ds.variables:
             ds = self.ds.drop_vars("z_t")
-            object.__setattr__(self, "ds", ds)
+            self.ds = ds
 
         mask = xr.where(
             self.ds[self.var_names["pco2_air"]]
@@ -1271,7 +1267,7 @@ class CESMBGCSurfaceForcingDataset(CESMDataset):
         self.ds["mask"] = mask
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class ERA5Dataset(Dataset):
     """Represents ERA5 data on original grid.
 
@@ -1371,27 +1367,27 @@ class ERA5Dataset(Dataset):
             ds["qair"].attrs["long_name"] = "Absolute humidity at 2m"
             ds["qair"].attrs["units"] = "kg/kg"
             ds = ds.drop_vars([self.var_names["d2m"]])
-            object.__setattr__(self, "ds", ds)
+            self.ds = ds
 
             # Update var_names dictionary
             var_names = {**self.var_names, "qair": "qair"}
             var_names.pop("d2m")
-            object.__setattr__(self, "var_names", var_names)
+            self.var_names = var_names
 
         if "mask" in self.var_names.keys():
             ds = self.ds
             mask = xr.where(self.ds[self.var_names["mask"]].isel(time=0).isnull(), 0, 1)
             ds["mask"] = mask
             ds = ds.drop_vars([self.var_names["mask"]])
-            object.__setattr__(self, "ds", ds)
+            self.ds = ds
 
             # Remove mask from var_names dictionary
             var_names = self.var_names
             var_names.pop("mask")
-            object.__setattr__(self, "var_names", var_names)
+            self.var_names = var_names
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class ERA5Correction(Dataset):
     """Global dataset to correct ERA5 radiation. The dataset contains multiplicative
     correction factors for the ERA5 shortwave radiation, obtained by comparing the
@@ -1493,10 +1489,10 @@ class ERA5Correction(Dataset):
             raise ValueError(
                 "The correction dataset does not contain all specified longitude values."
             )
-        object.__setattr__(self, "ds", subdomain)
+        self.ds = subdomain
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class ETOPO5Dataset(Dataset):
     """Represents topography data on the original grid from the ETOPO5 dataset.
 
@@ -1553,7 +1549,7 @@ class ETOPO5Dataset(Dataset):
         return ds
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class SRTM15Dataset(Dataset):
     """Represents topography data on the original grid from the SRTM15 dataset.
 
@@ -1589,7 +1585,7 @@ class SRTM15Dataset(Dataset):
 
 
 # river datasets
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class RiverDataset:
     """Represents river data.
 
@@ -1647,7 +1643,7 @@ class RiverDataset:
 
         # Select relevant times
         ds = self.add_time_info(ds)
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def load_data(self) -> xr.Dataset:
         """Load dataset from the specified file.
@@ -1785,13 +1781,13 @@ class RiverDataset:
 
         ds = assign_dates_to_climatology(self.ds, "month")
         ds = ds.swap_dims({"month": "time"})
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
         updated_dim_names = {**self.dim_names}
         updated_dim_names["time"] = "time"
-        object.__setattr__(self, "dim_names", updated_dim_names)
+        self.dim_names = updated_dim_names
 
-        object.__setattr__(self, "climatology", True)
+        self.climatology = True
 
     def sort_by_river_volume(self, ds: xr.Dataset) -> xr.Dataset:
         """Sorts the dataset by river volume in descending order (largest rivers first),
@@ -1911,7 +1907,7 @@ class RiverDataset:
             ds = xr.Dataset()
             river_indices = {}
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
         return river_indices
 
@@ -1961,10 +1957,10 @@ class RiverDataset:
             )
 
         # Set the filtered dataset as the new `ds`
-        object.__setattr__(self, "ds", ds_filtered)
+        self.ds = ds_filtered
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class DaiRiverDataset(RiverDataset):
     """Represents river data from the Dai river dataset.