How to best add or remove levels from a multiindex?

[aspitarl]

I often find myself with a DataArray with a mutlindexed dimension that I want to add and/or remove levels from before unstacking. I have developed some ways to do this, but they are somewhat complicated. I wanted to make sure there wasn't a better way to do this.

# Generate a multindex data array
# http://xarray.pydata.org/en/stable/user-guide/indexing.html

midx = pd.MultiIndex.from_product([list("abc"), [0, 1]], names=("one", "two"))
mda = xr.DataArray(np.random.rand(6, 3), [("x", midx), ("y", range(3))])

 <xarray.DataArray (x: 6, y: 3)>
array([[0.89409303, 0.47413507, 0.54185886],
       [0.71962262, 0.70948541, 0.82275026],
       [0.2616073 , 0.19551461, 0.91624799],
       [0.93271297, 0.46344353, 0.3927474 ],
       [0.98137793, 0.62304479, 0.7173114 ],
       [0.87360386, 0.72348197, 0.01473442]])
Coordinates:
  * x        (x) MultiIndex
  - one      (x) object 'a' 'a' 'b' 'b' 'c' 'c'
  - two      (x) int64 0 1 0 1 0 1
  * y        (y) int32 0 1 2

Here is the method I use to add a level to this multindex.

mda_new = mda.reset_index('x').assign_coords(three=('x', [1,2,3,4,5,6])).set_index(x = ['one','two', 'three'])

<xarray.DataArray (x: 6, y: 3)>
array([[0.89409303, 0.47413507, 0.54185886],
       [0.71962262, 0.70948541, 0.82275026],
       [0.2616073 , 0.19551461, 0.91624799],
       [0.93271297, 0.46344353, 0.3927474 ],
       [0.98137793, 0.62304479, 0.7173114 ],
       [0.87360386, 0.72348197, 0.01473442]])
Coordinates:
  * y        (y) int32 0 1 2
  * x        (x) MultiIndex
  - one      (x) object 'a' 'a' 'b' 'b' 'c' 'c'
  - two      (x) int64 0 1 0 1 0 1
  - three    (x) int64 1 2 3 4 5 6

I can also remove a level, but It's the same idea.

mda_new.reset_index('x').set_index(x=['two','three']).drop('one')

<xarray.DataArray (x: 6, y: 3)>
array([[0.89409303, 0.47413507, 0.54185886],
       [0.71962262, 0.70948541, 0.82275026],
       [0.2616073 , 0.19551461, 0.91624799],
       [0.93271297, 0.46344353, 0.3927474 ],
       [0.98137793, 0.62304479, 0.7173114 ],
       [0.87360386, 0.72348197, 0.01473442]])
Coordinates:
  * y        (y) int32 0 1 2
  * x        (x) MultiIndex
  - two      (x) int64 0 1 0 1 0 1
  - three    (x) int64 1 2 3 4 5 6

I'm using xarray = 0.20.1

It feels this should be all just handled by the 'assign_coords' function, but it appears I have to set and reset the index. This is OK but somewhat cumbersome because I have to explicitly write out the level names. perhaps there could be add_level or remove_level methods added?

[benbovy]

You actually don't need to first reset the index.

You can add a level with set_index(..., append=True):

mda_new = mda.assign_coords(three=('x', [1,2,3,4,5,6])).set_index(x='three', append=True)

Or remove a level with mda_new.reset_index(x="one", drop=True)

However, I wouldn't rely too much on that. Xarray indexes are currently being refactored and the API above might be depreciated eventually.