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

roms_tools/setup/surface_forcing.py

@@ -3,14 +3,12 @@ import logging
 from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Dict, List, Optional, Union
 
-import matplotlib.pyplot as plt
 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.setup.datasets import (
     CESMBGCSurfaceForcingDataset,
@@ -21,11 +19,9 @@ from roms_tools.setup.datasets import (
     UnifiedBGCSurfaceDataset,
 )
 from roms_tools.setup.utils import (
-    _from_yaml,
-    _to_dict,
-    _write_to_yaml,
     add_time_info_to_ds,
     compute_missing_surface_bgc_variables,
+    from_yaml,
     get_target_coords,
     get_variable_metadata,
     group_dataset,
@@ -34,6 +30,8 @@ from roms_tools.setup.utils import (
     nan_check,
     rotate_velocities,
     substitute_nans_by_fillvalue,
+    to_dict,
+    write_to_yaml,
 )
 from roms_tools.utils import save_datasets, transpose_dimensions
 
@@ -113,21 +111,36 @@ class SurfaceForcing:
     """
 
     grid: Grid
-    start_time: Optional[datetime] = None
-    end_time: Optional[datetime] = None
-    source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
+    """Object representing the grid information."""
+    start_time: datetime | None = None
+    """The start time of the desired surface forcing data."""
+    end_time: datetime | None = None
+    """The end time of the desired surface forcing data."""
+    source: dict[str, str | Path | list[str | Path]]
+    """Dictionary specifying the source of the surface forcing data."""
     type: str = "physics"
+    """Specifies the type of forcing data ("physics", "bgc")."""
     correct_radiation: bool = True
+    """Whether to correct shortwave radiation."""
     wind_dropoff: bool = False
+    """Whether to apply a coastal wind speed reduction to mimic nearshore wind drop-
+    off."""
     coarse_grid_mode: str = "auto"
+    """Specifies whether to interpolate onto grid coarsened by a factor of two."""
     model_reference_date: datetime = datetime(2000, 1, 1)
+    """Reference date for the model."""
     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."""
+    use_coarse_grid: bool = field(init=False, repr=False)
+    """Whether data is interpolated onto grid coarsened by factor 2."""
 
     def __post_init__(self):
-
         self._input_checks()
         data = self._get_data()
 
@@ -292,7 +305,6 @@ class SurfaceForcing:
         return use_coarse_grid
 
     def _get_data(self):
-
         data_dict = {
             "filename": self.source["path"],
             "start_time": self.start_time,
@@ -331,7 +343,6 @@ class SurfaceForcing:
         return data
 
     def _get_correction_data(self):
-
         if self.source["name"] == "ERA5":
             correction_data = ERA5Correction(use_dask=self.use_dask)
         else:
@@ -497,7 +508,6 @@ class SurfaceForcing:
         vwnd_corrected : xr.DataArray
             Corrected meridional wind component with reduced coastal values.
         """
-
         # calculate the distance from each ocean point to the closest land point
         cdist = min_dist_to_land(
             self.target_coords["lon"].values,
@@ -518,7 +528,6 @@ class SurfaceForcing:
         return uwnd_corrected, vwnd_corrected
 
     def _write_into_dataset(self, processed_fields, data, d_meta):
-
         # save in new dataset
         ds = xr.Dataset()
 
@@ -576,7 +585,6 @@ class SurfaceForcing:
         -----
         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:
             if self.variable_info[var_name]["validate"]:
                 # all variables are at rho-points
@@ -584,7 +592,6 @@ class SurfaceForcing:
                 nan_check(ds[var_name].isel(time=0), mask)
 
     def _add_global_metadata(self, ds=None):
-
         if ds is None:
             ds = xr.Dataset()
         ds.attrs["title"] = "ROMS surface forcing file created by ROMS-Tools"
@@ -607,7 +614,12 @@ class SurfaceForcing:
 
         return ds
 
-    def plot(self, var_name, time=0) -> None:
+    def plot(
+        self,
+        var_name: str,
+        time: int = 0,
+        save_path: str | None = None,
+    ) -> None:
         """Plot the specified surface forcing field for a given time slice.
 
         Parameters
@@ -633,6 +645,10 @@ class SurfaceForcing:
             The time index to plot. Default is 0, which corresponds to the first
             time slice.
 
+        save_path : str, optional
+            Path to save the generated plot. If None, the plot is shown interactively.
+            Default is None.
+
         Returns
         -------
         None
@@ -648,7 +664,6 @@ class SurfaceForcing:
         --------
         >>> atm_forcing.plot("uwnd", time=0)
         """
-
         if var_name not in self.ds:
             raise ValueError(f"Variable '{var_name}' is not found in dataset.")
 
@@ -660,41 +675,24 @@ class SurfaceForcing:
             with ProgressBar():
                 field = field.load()
 
-        field = field.where(self.target_coords["mask"])
-
-        lon_deg = self.target_coords["lon"]
-        lat_deg = self.target_coords["lat"]
-        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 = field.long_name
-
         if var_name in ["uwnd", "vwnd"]:
-            vmax = max(field.max().values, -field.min().values)
-            vmin = -vmax
-            cmap = plt.colormaps.get_cmap("RdBu_r")
+            cmap_name = "RdBu_r"
+        elif var_name in ["swrad", "lwrad", "Tair", "qair"]:
+            cmap_name = "YlOrRd"
         else:
-            vmax = field.max().values
-            vmin = field.min().values
-            if var_name in ["swrad", "lwrad", "Tair", "qair"]:
-                cmap = plt.colormaps.get_cmap("YlOrRd")
-            else:
-                cmap = plt.colormaps.get_cmap("YlGnBu")
-        cmap.set_bad(color="gray")
-
-        kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
+            cmap_name = "YlGnBu"
 
-        _plot(
+        plot(
             field=field,
-            title=title,
-            c="g",
-            kwargs=kwargs,
+            grid_ds=self.grid.ds,
+            use_coarse_grid=self.use_coarse_grid,
+            save_path=save_path,
+            cmap_name=cmap_name,
         )
 
     def save(
         self,
-        filepath: Union[str, Path],
+        filepath: str | Path,
         group: bool = True,
     ) -> None:
         """Save the surface forcing fields to one or more netCDF4 files.
@@ -716,7 +714,6 @@ class SurfaceForcing:
         List[Path]
             A list of `Path` objects representing the filenames of the saved file(s).
         """
-
         # Ensure filepath is a Path object
         filepath = Path(filepath)
 
@@ -736,7 +733,7 @@ class SurfaceForcing:
 
         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.
 
@@ -745,14 +742,13 @@ class SurfaceForcing:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-
-        forcing_dict = _to_dict(self)
-        _write_to_yaml(forcing_dict, filepath)
+        forcing_dict = to_dict(self, exclude=["use_dask", "use_coarse_grid"])
+        write_to_yaml(forcing_dict, filepath)
 
     @classmethod
     def from_yaml(
         cls,
-        filepath: Union[str, Path],
+        filepath: str | Path,
         use_dask: bool = False,
     ) -> "SurfaceForcing":
         """Create an instance of the SurfaceForcing class from a YAML file.
@@ -772,6 +768,6 @@ class SurfaceForcing:
         filepath = Path(filepath)
 
         grid = Grid.from_yaml(filepath)
-        params = _from_yaml(cls, filepath)
+        params = from_yaml(cls, filepath)
 
         return cls(grid=grid, **params, use_dask=use_dask)

roms_tools/setup/tides.py

@@ -2,20 +2,16 @@ import importlib.metadata
 from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Dict, List, Union
 
-import matplotlib.pyplot as plt
 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.setup.datasets import TPXOManager
 from roms_tools.setup.utils import (
-    _from_yaml,
-    _to_dict,
-    _write_to_yaml,
+    from_yaml,
     get_target_coords,
     get_variable_metadata,
     get_vector_pairs,
@@ -24,6 +20,8 @@ from roms_tools.setup.utils import (
     nan_check,
     rotate_velocities,
     substitute_nans_by_fillvalue,
+    to_dict,
+    write_to_yaml,
 )
 from roms_tools.utils import save_datasets
 
@@ -35,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:
 
@@ -81,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)
 
@@ -168,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():
@@ -178,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 = {
@@ -187,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"],
@@ -263,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()
 
@@ -277,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:
@@ -325,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
@@ -346,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
@@ -360,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:
@@ -373,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
@@ -429,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)
 
@@ -446,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.
 
@@ -455,14 +429,13 @@ class TidalForcing:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-
-        forcing_dict = _to_dict(self)
-        _write_to_yaml(forcing_dict, 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.
@@ -482,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

@@ -17,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:
@@ -39,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
@@ -52,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)
@@ -68,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
@@ -87,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)
@@ -165,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
@@ -187,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.
@@ -194,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(