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

roms_tools/setup/river_forcing.py

@@ -1,26 +1,31 @@
-import xarray as xr
-import numpy as np
 import logging
 from dataclasses import dataclass, field
-import cartopy.crs as ccrs
 from datetime import datetime
-from typing import Dict, Union, List, Optional
 from pathlib import Path
-import matplotlib.pyplot as plt
+from typing import Dict, List, Optional, Union
+
+import cartopy.crs as ccrs
 import matplotlib.cm as cm
+import matplotlib.pyplot as plt
+import numpy as np
+import xarray as xr
+
 from roms_tools import Grid
-from roms_tools.plot import _get_projection, _add_field_to_ax
-from roms_tools.utils import save_datasets
+from roms_tools.plot import _get_projection, _plot
 from roms_tools.setup.datasets import DaiRiverDataset
 from roms_tools.setup.utils import (
-    get_target_coords,
-    gc_dist,
-    substitute_nans_by_fillvalue,
-    convert_to_roms_time,
-    _to_yaml,
     _from_yaml,
+    _to_dict,
+    _write_to_yaml,
+    add_time_info_to_ds,
+    add_tracer_metadata_to_ds,
+    gc_dist,
+    get_target_coords,
+    get_tracer_defaults,
     get_variable_metadata,
+    substitute_nans_by_fillvalue,
 )
+from roms_tools.utils import save_datasets
 
 
 @dataclass(kw_only=True)
@@ -275,93 +280,50 @@ class RiverForcing:
 
         ds = xr.Dataset()
 
+        # Tracer metadata
+        ds = add_tracer_metadata_to_ds(ds, self.include_bgc)
+
+        # River volume
         river_volume = (
             data.ds[data.var_names["flux"]] * data.ds[data.var_names["ratio"]]
         ).astype(np.float32)
-        river_volume.attrs["long_name"] = "River volume flux"
-        river_volume.attrs["units"] = "m^3/s"
         river_volume = river_volume.rename(
             {data.dim_names["time"]: "river_time", data.dim_names["station"]: "nriver"}
         )
+        ds["river_volume"] = river_volume
+        ds["river_volume"].attrs["long_name"] = "River volume flux"
+        ds["river_volume"].attrs["units"] = "m^3/s"
 
