roms-tools 1.5.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

roms_tools/setup/regrid.py

@@ -2,45 +2,38 @@ import xarray as xr
 
 
 class LateralRegrid:
-    """Applies lateral fill and regridding to data.
+    """Handles lateral regridding of data onto a new spatial grid."""
 
-    This class fills missing values in ocean data and interpolates it onto a new grid
-    defined by the provided longitude and latitude.
-
-    Parameters
-    ----------
-    data : DataContainer
-        Container with variables to be interpolated, including a `mask` and dimension names.
-    lon : xarray.DataArray
-        Target longitude coordinates.
-    lat : xarray.DataArray
-        Target latitude coordinates.
-    """
-
-    def __init__(self, data, lon, lat):
-        """Initializes the lateral fill and target grid coordinates.
+    def __init__(self, target_coords, source_dim_names):
+        """Initialize target grid coordinates and names for grid dimensions.
 
         Parameters
         ----------
-        data : DataContainer
-            Data with dimensions and mask for filling.
-        lon : xarray.DataArray
-            Longitude for new grid.
-        lat : xarray.DataArray
-            Latitude for new grid.
+        target_coords : dict
+            Dictionary containing 'lon' and 'lat' as xarray.DataArrays representing
+            the longitude and latitude values of the target grid.
+        source_dim_names : dict
+            Dictionary specifying names for the latitude and longitude dimensions,
+            typically using keys like "latitude" and "longitude" to align with the dataset conventions.
+
+        Attributes
+        ----------
+        coords : dict
+            Maps the dimension names to the corresponding latitude and longitude
+            DataArrays, providing easy access to target grid coordinates.
         """
 
         self.coords = {
-            data.dim_names["latitude"]: lat,
-            data.dim_names["longitude"]: lon,
+            source_dim_names["latitude"]: target_coords["lat"],
+            source_dim_names["longitude"]: target_coords["lon"],
         }
 
-    def apply(self, var):
+    def apply(self, da):
         """Fills missing values and regrids the variable.
 
         Parameters
         ----------
-        var : xarray.DataArray
+        da : xarray.DataArray
             Input data to fill and regrid.
 
         Returns
@@ -48,38 +41,48 @@ class LateralRegrid:
         xarray.DataArray
             Regridded data with filled values.
         """
-        regridded = var.interp(self.coords, method="linear").drop_vars(
+        regridded = da.interp(self.coords, method="linear").drop_vars(
             list(self.coords.keys())
         )
         return regridded
 
 
 class VerticalRegrid:
-    """Performs vertical interpolation of data onto new depth coordinates.
+    """Interpolates data onto new vertical (depth) coordinates.
 
     Parameters
     ----------
-    data : DataContainer
-        Container holding the data to be regridded, with relevant dimension names.
-    target_depth : xarray.DataArray
-        Target depth coordinates for interpolation.
+    target_depth_coords : xarray.DataArray
+        Depth coordinates for the target grid.
+    source_depth_coords : xarray.DataArray
+        Depth coordinates for the source grid.
     """
 
-    def __init__(self, data, target_depth):
-        """Initializes vertical regridding with specified depth coordinates.
+    def __init__(self, target_depth_coords, source_depth_coords):
+        """Initialize regridding factors for interpolation.
 
         Parameters
         ----------
