Wrapper to match grid_labels (#161)

* First tests working

* Added whats new

* Add xesmf to ci envs

* Dont use pip for xesmf. Bad idea

* added docs

* some docs typos

* some minor coverage fix

* fix doc headings

* fix doc headings again

* Delete Untitled.ipynb

* Update whats-new.rst

* Clean up

.gitignore

@@ -65,6 +65,7 @@ target/
 .ipynb_checkpoints
 
 **/dask-worker-space/
+mydask.png
 
 .vscode
 .mypy_cache

ci/environment-upstream-dev.yml

@@ -5,7 +5,8 @@ dependencies:
   - cftime
   - dask
   - pip
-  - cartopy
+  - cartopy #installing this without conda is a nightmare, so ill leave it here
+  - xesmf # same here
   - pip:
     - codecov
     - pytest-cov

ci/environment.yml

@@ -9,6 +9,7 @@ dependencies:
   - xgcm
   - cftime
   - regionmask
+  - xesmf
   - xarrayutils
   - pytest-cov
   - pytest-xdist

cmip6_preprocessing/postprocessing.py

@@ -8,6 +8,12 @@
 from cmip6_preprocessing.utils import _key_from_attrs, cmip6_dataset_id
 
 
+try:
+    import xesmf
+except ImportError:
+    xesmf = None
+
+
 # define the attrs that are needed to get an 'exact' match
 exact_attrs = [
     "source_id",
@@ -217,6 +223,200 @@ def concat_members(
     )
 
 
+### Matching wrapper specific to combining grid labels via interpolation with xesmf
+def requires_xesmf(func):
+    @functools.wraps(func)
+    def wrapper_requires_xesmf(*args, **kwargs):
+
+        # Check if xesmf is installed
+        if xesmf is None:
+            raise ValueError(
+                "This function needs the optional package xesmf. Please install with `conda install -c conda-forge xesmf`."
+            )
+        return func(*args, **kwargs)
+
+    return wrapper_requires_xesmf
+
+
+def _pick_grid_label(ds_list, preferred_grid_label):
+    """From a list of datasets this picks the dataset with `ds.attrs['grid_label']==preferred_label`
+    if available"""
+    # TODO: Implement an option to pick the highest/lowest resolution if grid label dows not match.
+    matches = [ds for ds in ds_list if ds.attrs["grid_label"] == preferred_grid_label]
+    if len(matches) > 0:
+        return matches[0]
+    else:
+        return ds_list[0]
+
+
+def _drop_duplicate_grid_labels(ddict, preferred_grid_label):
+    """Eliminate multiple grid labels for the same dataset, by dropping multiples
+    with a preference for `preferred_grid_label`"""
+    match_attrs = [ma for ma in exact_attrs if ma not in ["version", "grid_label"]] + [
+        "variable_id"
+    ]
+    return combine_datasets(
+        ddict,
+        _pick_grid_label,
+        combine_func_args=(preferred_grid_label,),
+        match_attrs=match_attrs,
+    )
+
+
+def _clean_regridder(ds_source, ds_target, method, **xesmf_kwargs):
+    def _clean(ds):
+        # remove all unnecessary stuff for the regridding
+        ds = ds.isel(time=0, lev=0, rho=0, missing_dims="ignore")
+        for coord in [
+            co for co in ds.coords if co not in ["lon", "lat"]
+        ]:  # , "lat_bounds""lon_bounds",
+            ds = ds.drop_vars(coord)
+
+        # Ugly Hack to elminate 'seam' when regridding from gr to native grids
+        # There is something in xesmf that causes problems with broadcasted regular lon/lat values
+        if "gr" in ds.attrs["grid_label"]:
+            # actually revert the convention and return the 1d coordinates. ?Should I change that behavior in general?
+            ds = ds.assign(lon=ds.lon.isel(y=0))
+            ds = ds.assign(lat=ds.lat.isel(x=0))
+
+        # for now just eliminate the attrs here
+        # I can solve this more elegantly when I parse proper cf-attributes
+        for coord in ds.coords:
+            ds[coord].attrs = {}
+        #         ds = ds.rename({'lon_bounds':'lon_b', 'lat_bounds':'lat_b'})
+
+        # TODO: Make this work out of the box with lon/lat bounds and method='conservative'
+        # TODO: Maybe erase the need for this completely with cf-xarray
+        return ds
+
+    ds_source = _clean(ds_source)
+    ds_target = _clean(ds_target)
+    return xesmf.Regridder(ds_source, ds_target, method, **xesmf_kwargs)
+
+
+def _regrid_to_target(ds_source, ds_target, regridder):
+    ds_regridded = regridder(ds_source, keep_attrs=True)
+
+    # remove all coordinates that involve x and y (these will be merged from the native grid dataset)
+    ds_regridded = ds_regridded.reset_coords(drop=True)
+
+    # modify attributes
+    ds_regridded.attrs["grid_label"] = ds_target.attrs["grid_label"]
+    # identify the variables that are regridded
+    for var in ds_regridded.data_vars:
+        ds_regridded[var].attrs[
+            "cmip6_preprocessing_regrid_method"
+        ] = ds_regridded.attrs["regrid_method"]
+
+    return ds_regridded
+
+
+@requires_xesmf
+def interpolate_grid_label(
+    ds_dict,
+    target_grid_label="gn",
+    method="bilinear",
+    xesmf_kwargs={},
+    merge_kwargs={},
+    verbose=False,
+):
+    """Combines different grid labels via interpolation with xesmf
+
+    Parameters
+    ----------
+    ds_dict : dict
+        dictonary of input datasets
+    target_grid_label : str, optional
+        preferred grid_label value. If at least one dataset has this grid_label, otherse are interpolated to it.
+        Dataset with this grid label are not modified, by default "gn"
+    method : str, optional
+        interpolation method for xesmf, by default "bilinear"
+    xesmf_kwargs : dict, optional
+        optional arguments for building xesmf regridder, by default {}
+    merge_kwargs : dict, optional
+        optional arguments for the merging of interpolated datasets, by default {}
+    verbose : bool, optional
+        print output while creating regridder, by default False
+
+    Returns
+    -------
+    dict
+        dictionary of combined datasets (usually will combine across different variable ids)
+    """
+    match_attrs = [
+        ma for ma in exact_attrs if ma not in ["grid_label", "version"]
+    ]  # does this need to be more flexible?
+
+    xesmf_kwargs.setdefault("ignore_degenerate", True)
+    xesmf_kwargs.setdefault("periodic", True)
+
+    merge_kwargs.setdefault("combine_attrs", "drop_conflicts")
+
+    # first drop the datasets that might have both the target and another grid_label present
+    ds_dict = _drop_duplicate_grid_labels(ds_dict, target_grid_label)
+
+    def combine_func(ds_list, **kwargs):
+        grid_labels = np.unique([dss.attrs["grid_label"] for dss in ds_list])
+        target_grid = [
+            dss for dss in ds_list if target_grid_label == dss.attrs["grid_label"]
+        ]
+        if len(target_grid) < 1:
+            raise ValueError(
+                f"Could not find any variable with the target_grid_label{target_grid_label}. Found these instead: {grid_labels}"
+            )
+        else:
+            # just take the first one with a matching grid?
+            target_grid = target_grid[0]
+
+        # Construct a regridder for each other grid_label
+        regridder_dict = {}
+        if verbose:
+            print(
+                f'Constructing regridders for source_id {target_grid.attrs["source_id"]} ...'
+            )
+        for gl in grid_labels:
+            if gl != target_grid_label:
+                if verbose:
+                    print(gl)
+                source_grid = [dss for dss in ds_list if dss.attrs["grid_label"] == gl][
+                    0
+                ]  # again just take the first one available
+                regridder_dict[gl] = _clean_regridder(
+                    source_grid, target_grid, method, **xesmf_kwargs
+                )
+        if verbose:
+            print("FINISHED")
+
+        # Now regrid all datasets in the list (dont do anything if already on the target_grid)
+        ds_list_new = []
+        for ds_raw in ds_list:
+            if ds_raw.attrs["grid_label"] != target_grid_label:
+                if verbose:
+                    print(f"regridding {cmip6_dataset_id(ds_raw)}")
+                ds_regridded = _regrid_to_target(
+                    ds_raw, target_grid, regridder_dict[ds_raw.attrs["grid_label"]]
+                )
+                ds_list_new.append(ds_regridded)
+            else:
+                ds_list_new.append(ds_raw)
+
+        # check that horizontal dimensions are compatible.
+        xy_dimensions = [(len(ds.x), len(ds.y)) for ds in ds_list_new]
+        if not all([xy == xy_dimensions[0] for xy in xy_dimensions]):
+            raise ValueError(
+                f"Regridded datasets do not have the same dimensions. Found({xy_dimensions}). This will cause broadcasting problems during merge."
+            )
+        if verbose:
+            print("Merging regridded files")
+        return xr.merge(ds_list_new, **merge_kwargs)
+
+    return combine_datasets(
+        ds_dict,
+        combine_func,
+        match_attrs=match_attrs,
+    )
+
+
 ### Matching wrapper specific to metric datasets