Repository: pydata/xarray

Issue:
Set `allow_rechunk=True` in `apply_ufunc`
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports: http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples: https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**What happened**:

`blockwise` calls `unify_chunks` by default but `apply_gufunc` does not; so we have a regression in `apply_ufunc` now that we've switched from `blockwise` to `apply_gufunc`.

**Minimal Complete Verifiable Example**:

```python
import operator

a = xr.DataArray(np.arange(10), dims=("a")).chunk({"a": 2})
b = xr.DataArray(np.arange(10), dims=("a")).chunk({"a": 4})

xr.apply_ufunc(operator.add, a, b, dask="parallelized", output_dtypes=[a.dtype]).compute()
```

raises
```
ValueError: Dimension `'__loopdim0__'` with different chunksize present
```

on master but works with 0.16.0

I think we need to do `dask_gufunc_kwargs.setdefault("allow_rechunk", True)`

If we want to avoid that, we'll need to go through a deprecation cycle.

Discussion:
kmuehlbauer: As laid out in #4372 the behaviour before #4060 was to run `blockwise` with automatically aligning non-core dimensions (permit rechunking). For multiple chunks in core dimensions a ValueError would have been thrown. 

If I understand the comments of @dcherian and @shoyer in #4372 correctly `allow_rechunk` should be set `True` in the first case but not in the second. 

This PR just checks for the dask ValueError which is raised for non-core dimension chunk mismatch, issues a warning, that the behaviour will change in the future (not sure if this is intended, please recommend otherwise) and reruns `apply_gufunc` with `allow_rechunk=True`. For the second case the dask ValueError is raised. 

kmuehlbauer: @dcherian I think, that this is almost good to go. Any rework on the warning message needed?
kmuehlbauer: I think my latest changes according to @dcherian's comments are now backwards compatible. Short explanation:

