issue_comments

almost no attention is paid to minimizing memory consumption (whether through in-place operations, or more generally minimizing temporary memory usage).

I think we'd be open to fixing this when it doesn't compromise readability. Can you open a new issue with some particularly bad examples?

I wouldn't necessarily say that it's particularly bad, but see the discussion following https://github.com/pydata/xarray/pull/2922#issuecomment-601496897.

I'm not familiar at all with the various numpy interfaces, so I can't offer any input implementation-wise. But as a user, being able to do operations in place (via out or otherwise) is extremely useful when dealing with large arrays under memory constraints. In fact my one "philosophical" beef with xarray is that it seems almost no attention is paid to minimizing memory consumption (whether through in-place operations, or more generally minimizing temporary memory usage).

What is the case for having out kwargs?

It lets you reuse memory you already have. In particular for a simple operation like clip, you can do it in-place: da.clip(..., out=da.values). Very useful if you deal with lots of data and memory is a concern.

I'm wondering if one could just have a generic implementation of DataArray.func(*args, **kwargs) for any unrecognized "func" that calls np.func(self, *args, **kwargs) (perhaps conditional on some global numpy_fallthrough option)? And similarly for Dataset.func(*args, **kwargs).

Indeed, this is reported in https://github.com/pandas-dev/pandas/issues/35466#issuecomment-678407125 and https://github.com/pandas-dev/pandas/issues/35830. Also https://github.com/pandas-dev/pandas/pull/35478.

Does this fix #4363?

@TomNicholas , yes, thank you.

All good points:

What could be done, though is to only do da = da.fillna(0.0) if da contains NaNs.

Good idea, though I don't know what the performance hit would be of the extra check (in the case that da does contain NaNs, so the check is for naught).

I assume so. I don't know what kind of temporary variables np.einsum creates. Also np.einsum is wrapped in xr.apply_ufunc so all kinds of magic is going on.

Well, (da * weights) will be at least as large as da. I'm not certain, but I don't think np.einsum creates huge temporary arrays.

Do you want to leave it away for performance reasons? Because it was a deliberate decision to not support NaNs in the weights and I don't think this is going to change.

Yes. You can continue not supporting NaNs in the weights, yet not explicitly check that there are no NaNs (optionally, if the caller assures you that there are no NaNs).

None of your suggested functions support NaNs so they won't work.

Correct. These have nothing to do with the NaNs issue.

For profiling memory usage, I use psutil.Process(os.getpid()).memory_info().rss for current usage and resource.getusage(resource.RUSAGE_SElF).ru_maxrss for peak usage (on linux).

