Repository: pangeo-data/xESMF

Issue:
Regridding dataset with dask arrays fails
MWE
```python3
import xesmf as xe

ds = xe.util.grid_global(1, 1)
ds = ds.assign(data=ds.lon + ds.lat, data2=ds.lon * ds.lat)

reg = xe.Regridder(ds, ds, 'bilinear')
reg(ds)  # ok

reg(ds.chunk())  # Fails
```
Fails with:
```
TypeError: Field elements must be 2- or 3-tuples, got 'dtype('float64')'
```
in `/dask/array/utils.py in meta_from_array(x, ndim, dtype)`

This is due to a misuse of `xr.apply_ufunc` in the dataset case.  The  `output_dtypes` argument of `apply_ufunc` is only used if data uses the dask backend or if `vectorize=True` (not our case).  It is a list of dtypes _the same length as the number of outputs_.  When encountering a dataset, `xarray` will iteratively apply the `ufunc` for each variable. Thus, in each case there is only 1 output. Passing a list to `output_dtypes` has no meaning.

The best fix would be that xarray could accept some sort of mapping? I'm not sure how complicated that would be, probably not trivial.

For xesmf, we could:
- Call `astype` ourselves on the output dataset.
- Go back to float output. 
- Midway : force float32 output when all input variables are float32, use float64 otherwise. 

(When regridding, integer dtypes don't really make sense, no?)



Discussion:
aulemahal: @raphaeldussin It shouldn't :  As I pointed out in #121, passing `output_dtypes` as a list doesn't have the effect it was thought to have. 
Explanation 1 : The output of `xarray.apply_ufunc` called on a Dataset is a Dataset. A Dataset is mapping, how should xarray interpret a `output_dtypes` list to cast the different _variables_ correctly? In some python versions, dictionaries have inconsistent order.

Explanation 2 : As stated in the doc string of `apply_ufunc` : 
> output_dtypes (list of dtype, optional) – Optional list of output dtypes. Only used if dask='parallelized' or vectorize=True. 
Looking at the code, I can further precise that the argument is used ONLY when applying the func on dask-backed variables.

Thus, my change do not affect the behaviour wanted in #102, i.e.: conserving dtypes on numpy-backed variables. The changes only affect lines that are ignored when passing numpy-backed variables.

Also, because of how xarray applies  `apply_ufunc` on Datasets and how it understands `output_dtypes`, there is no way for us to coerce the dtype of the different variables on the output. Except by calling `astype` ourselves _afterwards_, but that adds dask tasks and I don't feel it makes much sense (but I'm open to discuss).
raphaeldussin: @huard I will do my best to get to this by next week 
raphaeldussin: @aulemahal @huard I have just a few comments but I believe this can be dealt with quickly
huard: @raphaeldussin Ready for your re-review. 

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 3b39db7..315c0c8 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -6,6 +6,7 @@ What's new
 
 Bug fixes
 ~~~~~~~~~
+- Regridding datasets with dask-backed variables is fixed. Dtype of the outputs is changed for this specific case. By `Pascal Bourgault <https://github.com/aulemahal>`_
 - SpatialAverager did not compute the same weights as Regridder when source cell areas were not uniform (:pull:`128`). By `David Huard <https://github.com/huard>`_
 
 
diff --git a/doc/conf.py b/doc/conf.py
index 094ed8e..dc53d12 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -58,6 +58,7 @@ autodoc_member_order = 'bysource'
 
 # To avoid installing xESMF and all its dependencies when building doc
 # https://stackoverflow.com/a/15912502/8729698
+
 autodoc_mock_imports = ['numpy', 'xarray', 'scipy', 'ESMF', 'dask', 'sparse', 'numba']
 
 # avoid automatic execution for notebooks
diff --git a/xesmf/frontend.py b/xesmf/frontend.py
index 764f7ca..0414172 100644
--- a/xesmf/frontend.py
+++ b/xesmf/frontend.py
@@ -432,7 +432,7 @@ class BaseRegridder(object):
 
         Returns
         -------
