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

roms_tools/setup/surface_forcing.py

@@ -1,34 +1,42 @@
-import xarray as xr
 import importlib.metadata
+import logging
 from dataclasses import dataclass, field
 from datetime import datetime
-import numpy as np
-import matplotlib.pyplot as plt
 from pathlib import Path
-import logging
-from typing import Dict, Union, List, Optional
+
+import numpy as np
+import xarray as xr
+
 from roms_tools import Grid
-from roms_tools.utils import save_datasets, transpose_dimensions
+from roms_tools.plot import plot
 from roms_tools.regrid import LateralRegridToROMS
-from roms_tools.plot import _plot
 from roms_tools.setup.datasets import (
-    ERA5Dataset,
-    ERA5Correction,
     CESMBGCSurfaceForcingDataset,
+    Dataset,
+    ERA5ARCODataset,
+    ERA5Correction,
+    ERA5Dataset,
     UnifiedBGCSurfaceDataset,
 )
 from roms_tools.setup.utils import (
+    add_time_info_to_ds,
+    compute_missing_surface_bgc_variables,
+    from_yaml,
     get_target_coords,
-    nan_check,
-    substitute_nans_by_fillvalue,
-    interpolate_from_climatology,
     get_variable_metadata,
     group_dataset,
+    interpolate_from_climatology,
+    min_dist_to_land,
+    nan_check,
     rotate_velocities,
-    compute_missing_surface_bgc_variables,
-    convert_to_roms_time,
-    _to_yaml,
-    _from_yaml,
+    substitute_nans_by_fillvalue,
+    to_dict,
+    write_to_yaml,
+)
+from roms_tools.utils import save_datasets, transpose_dimensions
+
+DEFAULT_ERA5_ARCO_PATH = (
+    "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
 )
 
 
@@ -51,12 +59,13 @@ class SurfaceForcing:
     source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool]
         Dictionary specifying the source of the surface forcing data. Keys include:
 
-          - "name" (str): Name of the data source (e.g., "ERA5").
-          - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file(s). This can be:
+          - "name" (str): Name of the data source. Currently supported: "ERA5"
+          - "path" (optional; Union[str, Path, List[Union[str, Path]]]): Path(s) to the raw data file(s). Accepted formats:
 
-            - A single string (with or without wildcards).
-            - A single Path object.
-            - A list of strings or Path objects containing multiple files.
+            - A single string (supports wildcards),
+            - A single Path object,
+            - A list of strings or Path objects.
+            If omitted or set to the ARCO URL, the data will be streamed from the cloud.
           - "climatology" (bool): Indicates if the data is climatology data. Defaults to False.
 
     type : str
@@ -67,6 +76,12 @@ class SurfaceForcing:
 
     correct_radiation : bool
         Whether to correct shortwave radiation. Default is True.
+
+    wind_dropoff : bool, optional
+        Whether to apply a coastal wind speed reduction to mimic nearshore wind drop-off.
+        This applies an exponential decay to wind magnitude near the coast, based on
+        a 12.5 km e-folding scale, with up to 40% reduction at the coastline. Default is False.
+
     coarse_grid_mode : str, optional
         Specifies whether to interpolate onto grid coarsened by a factor of two. Options are:
 
@@ -96,20 +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()
 
@@ -138,19 +169,23 @@ class SurfaceForcing:
         data.apply_lateral_fill()
 
         self._set_variable_info(data)
-        var_names = self.variable_info.keys()
+        var_names = {
+            var: {"name": name}
+            for d in [data.var_names, data.opt_var_names]
+            for var, name in d.items()
+            if name in data.ds.data_vars
+        }
 
         processed_fields = {}
         # lateral regridding
         lateral_regrid = LateralRegridToROMS(target_coords, data.dim_names)
         for var_name in var_names:
-            if var_name in data.var_names.keys():
-                processed_fields[var_name] = lateral_regrid.apply(
-                    data.ds[data.var_names[var_name]]
-                )
+            processed_fields[var_name] = lateral_regrid.apply(
+                data.ds[var_names[var_name]["name"]]
+            )
 
         # rotation of velocities
