Python Xarray, sort by index or dimension?

Question:

Is there a sort_index or sort_by_dimension method of some kind in xarray, much like pandas.DataFrame.sort_index(), where I can sort a xarray.DataArray object by one of its dimensions? In terms of usage, I’m thinking of something like data_array.sort(dim="dimension_name").

Asked By: Ray

||

Answers:

Do you mean rearrange the dimensions?
For example, with .transpose() you can reorder the dimensions of both a whole dataset or specific dataarray, e.g.

ds.transpose('lat','lon','time')
Answered By: dreab

I couldn’t find a good method to do this built into xarray, so I made a new array by taking a slice with the sorted values from the coordinate I wanted to sort: da_sorted=da.loc[{'lon':sorted(da.coords['lon'].values)}]

Here’s a larger example with test data showing what it looks like in practice:

import numpy as np
import xarray as xr

data = np.random.rand(2, 3)

lat = [47, 45]

lon = [-118, -122, -120]
da = xr.DataArray(data, coords=[lat, lon], dims=['lat', 'lon'])


>>>> da
<xarray.DataArray (lat: 2, lon: 3)>
array([[ 0.16118 ,  0.215621,  0.599749],
       [ 0.144778,  0.984167,  0.416469]])
Coordinates:
  * lat      (lat) int64 47 45
  * lon      (lon) int64 -118 -122 -120,


da_sorted = da.loc[{'lon':sorted(da.coords['lon'].values)}]
>>>> da_sorted
<xarray.DataArray (lat: 2, lon: 3)>
array([[ 0.215621,  0.599749,  0.16118 ],
       [ 0.984167,  0.416469,  0.144778]])
Coordinates:
  * lat      (lat) int64 47 45
  * lon      (lon) int64 -122 -120 -118
Answered By: Kyle Heuton

xarray has now a built-in method for this task, which can be performed easily using .sortby() https://xarray.pydata.org/en/stable/generated/xarray.DataArray.sortby.html

The code posted by @Kyle Heuton becomes now:

import numpy as np
import xarray as xr

data = np.random.rand(2, 3)

lat = [47, 45]

lon = [-118, -122, -120]
da = xr.DataArray(data, coords=[lat, lon], dims=['lat', 'lon'])


>>>> da
<xarray.DataArray (lat: 2, lon: 3)>
array([[ 0.16118 ,  0.215621,  0.599749],
       [ 0.144778,  0.984167,  0.416469]])
Coordinates:
  * lat      (lat) int64 47 45
  * lon      (lon) int64 -118 -122 -120,


da_sorted = da.sortby(da.lon)
>>>> da_sorted
<xarray.DataArray (lat: 2, lon: 3)>
array([[ 0.215621,  0.599749,  0.16118 ],
       [ 0.984167,  0.416469,  0.144778]])
Coordinates:
  * lat      (lat) int64 47 45
  * lon      (lon) int64 -122 -120 -118
Answered By: clausmichele
Categories: questions Tags: ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.