add workaround for dask.pad mode=mean that converts integers to floats, and add an additional check if the shape of output

Author: mark-boer

xarray/core/dask_array_compat.py

@@ -99,11 +99,43 @@ def meta_from_array(x, ndim=None, dtype=None):
         return meta
 
 
-# TODO figure out how Dask versioning works
-# if LooseVersion(dask_version) >= LooseVersion("1.7.0"):
-try:
-    pad = da.pad
-except AttributeError:
+def _validate_pad_output_shape(input_shape, pad_width, output_shape):
+    """ Dask.array.pad with mode='reflect' does not always return the correct output_shape. """
+    isint = lambda i: isinstance(i, int)
+
+    if isint(pad_width):
+        pass
+    elif len(pad_width) == 2 and all(map(isint, pad_width)):
+        pad_width = sum(pad_width)
+    elif (
+        len(pad_width) == len(input_shape)
+        and all(map(lambda x: len(x) == 2, pad_width))
+        and all((isint(i) for p in pad_width for i in p))
+    ):
+        pad_width = np.sum(pad_width, axis=1)
+    else:
+        return  # should be impossible
+
+    if not np.array_equal(np.array(input_shape) + pad_width, output_shape):
+        raise RuntimeError(
+            "There seems to be something wrong with the shape of the output of dask.array.pad, "
+            "try upgrading Dask, use a different pad mode e.g. mode='constant' or first convert "
+            "your DataArray/Dataset to one backed by a numpy array by calling the `compute()` method."
+        )
+
+
+if LooseVersion(dask_version) >= LooseVersion("0.18.1"):
+
+    def pad(array, pad_width, mode="constant", **kwargs):
+        padded = da.pad(array, pad_width, mode=mode, **kwargs)
+        # workaround for inconsistency between numpy and dask: https://github.com/dask/dask/issues/5303
+        if mode == "mean" and issubclass(array.dtype.type, np.integer):
+            return da.round(padded).astype(array.dtype)
+        _validate_pad_output_shape(array.shape, pad_width, padded.shape)
+        return padded
+
+
+else:
 
     def pad(array, pad_width, mode="constant", **kwargs):
         """

xarray/core/dataset.py

@@ -5820,8 +5820,8 @@ def pad(
                     reflect_type=reflect_type,
                 )
             else:
-                variables[name] = var.pad(  # type: ignore
-                    pad_width=var_pad_width, mode=coord_pad_mode, **coord_pad_options,
+                variables[name] = var.pad(
+                    pad_width=var_pad_width, mode=coord_pad_mode, **coord_pad_options,  # type: ignore
                 )
 
         return self._replace_vars_and_dims(variables)

xarray/core/duck_array_ops.py

@@ -11,7 +11,7 @@
 import numpy as np
 import pandas as pd
 
-from . import dask_array_ops, dask_array_compat, dtypes, npcompat, nputils
+from . import dask_array_ops, dtypes, npcompat, nputils
 from .nputils import nanfirst, nanlast
 from .pycompat import dask_array_type
 

xarray/tests/test_variable.py

@@ -788,7 +788,7 @@ def test_getitem_error(self):
     @pytest.mark.parametrize(
         "mode",
         [
-            pytest.param("mean", marks=pytest.mark.xfail),
+            "mean",
             pytest.param("median", marks=pytest.mark.xfail),
             pytest.param("reflect", marks=pytest.mark.xfail),
             "edge",
@@ -2070,7 +2070,7 @@ def test_getitem_uint(self):
     @pytest.mark.parametrize(
         "mode",
         [
-            pytest.param("mean", marks=pytest.mark.xfail),
+            "mean",
             pytest.param("median", marks=pytest.mark.xfail),
             pytest.param("reflect", marks=pytest.mark.xfail),
             "edge",