pydata / pydata/xarray

Sort the values of an nD array

Open
#8,129 11 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
4.2k
Forks
1.4k
Avg merge
2d 15h
Merged PRs (30d)
14

Description

Is your feature request related to a problem?

As far as I know, there is no straightforward API in xarray to do what np.sort or pandas.sort_values does. We have DataArray.sortby("x"), which will sort the array according to the coordinate itself. But if instead you want to sort the values of the array to be monotonic, you're on your own. There are probably a lot of ways we could do this, but I ended up with the couple line solution below after a little trial and error.

Describe the solution you'd like

Would there be interest in implementing a Dataset/DataArray.sort_values(dim="x") method?

Note: this 1D example is not really relevant, see the 2D version and more obvious implementation in comments below for what I really want.

def sort_values(self, dim: str):
    sort_idx = self.argsort(axis=self.get_axis_num(dim)).drop_vars(dim)
    return self.isel({dim: sort_idx}).drop_vars(dim).assign_coords({dim: self[dim]})

The goal is to handle arrays that we want to monotize like so:

da = xr.DataArray([1, 3, 2, 4], coords={"x": [1, 2, 3, 4]})
da.sort_values("x")

<xarray.DataArray (x: 4)>
array([1, 2, 3, 4])
Coordinates:
  * x        (x) int64 1 2 3 4

image

In addition to sortby which can deal with an array that is just unordered according to the coordinate:

da = xr.DataArray([1, 3, 2, 4], coords={"x": [1, 3, 2, 4]})
da.sortby("x")

<xarray.DataArray (x: 4)>
array([1, 2, 3, 4])
Coordinates:
  * x        (x) int64 1 2 3 4
Describe alternatives you've considered

I don't know if argsort is dask-enabled (the docs just point to the numpy function). Is there a more intelligent way to implement this with apply_ufunc and something else? I assume chunking in the sort dimension would be problematic.

Additional context

Some past related threads on this topic:
https://github.com/pydata/xarray/issues/3957
https://stackoverflow.com/questions/64518239/sorting-dataset-along-axis-with-dask

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading the existing DataArray.sortby, argsort, isel, and apply_ufunc entry points mentioned in the issue. Clarify the expected Dataset and DataArray behavior for multidimensional arrays and dask chunking. Done means a sort_values(dim) API handles the requested examples with documented behavior and tests.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, pandas, python
Domain
api, data
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.