Edit ArviZ data structure

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.

:warning: Use of positional indexing and multiple indexing aproaches :warning:

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.

2 Likes