Refactor lazy_regrid.py for wflow diagnostics

esmvaltool/diag_scripts/hydrology/lazy_regrid.py

@@ -1,12 +1,10 @@
 """Lazy regridding, because this is not supported by iris (yet).
 
 Iris issue requesting the feature:
-https://github.com/SciTools/iris/issues/3700
+https://github.com/SciTools/iris/issues/3808
 """
-import copy
-
-import iris
 import numpy as np
+import iris
 
 HORIZONTAL_SCHEMES = {
     'linear': iris.analysis.Linear(extrapolation_mode='mask'),
@@ -46,58 +44,17 @@ def _compute_chunks(src, tgt):
         (src.coord('latitude').shape[0], ),
         (src.coord('longitude').shape[0], ),
     )
-    tgt_chunks = (
-        time_chunks,
-        (tgt_nlat, ),
-        (tgt_nlon, ),
-    )
 
-    return src_chunks, tgt_chunks
+    return src_chunks
 
 
-def _regrid_data(src, tgt, scheme):
-    """Regrid data from cube src onto grid of cube tgt."""
-    src_chunks, tgt_chunks = _compute_chunks(src, tgt)
+def lazy_regrid(src, tgt, scheme):
+    """Regrid cube src onto the grid of cube tgt."""
+    src_chunks = _compute_chunks(src, tgt)
 
-    # Define the block regrid function
     if scheme not in HORIZONTAL_SCHEMES:
         raise ValueError(f"Regridding scheme {scheme} not supported, "
                          f"choose from {HORIZONTAL_SCHEMES.keys()}.")
     regridder = HORIZONTAL_SCHEMES[scheme].regridder(src, tgt)
-
-    def regrid(block):
-        tlen = block.shape[0]
-        cube = src[:tlen].copy(block)
-        return regridder(cube).core_data()
-
-    # Regrid
-    data = src.core_data().rechunk(src_chunks).map_blocks(
-        regrid,
-        dtype=src.dtype,
-        chunks=tgt_chunks,
-    )
-
-    return data
-
-
-def lazy_regrid(src, tgt, scheme):
-    """Regrid cube src onto the grid of cube tgt."""
-    data = _regrid_data(src, tgt, scheme)
-
-    result = iris.cube.Cube(data)
-    result.metadata = copy.deepcopy(src.metadata)
-
-    def copy_coords(src_coords, add_method):
-        for coord in src_coords:
-            dims = src.coord_dims(coord)
-            if coord == src.coord('longitude'):
-                coord = tgt.coord('longitude')
-            elif coord == src.coord('latitude'):
-                coord = tgt.coord('latitude')
-            result_coord = coord.copy()
-            add_method(result_coord, dims)
-
-    copy_coords(src.dim_coords, result.add_dim_coord)
-    copy_coords(src.aux_coords, result.add_aux_coord)
-
-    return result
+    src.data = src.lazy_data().rechunk(src_chunks)
+    return regridder(src)

esmvaltool/diag_scripts/hydrology/wflow.py

@@ -2,9 +2,9 @@
 import logging
 from pathlib import Path
 
-import iris
 import numpy as np
 from osgeo import gdal
+import iris
 
 from esmvaltool.diag_scripts.hydrology.derive_evspsblpot import debruin_pet
 from esmvaltool.diag_scripts.hydrology.lazy_regrid import lazy_regrid

esmvaltool/recipes/hydrology/recipe_wflow.yml

@@ -4,17 +4,17 @@
 documentation:
   description: |
     Pre-processes climate data for the WFlow hydrological model.
-    
+
   authors:
     - kalverla_peter
     - camphuijsen_jaro
     - alidoost_sarah
     - aerts_jerom
     - andela_bouwe
-    
+
   projects:
     - ewatercycle
-  
+
   references:
     - acknow_project
 
@@ -40,7 +40,7 @@ diagnostics:
         mip: day
         preprocessor: rough_cutout
         start_year: 1990
-        end_year: 1990
+        end_year: 2001
       pr: *daily_var
       # evspsblpot:  # doesn't exist for ERA-Interim.
       # Reconstruct evspsblpot using:
@@ -53,5 +53,5 @@ diagnostics:
       script:
         script: hydrology/wflow.py
         basin: Meuse
-        dem_file: 'wflow/wflow_dem_Meuse.nc'
+        dem_file: 'wflow_parameterset/meuse/staticmaps/wflow_dem.map'
         regrid: area_weighted