Adding Diganostics to Ecdf Plots in the spirit of TARP

[arrjon]

This pull request focuses on enhancing the functionality of the ECDF plotting functions plot_sbc_ecdf. The changes include adding new rank computation methods in the spirit of Lemos, Pablo, et al. "Sampling-based accuracy testing of posterior estimators for general inference.". Feedback is welcome!

The idea is to compute ranks based on distances to the origin or a random point instead of the marginal fractional ranks. This way one can analyse the joint calibration of the model instead of looking at single parameters. If the random reference has some dependency on the data, this can reveal biases which would not be detected by the standard ECDF plots (e.g. when the prior equals the posterior the ECDF indicate a well calibrated model; when we have a tutorial on the diagnostics as discussed in #236, I would add an examples there). One can also pass a distance function as an argument, which could be for instance a distance based on the log-probability of the approximator.

Furthermore, I fixed some minor bugs in the function plot_posterior_2d.

[paul-buerkner]

That looks cool, thank you! @jerrymhuang would you mind reviewing this PR since you have been working on the diagnostics module previously?

[stefanradev93]

@jerrymhuang Can you fix the conflicts due to the very recent changes in dev diagnostics?
@paul-buerkner As an author of the original ECDF paper, can you approve the new interface and functionality?

[jerrymhuang]

Nice changes. Thank you @arrjon !

@stefanradev93 @paul-buerkner The conflict is fixed to align with the recent partitioning scheme for the diagnostics module. The distance rank is also tested in the Linear Regression notebook.

[paul-buerkner]

Thank you for adding this feature! I had only a little time to check the paper but from what I understand, your implementation matches what they propose. I would have to investigate the properties of the method in more detail to fully understand it. But since we are currently having the new feature only optional with the standard SBC as default, I think we can safely merge this PR and improve the feature later on if we figure out that is needed for some yet unknown reason.

[arrjon]

Before implementing the feature, I considered the following: The key requirement is that the rank of the true parameter with respect to samples from the posterior distribution is uniformly distributed, regardless of how the parameters are transformed (e.g. just looking at the marginals or apply a norm). With this condition satisfied, everything, including the simultaneous confidence bands, works out fine.

In the paper by Lemos et al., they basically showed that you can even use a rank based on a distance to random reference point to compute valid coverage statistics. So instead of ECDFs (as we do here) they computed expected coverages based on these ranks.

[paul-buerkner]

That makes sense, thank you! We could later on generalize or extend this feature to allow for other kind of metrics that a distance, including the (log) likelihood density as target, which we suggested in https://arxiv.org/abs/2211.02383. But that is subject to a separate PR.

[arrjon]

That makes sense, thank you! We could later on generalize or extend this feature to allow for other kind of metrics that a distance, including the (log) likelihood density as target, which we suggested in https://arxiv.org/abs/2211.02383. But that is subject to a separate PR.

This is already possible by constructing a distance based on the likelihood density and passing it with the argument distance to ranks_kwargs. But we could certainly implement it in way, so the user can just apply it!

[paul-buerkner]

Good point! And I agree, some more convenience functionaly for common use cases would be cool down the line.

[paul-buerkner]

@stefanradev93 you can merge if you are happy with the PR as well.

[stefanradev93]

Thank you all!