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

roms_tools/setup/tides.py

@@ -16,26 +16,34 @@ 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.fill import _lateral_fill
+from roms_tools.setup.regrid import _lateral_regrid
 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
     ----------
     grid : Grid
         The grid object representing the ROMS grid associated with the tidal forcing data.
     source : Dict[str, Union[str, Path, List[Union[str, Path]]]]
-        Dictionary specifying the source of the tidal data:
-        - "name" (str): Name of the data source (e.g., "TPXO").
-        - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file(s). Can be a single string (with or without wildcards),
-          a single Path object, or a list of strings or Path objects containing multiple files.
+        Dictionary specifying the source of the tidal data. Keys include:
+
+          - "name" (str): Name of the data source (e.g., "TPXO").
+          - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file(s). This can be:
+
+            - A single string (with or without wildcards).
+            - A single Path object.
+            - A list of strings or Path objects containing multiple files.
+
     ntides : int, optional
         Number of constituents to consider. Maximum number is 14. Default is 10.
     allan_factor : float, optional
@@ -45,11 +53,6 @@ class TidalForcing(ROMSToolsMixins):
     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.
 
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The xarray Dataset containing the tidal forcing data.
-
     Examples
     --------
     >>> tidal_forcing = TidalForcing(
@@ -69,49 +72,58 @@ 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],
+            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=straddle,
+            straddle=target_coords["straddle"],
         )
-
         # 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
+        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_vars = _lateral_fill(data_vars, data)
+
+        # lateral regridding
+        var_names = variable_info.keys()
+        data_vars = _lateral_regrid(
+            data, target_coords["lon"], target_coords["lat"], data_vars, var_names
         )
 
-        # Convert to barotropic velocity
+        # 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],
+                    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
 
-        # Interpolate from rho- to velocity points
+        # interpolate from rho- to velocity points
         for uname in ["u_Re", "u_Im"]:
             data_vars[uname] = interpolate_from_rho_to_u(data_vars[uname])
         for vname in ["v_Re", "v_Im"]:
@@ -148,6 +160,62 @@ class TidalForcing(ROMSToolsMixins):
             raise ValueError('Only "TPXO" is a valid option for source["name"].')
         return data
 
+    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, data_vars, d_meta):
 
         # save in new dataset
@@ -180,13 +248,13 @@ class TidalForcing(ROMSToolsMixins):
         return ds
 
     def plot(self, varname, ntides=0) -> None:
-        """
-        Plot the specified tidal forcing variable for a given tidal constituent.
+        """Plot the specified tidal forcing variable for a given tidal constituent.
 
         Parameters
         ----------
         varname : str
             The tidal forcing variable to plot. Options include:
+
             - "ssh_Re": Real part of tidal elevation.
             - "ssh_Im": Imaginary part of tidal elevation.
             - "pot_Re": Real part of tidal potential.
@@ -195,6 +263,7 @@ class TidalForcing(ROMSToolsMixins):
             - "u_Im": Imaginary part of tidal velocity in the x-direction.
             - "v_Re": Real part of tidal velocity in the y-direction.
             - "v_Im": Imaginary part of tidal velocity in the y-direction.
+
         ntides : int, optional
             The index of the tidal constituent to plot. Default is 0, which corresponds
             to the first constituent.
@@ -258,17 +327,19 @@ class TidalForcing(ROMSToolsMixins):
     def save(
         self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
     ) -> None:
-        """
-        Save the tidal forcing information to a netCDF4 file.
+        """Save the tidal forcing information to a netCDF4 file.
 
         This method supports saving the dataset in two modes:
 
-        1. **Single File Mode (default)**:
-           - If both `np_eta` and `np_xi` are `None`, the entire dataset is saved as a single file at the specified `filepath.nc`.
+          1. **Single File Mode (default)**:
+
+            If both `np_eta` and `np_xi` are `None`, the entire dataset is saved as a single netCDF4 file
+            with the base filename specified by `filepath.nc`.
 
-        2. **Partitioned Mode**:
-           - If either `np_eta` or `np_xi` is specified, the dataset is divided into spatial tiles along the eta-axis and xi-axis.
-           - The files are saved as `filepath.0.nc`, `filepath.1.nc`, ..., where the numbering corresponds to the partition index.
+          2. **Partitioned Mode**:
+
+            - If either `np_eta` or `np_xi` is specified, the dataset is divided into spatial tiles along the eta-axis and xi-axis.
+            - Each spatial tile is saved as a separate netCDF4 file.
 
         Parameters
         ----------
@@ -302,8 +373,8 @@ class TidalForcing(ROMSToolsMixins):
         return saved_filenames
 
     def to_yaml(self, filepath: Union[str, Path]) -> None:
-        """
-        Export the parameters of the class to a YAML file, including the version of roms-tools.
+        """Export the parameters of the class to a YAML file, including the version of
+        roms-tools.
 
         Parameters
         ----------
@@ -351,8 +422,7 @@ class TidalForcing(ROMSToolsMixins):
     def from_yaml(
         cls, filepath: Union[str, Path], use_dask: bool = False
     ) -> "TidalForcing":
-        """
-        Create an instance of the TidalForcing class from a YAML file.
+        """Create an instance of the TidalForcing class from a YAML file.
 
         Parameters
         ----------
