roms-tools 1.2.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

roms_tools/setup/tides.py

@@ -3,20 +3,23 @@ import xarray as xr
 import numpy as np
 import yaml
 import importlib.metadata
-from typing import Dict, Union
+from typing import Dict, Union, List
 
 from dataclasses import dataclass, field, asdict
 from roms_tools.setup.grid import Grid
 from roms_tools.setup.plot import _plot
-from roms_tools.setup.fill import fill_and_interpolate
 from roms_tools.setup.datasets import TPXODataset
 from roms_tools.setup.utils import (
     nan_check,
+    substitute_nans_by_fillvalue,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
+    get_variable_metadata,
+    save_datasets,
 )
 from roms_tools.setup.mixins import ROMSToolsMixins
 import matplotlib.pyplot as plt
+from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -28,16 +31,19 @@ class TidalForcing(ROMSToolsMixins):
     ----------
     grid : Grid
         The grid object representing the ROMS grid associated with the tidal forcing data.
-    source : Dict[str, Union[str, None]]
+    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" (str): Path to the tidal data file. Can contain wildcards.
+        - "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.
     ntides : int, optional
         Number of constituents to consider. Maximum number is 14. Default is 10.
     allan_factor : float, optional
         The Allan factor used in tidal model computation. Default is 2.0.
     model_reference_date : datetime, optional
         The reference date for the ROMS simulation. Default is datetime(2000, 1, 1).
+    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
     ----------
@@ -52,10 +58,11 @@ class TidalForcing(ROMSToolsMixins):
     """
 
     grid: Grid
-    source: Dict[str, Union[str, None]]
+    source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
     ntides: int = 10
     allan_factor: float = 2.0
     model_reference_date: datetime = datetime(2000, 1, 1)
+    use_dask: bool = False
 
     ds: xr.Dataset = field(init=False, repr=False)
 
@@ -74,17 +81,12 @@ class TidalForcing(ROMSToolsMixins):
             straddle=straddle,
         )
 
-        tides = self._get_corrected_tides(data)
-
         # select desired number of constituents
-        for k in tides.keys():
-            tides[k] = tides[k].isel(ntides=slice(None, self.ntides))
+        object.__setattr__(data, "ds", data.ds.isel(ntides=slice(None, self.ntides)))
 
-        # interpolate onto desired grid
-        coords = {"latitude": lat, "longitude": lon}
-        mask = xr.where(data.ds.depth > 0, 1, 0)
+        self._correct_tides(data)
 
-        varnames = [
+        vars_2d = [
             "ssh_Re",
             "ssh_Im",
             "pot_Re",
@@ -94,16 +96,9 @@ class TidalForcing(ROMSToolsMixins):
             "v_Re",
             "v_Im",
         ]
-        data_vars = {}
-
-        for var in varnames:
-            data_vars[var] = fill_and_interpolate(
-                tides[var],
-                mask,
-                list(coords.keys()),
-                coords,
-                method="linear",
-            )
+        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
@@ -122,15 +117,20 @@ class TidalForcing(ROMSToolsMixins):
         for vname in ["v_Re", "v_Im"]:
             data_vars[vname] = interpolate_from_rho_to_v(data_vars[vname])
 
-        d_meta = super().get_variable_metadata()
+        d_meta = get_variable_metadata()
         ds = self._write_into_dataset(data_vars, d_meta)
-        ds["omega"] = tides["omega"]
+        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)
 
+        # 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])
+
         object.__setattr__(self, "ds", ds)
 
     def _input_checks(self):
@@ -143,7 +143,7 @@ class TidalForcing(ROMSToolsMixins):
     def _get_data(self):
 
         if self.source["name"] == "TPXO":
-            data = TPXODataset(filename=self.source["path"])
+            data = TPXODataset(filename=self.source["path"], use_dask=self.use_dask)
         else:
             raise ValueError('Only "TPXO" is a valid option for source["name"].')
         return data
@@ -255,25 +255,63 @@ class TidalForcing(ROMSToolsMixins):
             title=title,
         )
 
-    def save(self, filepath: str) -> None:
+    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.
 
+        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`.
+
+        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.
+
         Parameters
         ----------
-        filepath
+        filepath : Union[str, Path]
+            The base path or filename where the dataset should be saved.
+        np_eta : int, optional
+            The number of partitions along the `eta` direction. If `None`, no spatial partitioning is performed.
+        np_xi : int, optional
+            The number of partitions along the `xi` direction. If `None`, no spatial partitioning is performed.
+
+        Returns
+        -------
+        List[Path]
+            A list of Path objects for the filenames that were saved.
         """