-        if "uwnd" in self.variable_info and "vwnd" in self.variable_info:
+        if "uwnd" in processed_fields and "vwnd" in processed_fields:
             processed_fields["uwnd"], processed_fields["vwnd"] = rotate_velocities(
                 processed_fields["uwnd"],
                 processed_fields["vwnd"],
@@ -158,9 +193,18 @@ class SurfaceForcing:
                 interpolate=False,
             )
 
-        # correct radiation
-        if self.type == "physics" and self.correct_radiation:
-            processed_fields = self._apply_correction(processed_fields, data)
+        if self.type == "physics":
+            if self.correct_radiation:
+                processed_fields["swrad"] = self._apply_radiation_correction(
+                    processed_fields["swrad"], data
+                )
+            if self.wind_dropoff:
+                (
+                    processed_fields["uwnd"],
+                    processed_fields["vwnd"],
+                ) = self._apply_wind_correction(
+                    processed_fields["uwnd"], processed_fields["vwnd"]
+                )
 
         if self.type == "bgc":
             processed_fields = compute_missing_surface_bgc_variables(processed_fields)
@@ -205,7 +249,13 @@ class SurfaceForcing:
         if "name" not in self.source:
             raise ValueError("`source` must include a 'name'.")
         if "path" not in self.source:
