Add pairs plot for arbitrary quantities

[vpratz]

This PR adds a pairs_quantity plotting function that provides a pairs plot of an arbitrary quantity.

For a trained BasicWorkflow, one could use it like that:

def identity(x, *args, **kwargs):
    return x

test_data = simulator.sample(500)
samples = workflow.sample(num_samples=10, conditions=test_data)
values = bf.diagnostics.metrics.posterior_contraction(samples, test_data, aggregation=identity)
fig = bf.diagnostics.pairs_quantity(values["values"], test_data)

The current implementation only features structures as returned by the values field of posterior_contraction, but should be generalizable to other settings. A similar non-paired plot (containing only the diagonal) could be nice as well.

Below is an example of the plot:

[codecov]

Codecov Report

❌ Patch coverage is 96.68874% with 5 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
bayesflow/diagnostics/plots/plot_quantity.py	92.06%	5 Missing ⚠️

Files with missing lines	Coverage Δ
bayesflow/diagnostics/__init__.py	100.00% <ø> (ø)
...sflow/diagnostics/metrics/posterior_contraction.py	100.00% <100.00%> (ø)
bayesflow/diagnostics/plots/__init__.py	100.00% <100.00%> (ø)
bayesflow/diagnostics/plots/calibration_ecdf.py	72.72% <100.00%> (-4.89%)	⬇️
bayesflow/diagnostics/plots/pairs_quantity.py	100.00% <100.00%> (ø)
bayesflow/utils/dict_utils.py	69.81% <100.00%> (+6.30%)	⬆️
bayesflow/diagnostics/plots/plot_quantity.py	92.06% <92.06%> (ø)

[paul-buerkner]

Thank you! That looks good to be as a basic implementation. The documentation doesn't mention the required dimensionality of values. Could this be clarified?

[elseml]

Nice! I'd add a "Contraction" label to the colorbar and the y-axis of the diagonal for consistency with other plots and also to make the equivalence between the colorbar and the right-hand side y-axes clear. The cleanest solution would imo be to place them at the colorbar + at the right side of only the last diagonal element's y-axis, since it prevents cluttering the main plot.

[paul-buerkner]

We cannot provide a default metric name if the user decides, via the values argument, which metric is shown. But of course, we could add convenience funcationality allowing "values" to also be a string naming built-in supported metrics "contraction" or "z-score" for example.

[vpratz]

@elseml Thanks for the feedback! The colorbar can already be labeled by providing the label parameter, but I like the idea with the last diagonal element and will add it there as well.
@paul-buerkner Exactly. I think the main remaining task is to decide what values can be, to implement the different options and to document it properly... Currently I think the most natural would be:

• a numpy array (optionally a VariableArray from dict_utils)
• The output dict from a metric function, which has to be called without aggregating

We could think about renaming values to quantity and allowing to pass a function/function name there as well. But then we would also need the parameter estimates, and it might be confusing to have them both.
What do you think?

[paul-buerkner]

I have thought about the same of whether we should allow functions too. If we aim to also support test_quantities similar to the SBC diagnostics, we would have to provide estimates anyway. If quantity (or however called) is an an array or dict of array, ignore estimates and test_quantites. If quantity is a function then require estimates and support test_quantities. We can of course choose to support test_quantities only later, since it is not high priority. However, we should think about it already now because quantity and test_quantities will have clashing terminology. Perhaps values is not too bad after all even if it can optionally be a function.

[vpratz]

@paul-buerkner Thanks a lot for the input. The many different options make it a bit convoluted, but I have now implemented the following:

• values as array
• values as dict from a metric function with aggregation=None
• values as function that has to return one of the above

If have also included the test_quantities, here as well we have two cases:

1. the quantity is scalar, so we do not require estimates as we do not have to calculate it per quantity, they are only different axes
2. the quantity is a vector with one entry per variable, we require estimates and a function to calculate the values

Here we could skip 1. to simplify the code, this would then correspond to the description in your comment, but I think having the more general implementation does not hurt. I have created helper functions to avoid code duplication, I'm not sure yet where we want to put them.

Let me know what you think.

Example

fig = bf.diagnostics.pairs_quantity(
    values=bf.diagnostics.posterior_contraction,
    targets=test_data,
    estimates=samples,
    test_quantities=test_quantities,
    variable_names=[r"$\theta_1$", r"$\theta_2$"],
)

[paul-buerkner]

This looks very cool!

Two quick questions:

• Is aggregation = None equal to aggregation = identity (with defined identity function). I assume we would definitely want the None feature so users don't need to define identy first.
• I don't quite get the different options concerning test_quantities. Why do we have to distinquish the two cases here but not(?) in the SBC plots?

[vpratz]

Thanks for taking a look. Regarding your questions:

• Yes, aggregation=None should behave like an identity, i.e., not aggregate over datasets but instead return the individual values, as we require them for plotting. If I read it correctly, for now this only makes sense for posterior_contraction. It would be viable for example for z scores as well, but as far as I can tell we have not implemented that as a metric yet.
• I just discussed with @han-ol, and it probably makes sense to remove option 1. The idea was basically, if we have a value that is a property of the dataset, e.g. its log-prob, we do not need to recalculate it even if test_quantities is present, which might be nice if the quantity is costly or cannot easily be recalculated. test_quantities would then just supply new axes to plot against. This can become muddy quite quickly though (imagine the user wants to calculate the mean posterior contraction over all variables, which would be a scalar value that should change), and it is an edge case that probably no one will ever use, so for the sake of maintainability and clarity I would remove it. Do you agree?

[paul-buerkner]

• Okay. I think we should add a z-score metric too but of course this will be a different issue.
• Agreed.

[vpratz]

Thanks a lot for your input. I think the functions are now ready for a final review, @paul-buerkner.
I added a few tests and updated a few details to incorporate the changes discussed above.

[paul-buerkner]

Looks good to me, thank you! Just found a typo in the doc, I think. After fixing it, feel free to merge :-)