- `allow_rechunk` is None (default)
    - if core dimension chunks > 1 raise ValueError (to keep backward compatibility concerning core dimension chunking)
    - else sets `allow_rechunk=True` (to keep backward compatibility concerning chunking mismatch in non-core dimensions
- `allow_rechunk` == True  (user need to explicitoly set this in the `dask_gufunc_kwargs`)
    - just call `apply_gufunc`
- `allow_rechunk` == False (user need to explicitoly set this in the `dask_gufunc_kwargs`)
    - just call `apply_gufunc`




Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 7265bbac..16165288 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -6,7 +6,7 @@ repos:
     hooks:
       - id: isort
   # https://github.com/python/black#version-control-integration
-  - repo: https://github.com/python/black
+  - repo: https://github.com/psf/black
     rev: 20.8b1
     hooks:
       - id: black
diff --git a/doc/examples/area_weighted_temperature.ipynb b/doc/examples/area_weighted_temperature.ipynb
index 72876e3f..de705966 100644
--- a/doc/examples/area_weighted_temperature.ipynb
+++ b/doc/examples/area_weighted_temperature.ipynb
@@ -106,7 +106,7 @@
    "source": [
     "### Creating weights\n",
     "\n",
-    "For a for a rectangular grid the cosine of the latitude is proportional to the grid cell area."
+    "For a rectangular grid the cosine of the latitude is proportional to the grid cell area."
    ]
   },
   {
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index eda5c9df..f367286a 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -77,6 +77,8 @@ Bug fixes
   and :py:meth:`DataArray.str.wrap` (:issue:`4334`). By `Mathias Hauser <https://github.com/mathause>`_.
 - Fixed overflow issue causing incorrect results in computing means of :py:class:`cftime.datetime`
   arrays (:issue:`4341`). By `Spencer Clark <https://github.com/spencerkclark>`_.
+- Fixed :py:meth:`Dataset.coarsen`, :py:meth:`DataArray.coarsen` dropping attributes on original object (:issue:`4120`, :pull:`4360`). by `Julia Kent <https://github.com/jukent>`_.
+- fix the signature of the plot methods. (:pull:`4359`) By `Justus Magin <https://github.com/keewis>`_.
 - Fix :py:func:`xarray.apply_ufunc` with ``vectorize=True`` and ``exclude_dims`` (:issue:`3890`).
   By `Mathias Hauser <https://github.com/mathause>`_.
 - Fix `KeyError` when doing linear interpolation to an nd `DataArray`
@@ -96,7 +98,7 @@ Documentation
 Internal Changes
 ~~~~~~~~~~~~~~~~
 - Use :py:func:`dask.array.apply_gufunc` instead of :py:func:`dask.array.blockwise` in
-  :py:func:`xarray.apply_ufunc` when using ``dask='parallelized'``. (:pull:`4060`, :pull:`4391`)
+  :py:func:`xarray.apply_ufunc` when using ``dask='parallelized'``. (:pull:`4060`, :pull:`4391`, :pull:`4392`)
 - Fix ``pip install .`` when no ``.git`` directory exists; namely when the xarray source
   directory has been rsync'ed by PyCharm Professional for a remote deployment over SSH.
   By `Guido Imperiale <https://github.com/crusaderky>`_
diff --git a/xarray/core/computation.py b/xarray/core/computation.py
index 507f12fe..db4a9b46 100644
--- a/xarray/core/computation.py
+++ b/xarray/core/computation.py
@@ -625,6 +625,26 @@ def apply_variable_ufunc(
             if dask_gufunc_kwargs is None:
                 dask_gufunc_kwargs = {}
 
+            allow_rechunk = dask_gufunc_kwargs.get("allow_rechunk", None)
+            if allow_rechunk is None:
+                for n, (data, core_dims) in enumerate(
+                    zip(input_data, signature.input_core_dims)
+                ):
+                    if is_duck_dask_array(data):
+                        # core dimensions cannot span multiple chunks
+                        for axis, dim in enumerate(core_dims, start=-len(core_dims)):
+                            if len(data.chunks[axis]) != 1:
+                                raise ValueError(
+                                    f"dimension {dim} on {n}th function argument to "
+                                    "apply_ufunc with dask='parallelized' consists of "
+                                    "multiple chunks, but is also a core dimension. To "
+                                    "fix, either rechunk into a single dask array chunk along "
+                                    f"this dimension, i.e., ``.chunk({dim}: -1)``, or "
+                                    "pass ``allow_rechunk=True`` in ``dask_gufunc_kwargs`` "
+                                    "but beware that this may significantly increase memory usage."
+                                )
+                dask_gufunc_kwargs["allow_rechunk"] = True
+
             output_sizes = dask_gufunc_kwargs.pop("output_sizes", {})
             if output_sizes:
                 output_sizes_renamed = {}
diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py
index 0c4614e0..3d7be955 100644
--- a/xarray/core/rolling.py
+++ b/xarray/core/rolling.py
@@ -690,7 +690,7 @@ class DataArrayCoarsen(Coarsen):
             from .dataarray import DataArray
 
             reduced = self.obj.variable.coarsen(
-                self.windows, func, self.boundary, self.side, **kwargs
+                self.windows, func, self.boundary, self.side, self.keep_attrs, **kwargs
             )
             coords = {}
             for c, v in self.obj.coords.items():
@@ -703,6 +703,7 @@ class DataArrayCoarsen(Coarsen):
                             self.coord_func[c],
                             self.boundary,
                             self.side,
+                            self.keep_attrs,
                             **kwargs,
                         )
                     else:
diff --git a/xarray/core/variable.py b/xarray/core/variable.py
index ce6df528..075d7904 100644
--- a/xarray/core/variable.py
+++ b/xarray/core/variable.py
@@ -1993,7 +1993,9 @@ class Variable(
             ),
         )
 
-    def coarsen(self, windows, func, boundary="exact", side="left", **kwargs):
+    def coarsen(
+        self, windows, func, boundary="exact", side="left", keep_attrs=None, **kwargs
+    ):
         """
         Apply reduction function.
         """