-        name = data.ds[data.var_names["name"]].rename(
-            {data.dim_names["station"]: "nriver"}
+        # River tracers
+        ds["river_tracer"] = xr.zeros_like(
+            ds.river_time.astype(np.float32) * ds.ntracers * ds.nriver, dtype=np.float32
         )
-        name.attrs["long_name"] = "River name"
-        river_volume.coords["river_name"] = name
+        ds["river_tracer"].attrs["long_name"] = "River tracer data"
 
-        ds["river_volume"] = river_volume
+        defaults = get_tracer_defaults()
+        for ntracer in range(ds.ntracers.size):
+            tracer_name = ds.tracer_name[ntracer].item()
+            ds["river_tracer"].loc[{"ntracers": ntracer}] = defaults[tracer_name]
+
+        # River names
+        river_names = data.ds[data.var_names["name"]].rename(
+            {data.dim_names["station"]: "nriver"}
+        )
+        river_names.attrs["long_name"] = "River name"
+        ds = ds.assign_coords({"river_name": river_names})
 
+        # River IDs
         nriver = xr.DataArray(np.arange(1, len(ds.nriver) + 1), dims="nriver")
         nriver.attrs["long_name"] = "River ID (1-based Fortran indexing)"
         ds = ds.assign_coords({"nriver": nriver})
 
-        if self.include_bgc:
-            ntracers = 2 + 32
-        else:
-            ntracers = 2
-        tracer_data = np.zeros(
-            (len(ds.river_time), ntracers, len(ds.nriver)), dtype=np.float32
-        )
-        tracer_data[:, 0, :] = 17.0
-        tracer_data[:, 1, :] = 1.0
-        tracer_data[:, 2:, :] = 0.0
-
-        river_tracer = xr.DataArray(
-            tracer_data, dims=("river_time", "ntracers", "nriver")
+        # River time
+        ds["river_time"] = data.ds[data.dim_names["time"]].rename(
+            {"time": "river_time"}
         )
-        river_tracer.attrs["long_name"] = "River tracer data"
-
-        if self.include_bgc:
-            tracer_names = xr.DataArray(
-                [
-                    "temp",
-                    "salt",
-                    "PO4",
-                    "NO3",
-                    "SiO3",
-                    "NH4",
-                    "Fe",
-                    "Lig",
-                    "O2",
-                    "DIC",
-                    "DIC_ALT_CO2",
-                    "ALK",
-                    "ALK_ALT_CO2",
-                    "DOC",
-                    "DON",
-                    "DOP",
-                    "DOPr",
-                    "DONr",
-                    "DOCr",
-                    "zooC",
-                    "spChl",
-                    "spC",
-                    "spP",
-                    "spFe",
-                    "spCaCO3",
-                    "diatChl",
-                    "diatC",
-                    "diatP",
-                    "diatFe",
-                    "diatSi",
-                    "diazChl",
-                    "diazC",
-                    "diazP",
-                    "diazFe",
-                ],
-                dims="ntracers",
-            )
-        else:
-            tracer_names = xr.DataArray(["temp", "salt"], dims="ntracers")
-        tracer_names.attrs["long_name"] = "Tracer name"
-        river_tracer.coords["tracer_name"] = tracer_names
-        ds["river_tracer"] = river_tracer
-
-        ds, time = convert_to_roms_time(
+        ds, time = add_time_info_to_ds(
             ds, self.model_reference_date, self.climatology, time_name="river_time"
         )
-
         ds = ds.assign_coords({"river_time": time})
 
         return ds
@@ -459,8 +421,8 @@ class RiverForcing:
             The modified dataset with the "river_index" and "river_fraction" variables added.
         """
 
-        river_index = xr.zeros_like(self.grid.ds.h)
-        river_fraction = xr.zeros_like(self.grid.ds.h)
+        river_index = xr.zeros_like(self.grid.ds.h, dtype=np.float32)
+        river_fraction = xr.zeros_like(self.grid.ds.h, dtype=np.float32)
 
         for nriver in ds.nriver:
             river_name = str(ds.river_name.sel(nriver=nriver).values)
@@ -513,56 +475,25 @@ class RiverForcing:
         """Plots the original and updated river locations on a map projection."""
 
         field = self.grid.ds.mask_rho
-        vmax = 3
-        vmin = 0
-        cmap = plt.colormaps.get_cmap("Blues")
-        kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
-
         lon_deg = self.grid.ds.lon_rho
         lat_deg = self.grid.ds.lat_rho
-
-        # check if North or South pole are in domain
-        if lat_deg.max().values > 89 or lat_deg.min().values < -89:
-            raise NotImplementedError(
-                "Plotting is not implemented for the case that the domain contains the North or South pole."
-            )
-
         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})
 
-        trans = _get_projection(lon_deg, lat_deg)
+        vmax = 6
+        vmin = 0
+        cmap = plt.colormaps.get_cmap("Blues")
+        kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
-        lon_deg = lon_deg.values
-        lat_deg = lat_deg.values
+        trans = _get_projection(lon_deg, lat_deg)
 
         fig, axs = plt.subplots(
             1, 2, figsize=(13, 13), subplot_kw={"projection": trans}
         )
 
         for ax in axs:
-            _add_field_to_ax(
-                ax,
-                lon_deg,
-                lat_deg,
-                field,
-                add_colorbar=False,
-                kwargs=kwargs,
-            )
-            # Add gridlines with labels for latitude and longitude
-            gridlines = ax.gridlines(
-                draw_labels=True, linewidth=0.5, color="gray", alpha=0.7, linestyle="--"
-            )
-            gridlines.top_labels = False  # Hide top labels
-            gridlines.right_labels = False  # Hide right labels
-            gridlines.xlabel_style = {
-                "size": 10,
-                "color": "black",
-            }  # Customize longitude label style
-            gridlines.ylabel_style = {
-                "size": 10,
-                "color": "black",
-            }  # Customize latitude label style
+            _plot(field, kwargs=kwargs, ax=ax, c=None, add_colorbar=False)
 
         proj = ccrs.PlateCarree()
 
@@ -747,7 +678,16 @@ class RiverForcing:
             The path to the YAML file where the parameters will be saved.
         """
 
-        _to_yaml(self, filepath)
+        forcing_dict = _to_dict(self)
+
+        # Convert indices format
+        indices_data = forcing_dict["RiverForcing"]["indices"]
+        serialized_indices = {"_convention": "eta_rho, xi_rho"}
+        for key, value in indices_data.items():
+            serialized_indices[key] = [f"{tup[0]}, {tup[1]}" for tup in value]
+        forcing_dict["RiverForcing"]["indices"] = serialized_indices
+
+        _write_to_yaml(forcing_dict, filepath)
 
     @classmethod
     def from_yaml(cls, filepath: Union[str, Path]) -> "RiverForcing":
@@ -769,14 +709,11 @@ class RiverForcing:
         params = _from_yaml(cls, filepath)
 
         def convert_indices_format(indices):
-            # Remove the '_convention' key from the dictionary if present
             indices = {
                 key: value for key, value in indices.items() if key != "_convention"
             }
 
-            # Convert the string of indices into tuples
             for river, index_list in indices.items():
-                # Split the string by ',' and convert to tuples of integers
                 indices[river] = [tuple(map(int, idx.split(","))) for idx in index_list]
 
             return indices

roms_tools/setup/surface_forcing.py

@@ -1,34 +1,44 @@
-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
+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.utils import save_datasets, transpose_dimensions
-from roms_tools.regrid import LateralRegridToROMS
 from roms_tools.plot import _plot
+from roms_tools.regrid import LateralRegridToROMS
 from roms_tools.setup.datasets import (
-    ERA5Dataset,
-    ERA5Correction,
     CESMBGCSurfaceForcingDataset,
+    Dataset,
+    ERA5ARCODataset,
+    ERA5Correction,
+    ERA5Dataset,
     UnifiedBGCSurfaceDataset,
 )
 from roms_tools.setup.utils import (
+    _from_yaml,
+    _to_dict,
+    _write_to_yaml,
+    add_time_info_to_ds,
+    compute_missing_surface_bgc_variables,
     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,
+)
+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 +61,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 +78,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:
 
@@ -101,6 +118,7 @@ class SurfaceForcing:
     source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
     type: str = "physics"
     correct_radiation: bool = True
+    wind_dropoff: bool = False
     coarse_grid_mode: str = "auto"
     model_reference_date: datetime = datetime(2000, 1, 1)
     use_dask: bool = False
@@ -138,19 +156,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 +180,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 +236,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 = {
@@ -266,7 +303,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 +320,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(
@@ -356,8 +400,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 +448,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 +457,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,9 +466,56 @@ class SurfaceForcing:
         )
         corr_factor = lateral_regrid.apply(corr_factor)
 
-        processed_fields["swrad"] = processed_fields["swrad"] * corr_factor
+        radiation_corrected = radiation * corr_factor
+
+        return radiation_corrected
+
+    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 processed_fields
+        return uwnd_corrected, vwnd_corrected
 
     def _write_into_dataset(self, processed_fields, data, d_meta):
 
@@ -420,7 +531,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
         )
 
@@ -487,6 +598,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)
 
@@ -634,7 +746,8 @@ class SurfaceForcing:
             The path to the YAML file where the parameters will be saved.
         """
 
-        _to_yaml(self, filepath)
+        forcing_dict = _to_dict(self)
+        _write_to_yaml(forcing_dict, filepath)
 
     @classmethod
     def from_yaml(