Implement GLOFAS dataset

.gitignore

@@ -29,4 +29,4 @@ paper
 .idea/
 .vscode
 AGENTS.md
-uv.lock
+uv.lock

RELEASE_NOTES.rst

@@ -24,6 +24,10 @@ Upcoming Release
   buffer that computes a geometrically accurate Euclidean buffer, correct in
   diagonal directions.
 
+* Implement glofas dataset which contains daily river discharge. ``cutout.hydro()`` now 
+  returns discharge if ``cutout.module`` contains ``"glofas"``.
+  (https://github.com/PyPSA/atlite/pull/498)
+
 **Bug fixes**
 
 * Fix ``Cutout.line_rating`` passing line azimuth in radians while

atlite/convert.py

@@ -1595,14 +1595,20 @@ def runoff(
 def hydro(
     cutout,
     plants,
-    hydrobasins,
+    hydrobasins=None,
     flowspeed=1,
     weight_with_height=False,
     show_progress=False,
+    *,
+    module="auto",
+    time=None,
     **kwargs,
 ):
     """
-    Compute inflow time series for plants by aggregating over catchment basins.
+    Get inflow time-series for `plants` from discharge or runoff data.
+
+    Either extracts the discharge time series for the nearest grid points or
+    computes runoff-based inflow time series.
 
     Parameters
     ----------
@@ -1611,58 +1617,63 @@ def hydro(
     plants : pd.DataFrame
         Run-of-river plants or dams with lon, lat columns.
     hydrobasins : str|gpd.GeoDataFrame
-        Filename or GeoDataFrame of one level of the HydroBASINS dataset.
+        Filename or GeoDataFrame of one level of the HydroBASINS dataset. Only required
+        for runoff-based computation.
     flowspeed : float
         Average speed of water flows to estimate the water travel time from
-        basin to plant (default: 1 m/s).
+        basin to plant (default: 1 m/s). Only relevant for runoff-based computation.
     weight_with_height : bool
         Whether surface runoff should be weighted by potential height (probably
-        better for coarser resolution).
+        better for coarser resolution). Only relevant for runoff-based computation.
     show_progress : bool
-        Whether to display progressbars.
+        Whether to display progressbars. Only relevant for runoff-based computation.
+    module : str
+        The method to compute hydro time series. "auto" will prefer discharge but fall
+        back to runoff-based computation, "glofas" uses discharge directly, "era5" uses
+        runoff-based computation.
+    time : pd.DatetimeIndex, optional
+        Time index to interpolate the plant inflow onto. Only relevant for
+        discharge-based computation. Defaults to the cutout's own time index.
     **kwargs
-        Additional keyword arguments passed to `convert_and_aggregate`.
+        Additional arguments for runoff-based computation.
 
     Returns
     -------
     xr.DataArray
         Inflow time-series for each plant.
 
-    References
-    ----------
-    [1] Liu, Hailiang, et al. "A validated high-resolution hydro power
-    time-series model for energy systems analysis." arXiv preprint
-    arXiv:1901.08476 (2019).
-
-    [2] Lehner, B., Grill G. (2013): Global river hydrography and network
-    routing: baseline data and new approaches to study the world’s large river
-    systems. Hydrological Processes, 27(15): 2171–2186. Data is available at
-    www.hydrosheds.org.
-
+    Raises
+    ------
+    ValueError
+        If required data for the selected module is missing or the module is unknown.
     """
-    basins = hydrom.determine_basins(plants, hydrobasins, show_progress=show_progress)
+    module = module.lower()
+    if module == "auto":
+        module = "glofas" if "discharge" in cutout.data.data_vars else "era5"
 
-    matrix = cutout.indicatormatrix(basins.shapes)
-    # compute the average surface runoff in each basin
-    # Fix NaN and Inf values to 0.0 to avoid numerical issues
-    matrix_normalized = np.nan_to_num(
-        matrix / matrix.sum(axis=1), nan=0.0, posinf=0.0, neginf=0.0
-    )
-    runoff = cutout.runoff(
-        matrix=matrix_normalized,
-        index=basins.shapes.index,
-        weight_with_height=weight_with_height,
-        show_progress=show_progress,
-        **kwargs,
-    )
-    # The hydrological parameters are in units of "m of water per day" and so
-    # they should be multiplied by 1000 and the basin area to convert to m3
-    # d-1 = m3 h-1 / 24
-    runoff *= xr.DataArray(basins.shapes.to_crs({"proj": "cea"}).area)
+    if module == "glofas":
+        if "discharge" not in cutout.data.data_vars:
+            raise ValueError(
+                "For GloFAS-based hydro time series, the cutout must include discharge data."
+            )
+        return hydrom._hydro_from_discharge(cutout, plants, time=time)
 
-    return hydrom.shift_and_aggregate_runoff_for_plants(
-        basins, runoff, flowspeed, show_progress
-    )
+    if module == "era5":
+        if hydrobasins is None:
+            raise ValueError(
+                "For ERA5-based hydro time series, the hydrobasins dataset must be provided."
+            )
+        return hydrom._hydro_from_runoff(
+            cutout,
+            plants,
+            hydrobasins,
+            flowspeed=flowspeed,
+            weight_with_height=weight_with_height,
+            show_progress=show_progress,
+            **kwargs,
+        )
+
+    raise ValueError(f'Unknown hydro module option "{module}".')
 
 
 def convert_line_rating(

atlite/datasets/__init__.py

@@ -4,6 +4,6 @@
 
 """atlite datasets."""
 
-from atlite.datasets import era5, gebco, sarah
+from atlite.datasets import era5, gebco, glofas, sarah
 
-modules = {"era5": era5, "sarah": sarah, "gebco": gebco}
+modules = {"era5": era5, "sarah": sarah, "gebco": gebco, "glofas": glofas}

atlite/datasets/cds_helper.py

@@ -0,0 +1,188 @@
+# SPDX-FileCopyrightText: Contributors to atlite <https://github.com/pypsa/atlite>
+#
+# SPDX-License-Identifier: MIT
+"""
+Shared helpers for datasets downloaded via the Climate Data Store (CDS).
+
+Used by both the era5 and glofas dataset modules.
+"""
+
+from __future__ import annotations
+
+import logging
+import weakref
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+import xarray as xr
+
+if TYPE_CHECKING:
+    from atlite._types import PathLike
+
+logger = logging.getLogger(__name__)
+
+
+def _area(coords: dict[str, xr.DataArray]) -> list[float]:
+    """
+    Extract CDS API bounding box from coordinates.
+
+    Parameters
+    ----------
+    coords : dict[str, xr.DataArray]
+        Coordinate arrays with 'x' (longitude) and 'y' (latitude).
+
+    Returns
+    -------
+    list[float]
+        Bounding box as [north, west, south, east].
+    """
+    x0, x1 = coords["x"].min().item(), coords["x"].max().item()
+    y0, y1 = coords["y"].min().item(), coords["y"].max().item()
+    return [y1, x0, y0, x1]
+
+
+def noisy_unlink(path: PathLike) -> None:
+    """
+    Remove a file with debug logging, handling PermissionError gracefully.
+
+    Parameters
+    ----------
+    path : PathLike
+        Path to the file to delete.
+    """
+    logger.debug("Deleting file %s", path)
+    try:
+        Path(path).unlink()
+    except PermissionError:
+        logger.error("Unable to delete file %s, as it is still in use.", path)
+
+
+def add_finalizer(ds: xr.Dataset, target: PathLike) -> None:
+    """
+    Register a weak-reference callback to delete a temp file on garbage collection.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Dataset whose lifetime controls the temp file.
+    target : PathLike
+        Path to the temporary file to clean up.
+    """
+    logger.debug("Adding finalizer for %s", target)
+    assert ds._close is not None
+    weakref.finalize(cast("Any", ds._close).__self__.ds, noisy_unlink, target)
+
+
+def sanitize_chunks(chunks: Any, **dim_mapping: str) -> Any:
+    """
+    Remap internal dimension names to CDS dimension names in chunk specs.
+
+    Translates atlite dimension names (time, x, y) to the corresponding
+    CDS names (valid_time, longitude, latitude).
+
+    Parameters
+    ----------
+    chunks : Any
+        Chunk specification. If not a dict, returned as-is.
+    **dim_mapping : str
+        Additional or override dimension name mappings.
+
+    Returns
+    -------
+    Any
+        Remapped chunk dict, or original value if not a dict.
+    """
+    dim_mapping = {
+        "time": "valid_time",
+        "x": "longitude",
+        "y": "latitude",
+    } | dim_mapping
+    if not isinstance(chunks, dict):
+        return chunks
+
+    return {
+        extname: chunks[intname]
+        for intname, extname in dim_mapping.items()
+        if intname in chunks
+    }
+
+
+def open_with_grib_conventions(
+    grib_file: PathLike,
+    chunks: dict[str, int] | None = None,
+    tmpdir: PathLike | None = None,
+) -> xr.Dataset:
+    """
+    Open a GRIB file using cfgrib with standardized coordinate conventions.
+
+    Performs the same conversion as the CDS backend, but locally.
+    Based on the documentation at
+    https://confluence.ecmwf.int/display/CKB/GRIB+to+netCDF+conversion+on+new+CDS+and+ADS+systems
+
+    Parameters
+    ----------
+    grib_file : PathLike
+        Path to the GRIB file.
+    chunks : dict[str, int] or None, optional
+        Dask chunk specification for lazy loading.
+    tmpdir : PathLike or None, optional
+        If set, the file is kept (managed externally).
+
+    Returns
+    -------
+    xr.Dataset
+        Opened dataset with standardized dimensions.
+    """
+    # Open grib file as dataset.
+    # Options below normalize different grib variants into consistent
+    # netCDF-compatible hypercubes. Options relevant only to e.g. wave-model
+    # data have been removed to keep this routine focused on the products we use.
+    ds = xr.open_dataset(
+        grib_file,
+        engine="cfgrib",
+        time_dims=["valid_time"],
+        ignore_keys=["edition"],
+        coords_as_attributes=[
+            "surface",
+            "depthBelowLandLayer",
+            "entireAtmosphere",
+            "heightAboveGround",
+            "meanSea",
+        ],
+        chunks=sanitize_chunks(chunks),
+    )
+    if tmpdir is None:
+        add_finalizer(ds, grib_file)
+
+    def safely_expand_dims(dataset: xr.Dataset, expand_dims: list[str]) -> xr.Dataset:
+        """Expand missing dimensions while preserving their original order.
+
+        Returns
+        -------
+        xr.Dataset
+            Dataset with the requested dimensions present.
+        """
+        dims_required = [
+            c for c in dataset.coords if c in expand_dims + list(dataset.dims)
+        ]
+        dims_missing = [
+            (c, i) for i, c in enumerate(dims_required) if c not in dataset.dims
+        ]
+        dataset = dataset.expand_dims(
+            dim=[x[0] for x in dims_missing], axis=[x[1] for x in dims_missing]
+        )
+        return dataset
+
+    logger.debug("Converting grib file to netcdf format")
+    rename_vars = {
+        "time": "forecast_reference_time",
+        "step": "forecast_period",
+        "isobaricInhPa": "pressure_level",
+        "hybrid": "model_level",
+    }
+    rename_vars = {k: v for k, v in rename_vars.items() if k in ds}
+    ds = ds.rename(rename_vars)
+
+    ds = safely_expand_dims(ds, ["valid_time", "pressure_level", "model_level"])
+
+    return ds