@@ -2001,13 +2003,22 @@ class Variable(
         if not windows:
             return self.copy()
 
+        if keep_attrs is None:
+            keep_attrs = _get_keep_attrs(default=False)
+
+        if keep_attrs:
+            _attrs = self.attrs
+        else:
+            _attrs = None
+
         reshaped, axes = self._coarsen_reshape(windows, boundary, side)
         if isinstance(func, str):
             name = func
             func = getattr(duck_array_ops, name, None)
             if func is None:
                 raise NameError(f"{name} is not a valid method.")
-        return self._replace(data=func(reshaped, axis=axes, **kwargs))
+
+        return self._replace(data=func(reshaped, axis=axes, **kwargs), attrs=_attrs)
 
     def _coarsen_reshape(self, windows, boundary, side):
         """
@@ -2072,9 +2083,6 @@ class Variable(
             else:
                 shape.append(variable.shape[i])
 
-        keep_attrs = _get_keep_attrs(default=False)
-        variable.attrs = variable._attrs if keep_attrs else {}
-
         return variable.data.reshape(shape), tuple(axes)
 
     @property
diff --git a/xarray/plot/plot.py b/xarray/plot/plot.py
index 305405d4..0d44d5fd 100644
--- a/xarray/plot/plot.py
+++ b/xarray/plot/plot.py
@@ -463,6 +463,15 @@ class _PlotMethods:
         return step(self._da, *args, **kwargs)
 
 
+def override_signature(f):
+    def wrapper(func):
+        func.__wrapped__ = f
+
+        return func
+
+    return wrapper
+
+
 def _plot2d(plotfunc):
     """
     Decorator for common 2d plotting logic
@@ -572,6 +581,16 @@ def _plot2d(plotfunc):
     # Build on the original docstring
     plotfunc.__doc__ = f"{plotfunc.__doc__}\n{commondoc}"
 
+    # plotfunc and newplotfunc have different signatures:
+    # - plotfunc: (x, y, z, ax, **kwargs)
+    # - newplotfunc: (darray, x, y, **kwargs)
+    # where plotfunc accepts numpy arrays, while newplotfunc accepts a DataArray
+    # and variable names. newplotfunc also explicitly lists most kwargs, so we
+    # need to shorten it
+    def signature(darray, x, y, **kwargs):
+        pass
+
+    @override_signature(signature)
     @functools.wraps(plotfunc)
     def newplotfunc(
         darray,


Test patch:
diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py
index 7cb755b6..63bedfaf 100644
--- a/xarray/tests/test_computation.py
+++ b/xarray/tests/test_computation.py
@@ -695,8 +695,7 @@ def test_apply_dask_parallelized_two_args():
     check(data_array, 0 * data_array)
     check(data_array, 0 * data_array[0])
     check(data_array[:, 0], 0 * data_array[0])
-    with raises_regex(ValueError, "with different chunksize present"):
-        check(data_array, 0 * data_array.compute())
+    check(data_array, 0 * data_array.compute())
 
 
 @requires_dask
@@ -710,7 +709,7 @@ def test_apply_dask_parallelized_errors():
     with raises_regex(ValueError, "at least one input is an xarray object"):
         apply_ufunc(identity, array, dask="parallelized")
 
-    # formerly from _apply_blockwise, now from dask.array.apply_gufunc
+    # formerly from _apply_blockwise, now from apply_variable_ufunc
     with raises_regex(ValueError, "consists of multiple chunks"):
         apply_ufunc(
             identity,
diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
index 8fd56446..a22ed58c 100644
--- a/xarray/tests/test_dataarray.py
+++ b/xarray/tests/test_dataarray.py
@@ -6186,6 +6186,35 @@ def test_isin(da):
     assert_equal(result, expected)
 
 
+def test_coarsen_keep_attrs():
+    _attrs = {"units": "test", "long_name": "testing"}
+
+    da = xr.DataArray(
+        np.linspace(0, 364, num=364),
+        dims="time",
+        coords={"time": pd.date_range("15/12/1999", periods=364)},
+        attrs=_attrs,
+    )
+
+    da2 = da.copy(deep=True)
+
+    # Test dropped attrs
+    dat = da.coarsen(time=3, boundary="trim").mean()
+    assert dat.attrs == {}
+
+    # Test kept attrs using dataset keyword
+    dat = da.coarsen(time=3, boundary="trim", keep_attrs=True).mean()
+    assert dat.attrs == _attrs
+
+    # Test kept attrs using global option
+    with xr.set_options(keep_attrs=True):
+        dat = da.coarsen(time=3, boundary="trim").mean()
+    assert dat.attrs == _attrs
+
+    # Test kept attrs in original object
+    xr.testing.assert_identical(da, da2)
+
+
 @pytest.mark.filterwarnings("error:Mean of empty slice")
 @pytest.mark.parametrize("da", (1, 2), indirect=True)
 def test_rolling_iter(da):
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index 16b1cc33..40e2bdfc 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -5956,6 +5956,8 @@ def test_coarsen_keep_attrs():
         attrs=_attrs,
     )
 
+    ds2 = ds.copy(deep=True)
+
     # Test dropped attrs
     dat = ds.coarsen(coord=5).mean()
     assert dat.attrs == {}
@@ -5969,6 +5971,9 @@ def test_coarsen_keep_attrs():
         dat = ds.coarsen(coord=5).mean()
     assert dat.attrs == _attrs
 
+    # Test kept attrs in original object
+    xr.testing.assert_identical(ds, ds2)
+
 
 def test_rolling_keep_attrs():
     _attrs = {"units": "test", "long_name": "testing"}
diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py
index 169105bb..efebe09e 100644
--- a/xarray/tests/test_variable.py
+++ b/xarray/tests/test_variable.py
@@ -1949,7 +1949,10 @@ class TestVariable(VariableSubclassobjects):
         # Test kept attrs
         with set_options(keep_attrs=True):
             new = Variable(["coord"], np.linspace(1, 10, 100), attrs=_attrs).coarsen(
-                windows={"coord": 1}, func=test_func, boundary="exact", side="left"
+                windows={"coord": 1},
+                func=test_func,
+                boundary="exact",
+                side="left",
             )
         assert new.attrs == _attrs