roms-tools 1.1.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

roms_tools/setup/tides.py

@@ -3,7 +3,7 @@ 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
@@ -12,14 +12,19 @@ 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)
-class TidalForcing:
+class TidalForcing(ROMSToolsMixins):
     """
     Represents tidal forcing data used in ocean modeling.
 
@@ -27,16 +32,19 @@ class TidalForcing:
     ----------
     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
     ----------
@@ -51,38 +59,22 @@ class TidalForcing:
     """
 
     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)
 
     def __post_init__(self):
-        if "name" not in self.source.keys():
-            raise ValueError("`source` must include a 'name'.")
-        if "path" not in self.source.keys():
-            raise ValueError("`source` must include a 'path'.")
-        if self.source["name"] == "TPXO":
-            data = TPXODataset(filename=self.source["path"])
-        else:
-            raise ValueError('Only "TPXO" is a valid option for source["name"].')
+
+        self._input_checks()
+        lon, lat, angle, straddle = super().get_target_lon_lat()
+
+        data = self._get_data()
 
         data.check_number_constituents(self.ntides)
-        # operate on longitudes between -180 and 180 unless ROMS domain lies at least 5 degrees in longitude away from Greenwich meridian
-        lon = self.grid.ds.lon_rho
-        lat = self.grid.ds.lat_rho
-        angle = self.grid.ds.angle
-
-        lon = xr.where(lon > 180, lon - 360, lon)
-        straddle = True
-        if not self.grid.straddle and abs(lon).min() > 5:
-            lon = xr.where(lon < 0, lon + 360, lon)
-            straddle = False
-
-        # Restrict data to relevant subdomain to achieve better performance and to avoid discontinuous longitudes introduced by converting
-        # to a different longitude range (+- 360 degrees). Discontinues longitudes can lead to artifacts in the interpolation process that
-        # would not be detected by the nan_check function.
         data.choose_subdomain(
             latitude_range=[lat.min().values, lat.max().values],
             longitude_range=[lon.min().values, lon.max().values],
@@ -121,56 +113,69 @@ class TidalForcing:
                 method="linear",
             )
 
-        # Rotate to grid orientation
-        u_Re = data_vars["u_Re"] * np.cos(angle) + data_vars["v_Re"] * np.sin(angle)
-        v_Re = data_vars["v_Re"] * np.cos(angle) - data_vars["u_Re"] * np.sin(angle)
-        u_Im = data_vars["u_Im"] * np.cos(angle) + data_vars["v_Im"] * np.sin(angle)
-        v_Im = data_vars["v_Im"] * np.cos(angle) - data_vars["u_Im"] * np.sin(angle)
+        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
-        u_Re = u_Re / self.grid.ds.h
-        v_Re = v_Re / self.grid.ds.h
-        u_Im = u_Im / self.grid.ds.h
-        v_Im = v_Im / self.grid.ds.h
+        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
-        u_Re = interpolate_from_rho_to_u(u_Re)
-        v_Re = interpolate_from_rho_to_v(v_Re)
-        u_Im = interpolate_from_rho_to_u(u_Im)
-        v_Im = interpolate_from_rho_to_v(v_Im)
+        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"]:
+            data_vars[vname] = interpolate_from_rho_to_v(data_vars[vname])
+
+        d_meta = get_variable_metadata()
+        ds = self._write_into_dataset(data_vars, d_meta)
+        ds["omega"] = tides["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):
+
+        if "name" not in self.source.keys():
+            raise ValueError("`source` must include a 'name'.")
+        if "path" not in self.source.keys():
+            raise ValueError("`source` must include a 'path'.")
+
+    def _get_data(self):
+
+        if self.source["name"] == "TPXO":
+            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
+
+    def _write_into_dataset(self, data_vars, d_meta):
 
         # save in new dataset
         ds = xr.Dataset()
 
-        # ds["omega"] = tides["omega"]
-
-        ds["ssh_Re"] = data_vars["ssh_Re"].astype(np.float32)
-        ds["ssh_Im"] = data_vars["ssh_Im"].astype(np.float32)
-        ds["ssh_Re"].attrs["long_name"] = "Tidal elevation, real part"
-        ds["ssh_Im"].attrs["long_name"] = "Tidal elevation, complex part"
-        ds["ssh_Re"].attrs["units"] = "m"
-        ds["ssh_Im"].attrs["units"] = "m"
-
-        ds["pot_Re"] = data_vars["pot_Re"].astype(np.float32)
-        ds["pot_Im"] = data_vars["pot_Im"].astype(np.float32)
-        ds["pot_Re"].attrs["long_name"] = "Tidal potential, real part"
-        ds["pot_Im"].attrs["long_name"] = "Tidal potential, complex part"
-        ds["pot_Re"].attrs["units"] = "m"
-        ds["pot_Im"].attrs["units"] = "m"
-
-        ds["u_Re"] = u_Re.astype(np.float32)
-        ds["u_Im"] = u_Im.astype(np.float32)
-        ds["u_Re"].attrs["long_name"] = "Tidal velocity in x-direction, real part"
-        ds["u_Im"].attrs["long_name"] = "Tidal velocity in x-direction, complex part"
-        ds["u_Re"].attrs["units"] = "m/s"
-        ds["u_Im"].attrs["units"] = "m/s"
-
-        ds["v_Re"] = v_Re.astype(np.float32)
-        ds["v_Im"] = v_Im.astype(np.float32)
-        ds["v_Re"].attrs["long_name"] = "Tidal velocity in y-direction, real part"
-        ds["v_Im"].attrs["long_name"] = "Tidal velocity in y-direction, complex part"
-        ds["v_Re"].attrs["units"] = "m/s"
-        ds["v_Im"].attrs["units"] = "m/s"
+        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"]
+
+        ds = ds.drop_vars(["lat_rho", "lon_rho"])
+
+        return ds
+
+    def _add_global_metadata(self, ds):
 
         ds.attrs["title"] = "ROMS tidal forcing created by ROMS-Tools"
         # Include the version of roms-tools
@@ -185,10 +190,7 @@ class TidalForcing:
         ds.attrs["model_reference_date"] = str(self.model_reference_date)
         ds.attrs["allan_factor"] = self.allan_factor
 
-        object.__setattr__(self, "ds", ds)
-
-        for var in ["ssh_Re", "u_Re", "v_Im"]:
-            nan_check(self.ds[var].isel(ntides=0), self.grid.ds.mask_rho)
+        return ds
 
     def plot(self, varname, ntides=0) -> None:
         """
@@ -228,6 +230,25 @@ class TidalForcing:
         """
 
         field = self.ds[varname].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(
+                {"lon": self.grid.ds.lon_rho, "lat": self.grid.ds.lat_rho}
+            )
+
+        elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
+            field = field.where(self.grid.ds.mask_u)
+            field = field.assign_coords(
+                {"lon": self.grid.ds.lon_u, "lat": self.grid.ds.lat_u}
+            )
+
+        elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
+            field = field.where(self.grid.ds.mask_v)
+            field = field.assign_coords(
+                {"lon": self.grid.ds.lon_v, "lat": self.grid.ds.lat_v}
+            )
+        else:
+            ValueError("provided field does not have two horizontal dimension")
 
         title = "%s, ntides = %i" % (field.long_name, self.ds[varname].ntides[ntides])
 
@@ -247,25 +268,63 @@ class TidalForcing:
             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)
@@ -295,29 +354,34 @@ class TidalForcing:
         # 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
@@ -346,7 +410,7 @@ class TidalForcing:
         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):
 

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"
     )