roms-tools 2.7.0__py3-none-any.whl → 3.1.0__py3-none-any.whl

roms_tools/setup/tides.py

@@ -1,28 +1,29 @@
-from datetime import datetime
-import xarray as xr
-import numpy as np
-from typing import Dict, Union, List
 import importlib.metadata
-import matplotlib.pyplot as plt
-from pathlib import Path
 from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import numpy as np
+import xarray as xr
+
 from roms_tools import Grid
-from roms_tools.plot import _plot
+from roms_tools.plot import plot
 from roms_tools.regrid import LateralRegridToROMS
-from roms_tools.utils import save_datasets
 from roms_tools.setup.datasets import TPXOManager
 from roms_tools.setup.utils import (
-    nan_check,
-    substitute_nans_by_fillvalue,
+    from_yaml,
+    get_target_coords,
+    get_variable_metadata,
+    get_vector_pairs,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
-    get_variable_metadata,
-    get_target_coords,
+    nan_check,
     rotate_velocities,
-    get_vector_pairs,
-    _to_yaml,
-    _from_yaml,
+    substitute_nans_by_fillvalue,
+    to_dict,
+    write_to_yaml,
 )
+from roms_tools.utils import save_datasets
 
 
 @dataclass(kw_only=True)
@@ -32,7 +33,7 @@ class TidalForcing:
     Parameters
     ----------
     grid : Grid
-        The grid object representing the ROMS grid associated with the tidal forcing data.
+        Object representing the grid information.
     source : Dict[str, Union[str, Path, Dict[str, Union[str, Path]]]]
         Dictionary specifying the source of the tidal data. Keys include:
 
@@ -78,16 +79,23 @@ class TidalForcing:
     """
 
     grid: Grid
-    source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
+    """Object representing the grid information."""
+    source: dict[str, str | Path | list[str | Path]]
+    """Dictionary specifying the source of the tidal data."""
     ntides: int = 10
+    """Number of constituents to consider."""
     model_reference_date: datetime = datetime(2000, 1, 1)
+    """The reference date for the ROMS simulation."""
     use_dask: bool = False
+    """Whether to use dask for processing."""
     bypass_validation: bool = False
+    """Whether to skip validation checks in the processed data."""
 
     ds: xr.Dataset = field(init=False, repr=False)
+    """An xarray Dataset containing post-processed variables ready for input into
+    ROMS."""
 
     def __post_init__(self):
-
         self._input_checks()
         target_coords = get_target_coords(self.grid)
 
@@ -165,7 +173,6 @@ class TidalForcing:
         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():
@@ -175,7 +182,6 @@ class TidalForcing:
 
     def _get_data(self):
         """Loads tidal forcing data based on the specified source."""
-
         if self.source["name"] == "TPXO":
             if isinstance(self.source["path"], dict):
                 fname_dict = {
@@ -184,7 +190,7 @@ class TidalForcing:
                     "u": self.source["path"]["u"],
                 }
 
-            elif isinstance(self.source["path"], (str, Path)):
+            elif isinstance(self.source["path"], str | Path):
                 fname_dict = {
                     "grid": self.source["path"],
                     "h": self.source["path"],
@@ -260,7 +266,6 @@ class TidalForcing:
         self.variable_info = variable_info
 
     def _write_into_dataset(self, processed_fields, d_meta):
-
         # save in new dataset
         ds = xr.Dataset()
 
@@ -274,7 +279,6 @@ class TidalForcing:
         return ds
 
     def _add_global_metadata(self, ds):
-
         ds.attrs["title"] = "ROMS tidal forcing created by ROMS-Tools"
         # Include the version of roms-tools
         try:
@@ -322,7 +326,12 @@ class TidalForcing:
                 da = ds[var_name].isel(ntides=0)
                 nan_check(da, mask)
 
-    def plot(self, var_name, ntides=0) -> None:
+    def plot(
+        self,
+        var_name: str,
+        ntides: int = 0,
+        save_path: str | None = None,
+    ) -> None:
         """Plot the specified tidal forcing variable for a given tidal constituent.
 
         Parameters
@@ -343,6 +352,10 @@ class TidalForcing:
             The index of the tidal constituent to plot. Default is 0, which corresponds
             to the first constituent.
 