-            raise ValueError("`source` must include a 'path'.")
+            if self.source["name"] == "ERA5":
+                logging.info(
+                    "No path specified for ERA5 source; defaulting to ARCO ERA5 dataset on Google Cloud."
+                )
+                self.source["path"] = DEFAULT_ERA5_ARCO_PATH
+            else:
+                raise ValueError("`source` must include a 'path'.")
 
         # Set 'climatology' to False if not provided in 'source'
         self.source = {
@@ -255,7 +305,6 @@ class SurfaceForcing:
         return use_coarse_grid
 
     def _get_data(self):
-
         data_dict = {
             "filename": self.source["path"],
             "start_time": self.start_time,
@@ -266,7 +315,16 @@ class SurfaceForcing:
 
         if self.type == "physics":
             if self.source["name"] == "ERA5":
-                data = ERA5Dataset(**data_dict)
+                if str(self.source["path"]).startswith("gs://") or str(
+                    self.source["path"]
+                ).startswith("gcs://"):
+                    if not self.use_dask:
+                        raise ValueError(
+                            "Cloud-based ERA5 access requires `use_dask=True`. Please enable Dask by setting `use_dask=True`."
+                        )
+                    data = ERA5ARCODataset(**data_dict)
+                else:
+                    data = ERA5Dataset(**data_dict)
             else:
                 raise ValueError(
                     'Only "ERA5" is a valid option for source["name"] when type is "physics".'
@@ -274,10 +332,8 @@ class SurfaceForcing:
 
         elif self.type == "bgc":
             if self.source["name"] == "CESM_REGRIDDED":
-
                 data = CESMBGCSurfaceForcingDataset(**data_dict)
             elif self.source["name"] == "UNIFIED":
-
                 data = UnifiedBGCSurfaceDataset(**data_dict)
             else:
                 raise ValueError(
@@ -287,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:
@@ -356,8 +411,28 @@ class SurfaceForcing:
 
         self.variable_info = variable_info
 
-    def _apply_correction(self, processed_fields, data):
+    def _apply_radiation_correction(
+        self, radiation: xr.DataArray, data: Dataset
+    ) -> xr.DataArray:
+        """Apply a climatological correction to shortwave radiation.
+
+        This method scales the input `radiation` field using a correction factor
+        derived from climatological data, interpolated in time and regridded
+        to the ROMS domain.
+
+        Parameters
+        ----------
+        radiation : xr.DataArray
+            Shortwave radiation field to be corrected. Must include a `time` coordinate.
+
+        data : Dataset
+            Dataset containing ROMS grid and mask information used to align correction data.
 
+        Returns
+        -------
+        radiation_corrected : xr.DataArray
+            Radiation field scaled by the correction factor, with original coordinates.
+        """
         correction_data = self._get_correction_data()
         # Match subdomain to forcing data to reuse the mask
         coords_correction = {
@@ -384,7 +459,7 @@ class SurfaceForcing:
                         correction_data.dim_names["time"],
                         time=time,
                     )
-                    for time in processed_fields["swrad"].time
+                    for time in radiation.time
                 ],
                 dim="time",
             )
@@ -393,7 +468,7 @@ class SurfaceForcing:
             corr_factor = interpolate_from_climatology(
                 correction_data.ds[correction_data.var_names["swr_corr"]],
                 correction_data.dim_names["time"],
-                time=processed_fields["swrad"].time,
+                time=radiation.time,
             )
 
         # Spatial regridding
@@ -402,12 +477,57 @@ class SurfaceForcing:
         )
         corr_factor = lateral_regrid.apply(corr_factor)
 
-        processed_fields["swrad"] = processed_fields["swrad"] * corr_factor
+        radiation_corrected = radiation * corr_factor
 
-        return processed_fields
+        return radiation_corrected
 
-    def _write_into_dataset(self, processed_fields, data, d_meta):
+    def _apply_wind_correction(
+        self, uwnd: xr.DataArray, vwnd: xr.DataArray
+    ) -> tuple[xr.DataArray, xr.DataArray]:
+        """Apply coastal wind drop-off correction to wind components.
+
+        This correction reduces wind speed near the coastline by up to 40%,
+        transitioning smoothly from full magnitude offshore using an
+        exponential decay with an e-folding scale of 12.5 km.
+
+        Reanalysis wind products often lack sufficient resolution to capture
+        sharp coastal wind gradients caused by orography and land-sea contrasts.
+        This method adjusts wind magnitude to better reflect these coastal effects.
+
+        Parameters
+        ----------
+        uwnd : xr.DataArray
+            Zonal (east-west) wind component on the ROMS grid.
+        vwnd : xr.DataArray
+            Meridional (north-south) wind component on the ROMS grid.
 
+        Returns
+        -------
+        uwnd_corrected : xr.DataArray
+            Corrected zonal wind component with reduced coastal values.
+        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,
+            self.target_coords["lat"].values,
+            self.target_coords["mask"].values,
+        )
+
+        # Compute a spatially varying scaling factor to reduce wind near the coast.
+        # This uses an exponential decay with a 12.5 km e-folding scale,
+        # reducing wind magnitude by up to 40% at the coastline.
+        mult = 1 - 0.4 * np.exp(-0.08 * cdist / 1000)
+
+        mult = xr.DataArray(data=mult, dims=["eta_rho", "xi_rho"])
+
+        uwnd_corrected = mult * uwnd
+        vwnd_corrected = mult * vwnd
+
+        return uwnd_corrected, vwnd_corrected
+
+    def _write_into_dataset(self, processed_fields, data, d_meta):
         # save in new dataset
         ds = xr.Dataset()
 
@@ -420,7 +540,7 @@ class SurfaceForcing:
         ds = self._add_global_metadata(ds)
 
         # Convert the time coordinate to the format expected by ROMS
-        ds, sfc_time = convert_to_roms_time(
+        ds, sfc_time = add_time_info_to_ds(
             ds, self.model_reference_date, data.climatology
         )
 
@@ -465,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
@@ -473,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"
@@ -487,6 +605,7 @@ class SurfaceForcing:
         ds.attrs["end_time"] = str(self.end_time)
         ds.attrs["source"] = self.source["name"]
         ds.attrs["correct_radiation"] = str(self.correct_radiation)
+        ds.attrs["wind_dropoff"] = str(self.wind_dropoff)
         ds.attrs["use_coarse_grid"] = str(self.use_coarse_grid)
         ds.attrs["model_reference_date"] = str(self.model_reference_date)
 
@@ -495,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
@@ -521,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
@@ -536,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.")
 
@@ -548,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.
@@ -604,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)
 
@@ -624,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.
 
@@ -633,13 +742,13 @@ class SurfaceForcing:
         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", "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.
@@ -659,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)