Convert an existing xarray dimension into a MultiIndex

[shoyer]

Suppose I have two xarray Datasets, each defined along the 'x' axis:

ds1 = xarray.Dataset({'foo': (('x',), [1, 2, 3])}, {'x': [1, 2, 3], 'y': 'a'})
ds2 = xarray.Dataset({'foo': (('x',), [4, 5, 6])}, {'x': [1, 2, 3], 'y': 'b'})

<xarray.Dataset>
Dimensions:  (x: 3)
Coordinates:
    y        |S1 'a'
  * x        (x) int64 1 2 3
Data variables:
    foo      (x) int64 1 2 3

Now I want to stack them along a new MultiIndex 'yx' that consists of y and x levels:

desired = xarray.Dataset({'foo': (('yx',), [1, 2, 3, 4, 5, 6])},
                         {'yx': pandas.MultiIndex.from_product([['a', 'b'], [1, 2, 3]], names=['y', 'x'])})

<xarray.Dataset>
Dimensions:  (yx: 6)
Coordinates:
  * yx       (yx) MultiIndex
  - y        (yx) object 'a' 'a' 'a' 'b' 'b' 'b'
  - x        (yx) int64 1 2 3 1 2 3
Data variables:
    foo      (yx) int64 1 2 3 4 5 6

How can this be achieved with the minimum effort? What is the missing utility function that we need? I attempted to use set_index and swap_dims but so far have been unsuccessful.

@benbovy any ideas? I think something similar may have come up when we were discussing your set_index PR.

[shoyer]

Currently the best I have is:

(ds1.reset_index('x')
 .rename({'x': 'yx', 'x_': 'x'})
 .assign_coords(y=lambda ds: (('yx',), [ds.y.item()] * 3))
 .set_index(yx=['y', 'x']))

<xarray.Dataset>
Dimensions:  (yx: 3)
Coordinates:
  * yx       (yx) MultiIndex
  - y        (yx) object 'a' 'a' 'a'
  - x        (yx) int64 1 2 3
Data variables:
    foo      (yx) int64 1 2 3

which is rather un-intuitive.

It occurs to me that maybe this is just a good use case for expand_dims (#1326), which could make this as simple as ds1.expand_dims('y').stack(yx=['y', 'x']))

[benbovy]

Given your two Datasets and your desired result, why not simply using concat?

xr.concat([ds1, ds2], dim='y').stack(yx=['y', 'x'])

<xarray.Dataset>
Dimensions:  (yx: 6)
Coordinates:
  * yx       (yx) MultiIndex
  - y        (yx) object 'a' 'a' 'a' 'b' 'b' 'b'
  - x        (yx) int64 1 2 3 1 2 3
Data variables:
    foo      (yx) int64 1 2 3 4 5 6

I don't really understand your last comment (only using ds1). Am I missing something?

[shoyer]

Given your two Datasets and your desired result, why not simply using concat?

OK, I feel pretty dumb now. Yes, that's totally the right way to do this. For some reason I was fixated on converting each dataset individually to have a MultiIndex before combining them.