GroupBy.map with keep_attrs=True gives error

[mktippett]

I wanted to apply a function (not as simple as np.mean in the example below) using xarray.map and keep the coordinate (lat) attributes. I got an error with keep_attrs=True that I didn't understand.

import numpy as np
import xarray as xr
ds =  xr.tutorial.open_dataset("air_temperature")
test_map = ds.groupby('lat').map(np.mean, keep_attrs=True)
test_map.lat

Error

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-26-e572326fea03> in <module>
      2 import xarray as xr
      3 ds =  xr.tutorial.open_dataset("air_temperature")
----> 4 test_map = ds.groupby('lat').map(np.mean, keep_attrs=True)
      5 test_map.lat

/opt/conda/lib/python3.8/site-packages/xarray/core/groupby.py in map(self, func, args, shortcut, **kwargs)
    921         # ignore shortcut if set (for now)
    922         applied = (func(ds, *args, **kwargs) for ds in self._iter_grouped())
--> 923         return self._combine(applied)
    924 
    925     def apply(self, func, args=(), shortcut=None, **kwargs):

/opt/conda/lib/python3.8/site-packages/xarray/core/groupby.py in _combine(self, applied)
    941     def _combine(self, applied):
    942         """Recombine the applied objects like the original."""
--> 943         applied_example, applied = peek_at(applied)
    944         coord, dim, positions = self._infer_concat_args(applied_example)
    945         combined = concat(applied, dim)

/opt/conda/lib/python3.8/site-packages/xarray/core/utils.py in peek_at(iterable)
    181     """
    182     gen = iter(iterable)
--> 183     peek = next(gen)
    184     return peek, itertools.chain([peek], gen)
    185 

/opt/conda/lib/python3.8/site-packages/xarray/core/groupby.py in <genexpr>(.0)
    920         """
    921         # ignore shortcut if set (for now)
--> 922         applied = (func(ds, *args, **kwargs) for ds in self._iter_grouped())
    923         return self._combine(applied)
    924 

<__array_function__ internals> in mean(*args, **kwargs)

TypeError: _mean_dispatcher() got an unexpected keyword argument 'keep_attrs'

Output of xr.show_versions()

``` INSTALLED VERSIONS ------------------ commit: None python: 3.8.5 | packaged by conda-forge | (default, Aug 21 2020, 18:21:27) [GCC 7.5.0] python-bits: 64 OS: Linux OS-release: 4.19.112+ machine: x86_64 processor: x86_64 byteorder: little LC_ALL: en_US.UTF-8 LANG: en_US.UTF-8 LOCALE: en_US.UTF-8 libhdf5: 1.10.6 libnetcdf: 4.7.4

xarray: 0.16.0
pandas: 1.1.1
numpy: 1.19.1
scipy: 1.5.2
netCDF4: 1.5.4
pydap: None
h5netcdf: None
h5py: 2.10.0
Nio: None
zarr: 2.4.0
cftime: 1.2.1
nc_time_axis: 1.2.0
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: 1.3.2
dask: 2.20.0
distributed: 2.24.0
matplotlib: 3.2.2
cartopy: 0.18.0
seaborn: 0.10.1
numbagg: None
pint: None
setuptools: 49.6.0.post20200814
pip: 20.2.2
conda: 4.8.3
pytest: None
IPython: 7.17.0
sphinx: None

</details>

[dcherian]

That code passes keep_attrs to np.mean so it won't work. we should add keep_attrs as a kwarg to map

you'll have to do it like this

import numpy as np
import xarray as xr
ds =  xr.tutorial.open_dataset("air_temperature")
with xr.set_options(keep_attrs=True):
    test_map = ds.groupby('lat').map(np.mean)
test_map.lat

but it doesn't work for the coordinate variables. We'll have to fix it. Right now, you'll have to propagate attrs manually

[mathause]

Implementing keep_attrs as a kwarg to map and

with xr.set_options(keep_attrs=True):
    test_map = ds.groupby('lat').map(np.mean)

not working for coordinate variables are two different issues, right?

[dcherian]

Yes.

Somewhat related: I'm always confused about whether ds.groupby(..., keep_attrs=True).mean() or ds.groupby(...).mean(keep_attrs=True) is correct. (similarly for rolling, coarsen etc.)