Allow `.groupby().map()` to return scalars?

[max-sixty]

Is your feature request related to a problem?

I'm trying to get a count of unique values along a dimension. It's not so easy, unless I'm missing something.

One approach is:

da = xr.tutorial.load_dataset('air_temperature').to_dataarray()

xr.apply_ufunc(lambda x: len(np.unique(x)), da, input_core_dims=[["lat","lon"]], vectorize=True)  # NB: requires vectorize to work!

but apply_ufunc is generally too complex for normal users to use, I think.

Another approach could be

da.groupby('time').map(lambda x: len(np.unique(x)))

But this raises:

AttributeError: 'int' object has no attribute 'dims'

Instead, surrounding the expression with DataArray makes it work:

da.groupby('uptr').map(lambda x: xr.DataArray(len(np.unique(x))))

<xarray.DataArray (time: 2920)> Size: 23kB
array([546, 547, 545, ..., 555, 558, 566])
Coordinates:
  * time     (time) datetime64[ns] 23kB 2013-01-01 ... 2014-12-31T18:00:00

Describe the solution you'd like

Should we allow returning scalars from .groupby().map()?

I don't think there can be any ambiguity on what the result should be...

Describe alternatives you've considered

No response

Additional context

No response

[dcherian]

map expects xarray objects to be returned. We should improve the error message to point to reduce which should work for you

[dcherian]

Though perhaps we should add nunique along a dimension implemented as sort along axis, succeeding-elements-are-not-equal along axis handling NaNs, then sum along axis.

xref https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.nunique.html

[max-sixty]

Doesn't .reduce have similar restrictions though?

If we return a scalar then we get:

da.reduce(lambda x, axis: len(np.unique(x, axis=axis)), dim=['lat','lon'])

...


File ~/workspace/xarray/xarray/namedarray/core.py:910, in NamedArray.reduce(self, func, dim, axis, keepdims, **kwargs)
    906     if isinstance(axis, tuple) and len(axis) == 1:
    907         # unpack axis for the benefit of functions
    908         # like np.argmin which can't handle tuple arguments
    909         axis = axis[0]
--> 910     data = func(self.data, axis=axis, **kwargs)
    911 else:
    912     data = func(self.data, **kwargs)

Cell In[15], line 1, in <lambda>(x, axis)
----> 1 da.reduce(lambda x, axis: len(np.unique(x, axis=axis)), dim=['lat','lon'])

File ~/workspace/xarray/.venv/lib/python3.11/site-packages/numpy/lib/_arraysetops_impl.py:295, in unique(ar, return_index, return_inverse, return_counts, axis, equal_nan)
    293 # axis was specified and not None
    294 try:
--> 295     ar = np.moveaxis(ar, axis, 0)
    296 except np.exceptions.AxisError:
    297     # this removes the "axis1" or "axis2" prefix from the error message
    298     raise np.exceptions.AxisError(axis, ar.ndim) from None

File ~/workspace/xarray/.venv/lib/python3.11/site-packages/numpy/_core/numeric.py:1496, in moveaxis(a, source, destination)
   1494 destination = normalize_axis_tuple(destination, a.ndim, 'destination')
   1495 if len(source) != len(destination):
-> 1496     raise ValueError('`source` and `destination` arguments must have '
   1497                      'the same number of elements')
   1499 order = [n for n in range(a.ndim) if n not in source]
   1501 for dest, src in sorted(zip(destination, source)):

Or am I making some basic mistake?

[dcherian]

I think that's an error in using unique? I don't think it supports multiple axes. Also not sure the single axis version does what you want it to

[max-sixty]

OK, right, yes I think we need our own method to do this, agree with:

implemented as sort along axis, succeeding-elements-are-not-equal along axis handling NaNs, then sum along axis.

---

...though my claim stands re allowing scalars in groupby().map() seems reasonable, without having thought about the downsides?

[dcherian]

I like the invariant that map expects xarray objects and reduce expects bare arrays. If the error pointed this out, it would be more helpful

da.groupby('time').map(lambda x: xr.DataArray(len(np.unique(x))))

It'd be nice to make this work, but it looks quite broken?

da.groupby('time').reduce(lambda x: len(np.unique(x)))

[max-sixty]

Ah great, I hadn't realized you meant .reduce on the .groupby rather than the array — this actually works pretty well with the dims specified: da.groupby('time').reduce(lambda x: len(np.unique(x)), dim=...)

(I still think not quite well enough for an ordinary user to figure it out, but not sure that's feasible here)

I'll close this, thanks for the discussion.

I think nunique could be a worthwhile method to add at some point.

[dcherian]

I still think not quite well enough for an ordinary user to figure it out, but not sure that's feasible here

I think better error message would help.