Sort the values of an nD array
Nobody has claimed this yet.
- 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
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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