add setup_pet_forcing method (#257)

* add setup_pet_forcing method

* update changelog

docs/api.rst

@@ -45,6 +45,7 @@ Setup components
    WflowModel.setup_config_output_timeseries
    WflowModel.setup_precip_forcing
    WflowModel.setup_temp_pet_forcing
+   WflowModel.setup_pet_forcing
    WflowModel.setup_constant_pars
    WflowModel.setup_1dmodel_connection
    WflowModel.setup_grid_from_raster
@@ -242,6 +243,7 @@ Wflow workflows
    workflows.topography
    workflows.river
    workflows.river_bathymetry
+   workflows.pet
    workflows.landuse
    workflows.ksathorfrac
    workflows.soilgrids

docs/changelog.rst

@@ -13,6 +13,7 @@ Added
 -----
 - New function **setup_1dmodel_connection** to connect wflow to 1D river model (eg Delft3D FM 1D, HEC-RAS, etc.) `PR #210 <https://github.com/Deltares/hydromt_wflow/pull/210>`_
 - New setup method for the **KsatHorFrac** parameter **setup_ksathorfarc** to up-downscale existing ksathorfrac maps. `PR #255 <https://github.com/Deltares/hydromt_wflow/pull/255>`_
+- new function **setup_pet_forcing** to reproject existing pet data rather than computing from other meteo data. PR #257
 - Workflow to compute brooks corey c for the wflow layers based on soilgrids data, soilgrids_brooks_corey. PR #242
 
 Changed

docs/user_guide/wflow_model_setup.rst

@@ -67,6 +67,8 @@ a specific method see its documentation.
     * - :py:func:`~WflowModel.setup_precip_forcing`
       -  Setup gridded precipitation forcing at model resolution.
     * - :py:func:`~WflowModel.setup_temp_pet_forcing`
+      -  Setup gridded temperature and optionally compute reference evapotranspiration forcing at model resolution.
+    * - :py:func:`~WflowModel.setup_pet_forcing`
       -  Setup gridded reference evapotranspiration forcing at model resolution.
     * - :py:func:`~WflowModel.setup_constant_pars`
       -  Setup constant parameter maps for all active model cells.

hydromt_wflow/wflow.py

@@ -2469,6 +2469,58 @@ def setup_temp_pet_forcing(
         temp_out.attrs.update(opt_attr)
         self.set_forcing(temp_out.where(mask), name="temp")
 
+    def setup_pet_forcing(
+        self,
+        pet_fn: Union[str, xr.DataArray],
+        chunksize: Optional[int] = None,
+    ):
+        """
+        Prepare PET forcing from existig PET data.
+
+        Adds model layer:
+
+        * **pet**: reference evapotranspiration [mm]
+
+        Parameters
+        ----------
+        pet_fn: str, xr.DataArray
+            RasterDataset source or data for PET to be resampled.
+
+            * Required variable: 'pet' [mm]
+
+        chunksize: int, optional
+            Chunksize on time dimension for processing data (not for saving to disk!).
+            If None the data chunksize is used, this can however be optimized for
+            large/small catchments. By default None.
+
+        """
+        self.logger.info("Preparing potential evapotranspiration forcing maps.")
+
+        starttime = self.get_config("starttime")
+        endtime = self.get_config("endtime")
+        freq = pd.to_timedelta(self.get_config("timestepsecs"), unit="s")
+
+        pet = self.data_catalog.get_rasterdataset(
+            pet_fn,
+            geom=self.region,
+            buffer=2,
+            variables=["pet"],
+            time_tuple=(starttime, endtime),
+        )
+
+        pet_out = workflows.forcing.pet(
+            pet=pet,
+            ds_like=self.grid,
+            freq=freq,
+            mask_name=self._MAPS["basins"],
+            chunksize=chunksize,
+            logger=self.logger,
+        )
+
+        # Update meta attributes (used for default output filename later)
+        pet_out.attrs.update({"pet_fn": pet_fn})
+        self.set_forcing(pet_out, name="pet")
+
     def setup_rootzoneclim(
         self,
         run_fn: Union[str, Path, xr.Dataset],

hydromt_wflow/workflows/__init__.py

@@ -2,6 +2,7 @@
 
 from .basemaps import *
 from .connect import *
+from .forcing import *
 from .gauges import *
 from .glaciers import *
 from .landuse import *

hydromt_wflow/workflows/forcing.py

@@ -0,0 +1,74 @@
+"""Forcing workflow for wflow."""
+
+import logging
+from typing import Optional
+
+import numpy as np
+import xarray as xr
+from hydromt.workflows.forcing import resample_time
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["pet"]
+
+
+def pet(
+    pet: xr.DataArray,
+    ds_like: xr.Dataset,
+    freq: str = "D",
+    mask_name: Optional[str] = None,
+    chunksize: Optional[int] = None,
+    logger: Optional[logging.Logger] = logger,
+) -> xr.DataArray:
+    """
+    Resample and reproject PET to the grid of ds_like.
+
+    Parameters
+    ----------
+    pet : xr.DataArray
+        PET data array with time as first dimension.
+    ds_like : xr.Dataset
+        Dataset with the grid to reproject to.
+    freq : str, optional
+        Resampling frequency, by default "D".
+    mask_name : str, optional
+        Name of the mask variable in ds_like, by default None.
+    chunksize : int, optional
+        Chunksize for the time dimension for resampling, by default None to use default
+        time chunk.
+
+    Returns
+    -------
+    pet_out : xr.DataArray
+        Resampled and reprojected PET data array.
+    """
+    if chunksize is not None:
+        pet = pet.chunk({"time": chunksize})
+
+    if pet.raster.dim0 != "time":
+        raise ValueError(f'First pet dim should be "time", not {pet.raster.dim0}')
+
+    # change nodata
+    pet = pet.where(pet != pet.raster.nodata, np.nan)
+    pet.raster.set_nodata(np.nan)
+
+    # Minimum is zero
+    pet_out = np.fmax(pet.raster.reproject_like(ds_like, method="nearest_index"), 0)
+
+    # resample time
+    resample_kwargs = dict(label="right", closed="right")
+    if freq is not None:
+        resample_kwargs.update(upsampling="bfill", downsampling="sum", logger=logger)
+
+        pet_out = resample_time(pet_out, freq, conserve_mass=True, **resample_kwargs)
+
+    # Mask
+    if mask_name is not None:
+        mask = ds_like[mask_name].values > 0
+        pet_out = pet_out.where(mask, pet_out.raster.nodata)
+
+    # Attributes
+    pet_out.name = "pet"
+    pet_out.attrs.update(unit="mm")
+
+    return pet_out

tests/conftest.py

@@ -94,3 +94,14 @@ def rivers1d():
         join(TESTDATADIR, "rivers.geojson"),
     )
     return data
+
+
+@pytest.fixture()
+def da_pet(example_wflow_model):
+    da = example_wflow_model.data_catalog.get_rasterdataset(
+        "era5", geom=example_wflow_model.region, buffer=2, variables=["temp"]
+    )
+    da = 0.5 * (0.45 * da + 8)  # simple pet from Bradley Criddle
+    da.name = "pet"
+
+    return da

tests/test_model_methods.py

@@ -594,6 +594,19 @@ def test_setup_floodplains_2d(elevtn_map, example_wflow_model, floodplain1d_test
     )
 
 
+def test_setup_pet_forcing(example_wflow_model, da_pet):
+    example_wflow_model.setup_pet_forcing(
+        pet_fn=da_pet,
+    )
+
+    assert "pet" in example_wflow_model.forcing
+    # used to be debruin before update
+    assert "pet_method" not in example_wflow_model.forcing["pet"].attrs
+    assert example_wflow_model.forcing["pet"].min().values == da_pet.min().values
+    mean_val = example_wflow_model.forcing["pet"].mean().values
+    assert int(mean_val * 1000) == 2984
+
+
 def test_setup_1dmodel_connection(example_wflow_model, rivers1d):
     # test subbasin_area method with river boundaries
     example_wflow_model.setup_1dmodel_connection(