Reshaping and reorganizing data

These methods allow you to reorganize

Reordering dimensions

To reorder dimensions on a DataArray or across all variables on a Dataset, use transpose() or the .T property:

In [1]: ds = xr.Dataset({'foo': (('x', 'y', 'z'), [[[42]]]), 'bar': (('y', 'z'), [[24]])})

In [2]: ds.transpose('y', 'z', 'x')
Out[2]: 
<xarray.Dataset>
Dimensions:  (x: 1, y: 1, z: 1)
Coordinates:
    *empty*
Unindexed dimensions:
    x, y, z
Data variables:
    foo      (y, z, x) int64 42
    bar      (y, z) int64 24

In [3]: ds.T
Out[3]: 
<xarray.Dataset>
Dimensions:  (x: 1, y: 1, z: 1)
Coordinates:
    *empty*
Unindexed dimensions:
    x, y, z
Data variables:
    foo      (z, y, x) int64 42
    bar      (z, y) int64 24

Converting between datasets and arrays

To convert from a Dataset to a DataArray, use to_array():

In [4]: arr = ds.to_array()

In [5]: arr
Out[5]: 
<xarray.DataArray (variable: 2, x: 1, y: 1, z: 1)>
array([[[[42]]],


       [[[24]]]])
Coordinates:
  * variable  (variable) <U3 'foo' 'bar'
Unindexed dimensions:
    x, y, z

This method broadcasts all data variables in the dataset against each other, then concatenates them along a new dimension into a new array while preserving coordinates.

To convert back from a DataArray to a Dataset, use to_dataset():

In [6]: arr.to_dataset(dim='variable')
Out[6]: 
<xarray.Dataset>
Dimensions:  (x: 1, y: 1, z: 1)
Coordinates:
    *empty*
Unindexed dimensions:
    x, y, z
Data variables:
    foo      (x, y, z) int64 42
    bar      (x, y, z) int64 24

The broadcasting behavior of to_array means that the resulting array includes the union of data variable dimensions:

In [7]: ds2 = xr.Dataset({'a': 0, 'b': ('x', [3, 4, 5])})

# the input dataset has 4 elements
In [8]: ds2
Out[8]: 
<xarray.Dataset>
Dimensions:  (x: 3)
Coordinates:
    *empty*
Unindexed dimensions:
    x
Data variables:
    a        int64 0
    b        (x) int64 3 4 5

# the resulting array has 6 elements
In [9]: ds2.to_array()
Out[9]: 
<xarray.DataArray (variable: 2, x: 3)>
array([[0, 0, 0],
       [3, 4, 5]])
Coordinates:
  * variable  (variable) <U1 'a' 'b'
Unindexed dimensions:
    x

Otherwise, the result could not be represented as an orthogonal array.

If you use to_dataset without supplying the dim argument, the DataArray will be converted into a Dataset of one variable:

In [10]: arr.to_dataset(name='combined')
Out[10]: 
<xarray.Dataset>
Dimensions:   (variable: 2, x: 1, y: 1, z: 1)
Coordinates:
  * variable  (variable) <U3 'foo' 'bar'
Unindexed dimensions:
    x, y, z
Data variables:
    combined  (variable, x, y, z) int64 42 24

Stack and unstack

As part of xarray’s nascent support for pandas.MultiIndex, we have implemented stack() and unstack() method, for combining or splitting dimensions:

In [11]: array = xr.DataArray(np.random.randn(2, 3),
   ....:                      coords=[('x', ['a', 'b']), ('y', [0, 1, 2])])
   ....: 

In [12]: stacked = array.stack(z=('x', 'y'))

In [13]: stacked
Out[13]: 
<xarray.DataArray (z: 6)>
array([ 0.469112, -0.282863, -1.509059, -1.135632,  1.212112, -0.173215])
Coordinates:
  * z        (z) MultiIndex
  - x        (z) object 'a' 'a' 'a' 'b' 'b' 'b'
  - y        (z) int64 0 1 2 0 1 2

In [14]: stacked.unstack('z')
Out[14]: 
<xarray.DataArray (x: 2, y: 3)>
array([[ 0.469112, -0.282863, -1.509059],
       [-1.135632,  1.212112, -0.173215]])
Coordinates:
  * x        (x) object 'a' 'b'
  * y        (y) int64 0 1 2

These methods are modeled on the pandas.DataFrame methods of the same name, although in xarray they always create new dimensions rather than adding to the existing index or columns.

Like DataFrame.unstack, xarray’s unstack always succeeds, even if the multi-index being unstacked does not contain all possible levels. Missing levels are filled in with NaN in the resulting object:

In [15]: stacked2 = stacked[::2]

In [16]: stacked2
Out[16]: 
<xarray.DataArray (z: 3)>
array([ 0.469112, -1.509059,  1.212112])
Coordinates:
  * z        (z) MultiIndex
  - x        (z) object 'a' 'a' 'b'
  - y        (z) int64 0 2 1