@mathause, have you considered using these functions? - np.average() to calculate weighted mean(). - np.cov() to calculate weighted cov(), var(), and std(). - sp.stats.cumfreq() to calculate weighted median() (I haven't thought this through). - sp.spatial.distance.correlation() to calculate weighted corrcoef(). (Of course one could also calculate this from weighted cov() (see above), but first need to mask the two arrays simultaneously.) - sklearn.utils.extmath.weighted_mode() to calculate weighted mode(). - gmisclib.weighted_percentile.{wp,wtd_median}() to calculate weighted quantile() and median().

@mathause, ideally dot() would support skipna, so you could eliminate the da = da.fillna(0.0) and pass the skipna down the line. But alas it doesn't...

(da * weights).sum(dim=dim, skipna=skipna) would likely make things worse, I think, as it would necessarily create a temporary array of sized at least da, no?

Either way, this only addresses the da = da.fillna(0.0), not the mask = da.notnull().

Also, perhaps the test if weights.isnull().any() in Weighted.__init__() should be optional?

Maybe I'm more sensitive to this than others, but I regularly deal with 10-100GB arrays.

@max-sixty, I wish I could, but I'm afraid that I cannot submit code due to employer limitations.

I realize this is a bit late, but I'm still concerned about memory usage, specifically in https://github.com/pydata/xarray/blob/master/xarray/core/weighted.py#L130 and https://github.com/pydata/xarray/blob/master/xarray/core/weighted.py#L143. If da.sizes = {'dim_0': 100000, 'dim_1': 100000}, the two lines above will cause da.weighted(weights).mean('dim_0') to create two simultaneous temporary 100000x100000 arrays, which could be problematic.

I would have implemented this using apply_ufunc, so that one creates these temporary variables only on as small an array as absolutely necessary -- in this case just of size sizes['dim_0'] = 100000. (Much as I would like to, I'm afraid I'm not able to contribute code.) Of course this won't help in the case one is summing over all dimensions, but might as well minimize memory usage in some cases even if not in all.

@keewis, thanks for the suggestions. Both seem reasonable.

In your first example, if you wanted to prohibit obj.weighted.sum(dim), you could just check for self._weight in sum(). Though I suppose it would be nice to be able to have the interpreter enforce the requirement and not have to do an explicit check in every method.

Note that inferring dimensions from coords when it is a list of tuples does still work (with no deprecation warning): ``` In [1]: import numpy as np, xarray as xr

In [2]: xr.DataArray(np.zeros((2, 2)), coords=[('x', [1, 2]), ('y', [1, 2])])
Out[2]: <xarray.DataArray (x: 2, y: 2)> array([[0., 0.], [0., 0.]]) Coordinates: * x (x) int64 1 2 * y (y) int64 1 2 ```

Note that with the apply_ufunc implementation we're only reshaping dims-sized ndarrays, not (necessarily) the whole DataArray, so maybe it's not too bad? It might be better to first sort dims to be in the same order as self.dims. i.e. dims = [dim_ for dim_ in self.dims if dim_ in dims]. But I'm just speculating.

A few minor tweaks needed: ``` In [20]: import bottleneck

In [21]: xr.apply_ufunc( ...: lambda x: bottleneck.rankdata(x).reshape(x.shape), ...: d, ...: input_core_dims=[['xyz', 'abc']], ...: output_core_dims=[['xyz', 'abc']], ...: vectorize=True ...: ).transpose(*d.dims)
Out[21]: <xarray.DataArray (abc: 4, xyz: 3)> array([[ 1., 2., 3.], [ 4., 5., 6.], [ 7., 8., 9.], [10., 11., 12.]]) Dimensions without coordinates: abc, xyz ```

Despite what the docs say, bottleneck.{nan}rankdata(a) returns a 1-dimensional ndarray, not an array with the same shape as a.

What's wrong with the following? (Still need to deal with pct and keep_attrs.) apply_ufunc( bottleneck.{nan}rankdata, self, input_core_dims=[dims], output_core_dims=[dims], vectorize=True )

Per https://kwgoodman.github.io/bottleneck-doc/reference.html#bottleneck.rankdata, "The default (axis=None) is to rank the elements of the flattened array."

Assuming dims is a non-empty list of dimensions, the following code seems to work: temp_dim = '__temp_dim__' return da.stack(**{temp_dim: dims}).\ rank(temp_dim, pct=pct, keep_attrs=keep_attrs).\ unstack(temp_dim).transpose(*da.dims).\ drop_vars([dim_ for dim_ in dims if dim_ not in da.coords])

I'm afraid I'm not able to submit a PR. Sorry.

I think the only necessary changes are (a) delete the if method != "__call__" check (https://github.com/pydata/xarray/blob/master/xarray/core/arithmetic.py#L49), and (b) in the apply_ufunc() call, replace ufunc with getattr(ufunc, method) (https://github.com/pydata/xarray/blob/master/xarray/core/arithmetic.py#L71).

This issue is still relevant.

This is fixed in Pandas 1.0.1.

Pandas bug: https://github.com/pandas-dev/pandas/issues/31501

I think this issue is still relevant.

I was thinking a tuple/list (corresponding to args) of dicts (dim -> value) containing the non-input_core_dims being evaluated. (If it weren't for exclude_dims, if I understand it correctly, I think one would need only a single dict (dim -> value).)

I would be fine with this being an optional kwarg to the actual func.

I think this is still an issue.

Also, the following code seems to accomplish the same as the above:

``` def apply_func_rolling(func, args, kwargs): # determine rolling parameters, and remove them from kwargs apply_func_kwargs = {'input_core_dims', 'output_core_dims', 'vectorize', 'join', 'dataset_join', 'keep_attrs', 'exclude_dims', 'dataset_fill_value', 'kwargs', 'dask', 'output_dtypes', 'output_sizes'} min_periods = kwargs.pop('min_periods', None) center = kwargs.pop('center', False) dim = xr.core.utils.either_dict_or_kwargs(kwargs.pop('dim', None), {k: v for k, v in kwargs.items() if k not in apply_func_kwargs}, 'apply_func_rolling') if len(dim) != 1: raise ValueError("precisely one rolling dimension must be specified") rolling_dim = list(dim.keys())[0] kwargs.pop(rolling_dim) temp_rolling_dim = 'temp{}__'.format(rolling_dim) # change input_core_dims rolling_dim values to temp_rolling_dim input_core_dims = kwargs.get('input_core_dims', None) if input_core_dims: kwargs['input_core_dims'] = [[(temp_rolling_dim if (dim_ == rolling_dim) else dim_) for dim_ in dims_] for dims_ in input_core_dims] # change exclude_dims rolling_dim values to temp_rolling_dim exclude_dims = kwargs.get('exclude_dims', None) if exclude_dims: kwargs['exclude_dims'] = [[(temp_rolling_dim if (dim_ == rolling_dim) else dim_) for dim_ in dims_] for dims_ in exclude_dims] # call apply_func() with rolling-constructed objects return xr.apply_ufunc(func, [(arg.rolling(dim=dim, min_periods=min_periods, center=center). construct(temp_rolling_dim) if (rolling_dim in arg.dims) else arg) for arg in args], **kwargs)

apply_func_rolling(lambda a, b, w: ..., variables, observations, weights, date=N, input_core_dims=[['date', 'dim1', 'dim2', 'var'], ['date', 'dim1', 'dim2'], ['date', 'dim1', 'dim2']], output_core_dims=[['var']], vectorize=True) ```

construct method does not allocate that large array in memory. It uses the strided trick and therefore consumers only the order of 1000x1000x1000.

Ah. I didn't realize that. Good to know.

What I'm actually looking to do is a rolling weighted regression. I have three DataArrays: - observations, dims=('date', 'dim1', 'dim2') - variables, dims=('date', 'dim1', 'dim2', 'var') - weights, dims=('date', 'dim1', 'dim2')

I want to calculate a regression_coefficients DataArray with dims=('date', 'var'), where for each date it has the weighted regression coefficients calculated over the trailing N dates (over 'dim1' and 'dim2'). One way would be to put the three DataArrays in a Dataset, and then use a newly-defined Dataset.rolling().apply(). Another way would be to use an enhanced version of apply_ufunc() that can take Rolling objects. But now that I know that DataArrayRolling.construct() won't kill my machine, I'll try apply_ufunc() with the three DataArrayRolling.construct() objects. I'd welcome other suggestions.

OK, I seem to have got my problem working using: apply_ufunc(lambda a, b, w: ..., variables.rolling(date=N).construct('temp_date'), observations.rolling(date=N).construct('temp_date'), weights.rolling(date=N).construct('temp_date'), input_core_dims=[['temp_date', 'dim1', 'dim2', 'var'], ['temp_date', 'dim1', 'dim2'], ['temp_date', 'dim1', 'dim2']], output_core_dims=[['var']], vectorize=True) Still, I wonder if there isn't a more "natural" way of accomplishing this.

I think there are actually a couple different ways Rolling.apply could work, but this seems like one possible way:

``` from xarray.core.utils import maybe_wrap_array from xarray.core.combine import concat

def rolling_apply(rolling, func, args, kwargs): applied = [maybe_wrap_array(label, func(arr, args, **kwargs)) for label, arr in rolling] combined = concat(applied, dim=rolling.obj.coords[rolling.dim]) return combined ```

The problem I have with Rolling.construct is the same that I have with Rolling.reduce: it's very (potentially) memory-inefficient. E.g. consider a 1000x1000x1000 array for which I want to apply rolling window of length 500 along the final dimension; I believe Rolling.construct/reduce will construct a 1000x1000x1000x500 array. This can quickly get out of hand.

Separately, maybe apply_ufunc should accepts Rolling objects?

I think what is needed are DataArrayRolling.apply and DatasetRolling.apply (like DataArrayGroupBy.apply and DatasetGroupBy.apply). The problem with the reduce methods is that they are memory-inefficient.

This code isn't particularly pretty, and I'm not sure if it handles all cases, but it enables serialization of MultiIndex indices by calling ds.mi.encode_multiindices() before serializing and ds.mi.decode_multiindices() after deserializing.

``` @xr.register_dataset_accessor('mi') class MiscDatasetAccessor(): def init(self, xarray_obj): self._obj = xarray_obj

def encode_multiindices(self):
    result = self._obj
    for name, index in list(self._obj.indexes.items()):
        if isinstance(index, pd.MultiIndex):
            temp_name = '__' + name
            new_coords = {'{}__{}'.format(temp_name, level_name): level_values.rename(None)
                          for level_name, level_values in zip(index.names, index.levels)}
            new_coords[temp_name] = xr.DataArray(index.labels,
                                                 dims=('{}__names__'.format(temp_name),
                                                       '{}__num__'.format(temp_name)),
                                                 coords={'{}__names__'.format(temp_name): index.names,
                                                         '{}__num__'.format(temp_name): list(range(len(index)))},
                                                 attrs={'__is_multiindex': 1})
            result = result.drop(name).assign_coords(**new_coords)
    return result

def decode_multiindices(self):
    result = self._obj
    for temp_name, da in list(self._obj.coords.items()):
        if temp_name.startswith('__') and da.attrs.get('__is_multiindex', False):
            name = temp_name[2:]
            level_names = da.coords['{}__names__'.format(temp_name)].values
            levels = [result.coords['{}__{}'.format(temp_name, level_name)].values for level_name in level_names]
            labels = da.values
            result = result.assign_coords(**{name: pd.MultiIndex(levels=levels, labels=labels, names=level_names)})
            result = result.drop(['{}__{}'.format(temp_name, level_name) for level_name in level_names] +
                                 list(da.dims) + [temp_name])
    return result

```

I second this request.

The following may not be optimal, but seems to work for me as a keepdims=True version of reduce(): def dim_preserving_reduce(self, func, dim=None, axis=None, label=None, keep_attrs=False, **kwargs): if axis is not None: dim = np.take(self._obj.dims, axis, mode='wrap') dims = dim if isinstance(dim, (list, tuple)) else [dim] dims_coords = {dim: [lab] for dim, lab in zip(dims, (label if isinstance(label, list) else [label]))} return self._obj.reduce(func, dim=dims, keep_attrs=keep_attrs, **kwargs). \ expand_dims(dims, axis=[self._obj.dims.index(dim) for dim in dims]). \ assign_coords(**dims_coords)

I added a note to whats-new.rst.

One slight oddity is that formatting.format_array_flat(np.arange(4), 0) returns 0 ... 3 even though 0 1 2 3 would take up the same number of characters. I'm not inclined to add a special case for this, but let me know if you think I should.

@shoyer, I think I've implemented all your suggestions. Let me know what you think. (I haven't yet updated whats-new.rst.)

I included sample output in https://github.com/pydata/xarray/pull/2293#issuecomment-405369643 and https://github.com/pydata/xarray/pull/2293/files#diff-f82411dbe6aa53e3b6a5d9c2b601094c.

Replaced with https://github.com/pydata/xarray/pull/2293.

Sample output:

``` (base) C:\Users\Seth\github\xarray>ipython Python 3.6.5 | packaged by conda-forge | (default, Apr 6 2018, 16:13:55) [MSC v.1900 64 bit (AMD64)] Type 'copyright', 'credits' or 'license' for more information IPython 6.4.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import xarray as xr

In [2]: words = "This is the time for all good men to come to the aid of their country".split(' ')

In [3]: for i in range(0, len(words) + 1): ...: print("-------------------------------------------------------------------------------") ...: print(xr.DataArray(words[:i], dims=('foo',), coords={'foo': words[:i]})) ...:

<xarray.DataArray (foo: 0)> array([], dtype=float64) Coordinates: * foo (foo) float64

<xarray.DataArray (foo: 1)> array(['This'], dtype='<U4') Coordinates: * foo (foo) <U4 'This'

<xarray.DataArray (foo: 2)> array(['This', 'is'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is'

<xarray.DataArray (foo: 3)> array(['This', 'is', 'the'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the'

<xarray.DataArray (foo: 4)> array(['This', 'is', 'the', 'time'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time'

<xarray.DataArray (foo: 5)> array(['This', 'is', 'the', 'time', 'for'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for'

<xarray.DataArray (foo: 6)> array(['This', 'is', 'the', 'time', 'for', 'all'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for' 'all'

<xarray.DataArray (foo: 7)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for' 'all' 'good'

<xarray.DataArray (foo: 8)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for' 'all' 'good' 'men'

<xarray.DataArray (foo: 9)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for' 'all' 'good' 'men' 'to'

<xarray.DataArray (foo: 10)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' ... 'good' 'men' 'to' 'come'

<xarray.DataArray (foo: 11)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come', 'to'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for' ... 'men' 'to' 'come' 'to'

<xarray.DataArray (foo: 12)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come', 'to', 'the'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for' ... 'to' 'come' 'to' 'the'

<xarray.DataArray (foo: 13)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come', 'to', 'the', 'aid'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' ... 'come' 'to' 'the' 'aid'

<xarray.DataArray (foo: 14)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come', 'to', 'the', 'aid', 'of'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for' ... 'to' 'the' 'aid' 'of'

<xarray.DataArray (foo: 15)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come', 'to', 'the', 'aid', 'of', 'their'], dtype='<U5') Coordinates: * foo (foo) <U5 'This' 'is' 'the' 'time' ... 'the' 'aid' 'of' 'their'

<xarray.DataArray (foo: 16)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come', 'to', 'the', 'aid', 'of', 'their', 'country'], dtype='<U7') Coordinates: * foo (foo) <U7 'This' 'is' 'the' 'time' ... 'aid' 'of' 'their' 'country' ```

For some reason (presumably due to Github's outage today), this PR isn't updating to reflect the latest commits in my branch. So I'm going to close this one and create a new one.

Sample output:

``` (base) C:\Users\Seth\github\xarray>ipython Python 3.6.5 | packaged by conda-forge | (default, Apr 6 2018, 16:13:55) [MSC v.1900 64 bit (AMD64)] Type 'copyright', 'credits' or 'license' for more information IPython 6.4.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import xarray as xr

In [2]: words = "This is the time for all good men to come to the aid of their country".split(' ')

In [3]: for i in range(0, len(words) + 1): ...: print("-------------------------------------------------------------------------------") ...: print(xr.DataArray(words[:i], dims=('foo',), coords={'foo': words[:i]})) ...:

<xarray.DataArray (foo: 0)> array([], dtype=float64) Coordinates: * foo (foo) float64

<xarray.DataArray (foo: 1)> array(['This'], dtype='<U4') Coordinates: * foo (foo) <U4 'This'

<xarray.DataArray (foo: 2)> array(['This', 'is'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is'

<xarray.DataArray (foo: 3)> array(['This', 'is', 'the'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the'

<xarray.DataArray (foo: 4)> array(['This', 'is', 'the', 'time'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time'

<xarray.DataArray (foo: 5)> array(['This', 'is', 'the', 'time', 'for'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for'

<xarray.DataArray (foo: 6)> array(['This', 'is', 'the', 'time', 'for', 'all'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for' 'all'

<xarray.DataArray (foo: 7)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for' 'all' 'good'

<xarray.DataArray (foo: 8)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for' 'all' 'good' 'men'

<xarray.DataArray (foo: 9)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' ... 'all' 'good' 'men' 'to'

<xarray.DataArray (foo: 10)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' ... 'good' 'men' 'to' 'come'

<xarray.DataArray (foo: 11)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come', 'to'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' ... 'men' 'to' 'come' 'to'

<xarray.DataArray (foo: 12)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come', 'to', 'the'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' ... 'to' 'come' 'to' 'the'

<xarray.DataArray (foo: 13)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come', 'to', 'the', 'aid'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' ... 'come' 'to' 'the' 'aid'

<xarray.DataArray (foo: 14)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come', 'to', 'the', 'aid', 'of'], dtype='<U4') Coordinates: * foo (foo) <U4 'This' 'is' 'the' 'time' 'for' ... 'to' 'the' 'aid' 'of'

<xarray.DataArray (foo: 15)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come', 'to', 'the', 'aid', 'of', 'their'], dtype='<U5') Coordinates: * foo (foo) <U5 'This' 'is' 'the' 'time' ... 'the' 'aid' 'of' 'their'

<xarray.DataArray (foo: 16)> array(['This', 'is', 'the', 'time', 'for', 'all', 'good', 'men', 'to', 'come', 'to', 'the', 'aid', 'of', 'their', 'country'], dtype='<U7') Coordinates: * foo (foo) <U7 'This' 'is' 'the' 'time' ... 'of' 'their' 'country' ```

I'm not a fan of auto-flattening either, but that's what nd.argsort() does...

One option is to have DataArray.arg{min,max,sort}() all take an optional flag argument specifying whether to return integer indices or index labels. But I think my preference would be be to have six separate functions: DataArray.{idx,}arg{min,max,sort}() (or some such nomenclature that includes arg in all six functions).

I think that makes sense, though I don't quite understand what would go in its place. Another possibility -- perhaps a bad one -- is to permute the values in the sorted dimension so that they match the resulting values (i.e. something like result.coords[dim] = np.take(da.coords[dim].values, result.values, axis=axis)).

Note that ndarray.argsort(axis=None) sorts the flattened array, so the returned DataArray should respect this

Alternative suggestion: have DataArray.argsort() return an ndarray filled with labels from the sorted dimension, i.e. something like: class DataArray: def argsort(self, **kwargs): # TODO: update kwargs['axis'] based 'axis' and 'dim', and remove 'dim' if kwargs['axis'] is None: kwargs['axis'] = -1 return self.stack(dim=self.dims).argsort(**kwargs) return np.take(self.coords[self.dims[kwargs['axis']].values, self.values.argsort(**kwargs))

BTW, I'm just thinking in terms of ndarrays. Someone more knowledgeable than me may want to consider how to make it work intelligently with dask.