+        save_path : str, optional
+            Path to save the generated plot. If None, the plot is shown interactively.
+            Default is None.
+
         Returns
         -------
         None
@@ -357,11 +370,11 @@ class TidalForcing:
         Examples
         --------
         >>> tidal_forcing = TidalForcing(grid)
-        >>> tidal_forcing.plot("ssh_Re", nc=0)
+        >>> tidal_forcing.plot("ssh_Re", ntides=0)
         """
-
         if var_name not in self.ds:
             raise ValueError(f"Variable '{var_name}' is not found in dataset.")
+
         field = self.ds[var_name].isel(ntides=ntides)
 
         if self.use_dask:
@@ -370,49 +383,14 @@ class TidalForcing:
             with ProgressBar():
                 field = field.load()
 
-        if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-            lon_deg = self.grid.ds["lon_rho"]
-            lat_deg = self.grid.ds["lat_rho"]
-            mask = self.grid.ds["mask_rho"]
-
-        elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-            lon_deg = self.grid.ds["lon_u"]
-            lat_deg = self.grid.ds["lat_u"]
-            mask = self.grid.ds["mask_u"]
-
-        elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-            lon_deg = self.grid.ds["lon_v"]
-            lat_deg = self.grid.ds["lat_v"]
-            mask = self.grid.ds["mask_v"]
-
-        else:
-            ValueError("provided field does not have two horizontal dimension")
-
-        field = field.where(mask)
-        if self.grid.straddle:
-            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
-        field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
-
-        title = "%s, constituent: %s" % (
-            field.long_name,
-            self.ds[var_name].ntides[ntides].values.item().decode("utf-8"),
-        )
-
-        vmax = max(field.max(), -field.min())
-        vmin = -vmax
-        cmap = plt.colormaps.get_cmap("RdBu_r")
-        cmap.set_bad(color="gray")
-
-        kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
-
-        _plot(
+        plot(
             field=field,
-            title=title,
-            c="g",
-            kwargs=kwargs,
+            grid_ds=self.grid.ds,
+            save_path=save_path,
+            cmap_name="RdBu_r",
         )
 
-    def save(self, filepath: Union[str, Path]) -> None:
+    def save(self, filepath: str | Path) -> None:
         """Save the tidal forcing information to a netCDF4 file.
 
         Parameters
@@ -426,7 +404,6 @@ class TidalForcing:
         Path
             A `Path` object representing the location of the saved file.
         """
-
         # Ensure filepath is a Path object
         filepath = Path(filepath)
 
@@ -443,7 +420,7 @@ class TidalForcing:
 
         return saved_filenames
 
-    def to_yaml(self, filepath: Union[str, Path]) -> None:
+    def to_yaml(self, filepath: str | Path) -> None:
         """Export the parameters of the class to a YAML file, including the version of
         roms-tools.
 
@@ -452,13 +429,13 @@ class TidalForcing:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-
-        _to_yaml(self, filepath)
+        forcing_dict = to_dict(self, exclude=["use_dask"])
+        write_to_yaml(forcing_dict, filepath)
 
     @classmethod
     def from_yaml(
         cls,
-        filepath: Union[str, Path],
+        filepath: str | Path,
         use_dask: bool = False,
     ) -> "TidalForcing":
         """Create an instance of the TidalForcing class from a YAML file.
@@ -478,7 +455,7 @@ class TidalForcing:
         filepath = Path(filepath)
 
         grid = Grid.from_yaml(filepath)
-        tidal_forcing_params = _from_yaml(cls, filepath)
+        tidal_forcing_params = from_yaml(cls, filepath)
         return cls(
             grid=grid,
             **tidal_forcing_params,

roms_tools/setup/topography.py

@@ -1,13 +1,15 @@
-import time
 import logging
-import xarray as xr
-import numpy as np
-import gcm_filters
-from roms_tools.setup.utils import handle_boundaries
+import time
 import warnings
 from itertools import count
+
+import gcm_filters
+import numpy as np
+import xarray as xr
+
 from roms_tools.regrid import LateralRegridToROMS
 from roms_tools.setup.datasets import ETOPO5Dataset, SRTM15Dataset
+from roms_tools.setup.utils import handle_boundaries
 
 
 def _add_topography(
@@ -15,7 +17,7 @@ def _add_topography(
     target_coords,
     topography_source,
     hmin,
-    smooth_factor=8.0,
+    smooth_factor=6.0,
     rmax=0.2,
     verbose=False,
 ) -> xr.Dataset:
@@ -37,7 +39,7 @@ def _add_topography(
     smooth_factor : float, optional
         The smoothing factor used in the domain-wide Gaussian smoothing of the
         topography. Smaller values result in less smoothing, while larger
-        values produce more smoothing. The default is 8.0.
+        values produce more smoothing. The default is 6.0.
     rmax : float, optional
         The maximum allowable steepness ratio for the topography smoothing.
         This parameter controls the local smoothing of the topography. Smaller values result in
@@ -50,7 +52,6 @@ def _add_topography(
     xr.Dataset
         Updated dataset with added topography and metadata.
     """