-        outdata : Data type is the same as input data type.
+        outdata : Data type is the same as input data type, except for datasets.
             On the same horizontal grid as ``ds_out``,
             with extra dims in ``dr_in``.
 
@@ -443,6 +443,10 @@ class BaseRegridder(object):
             - (n_time, n_lev, n_y_out, n_x_out), if ``dr_in`` has shape
               (n_time, n_lev, n_y, n_x)
 
+            Datasets with dask-backed variables will have modified dtypes.
+            If all input variables are 'float32', all output will be 'float32',
+            for any other case, all outputs will be 'float64'.
+
         """
         if isinstance(indata, np.ndarray):
             return self.regrid_numpy(indata, skipna=skipna, na_thres=na_thres)
@@ -560,7 +564,15 @@ class BaseRegridder(object):
         ]
         ds_in = ds_in.drop_vars(non_regriddable)
 
-        ds_dtypes = [d.dtype for d in ds_in.data_vars.values()]
+        # Select dtype (only relevant for dask-backed arrays, this has no effect otherwise)
+        dtypes = set(da.dtype for da in ds_in.data_vars.values())
+        if len(dtypes) == 1:
+            out_dtype = dtypes.pop()
+        else:
+            size = 1  # init to 1 byte
+            for d in dtypes:
+                size = max(size, d.itemsize)  # use max number of bytes found in dtypes
+                out_dtype = np.dtype(f'f{size}')
 
         ds_out = xr.apply_ufunc(
             self._regrid_array,
@@ -570,7 +582,7 @@ class BaseRegridder(object):
             input_core_dims=[input_horiz_dims, ('out_dim', 'in_dim')],
             output_core_dims=[temp_horiz_dims],
             dask='parallelized',
-            output_dtypes=ds_dtypes,
+            output_dtypes=[out_dtype],
             dask_gufunc_kwargs={
                 'output_sizes': {
                     temp_horiz_dims[0]: self.shape_out[0],
@@ -580,6 +592,13 @@ class BaseRegridder(object):
             keep_attrs=keep_attrs,
         )
 
+        # For dask-backed, we force the dtype to fit the input's
+        # (numpy-backed variables already have been handled in `apply_weights`)
+        for name, data in ds_out.data_vars.items():
+            indtype = ds_in[name].dtype
+            if xr.core.common.is_duck_dask_array(data.data) and data.dtype != indtype:
+                ds_out[name] = data.astype(indtype, keep_attrs=True)
+
         return self._format_xroutput(ds_out, temp_horiz_dims)
 
     def _parse_xrinput(self, dr_in):


Test patch:
diff --git a/xesmf/tests/test_frontend.py b/xesmf/tests/test_frontend.py
index b13c5a4..89502b7 100644
--- a/xesmf/tests/test_frontend.py
+++ b/xesmf/tests/test_frontend.py
@@ -575,6 +575,27 @@ def test_regrid_dataset():
     xr.testing.assert_identical(ds_result2, ds_result)
 
 
+@pytest.mark.parametrize('scheduler', dask_schedulers)
+def test_regrid_dataset_dask(request, scheduler):
+    scheduler = request.getfixturevalue(scheduler)
+    # xarray.Dataset containing dask array
+    regridder = xe.Regridder(ds_in, ds_out, 'conservative')
+
+    # `ds_out` already refers to output grid object
+    ds_result = regridder(ds_in.chunk())
+
+    # output should contain all data variables
+    assert set(ds_result.data_vars.keys()) == set(ds_in.data_vars.keys())
+    assert dask.is_dask_collection(ds_result)
+    assert ds_result.data.dtype == ds_in.data.dtype
+
+    ds_in_f4 = ds_in.copy()
+    ds_in_f4['data'] = ds_in_f4.data.astype('float32')
+    ds_in_f4['data4D'] = ds_in_f4.data4D.astype('float32')
+    ds_result = regridder(ds_in_f4.chunk())
+    assert ds_result.data.dtype == 'float32'
+
+
 def test_regrid_dataset_to_locstream():
     # xarray.Dataset containing in-memory numpy array