@@ -400,10 +470,10 @@ class TidalForcing(ROMSToolsMixins):
         return cls(grid=grid, **tidal_forcing_params, use_dask=use_dask)
 
     def _correct_tides(self, data):
-        """
-        Apply tidal corrections to the dataset.
-        This method corrects the dataset for equilibrium tides, self-attraction and loading (SAL) effects, and
-        adjusts phases and amplitudes of tidal elevations and transports using Egbert's correction.
+        """Apply tidal corrections to the dataset. This method corrects the dataset for
+        equilibrium tides, self-attraction and loading (SAL) effects, and adjusts phases
+        and amplitudes of tidal elevations and transports using Egbert's correction.
+
         Parameters
         ----------
         data : Dataset
@@ -456,12 +526,14 @@ 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)
 
 
 def modified_julian_days(year, month, day, hour=0):
-    """
-    Calculate the Modified Julian Day (MJD) for a given date and time.
+    """Calculate the Modified Julian Day (MJD) for a given date and time.
 
     The Modified Julian Day (MJD) is a modified Julian day count starting from
     November 17, 1858 AD. It is commonly used in astronomy and geodesy.
@@ -519,9 +591,8 @@ def modified_julian_days(year, month, day, hour=0):
 
 
 def egbert_correction(date):
-    """
-    Correct phases and amplitudes for real-time runs using parts of the
-    post-processing code from Egbert's & Erofeeva's (OSU) TPXO model.
+    """Correct phases and amplitudes for real-time runs using parts of the post-
+    processing code from Egbert's & Erofeeva's (OSU) TPXO model.
 
     Parameters
     ----------
@@ -541,7 +612,6 @@ def egbert_correction(date):
     ----------
     - Egbert, G.D., and S.Y. Erofeeva. "Efficient inverse modeling of barotropic ocean
       tides." Journal of Atmospheric and Oceanic Technology 19, no. 2 (2002): 183-204.
-
     """
 
     year = date.year
@@ -676,8 +746,7 @@ def egbert_correction(date):
 
 
 def compute_equilibrium_tide(lon, lat):
-    """
-    Compute equilibrium tide for given longitudes and latitudes.
+    """Compute equilibrium tide for given longitudes and latitudes.
 
     Parameters
     ----------
@@ -699,7 +768,6 @@ def compute_equilibrium_tide(lon, lat):
         - 2: semidiurnal
         - 1: diurnal
         - 0: long-term
-
     """
 
     # Amplitudes and elasticity factors for 15 tidal constituents

roms_tools/setup/topography.py

@@ -12,8 +12,8 @@ from itertools import count
 def _add_topography_and_mask(
     ds, topography_source, hmin, smooth_factor=8.0, rmax=0.2
 ) -> xr.Dataset:
-    """
-    Adds topography and a land/water mask to the dataset based on the provided topography source.
+    """Adds topography and a land/water mask to the dataset based on the provided
+    topography source.
 
     This function performs the following operations:
     1. Interpolates topography data onto the desired grid.
@@ -86,9 +86,8 @@ def _add_topography_and_mask(
 
 
 def _make_raw_topography(lon, lat, topography_source) -> np.ndarray:
-    """
-    Given a grid of (lon, lat) points, fetch the topography file and interpolate height values onto the desired grid.
-    """
+    """Given a grid of (lon, lat) points, fetch the topography file and interpolate
+    height values onto the desired grid."""
 
     topo_ds = fetch_topo(topography_source)
 
@@ -150,9 +149,7 @@ def _smooth_topography_globally(hraw, factor) -> xr.DataArray:
 
 
 def _fill_enclosed_basins(mask) -> np.ndarray:
-    """
-    Fills in enclosed basins with land
-    """
+    """Fills in enclosed basins with land."""
 
     # Label connected regions in the mask
     reg, nreg = label(mask)
@@ -174,9 +171,7 @@ def _fill_enclosed_basins(mask) -> np.ndarray:
 
 
 def _smooth_topography_locally(h, hmin=5, rmax=0.2):
-    """
-    Smoothes topography locally to satisfy r < rmax
-    """
+    """Smoothes topography locally to satisfy r < rmax."""
     # Compute rmax_log
     if rmax > 0.0:
         rmax_log = np.log((1.0 + rmax * 0.9) / (1.0 - rmax * 0.9))
@@ -254,8 +249,7 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
 
 
 def _handle_boundaries(field):
-    """
-    Adjust the boundaries of a 2D field by copying values from adjacent cells.
+    """Adjust the boundaries of a 2D field by copying values from adjacent cells.
 
     Parameters
     ----------
@@ -267,7 +261,6 @@ def _handle_boundaries(field):
     -------
     field : numpy.ndarray or xarray.DataArray
         The input field with adjusted boundary values.
-
     """
 
     field[0, :] = field[1, :]
@@ -279,9 +272,8 @@ def _handle_boundaries(field):
 
 
 def _compute_rfactor(h):
-    """
-    Computes slope parameter (or r-factor) r = |Delta h| / 2h in both horizontal grid directions.
-    """
+    """Computes slope parameter (or r-factor) r = |Delta h| / 2h in both horizontal grid
+    directions."""
     # compute r_{i-1/2} = |h_i - h_{i-1}| / (h_i + h_{i+1})
     r_eta = np.abs(h.diff("eta_rho")) / (h + h.shift(eta_rho=1)).isel(
         eta_rho=slice(1, None)