In [17]: stacked2.unstack('z')
Out[17]: 
<xarray.DataArray (x: 2, y: 3)>
array([[ 0.469112,       nan, -1.509059],
       [      nan,  1.212112,       nan]])
Coordinates:
  * x        (x) object 'a' 'b'
  * y        (y) int64 0 1 2

However, xarray’s stack has an important difference from pandas: unlike pandas, it does not automatically drop missing values. Compare:

In [18]: array = xr.DataArray([[np.nan, 1], [2, 3]], dims=['x', 'y'])

In [19]: array.stack(z=('x', 'y'))
Out[19]: 
<xarray.DataArray (z: 4)>
array([ nan,   1.,   2.,   3.])
Coordinates:
  * z        (z) MultiIndex
  - x        (z) int64 0 0 1 1
  - y        (z) int64 0 1 0 1

In [20]: array.to_pandas().stack()
Out[20]: 
x  y
0  1    1.0
1  0    2.0
   1    3.0
dtype: float64

We departed from pandas’s behavior here because predictable shapes for new array dimensions is necessary for Out of core computation with dask.

Set and reset index

Complementary to stack / unstack, xarray’s .set_index, .reset_index and .reorder_levels allow easy manipulation of DataArray or Dataset multi-indexes without modifying the data and its dimensions.

You can create a multi-index from several 1-dimensional variables and/or coordinates using set_index():

In [21]: da = xr.DataArray(np.random.rand(4),
   ....:                   coords={'band': ('x', ['a', 'a', 'b', 'b']),
   ....:                           'wavenumber': ('x', np.linspace(200, 400, 4))},
   ....:                   dims='x')
   ....: 

In [22]: da
Out[22]: 
<xarray.DataArray (x: 4)>
array([ 0.123102,  0.543026,  0.373012,  0.447997])
Coordinates:
    band        (x) <U1 'a' 'a' 'b' 'b'
    wavenumber  (x) float64 200.0 266.7 333.3 400.0
Unindexed dimensions:
    x

In [23]: mda = da.set_index(x=['band', 'wavenumber'])

In [24]: mda
Out[24]: 
<xarray.DataArray (x: 4)>
array([ 0.123102,  0.543026,  0.373012,  0.447997])
Coordinates:
  * x           (x) MultiIndex
  - band        (x) object 'a' 'a' 'b' 'b'
  - wavenumber  (x) float64 200.0 266.7 333.3 400.0

These coordinates can now be used for indexing, e.g.,

In [25]: mda.sel(band='a')
Out[25]: 
<xarray.DataArray (wavenumber: 2)>
array([ 0.123102,  0.543026])
Coordinates:
  * wavenumber  (wavenumber) float64 200.0 266.7

Conversely, you can use reset_index() to extract multi-index levels as coordinates (this is mainly useful for serialization):

In [26]: mda.reset_index('x')
Out[26]: 
<xarray.DataArray (x: 4)>
array([ 0.123102,  0.543026,  0.373012,  0.447997])
Coordinates:
    band        (x) object 'a' 'a' 'b' 'b'
    wavenumber  (x) float64 200.0 266.7 333.3 400.0
Unindexed dimensions:
    x

reorder_levels() allows changing the order of multi-index levels:

In [27]: mda.reorder_levels(x=['wavenumber', 'band'])
Out[27]: 
<xarray.DataArray (x: 4)>
array([ 0.123102,  0.543026,  0.373012,  0.447997])
Coordinates:
  * x           (x) MultiIndex
  - wavenumber  (x) float64 200.0 266.7 333.3 400.0
  - band        (x) object 'a' 'a' 'b' 'b'

As of xarray v0.9 coordinate labels for each dimension are optional. You can also use .set_index / .reset_index to add / remove labels for one or several dimensions:

In [28]: array = xr.DataArray([1, 2, 3], dims='x')

In [29]: array
Out[29]: 
<xarray.DataArray (x: 3)>
array([1, 2, 3])
Unindexed dimensions:
    x

In [30]: array['c'] = ('x', ['a', 'b', 'c'])

In [31]: array.set_index(x='c')
Out[31]: 
<xarray.DataArray (x: 3)>
array([1, 2, 3])
Coordinates:
  * x        (x) object 'a' 'b' 'c'

In [32]: array.set_index(x='c', inplace=True)

In [33]: array.reset_index('x', drop=True)
Out[33]: 
<xarray.DataArray (x: 3)>
array([1, 2, 3])
Unindexed dimensions:
    x

Shift and roll

To adjust coordinate labels, you can use the shift() and roll() methods:

In [34]: array = xr.DataArray([1, 2, 3, 4], dims='x')

In [35]: array.shift(x=2)
Out[35]: 
<xarray.DataArray (x: 4)>
array([ nan,  nan,   1.,   2.])
Unindexed dimensions:
    x

In [36]: array.roll(x=2)
Out[36]: 
<xarray.DataArray (x: 4)>
array([3, 4, 1, 2])
Unindexed dimensions:
    x