storing superchain_ids for rhat_nested

[n-kall]

I have been working a bit with many short chains, and using nested R-hat as a convergence diagnostic. rhat_nested has an additional argument superchain_ids which is a vector specifying which superchain each chain belongs to. My current workflow is to add the superchain_ids as an attribute to a posterior::draws object, and then do something like

summarise_draws(draws, rhat_nested = ~rhat_nested(.x, superchain_ids = attr(draws, "superchain_ids"))

This is a bit cumbersome, and maybe it could be improved with:

• a new function add_superchain_ids, which will then add the superchain_ids to the draws object (as an attribute?)
• have summarise_draws automatically use the stored superchain_ids for rhat_nested

Any thoughts on this? I could make a PR if this seems reasonable

[paul-buerkner]

I like the idea. Just one more thought: at least for the draws_df format, we could also add another protected column, .superchain say. For the other formats however, we would probably have not easy way beyond attribute to store superchain info so I am not sure we should make a special case for draws_df here. Any thoughts? @avehtari what do you think?

[avehtari]

I think it would make sense to add superchain_ids for each format in the same way as chain, iteration and draw are implemented in each format. Also then the logical name would be superchain and in case of draws_df the protected column is .superchain

If adding another protected column, we would need to do reverse dependency check, as it is possible that someone is making assumptions about the number of protected columns (they should not, but with so many packages using posterior the probability is non-zero).

[n-kall]

If we treat superchain info similarly to chain info in the different formats, we might want to think about how the user should add the information. At the moment, superchain_ids is expected to be a vector of length nchains. However, draws_matrix just stores nchains as an attribute, and it is implicit that the first niterations draws are from chain 1, and the next niterations draws are from chain 2, etc. So we could do the same thing with superchains and just store nsuperchains and implicitly assign each draw to a chain and a superchain based on its row in the matrix.

Then maybe all a user would need to do is specify how many superchains there are, rather than specifying which chains are part of which superchain. But I wonder if this is error prone, as it relies on the order of the draws/chains in the object matching the superchain configuration. It would then not be possible to have chains e.g. (1,5,9), (2, 6, 10), (3, 7, 11), (4, 8, 12) be the superchains as they would need to be (1,2,3), (4,5,6), (7,8,9), (10,11,12).

For reference, my current workflow is:

• define vector superchain_ids of length nchains (e.g. [1,1,2,2])
• generate inits for each superchain
• fit single chains, using init for corresponding superchain
• bind the draws from all fits along chains (this assumes that the order of the draws objects passed matches the superchain_ids order)
• add the superchain_ids attr to the bound draws object