How to stack several DataArrays of the same shape?

[johann-petrak]

I have a list of (thousands of) DataArrays which all have the same shape (r,c) same dimensions, and same coordinates for dimension 1, but different coordinates for dimension 0.

Is there a way/function to combine these into a single 3-dimensional DataArray, so that I can get back each of the original 2-dimensional arrays by indexing along the new dimension?

It would be especially useful if this method of combining the original DataArrays could also combine the coordinates for dimension 0 in such a way that I get those back properly as well.

[TomNicholas]

Hi Johann - have you tried xr.concat? It can create a new dimension for along which all your 2D arrays will be stacked, returning a 3D result. By default it should also stack your differing coordinates (extending them from 1D to 2D), and you can ensure this by passing coords='all'.

Our page on combining data will also be useful here.

[johann-petrak]

My own approach, outlined above seems to give a result where the underlying numpy values match for the arrays and coords (but direct comparison of the DataArray objects shows False for many values, weirdly):

concattime = np.stack([x.coords["time"].values for x in segments])
stackedsegments = np.stack([s.values for s in segments])
xsegments = xr.DataArray(
    stackedsegments, dims=["segment", "time0", "cols"],
    coords=dict(
        segment=np.arange(stackedsegments.shape[0]),
        time=(["segment", "time0"],  concattime),
        cols=segments[0].coords["cols"]
    )
)

So how to achieve this with xr.concat?

[TomNicholas]

Thanks for the example! Does this do what you want?

In [15]: concatenated = xr.concat(segments, dim='row', coords='all', join='override')

In [16]: concatenated
Out[16]: 
<xarray.DataArray (row: 7, time: 10, cols: 2)>
array([[[-4.16025006e-01,  2.56799683e-01],
        [-7.14348469e-01, -1.09972270e+00],
        [ 2.74527609e-01,  1.13956902e+00],
        [ 1.37386705e+00,  2.38558844e-01],
        [-6.77915330e-01, -4.37836623e-02],
        [ 6.29159848e-03,  5.83046584e-01],
        [ 4.78322945e-01, -1.66948143e+00],
        [ 1.24739476e+00, -2.90357793e+00],
        [ 3.83402084e-01, -1.92930193e-02],
        [-1.40905286e+00,  2.15999766e-01]],

       [[-3.74356776e-01, -8.06066149e-01],
        [ 1.34125014e+00, -2.28119875e+00],
        [ 5.18941543e-01,  5.13077071e-01],
        [-1.33885907e+00,  7.73590869e-01],
        [ 1.56885115e+00, -6.30148971e-02],
        [-1.45163180e+00,  5.67262288e-01],
        [ 5.79305613e-01,  2.30645952e+00],
        [-2.43251969e-01, -6.15721094e-01],
        [-7.42373559e-01, -3.15947717e+00],
...
        [ 5.89551799e-01, -4.52219479e-01],
        [ 5.65141246e-01, -2.48603264e-03],
        [ 1.22767296e+00,  6.66604660e-01],
        [ 1.78343468e+00, -1.38257620e+00],
        [ 7.31930872e-01,  9.15238994e-02],
        [ 2.93210215e+00,  1.32295396e-01],
        [-3.45927474e-01, -7.50466560e-01],
        [ 1.59752240e+00,  1.62683874e+00],
        [ 6.80008087e-01,  9.97619148e-01]],

       [[ 1.28863031e+00, -3.85879510e-01],
        [-1.05853853e+00,  6.59389647e-01],
        [-2.04081130e+00,  1.06280500e+00],
        [-4.33664791e-01,  3.85927319e-01],
        [-2.95199379e-01,  8.23573366e-03],
        [-1.16359325e+00, -1.34083027e-01],
        [-1.10795389e-01,  4.82629031e-01],
        [ 1.21437447e+00, -9.64967105e-01],
        [-7.83617158e-01, -4.43077900e-01],
        [-5.37524297e-01, -7.92694731e-01]]])
Coordinates:
  * time     (time) float64 0.04008 0.086 0.1103 0.2499 ... 0.632 0.7233 0.8759
  * cols     (cols) <U4 'col1' 'col2'
Dimensions without coordinates: row

In [17]: concatenated[0].equals(segments[0])
Out[17]: True

[johann-petrak]

Thank you for your help! I have to admit that I have difficulties to understand what exactly all those parameters will do exactly.

Unfortunately, this does not produce the correct result for me: if you compare e.g. the time coordinates of segments[1] with concatenated[1], they are different:

segments[1].coords["time"].values == concatenated[1].coords["time"].values

while they agree when I use my method:

segments[1].coords["time"].values == xsegments[1].coords["time"].values

The other arrays and coordinates match, but not the time coordinate when I test this.

UPDATE: on closer inspection, it turns out that the "time" coordinate has the same values for each of the 7 segments, so
concatendated[0].coords["time"] == concatednated[1].coords["time"] while the should be as in the original segments, different. With my method I could achieve this with the trick of first using "time0" as the identical coordinates and then providing "time" as a 2D coordinate when creating the DataArray. At least this is what I believe to understand. :)

[johann-petrak]

@TomNicholas is there a solution for this using xr.cocat ?

[TomNicholas]

Only by renaming your time coordinate, I believe.

So what we want to happen is for xr.concat to concatenate the coords on each segment, so that the time coord gets a new dimension on the result, without dropping any values.

You would reasonably expect this to happen by default, but it's complicated by the fact that your time (and cols) coordinates are backed by indexes (indicated by the asterisk in the repr, or bolding in the html repr). We say they are "dimension coordinates" - defined as coordinates with the same name as underlying dimensions.

If you rename your time coordinate to something else (or rename your time dimension), then you will get the behavior you expect.

In [88]: segments = []
    ...: for i in range(2):
    ...:     time = np.sort(np.random.random(4))
    ...:     da = xr.DataArray(
    ...:         np.random.randn(4,2),
    ...:         dims=["time", "cols"],
    ...:         coords=dict(t=('time', time), cols=["col1", "col2"]),
    ...:         )
    ...:     segments.append(da)
    ...: 

In [89]: segments
Out[89]: 
[<xarray.DataArray (time: 4, cols: 2)>
 array([[-0.19611065, -0.60953111],
        [-0.62044044, -0.66667808],
        [-1.82047546,  1.50008308],
        [ 0.40023352,  0.19898603]])
 Coordinates:
     t        (time) float64 0.6919 0.7359 0.9214 0.9453
   * cols     (cols) <U4 'col1' 'col2'
 Dimensions without coordinates: time,
 <xarray.DataArray (time: 4, cols: 2)>
 array([[-2.77324438, -1.36313211],
        [-0.01420488, -0.67174893],
        [ 1.90601694,  0.78128248],
        [ 1.10752682,  0.24384413]])
 Coordinates:
     t        (time) float64 0.06925 0.8341 0.8693 0.9828
   * cols     (cols) <U4 'col1' 'col2'
 Dimensions without coordinates: time]

In [90]: xr.concat(segments, dim='rows')
Out[90]: 
<xarray.DataArray (rows: 2, time: 4, cols: 2)>
array([[[-0.19611065, -0.60953111],
        [-0.62044044, -0.66667808],
        [-1.82047546,  1.50008308],
        [ 0.40023352,  0.19898603]],

       [[-2.77324438, -1.36313211],
        [-0.01420488, -0.67174893],
        [ 1.90601694,  0.78128248],
        [ 1.10752682,  0.24384413]]])
Coordinates:
    t        (rows, time) float64 0.6919 0.7359 0.9214 ... 0.8341 0.8693 0.9828
  * cols     (cols) <U4 'col1' 'col2'
Dimensions without coordinates: rows, time

This is definitely unintuitive, so I've documented it for discussion in #7539.