-        data : DataContainer
-            Container holding the data to be regridded, with relevant dimension names.
-        target_depth : xarray.DataArray
-            Target depth coordinates for interpolation.
+        target_depth_coords : xarray.DataArray
+            Depth coordinates for the target grid.
+        source_depth_coords : xarray.DataArray
+            Depth coordinates for the source grid.
+
+        Attributes
+        ----------
+        coeff : xarray.Dataset
+            Dataset containing:
+            - `is_below` : Boolean mask for depths just below target.
+            - `is_above` : Boolean mask for depths just above target.
+            - `upper_mask`, `lower_mask` : Masks for valid interpolation bounds.
+            - `factor` : Weight for blending values between levels.
         """
 
-        self.depth_dim = data.dim_names["depth"]
+        self.depth_dim = source_depth_coords.dims[0]
+        source_depth = source_depth_coords
         dims = {"dim": self.depth_dim}
 
-        dlev = data.ds[data.dim_names["depth"]] - target_depth
+        dlev = source_depth - target_depth_coords
         is_below = dlev == dlev.where(dlev >= 0).min(**dims)
         is_above = dlev == dlev.where(dlev <= 0).max(**dims)
         p_below = dlev.where(is_below).sum(**dims)
@@ -138,61 +141,3 @@ class VerticalRegrid:
             )
 
         return result
-
-
-def _lateral_regrid(data, lon, lat, data_vars, var_names):
-    """Laterally regrid specified variables onto new latitude and longitude coordinates.
-
-    Parameters
-    ----------
-    data : Dataset
-        Input data containing the information about source dimensions.
-    lon : xarray.DataArray
-        Target longitude coordinates.
-    lat : xarray.DataArray
-        Target latitude coordinates.
-    data_vars : dict of str : xarray.DataArray
-        Dictionary of variables to regrid.
-    var_names : list of str
-        Names of variables to regrid.
-
-    Returns
-    -------
-    dict of str : xarray.DataArray
-        Updated data_vars with regridded variables.
-    """
-    lateral_regrid = LateralRegrid(data, lon, lat)
-
-    for var_name in var_names:
-        if var_name in data_vars:
-            data_vars[var_name] = lateral_regrid.apply(data_vars[var_name])
-
-    return data_vars
-
-
-def _vertical_regrid(data, target_depth, data_vars, var_names):
-    """Vertically regrid specified variables onto new depth coordinates.
-
-    Parameters
-    ----------
-    data : Dataset
-        Input dataset containing the variables and source depth information.
-    target_depth : xarray.DataArray
-        Target depth coordinates for regridding.
-    data_vars : dict of str : xarray.DataArray
-        Dictionary of variables to be regridded.
-    var_names : list of str
-        Names of variables to regrid.
-
-    Returns
-    -------
-    dict of str : xarray.DataArray
-        Updated data_vars with variables regridded onto the target depth coordinates.
-    """
-    vertical_regrid = VerticalRegrid(data, target_depth)
-
-    for var_name in var_names:
-        if var_name in data_vars:
-            data_vars[var_name] = vertical_regrid.apply(data_vars[var_name])
-
-    return data_vars

roms_tools/setup/surface_forcing.py

@@ -7,8 +7,7 @@ from roms_tools.setup.grid import Grid
 from datetime import datetime
 import numpy as np
 from typing import Dict, Union, List
-from roms_tools.setup.fill import _lateral_fill, LateralFill
-from roms_tools.setup.regrid import _lateral_regrid, LateralRegrid
+from roms_tools.setup.regrid import LateralRegrid
 from roms_tools.setup.datasets import (
     ERA5Dataset,
     ERA5Correction,
@@ -99,51 +98,40 @@ class SurfaceForcing:
 
         data = self._get_data()
         data.choose_subdomain(
-            latitude_range=[
-                target_coords["lat"].min().values,
-                target_coords["lat"].max().values,
-            ],
-            longitude_range=[
-                target_coords["lon"].min().values,
-                target_coords["lon"].max().values,
-            ],
-            margin=2,
-            straddle=target_coords["straddle"],
+            target_coords,
+            buffer_points=20,  # lateral fill needs some buffer from data margin
         )
 
-        variable_info = self._set_variable_info(data)
-
-        data_vars = {}
-        for var_name in data.var_names:
-            if var_name != "mask":
-                data_vars[var_name] = data.ds[data.var_names[var_name]]
+        data.apply_lateral_fill()
 
-        data_vars = _lateral_fill(data_vars, data)
+        variable_info = self._set_variable_info(data)
+        var_names = variable_info.keys()
 
+        processed_fields = {}
         # lateral regridding
-        var_names = variable_info.keys()
-        data_vars = _lateral_regrid(
-            data, target_coords["lon"], target_coords["lat"], data_vars, var_names
-        )
+        lateral_regrid = LateralRegrid(target_coords, data.dim_names)
+        for var_name in var_names:
+            if var_name in data.var_names.keys():
+                processed_fields[var_name] = lateral_regrid.apply(
+                    data.ds[data.var_names[var_name]]
+                )
 
         # rotation of velocities and interpolation to u/v points
         if "uwnd" in variable_info and "vwnd" in variable_info:
-            data_vars["uwnd"], data_vars["vwnd"] = rotate_velocities(
-                data_vars["uwnd"],
-                data_vars["vwnd"],
+            processed_fields["uwnd"], processed_fields["vwnd"] = rotate_velocities(
+                processed_fields["uwnd"],
+                processed_fields["vwnd"],
                 target_coords["angle"],
                 interpolate=False,
             )
 
         # correct radiation
         if self.type == "physics" and self.correct_radiation:
-            data_vars = self._apply_correction(data_vars, data)
-
-        object.__setattr__(data, "data_vars", data_vars)
+            processed_fields = self._apply_correction(processed_fields, data)
 
         d_meta = get_variable_metadata()
 
-        ds = self._write_into_dataset(data, d_meta)
+        ds = self._write_into_dataset(processed_fields, data, d_meta)
 
         if self.use_coarse_grid:
             mask = self.grid.ds["mask_coarse"].rename(
@@ -152,13 +140,11 @@ class SurfaceForcing:
         else:
             mask = self.grid.ds["mask_rho"]
 
-        # NaN values at wet points indicate that the raw data did not cover the domain, and the following will raise a ValueError
-        for var in ds.data_vars:
-            nan_check(ds[var].isel(time=0), mask)
+        self._validate(ds, mask)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
-        for var in ds.data_vars:
-            ds[var] = substitute_nans_by_fillvalue(ds[var])
+        for var_name in ds.data_vars:
+            ds[var_name] = substitute_nans_by_fillvalue(ds[var_name])
 
         object.__setattr__(self, "ds", ds)
 
@@ -272,7 +258,7 @@ class SurfaceForcing:
 
         return variable_info
 
-    def _apply_correction(self, data_vars, data):
+    def _apply_correction(self, processed_fields, data):
 
         correction_data = self._get_correction_data()
         # choose same subdomain as forcing data so that we can use same mask
@@ -285,44 +271,34 @@ class SurfaceForcing:
         correction_data.choose_subdomain(
             coords_correction, straddle=self.target_coords["straddle"]
         )
-
+        correction_data.ds["mask"] = data.ds["mask"]  # use mask from ERA5 data
+        correction_data.apply_lateral_fill()
         # regrid
-        lateral_fill = LateralFill(
-            data.ds["mask"],  # use mask from ERA5 data
-            [
-                correction_data.dim_names["latitude"],
-                correction_data.dim_names["longitude"],
-            ],
-        )
-        lateral_regrid = LateralRegrid(
-            correction_data, self.target_coords["lon"], self.target_coords["lat"]
-        )
-
-        filled = lateral_fill.apply(
+        lateral_regrid = LateralRegrid(self.target_coords, correction_data.dim_names)
+        corr_factor = lateral_regrid.apply(
             correction_data.ds[correction_data.var_names["swr_corr"]]
         )
-        corr_factor = lateral_regrid.apply(filled)
 
         # temporal interpolation
         corr_factor = interpolate_from_climatology(
             corr_factor,
             correction_data.dim_names["time"],
-            time=data_vars["swrad"].time,
+            time=processed_fields["swrad"].time,
         )
 
-        data_vars["swrad"] = data_vars["swrad"] * corr_factor
+        processed_fields["swrad"] = processed_fields["swrad"] * corr_factor
 
-        return data_vars
+        return processed_fields
 
-    def _write_into_dataset(self, data, d_meta):
+    def _write_into_dataset(self, processed_fields, data, d_meta):
 
         # save in new dataset
         ds = xr.Dataset()
 
-        for var in data.data_vars.keys():
-            ds[var] = data.data_vars[var].astype(np.float32)
-            ds[var].attrs["long_name"] = d_meta[var]["long_name"]
-            ds[var].attrs["units"] = d_meta[var]["units"]
+        for var_name in processed_fields.keys():
+            ds[var_name] = processed_fields[var_name].astype(np.float32)
+            ds[var_name].attrs["long_name"] = d_meta[var_name]["long_name"]
+            ds[var_name].attrs["units"] = d_meta[var_name]["units"]
 
         if self.use_coarse_grid:
             ds = ds.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
@@ -388,11 +364,36 @@ class SurfaceForcing:
             ds = ds.drop_vars(["time"])
 
         variables_to_drop = ["lat_rho", "lon_rho", "lat_coarse", "lon_coarse"]
-        existing_vars = [var for var in variables_to_drop if var in ds]
+        existing_vars = [var_name for var_name in variables_to_drop if var_name in ds]
         ds = ds.drop_vars(existing_vars)
 
         return ds
 
+    def _validate(self, ds, mask):
+        """Validates the dataset by checking for NaN values at wet points, which would
+        indicate missing raw data coverage over the target domain.
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            The dataset to validate.
+        mask : xarray.DataArray
+            Land mask (1=ocean, 0=land) to determine wet points in the domain.
+
+        Raises
+        ------
+        ValueError
+            If NaN values are found in any of the specified variables at wet points,
+            indicating incomplete data coverage.
+
+        Notes
+        -----
+        This check is applied to the first time step (`time=0`) of each variable in the provided dataset.
+        """
+
+        for var_name in ds.data_vars:
+            nan_check(ds[var_name].isel(time=0), mask)
+
     def _add_global_metadata(self, ds=None):
 
         if ds is None:
@@ -416,12 +417,12 @@ class SurfaceForcing:
 
         return ds
 
-    def plot(self, varname, time=0) -> None:
+    def plot(self, var_name, time=0) -> None:
         """Plot the specified surface forcing field for a given time slice.
 
         Parameters
         ----------