-        self.ds.to_netcdf(filepath)
 
-    def to_yaml(self, filepath: str) -> None:
+        # Ensure filepath is a Path object
+        filepath = Path(filepath)
+
+        # Remove ".nc" suffix if present
+        if filepath.suffix == ".nc":
+            filepath = filepath.with_suffix("")
+
+        dataset_list = [self.ds.load()]
+        output_filenames = [str(filepath)]
+
+        saved_filenames = save_datasets(
+            dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
+        )
+
+        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.
 
         Parameters
         ----------
-        filepath : str
+        filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
+        filepath = Path(filepath)
+
         grid_data = asdict(self.grid)
         grid_data.pop("ds", None)  # Exclude non-serializable fields
         grid_data.pop("straddle", None)
@@ -303,29 +341,34 @@ class TidalForcing(ROMSToolsMixins):
         # Combine both sections
         yaml_data = {**grid_yaml_data, **tidal_forcing_data}
 
-        with open(filepath, "w") as file:
+        with filepath.open("w") as file:
             # Write header
             file.write(header)
             # Write YAML data
             yaml.dump(yaml_data, file, default_flow_style=False)
 
     @classmethod
-    def from_yaml(cls, filepath: str) -> "TidalForcing":
+    def from_yaml(
+        cls, filepath: Union[str, Path], use_dask: bool = False
+    ) -> "TidalForcing":
         """
         Create an instance of the TidalForcing class from a YAML file.
 
         Parameters
         ----------
-        filepath : str
+        filepath : Union[str, Path]
             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.
 
         Returns
         -------
         TidalForcing
             An instance of the TidalForcing class.
         """
+        filepath = Path(filepath)
         # Read the entire file content
-        with open(filepath, "r") as file:
+        with filepath.open("r") as file:
             file_content = file.read()
 
         # Split the content into YAML documents
@@ -354,15 +397,30 @@ class TidalForcing(ROMSToolsMixins):
         grid = Grid.from_yaml(filepath)
 
         # Create and return an instance of TidalForcing
-        return cls(grid=grid, **tidal_forcing_params)
+        return cls(grid=grid, **tidal_forcing_params, use_dask=use_dask)
 
