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

roms_tools/setup/surface_forcing.py

@@ -7,7 +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.mixins import ROMSToolsMixins
+from roms_tools.setup.regrid import LateralRegrid
 from roms_tools.setup.datasets import (
     ERA5Dataset,
     ERA5Correction,
@@ -20,6 +20,8 @@ from roms_tools.setup.utils import (
     get_variable_metadata,
     group_dataset,
     save_datasets,
+    get_target_coords,
+    rotate_velocities,
 )
 from roms_tools.setup.plot import _plot
 import matplotlib.pyplot as plt
@@ -27,7 +29,7 @@ from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
-class SurfaceForcing(ROMSToolsMixins):
+class SurfaceForcing:
     """Represents surface forcing input data for ROMS.
 
     Parameters
@@ -91,69 +93,45 @@ class SurfaceForcing(ROMSToolsMixins):
     def __post_init__(self):
 
         self._input_checks()
-        lon, lat, angle, straddle = super()._get_target_lon_lat(self.use_coarse_grid)
-        object.__setattr__(self, "target_lon", lon)
-        object.__setattr__(self, "target_lat", lat)
+        target_coords = get_target_coords(self.grid, self.use_coarse_grid)
+        object.__setattr__(self, "target_coords", target_coords)
 
         data = self._get_data()
         data.choose_subdomain(
-            latitude_range=[lat.min().values, lat.max().values],
-            longitude_range=[lon.min().values, lon.max().values],
-            margin=2,
-            straddle=straddle,
+            target_coords,
+            buffer_points=20,  # lateral fill needs some buffer from data margin
         )
-        if self.type == "physics":
-            vars_2d = ["uwnd", "vwnd", "swrad", "lwrad", "Tair", "qair", "rain"]
-        elif self.type == "bgc":
-            vars_2d = data.var_names.keys()
-        vars_3d = []
 
-        data_vars = super()._regrid_data(data, vars_2d, vars_3d, lon, lat)
+        data.apply_lateral_fill()
 
-        if self.type == "physics":
-            data_vars = super()._process_velocities(
-                data_vars, angle, "uwnd", "vwnd", interpolate=False
-            )
-            if self.correct_radiation:
-                correction_data = self._get_correction_data()
-                # choose same subdomain as forcing data so that we can use same mask
-                coords_correction = {
-                    correction_data.dim_names["latitude"]: data.ds[
-                        data.dim_names["latitude"]
-                    ],
-                    correction_data.dim_names["longitude"]: data.ds[
-                        data.dim_names["longitude"]
-                    ],
-                }
-                correction_data.choose_subdomain(coords_correction, straddle=straddle)
-                # apply mask from ERA5 data
-                if "mask" in data.var_names.keys():
-                    mask = data.ds["mask"]
-                    for var in correction_data.ds.data_vars:
-                        correction_data.ds[var] = xr.where(
-                            mask == 1, correction_data.ds[var], np.nan
-                        )
-                    correction_data.ds["mask"] = mask
-                vars_2d = ["swr_corr"]
-                vars_3d = []
-                # spatial interpolation
-                data_vars_corr = super()._regrid_data(
-                    correction_data, vars_2d, vars_3d, lon, lat
-                )
-                # temporal interpolation
-                corr_factor = interpolate_from_climatology(
-                    data_vars_corr["swr_corr"],
-                    correction_data.dim_names["time"],
-                    time=data_vars["swrad"].time,
+        variable_info = self._set_variable_info(data)
+        var_names = variable_info.keys()
+
+        processed_fields = {}
+        # lateral regridding
+        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]]
                 )
 
-                data_vars["swrad"] = data_vars["swrad"] * corr_factor
+        # rotation of velocities and interpolation to u/v points
+        if "uwnd" in variable_info and "vwnd" in variable_info:
+            processed_fields["uwnd"], processed_fields["vwnd"] = rotate_velocities(
+                processed_fields["uwnd"],
+                processed_fields["vwnd"],
+                target_coords["angle"],
+                interpolate=False,
+            )
 
-        object.__setattr__(data, "data_vars", data_vars)
+        # correct radiation
+        if self.type == "physics" and self.correct_radiation:
+            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(
@@ -162,13 +140,11 @@ class SurfaceForcing(ROMSToolsMixins):
         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)
 
@@ -231,15 +207,98 @@ class SurfaceForcing(ROMSToolsMixins):
 
         return correction_data
 
-    def _write_into_dataset(self, data, d_meta):
+    def _set_variable_info(self, data):
+        """Sets up a dictionary with metadata for variables based on the type of data
+        (physics or BGC).
+
+        The dictionary contains the following information:
+        - `location`: Where the variable resides in the grid (e.g., rho, u, or v points).
+        - `is_vector`: Whether the variable is part of a vector (True for velocity components like 'u' and 'v').
+        - `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').
+
+        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'.
+        """
+        default_info = {
+            "location": "rho",
+            "is_vector": False,
+            "vector_pair": None,
+            "is_3d": False,
+        }
+
+        # Define a dictionary for variable names and their associated information
+        if self.type == "physics":
+            variable_info = {
+                "swrad": default_info,
+                "lwrad": default_info,
+                "Tair": default_info,
+                "qair": default_info,
+                "rain": default_info,
+                "uwnd": {
+                    "location": "u",
+                    "is_vector": True,
+                    "vector_pair": "vwnd",
+                    "is_3d": False,
+                },
+                "vwnd": {
+                    "location": "v",
+                    "is_vector": True,
+                    "vector_pair": "uwnd",
+                    "is_3d": False,
+                },
+            }
+        elif self.type == "bgc":
+            variable_info = {}
+            for var in data.var_names.keys():
+                variable_info[var] = default_info
+
+        return variable_info
+
+    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
+        coords_correction = {
+            correction_data.dim_names["latitude"]: data.ds[data.dim_names["latitude"]],
+            correction_data.dim_names["longitude"]: data.ds[
+                data.dim_names["longitude"]
+            ],
+        }
+        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_regrid = LateralRegrid(self.target_coords, correction_data.dim_names)
+        corr_factor = lateral_regrid.apply(
+            correction_data.ds[correction_data.var_names["swr_corr"]]
+        )
+
+        # temporal interpolation
+        corr_factor = interpolate_from_climatology(
+            corr_factor,
+            correction_data.dim_names["time"],
+            time=processed_fields["swrad"].time,
+        )
+
+        processed_fields["swrad"] = processed_fields["swrad"] * corr_factor
+
+        return processed_fields
+
+    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"})
@@ -305,11 +364,36 @@ class SurfaceForcing(ROMSToolsMixins):
             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:
@@ -333,12 +417,12 @@ class SurfaceForcing(ROMSToolsMixins):
 
         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.
@@ -367,7 +451,7 @@ class SurfaceForcing(ROMSToolsMixins):
         Raises
         ------
         ValueError
-            If the specified varname is not found in dataset.
+            If the specified var_name is not found in dataset.
 
 
         Examples
@@ -375,10 +459,10 @@ class SurfaceForcing(ROMSToolsMixins):
         >>> 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
@@ -388,17 +472,19 @@ class SurfaceForcing(ROMSToolsMixins):
         else:
             field = field.where(self.grid.ds.mask_rho)
 
-        field = field.assign_coords({"lon": self.target_lon, "lat": self.target_lat})
+        field = field.assign_coords(
+            {"lon": self.target_coords["lon"], "lat": self.target_coords["lat"]}
+        )
 
         # 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

@@ -16,15 +16,18 @@ from roms_tools.setup.utils import (
     interpolate_from_rho_to_v,
     get_variable_metadata,
     save_datasets,
+    get_target_coords,
+    rotate_velocities,
+    get_vector_pairs,
 )
-from roms_tools.setup.mixins import ROMSToolsMixins
+from roms_tools.setup.regrid import LateralRegrid
 import matplotlib.pyplot as plt
 from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
-class TidalForcing(ROMSToolsMixins):
-    """Represents tidal forcing data used in ocean modeling.
+class TidalForcing:
+    """Represents tidal forcing for ROMS.
 
     Parameters
     ----------
@@ -68,67 +71,69 @@ class TidalForcing(ROMSToolsMixins):
     def __post_init__(self):
 
         self._input_checks()
-        lon, lat, angle, straddle = super()._get_target_lon_lat()
+        target_coords = get_target_coords(self.grid)
 
         data = self._get_data()
-
         data.check_number_constituents(self.ntides)
         data.choose_subdomain(
-            latitude_range=[lat.min().values, lat.max().values],
-            longitude_range=[lon.min().values, lon.max().values],
-            margin=2,
-            straddle=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)
 
-        vars_2d = [
-            "ssh_Re",
-            "ssh_Im",
-            "pot_Re",
-            "pot_Im",
-            "u_Re",
-            "u_Im",
-            "v_Re",
-            "v_Im",
-        ]
-        vars_3d = []
-
-        data_vars = super()._regrid_data(data, vars_2d, vars_3d, lon, lat)
-
-        data_vars = super()._process_velocities(
-            data_vars, angle, "u_Re", "v_Re", interpolate=False
-        )
-        data_vars = super()._process_velocities(
-            data_vars, angle, "u_Im", "v_Im", 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
-
-        # Interpolate from rho- to velocity points
+        data.apply_lateral_fill()
+
+        variable_info = self._set_variable_info()
+        var_names = variable_info.keys()
+
+        processed_fields = {}
+        # lateral regridding
+        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 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 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)
 
@@ -147,15 +152,71 @@ class TidalForcing(ROMSToolsMixins):
             raise ValueError('Only "TPXO" is a valid option for source["name"].')
         return data
 
-    def _write_into_dataset(self, data_vars, d_meta):
+    def _set_variable_info(self):
+        """Sets up a dictionary with metadata for variables based on the type.
+
+        The dictionary contains the following information:
+        - `location`: Where the variable resides in the grid (e.g., rho, u, or v points).
+        - `is_vector`: Whether the variable is part of a vector (True for velocity components like 'u' and 'v').
+        - `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').
+
+        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'.
+        """
+        default_info = {
+            "location": "rho",
+            "is_vector": False,
+            "vector_pair": None,
+            "is_3d": False,
+        }
+
+        # Define a dictionary for variable names and their associated information
+        variable_info = {
+            "ssh_Re": default_info,
+            "ssh_Im": default_info,
+            "pot_Re": default_info,
+            "pot_Im": default_info,
+            "u_Re": {
+                "location": "u",
+                "is_vector": True,
+                "vector_pair": "v_Re",
+                "is_3d": False,
+            },
+            "v_Re": {
+                "location": "v",
+                "is_vector": True,
+                "vector_pair": "u_Re",
+                "is_3d": False,
+            },
+            "u_Im": {
+                "location": "u",
+                "is_vector": True,
+                "vector_pair": "v_Im",
+                "is_3d": False,
+            },
+            "v_Im": {
+                "location": "v",
+                "is_vector": True,
+                "vector_pair": "u_Im",
+                "is_3d": False,
+            },
+        }
+
+        return variable_info
+
+    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"])
 
@@ -178,12 +239,35 @@ class TidalForcing(ROMSToolsMixins):
 
         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.
@@ -216,7 +300,7 @@ class TidalForcing(ROMSToolsMixins):
         >>> 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(
@@ -237,7 +321,7 @@ class TidalForcing(ROMSToolsMixins):
         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
@@ -457,6 +541,9 @@ class TidalForcing(ROMSToolsMixins):
 
         # Update var_names dictionary
         var_names = {**data.var_names, "pot_Re": "pot_Re", "pot_Im": "pot_Im"}
+        var_names.pop("sal_Re", None)  # Remove "sal_Re" if it exists
+        var_names.pop("sal_Im", None)  # Remove "sal_Im" if it exists
+
         object.__setattr__(data, "var_names", var_names)