-        varname : str
+        var_name : str
             The name of the surface forcing field to plot. Options include:
 
             - "uwnd": 10 meter wind in x-direction.
@@ -450,7 +451,7 @@ class SurfaceForcing:
         Raises
         ------
         ValueError
-            If the specified varname is not found in dataset.
+            If the specified var_name is not found in dataset.
 
 
         Examples
@@ -458,10 +459,10 @@ class SurfaceForcing:
         >>> atm_forcing.plot("uwnd", time=0)
         """
 
-        if varname not in self.ds:
-            raise ValueError(f"Variable '{varname}' is not found in dataset.")
+        if var_name not in self.ds:
+            raise ValueError(f"Variable '{var_name}' is not found in dataset.")
 
-        field = self.ds[varname].isel(time=time).load()
+        field = self.ds[var_name].isel(time=time).load()
         title = field.long_name
 
         # assign lat / lon
@@ -476,14 +477,14 @@ class SurfaceForcing:
         )
 
         # choose colorbar
-        if varname in ["uwnd", "vwnd"]:
+        if var_name in ["uwnd", "vwnd"]:
             vmax = max(field.max().values, -field.min().values)
             vmin = -vmax
             cmap = plt.colormaps.get_cmap("RdBu_r")
         else:
             vmax = field.max().values
             vmin = field.min().values
-            if varname in ["swrad", "lwrad", "Tair", "qair"]:
+            if var_name in ["swrad", "lwrad", "Tair", "qair"]:
                 cmap = plt.colormaps.get_cmap("YlOrRd")
             else:
                 cmap = plt.colormaps.get_cmap("YlGnBu")

roms_tools/setup/tides.py

@@ -20,8 +20,7 @@ from roms_tools.setup.utils import (
     rotate_velocities,
     get_vector_pairs,
 )
-from roms_tools.setup.fill import _lateral_fill
-from roms_tools.setup.regrid import _lateral_regrid
+from roms_tools.setup.regrid import LateralRegrid
 import matplotlib.pyplot as plt
 from pathlib import Path
 
@@ -77,71 +76,64 @@ class TidalForcing:
         data = self._get_data()
         data.check_number_constituents(self.ntides)
         data.choose_subdomain(
-            latitude_range=[
-                target_coords["lat"].min().values,
-                target_coords["lat"].max().values,
-            ],
-            longitude_range=[
-                target_coords["lon"].min().values,
-                target_coords["lon"].max().values,
-            ],
-            margin=2,
-            straddle=target_coords["straddle"],
+            target_coords,
+            buffer_points=20,
         )
         # select desired number of constituents
         object.__setattr__(data, "ds", data.ds.isel(ntides=slice(None, self.ntides)))
         self._correct_tides(data)
 
-        variable_info = self._set_variable_info()
-
-        data_vars = {}
-        for var_name in data.var_names:
-            data_vars[var_name] = data.ds[data.var_names[var_name]]
+        data.apply_lateral_fill()
 
-        data_vars = _lateral_fill(data_vars, data)
+        variable_info = self._set_variable_info()
+        var_names = variable_info.keys()
 
+        processed_fields = {}
         # lateral regridding
-        var_names = variable_info.keys()
-        data_vars = _lateral_regrid(
-            data, target_coords["lon"], target_coords["lat"], data_vars, var_names
-        )
+        lateral_regrid = LateralRegrid(target_coords, data.dim_names)
+        for var_name in var_names:
+            if var_name in data.var_names.keys():
+                processed_fields[var_name] = lateral_regrid.apply(
+                    data.ds[data.var_names[var_name]]
+                )
 
         # rotation of velocities and interpolation to u/v points
         vector_pairs = get_vector_pairs(variable_info)
         for pair in vector_pairs:
             u_component = pair[0]
             v_component = pair[1]
-            if u_component in data_vars and v_component in data_vars:
-                (data_vars[u_component], data_vars[v_component],) = rotate_velocities(
-                    data_vars[u_component],
-                    data_vars[v_component],
+            if u_component in processed_fields and v_component in processed_fields:
+                (
+                    processed_fields[u_component],
+                    processed_fields[v_component],
+                ) = rotate_velocities(
+                    processed_fields[u_component],
+                    processed_fields[v_component],
                     target_coords["angle"],
                     interpolate=False,
                 )
 
         # convert to barotropic velocity
-        for varname in ["u_Re", "v_Re", "u_Im", "v_Im"]:
-            data_vars[varname] = data_vars[varname] / self.grid.ds.h
+        for var_name in ["u_Re", "v_Re", "u_Im", "v_Im"]:
+            processed_fields[var_name] = processed_fields[var_name] / self.grid.ds.h
 
         # interpolate from rho- to velocity points
         for uname in ["u_Re", "u_Im"]:
-            data_vars[uname] = interpolate_from_rho_to_u(data_vars[uname])
+            processed_fields[uname] = interpolate_from_rho_to_u(processed_fields[uname])
         for vname in ["v_Re", "v_Im"]:
-            data_vars[vname] = interpolate_from_rho_to_v(data_vars[vname])
+            processed_fields[vname] = interpolate_from_rho_to_v(processed_fields[vname])
 
         d_meta = get_variable_metadata()
-        ds = self._write_into_dataset(data_vars, d_meta)
+        ds = self._write_into_dataset(processed_fields, d_meta)
         ds["omega"] = data.ds["omega"]
 
         ds = self._add_global_metadata(ds)
 
-        # NaN values at wet points indicate that the raw data did not cover the domain, and the following will raise a ValueError
-        for var in ["ssh_Re", "u_Re", "v_Im"]:
-            nan_check(ds[var].isel(ntides=0), self.grid.ds.mask_rho)
+        self._validate(ds)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
-        for var in ds.data_vars:
-            ds[var] = substitute_nans_by_fillvalue(ds[var])
+        for var_name in ds.data_vars:
+            ds[var_name] = substitute_nans_by_fillvalue(ds[var_name])
 
         object.__setattr__(self, "ds", ds)
 
@@ -216,15 +208,15 @@ class TidalForcing:
 
         return variable_info
 
-    def _write_into_dataset(self, data_vars, d_meta):
+    def _write_into_dataset(self, processed_fields, d_meta):
 
         # save in new dataset
         ds = xr.Dataset()
 
-        for var in data_vars.keys():
-            ds[var] = data_vars[var].astype(np.float32)
-            ds[var].attrs["long_name"] = d_meta[var]["long_name"]
-            ds[var].attrs["units"] = d_meta[var]["units"]
+        for var_name in processed_fields.keys():
+            ds[var_name] = processed_fields[var_name].astype(np.float32)
+            ds[var_name].attrs["long_name"] = d_meta[var_name]["long_name"]
+            ds[var_name].attrs["units"] = d_meta[var_name]["units"]
 
         ds = ds.drop_vars(["lat_rho", "lon_rho"])
 
@@ -247,12 +239,35 @@ class TidalForcing:
 
         return ds
 
-    def plot(self, varname, ntides=0) -> None:
+    def _validate(self, ds):
+        """Validates the dataset by checking for NaN values at wet points, which would
+        indicate missing raw data coverage over the target domain.
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            The dataset to validate, containing tidal variables and a mask for wet points.
+
+        Raises
+        ------
+        ValueError
+            If NaN values are found in any of the specified variables at wet points,
+            indicating incomplete data coverage.
+
+        Notes
+        -----
+        This check is applied to the first constituent (`ntides=0`) of each variable in the dataset.
+        The method utilizes `self.grid.ds.mask_rho` to determine the wet points in the domain.
+        """
+        for var_name in ds.data_vars:
+            nan_check(ds[var_name].isel(ntides=0), self.grid.ds.mask_rho)
+
+    def plot(self, var_name, ntides=0) -> None:
         """Plot the specified tidal forcing variable for a given tidal constituent.
 
         Parameters
         ----------
-        varname : str
+        var_name : str
             The tidal forcing variable to plot. Options include:
 
             - "ssh_Re": Real part of tidal elevation.
@@ -285,7 +300,7 @@ class TidalForcing:
         >>> tidal_forcing.plot("ssh_Re", nc=0)
         """
 
-        field = self.ds[varname].isel(ntides=ntides).compute()
+        field = self.ds[var_name].isel(ntides=ntides).compute()
         if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
             field = field.where(self.grid.ds.mask_rho)
             field = field.assign_coords(
@@ -306,7 +321,7 @@ class TidalForcing:
         else:
             ValueError("provided field does not have two horizontal dimension")
 
-        title = "%s, ntides = %i" % (field.long_name, self.ds[varname].ntides[ntides])
+        title = "%s, ntides = %i" % (field.long_name, self.ds[var_name].ntides[ntides])
 
         vmax = max(field.max(), -field.min())
         vmin = -vmax