-
     if verbose:
         start_time = time.time()
     data = _get_topography_data(topography_source)
@@ -66,10 +67,11 @@ def _add_topography(
     # smooth topography domain-wide with Gaussian kernel to avoid grid scale instabilities
     if verbose:
         start_time = time.time()
-    hraw = _smooth_topography_globally(hraw, smooth_factor)
+    area = 1 / ds["pm"] / ds["pn"]
+    hraw = _smooth_topography_globally(hraw, smooth_factor, area)
     if verbose:
         logging.info(
-            f"Smoothing the topography globally: {time.time() - start_time:.3f} seconds"
+            f"Domain-wide topography smoothing: {time.time() - start_time:.3f} seconds"
         )
 
     # smooth topography locally to satisfy r < rmax
@@ -85,7 +87,7 @@ def _add_topography(
     }
     if verbose:
         logging.info(
-            f"Smoothing the topography locally: {time.time() - start_time:.3f} seconds"
+            f"Local topography smoothing: {time.time() - start_time:.3f} seconds"
         )
 
     ds = _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax)
@@ -163,7 +165,7 @@ def _make_raw_topography(
     return hraw
 
 
-def _smooth_topography_globally(hraw, factor) -> xr.DataArray:
+def _smooth_topography_globally(hraw, factor, area) -> xr.DataArray:
     """Apply global smoothing to the topography using a Gaussian filter.
 
     Parameters
@@ -185,6 +187,10 @@ def _smooth_topography_globally(hraw, factor) -> xr.DataArray:
     margin_mask = xr.concat(
         [margin_mask, 0 * margin_mask.isel(xi_rho=-1)], dim="xi_rho"
     )
+    area_extended = xr.concat([area, area.isel(eta_rho=-1)], dim="eta_rho")
+    area_extended = xr.concat(
+        [area_extended, area_extended.isel(xi_rho=-1)], dim="xi_rho"
+    )
 
     # we choose a Gaussian filter kernel corresponding to a Gaussian with standard deviation factor/sqrt(12);
     # this standard deviation matches the standard deviation of a boxcar kernel with total width equal to factor.
@@ -192,8 +198,8 @@ def _smooth_topography_globally(hraw, factor) -> xr.DataArray:
         filter_scale=factor,
         dx_min=1,
         filter_shape=gcm_filters.FilterShape.GAUSSIAN,
-        grid_type=gcm_filters.GridType.REGULAR_WITH_LAND,
-        grid_vars={"wet_mask": margin_mask},
+        grid_type=gcm_filters.GridType.REGULAR_WITH_LAND_AREA_WEIGHTED,
+        grid_vars={"wet_mask": margin_mask, "area": area_extended},
     )
     hraw_extended = xr.concat([hraw, hraw.isel(eta_rho=-1)], dim="eta_rho")
     hraw_extended = xr.concat(
@@ -400,7 +406,9 @@ def _add_topography_metadata(ds, topography_source, smooth_factor, hmin, rmax):
     xarray.Dataset
         Updated dataset with added metadata.
     """
-    ds.attrs["topography_source"] = topography_source["name"]
+    ds.attrs["topography_source_name"] = topography_source["name"]
+    if "path" in topography_source:
+        ds.attrs["topography_source_path"] = topography_source["path"]
     ds.attrs["hmin"] = hmin
 
     return ds