I am trying to fix a label switching problem in a mixture model by post-processing the inference data I obtained during sampling and re-label some of the components. However, to do this, I need to edit my posterior trace. I want to do something like
I suspect that you will need to use the operations describe here. Note that xarray explicitly states that assignment using .sel() will fail silently, which is probably what is happening when you run your assignment code.
Commenting mostly to ratify the link provided by @cluhmann. I think in this specific situation you want .loc.
Extra tangential notes:
You can use post = trace.posterior as indicated in Working with InferenceData — ArviZ dev documentation to work with the posterior Dataset directly in a less verbose way. Python assigns by reference, so modifying post will modify trace.posterior unless you make a copy when defining post.
Use of positional indexing and multiple indexing aproaches
Combining sel(chain=1) with multiple instances of positional indexing is a recipe for disaster. In xarray the dimension order should be irrelevant, only the dimension name matters, and there are xarray functions that modify the dimension order. Here it looks like you are only using sel right after the trace is obtained, so nothing should go wrong, but it is bad practice to rely on dimensions being in a given order instead of using their names. Moreover, by using label indexing you won’t need to use : for dimensions you don’t want to index, simply skip it.
This is perfectly valid and recommended practice. But I’d recomment you annotate the dimension in the mu variable so they you can select the component of my with prior mean 1 with sel(dim="coord_for_1").
You can’t use .sel on the left hand side of the equal, because it then fails silently. @cluhmann mentioned this explicitly because that is what was happening in your original example. But before that he share a link to the assigning values with indexing section of the xarray docs.
There it is explained that to modify existing objects you need to use .loc or .where.
Assuming your mu has a single dimension dim of length 3 with coordinate values a, b, c, something like: