How to sort the coordinates and variables of an xr.Dataset

Question

I would like to sort the coordinates and variables of an xarray Dataset in alphabetical order. I have tried to do this using ds.transpose(*sorted(ds.dims)). This seems to sort the coordinates/dimensions of each DataArray in the Dataset, but not the coordinates of the Dataset itself.

Example:

>>> ds = xr.Dataset(
...     {
...         'z': (['c', 'a', 'b'], np.ones(shape=(2, 2, 2))),
...         'x': (['a', 'b', 'c'], np.zeros(shape=(2, 2, 2))),
...         'y': (['c'], [0, 1]),
...     },
...     coords={'c': [30, 31], 'a': [10, 11], 'b': [20, 21]}
... )
    
>>> ds.transpose('a', 'b', 'c')
<xarray.Dataset>
Dimensions:  (c: 2, a: 2, b: 2)
Coordinates:
  * c        (c) int64 30 31
  * a        (a) int64 10 11
  * b        (b) int64 20 21
Data variables:
    z        (a, b, c) float64 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0
    x        (a, b, c) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
    y        (c) int64 0 1

Expected behaviour is that the coordinates and dimensions of the entire xr.Dataset would be sorted as 'a', 'b', 'c'. However, you can see that only the dimensions of the data variables themselves are in this order.

How can I sort the coordinates and variables of a Dataset, not just the array dimensions?

Any help is deeply appreciated, thank you!

Answer 1

You can do this in xarray with xr.Dataset.reindex_like and xr.Dataset.reindex. The former is helpful if you have another object that you want to match.

ds = xr.Dataset(
data_vars={
'z': (['c', 'a', 'b'], np.ones(shape=(2, 2, 2))),
'x': (['a', 'b', 'c'], np.zeros(shape=(2, 2, 2))),
'y': (['c'], [0, 1]),
},
coords={'c': [30, 31], 'a': [10, 11], 'b': [20, 21]}
)

current_indexes = ds.indexes
desired_order = ['a', 'b', 'c']
reordered_indexes = {index_name: current_indexes[index_name] for index_name in desired_order}
reordered_ds = ds.reindex(reordered_indexes)

>>> ds
<xarray.Dataset>
Dimensions:  (c: 2, a: 2, b: 2)
Coordinates:
  * c        (c) int64 30 31
  * a        (a) int64 10 11
  * b        (b) int64 20 21
Data variables:
    z        (c, a, b) float64 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0
    x        (a, b, c) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
    y        (c) int64 0 1
>>> reordered_ds
<xarray.Dataset>
Dimensions:  (a: 2, b: 2, c: 2)
Coordinates:
  * a        (a) int64 10 11
  * b        (b) int64 20 21
  * c        (c) int64 30 31
Data variables:
    z        (c, a, b) float64 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0
    x        (a, b, c) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
    y        (c) int64 0 1