How to replace coordinate levels in a MultiIndex

[aspitarl]

I am interested in being able to replace the coordinate in a MultiIndex with new values after some mathematical operation. There appears to be some circumstances in which this works and doesn't. I'm aware there is a indexing refactor that is ongoing and not sure what is a bug or isn't so I figured I'd start a discussion instead of an issue. I previously started a similar discussion regarding adding or removing levels here, which was resolved and I'm using that method below.

I'm making a test mutindexed array like this:

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

<xarray.DataArray (x: 12, y: 3)>
array([[...]])
Coordinates:
  * x        (x) object MultiIndex
  * one      (x) int64 0 0 0 0 1 1 1 1 2 2 2 2
  * two      (x) int64 3 3 4 4 3 3 4 4 3 3 4 4
  * three    (x) int64 5 6 5 6 5 6 5 6 5 6 5 6
  * y        (y) int32 0 1 2

And want to replace the values in the level 'one' with new coordinates of the same length

new_coords = mda.coords['one'].values*2
mda.assign_coords(one=('x', new_coords))

This doesn't just work directly giving the following traceback.

ValueError                                Traceback (most recent call last)

c:\Users\aspit\anaconda3\envs\dataanalysis\lib\site-packages\xarray\core\common.py in assign_coords(self, coords, **coords_kwargs)
    590         data = self.copy(deep=False)
    591         results: dict[Hashable, Any] = self._calc_assign_results(coords_combined)
--> 592         data.coords.update(results)
    593         return data
    594 

c:\Users\aspit\anaconda3\envs\dataanalysis\lib\site-packages\xarray\core\coordinates.py in update(self, other)
    160         other_vars = getattr(other, "variables", other)
    161         self._maybe_drop_multiindex_coords(set(other_vars))
--> 162         coords, indexes = merge_coords(
    163             [self.variables, other_vars], priority_arg=1, indexes=self.xindexes
    164         )

c:\Users\aspit\anaconda3\envs\dataanalysis\lib\site-packages\xarray\core\merge.py in merge_coords(objects, compat, join, priority_arg, indexes, fill_value)
    564     collected = collect_variables_and_indexes(aligned)
...
    202                 f"cannot set or update variable(s) {common_names_str}, which would corrupt "
    203                 f"the following index built from coordinates {index_names_str}:\n"

ValueError: cannot set or update variable(s) 'one', which would corrupt the following index built from coordinates 'x', 'one', 'two', 'three':
<xarray.core.indexes.PandasMultiIndex object at 0x000001AF774C4510>

I guess my first question is: is there some other operation I should do to accomplish this coordinate reassignment?

I have been able to make this work by dropping and re-adding the index, which gives the expected result, but this doesn't seem like it should be necessary

mda.reset_index('one', drop=True).assign_coords(one= ('x',new_coords)).set_index(x='one',append=True)

<xarray.DataArray (x: 12, y: 3)>
array([[... ]])
Coordinates:
  * x        (x) object MultiIndex
  * two      (x) int64 3 3 4 4 3 3 4 4 3 3 4 4
  * three    (x) int64 5 6 5 6 5 6 5 6 5 6 5 6
  * y        (y) int32 0 1 2
  * one      (x) int64 0 0 0 0 2 2 2 2 4 4 4 4

However, this method doesn't work if you have a multiindex with one level, which I suspect is a bug. I'm accomplishing this by just dropping the level 'two' first.

mda.reset_index('two',drop=True).reset_index('one', drop=True).assign_coords(one= ('x',new_coords))

This runs into the same value error:

ValueError: cannot set or update variable(s) 'one', which would corrupt the following index built from coordinates 'x', 'one', 'three':
<xarray.core.indexes.PandasMultiIndex object at 0x000001AF774DD7B0>

[benbovy]

Hi @aspitarl

is there some other operation I should do to accomplish this coordinate reassignment?
I have been able to make this work by dropping and re-adding the index, which gives the expected result, but this doesn't seem like it should be necessary

Reassigning a coordinate that takes part in a multi-coordinate index indeed requires to drop and rebuild the index since the latter is immutable. It would certainly be possible to do this automatically, but IMHO it is probably better if it is done explicitly. Indexes vs. coordinates are often a source of confusion and handling them separately (as much as possible) in the API would make things clearer. Another reason is that building an index may be costly: it is not straightforward to figure out why updating a coordinate is slow if the index is rebuilt under the hood.

Regarding pandas.MultiIndex, there are convenient methods like droplevel or reorder_levels for creating new indexes from existing ones, which are reused in Xarray via set_index and reset_index. However, in my opinion we should probably depreciate those special cases eventually in favor of a more simple and predictable behavior for any kind of multi-coordinate index.

However, this method doesn't work if you have a multiindex with one level, which I suspect is a bug.

It looks like a bug indeed. Could you open a new issue, please?

[aspitarl]

Thanks for the reply. Keeping things explicit sounds good. Its just difficult to know if I'm missing a more built-in/elegant method or doing something that's not a best practice, which stems from the confusion of indices and coordinates as you mention.

I've started an issue at #6959