-    def _get_corrected_tides(self, data):
+    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.
+        Parameters
+        ----------
+        data : Dataset
+            The dataset containing tidal data, including variables for sea surface height (ssh), zonal and meridional
+            currents (u, v), and self-attraction and loading corrections (sal).
+        Returns
+        -------
+        None
+            The dataset is modified in-place with corrected real and imaginary components for ssh, u, v, and the
+            potential field ('pot_Re', 'pot_Im').
+        """
 
         # Get equilibrium tides
         tpc = compute_equilibrium_tide(
             data.ds[data.dim_names["longitude"]], data.ds[data.dim_names["latitude"]]
         )
-        tpc = tpc.isel(**{data.dim_names["ntides"]: data.ds[data.dim_names["ntides"]]})
+        tpc = tpc.isel(ntides=data.ds["ntides"])
         # Correct for SAL
         tsc = self.allan_factor * (
             data.ds[data.var_names["sal_Re"]] + 1j * data.ds[data.var_names["sal_Im"]]
@@ -376,9 +434,9 @@ class TidalForcing(ROMSToolsMixins):
 
         # Apply correction for phases and amplitudes
         pf, pu, aa = egbert_correction(self.model_reference_date)
-        pf = pf.isel(**{data.dim_names["ntides"]: data.ds[data.dim_names["ntides"]]})
-        pu = pu.isel(**{data.dim_names["ntides"]: data.ds[data.dim_names["ntides"]]})
-        aa = aa.isel(**{data.dim_names["ntides"]: data.ds[data.dim_names["ntides"]]})
+        pf = pf.isel(ntides=data.ds["ntides"])
+        pu = pu.isel(ntides=data.ds["ntides"])
+        aa = aa.isel(ntides=data.ds["ntides"])
 
         dt = (self.model_reference_date - data.reference_date).days * 3600 * 24
 
@@ -387,22 +445,18 @@ class TidalForcing(ROMSToolsMixins):
         tvc = pf * tvc * np.exp(1j * (data.ds["omega"] * dt + pu + aa))
         tpc = pf * tpc * np.exp(1j * (data.ds["omega"] * dt + pu + aa))
 
-        tides = {
-            "ssh_Re": thc.real,
-            "ssh_Im": thc.imag,
-            "u_Re": tuc.real,
-            "u_Im": tuc.imag,
-            "v_Re": tvc.real,
-            "v_Im": tvc.imag,
-            "pot_Re": tpc.real,
-            "pot_Im": tpc.imag,
-            "omega": data.ds["omega"],
-        }
-
-        for k in tides.keys():
-            tides[k] = tides[k].rename({data.dim_names["ntides"]: "ntides"})
+        data.ds[data.var_names["ssh_Re"]] = thc.real
+        data.ds[data.var_names["ssh_Im"]] = thc.imag
+        data.ds[data.var_names["u_Re"]] = tuc.real
+        data.ds[data.var_names["u_Im"]] = tuc.imag
+        data.ds[data.var_names["v_Re"]] = tvc.real
+        data.ds[data.var_names["v_Im"]] = tvc.imag
+        data.ds["pot_Re"] = tpc.real
+        data.ds["pot_Im"] = tpc.imag
 
-        return tides
+        # Update var_names dictionary
+        var_names = {**data.var_names, "pot_Re": "pot_Re", "pot_Im": "pot_Im"}
+        object.__setattr__(data, "var_names", var_names)
 
 
 def modified_julian_days(year, month, day, hour=0):
@@ -555,7 +609,7 @@ def egbert_correction(date):
     pf[12] = pftmp**2  # Ms4
     pf[13] = pftmp  # 2n2
     pf[14] = 1.0  # S1
-    pf = xr.DataArray(pf, dims="nc")
+    pf = xr.DataArray(pf, dims="ntides")
 
     putmp = (
         np.arctan(
@@ -591,7 +645,7 @@ def egbert_correction(date):
     pu[12] = putmp  # Ms4
     pu[13] = putmp  # 2n2
     pu[14] = 0.0  # S1
-    pu = xr.DataArray(pu, dims="nc")
+    pu = xr.DataArray(pu, dims="ntides")
     # convert from degrees to radians
     pu = pu * rad
 
@@ -615,7 +669,7 @@ def egbert_correction(date):
                 0.0,  # S1
             ]
         ),
-        dims="nc",
+        dims="ntides",
     )
 
     return pf, pu, aa
@@ -669,7 +723,7 @@ def compute_equilibrium_tide(lon, lat):
                 0.000764,  # S1
             ]
         ),
-        dims="nc",
+        dims="ntides",
     )
     B = xr.DataArray(
         data=np.array(
@@ -691,12 +745,12 @@ def compute_equilibrium_tide(lon, lat):
                 0.693,  # S1
             ]
         ),
-        dims="nc",
+        dims="ntides",
     )
 
     # types: 2 = semidiurnal, 1 = diurnal, 0 = long-term
     ityp = xr.DataArray(
-        data=np.array([2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0, 2, 1]), dims="nc"
+        data=np.array([2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0, 2, 1]), dims="ntides"
     )
 
     d2r = np.pi / 180

roms_tools/setup/topography.py

@@ -57,7 +57,7 @@ def _add_topography_and_mask(
     mask = xr.where(hraw > 0, 1.0, 0.0)
 
     # smooth topography domain-wide with Gaussian kernel to avoid grid scale instabilities
-    hraw = _smooth_topography_globally(hraw, mask, smooth_factor)
+    hraw = _smooth_topography_globally(hraw, smooth_factor)
 
     # fill enclosed basins with land
     mask = _fill_enclosed_basins(mask.values)
@@ -115,10 +115,12 @@ def _make_raw_topography(lon, lat, topography_source) -> np.ndarray:
     return hraw
 
 
-def _smooth_topography_globally(hraw, wet_mask, factor) -> xr.DataArray:
+def _smooth_topography_globally(hraw, factor) -> xr.DataArray:
     # since GCM-Filters assumes periodic domain, we extend the domain by one grid cell in each dimension
     # and set that margin to land
-    margin_mask = xr.concat([wet_mask, 0 * wet_mask.isel(eta_rho=-1)], dim="eta_rho")
+
+    mask = xr.ones_like(hraw)
+    margin_mask = xr.concat([mask, 0 * mask.isel(eta_rho=-1)], dim="eta_rho")
     margin_mask = xr.concat(
         [margin_mask, 0 * margin_mask.isel(xi_rho=-1)], dim="xi_rho"
     )