8

Is there an existing method or approach to expand the dimensions (and coordinates) of an xarray.DataArray object?

I would like to obtain something similar to np.expand_dims while at the same time defining a new dimension and coordinate variable for the new expanded DataArray.

Using DataArray.assign_coords() I can create a new coordinate variable but the array itself is not expanded with a new axis.

gerrit
  • 24,025
  • 17
  • 97
  • 170
rafa
  • 235
  • 1
  • 2
  • 10
  • What would be your expected output in *expand coordinates*? Unlike numpy, `pandas` and `xarray` (formerly `xray`) are nd-table oriented, there is no point to have an empty dimension in pandas/xray. – Imanol Luengo Jan 25 '16 at 09:40
  • I integrate some DataArray over some dimension and as a result this dimension is dropped but I need to keep (or recreate) that dimension (with length-1, not empty) in the integrated DataArray as it is expected by some other functions / methods I use the DataArray for. – rafa Jan 25 '16 at 10:23
  • 1
    Does any of [this](http://xarray.pydata.org/en/stable/reshaping.html) help you? (Maybe the `stack`/`unstack` bit) – Imanol Luengo Jan 25 '16 at 11:38

2 Answers2

10

In xarray v0.10.0, I use a combination of assign_coords() and expand_dims() to add a new dimension and coordinate variable.

For example:

import xarray as xr
import numpy as np
data = xr.DataArray([1, 2, 3], dims='x', coords={'x': [10, 20, 30]})
data_newcoord = data.assign_coords(y='coord_value')
data_expanded = data_newcoord.expand_dims('y')
print(data_expanded)
# <xarray.DataArray (y: 1, x: 3)>
# array([[1, 2, 3]])
# Coordinates:
#   * x        (x) int64 10 20 30
#   * y        (y) <U11 'coord_value'
crjones
  • 181
  • 1
  • 4
  • 1
    How can we add more than one coordinates to the `y` dimension, e.g. y = [coord_value1, coord_value2, ..., coord_value10]? – alextc Jul 16 '20 at 00:15
  • @alextc You can achieve this in one go by `data.expand_dims(dim=dict(y=[coord_value1, coord_value2,..., coord_value10]))`. – Fei Yao Feb 02 '21 at 10:17
5

I agree that some sort of method for doing this would be useful. It does not currently exist directly in xarray, but I would encourage you to file an issue on GitHub to discuss API for a new feature and/or make a pull request implementing it.

The new xarray.broadcast function contains some related functionality that may suffice for this purposes:

import xarray as xr
import numpy as np
data = xr.DataArray([1, 2, 3], dims='x')
other = xr.DataArray(np.zeros(4), coords=[('y', list('abcd'))])
data2, other2 = xr.broadcast(data, other)
print(data2)
# <xarray.DataArray (x: 3, y: 4)>
# array([[1, 1, 1, 1],
#       [2, 2, 2, 2],
#       [3, 3, 3, 3]])
# Coordinates:
#   * x        (x) int64 0 1 2
#   * y        (y) |S1 'a' 'b' 'c' 'd'
shoyer
  • 9,165
  • 1
  • 37
  • 55