Expanding plotting options of plot_posterior_predictive

[jsnyde0]

See this discussion

Goal is to extend the current function of plot_posterior_predictive with backward compatibility using individual private functions:

def _add_hdi_to_plot(ax, hdi_prob, color=None, alpha=0.3, **kwargs) -> plt.Axes:
    """Add HDI to existing plot."""

def _add_mean_to_plot(ax, color='red', linestyle='-', **kwargs) -> plt.Axes:
    """Add mean prediction to existing plot."""

def _add_gradient_to_plot(ax, cmap='Blues', n_percentiles=100, **kwargs) -> plt.Axes:
    """Add gradient representation to existing plot."""

This allows to add some plotting options like adding the mean or a gradient representation of the full posterior predictive distribution to plot_posterior_predictive as follows:

mmm.plot_posterior_predictive( . . . , add_mean: bool, ....):
    . . .
    if add_mean:
        _add_mean_to_plot( . . . )
    . . .

[jsnyde0]

@juanitorduz this is my very first contribution to someone else's open-source project. I've created a draft PR here: #1058

I went through all the checklists. One test failed but I don't see how that's related?

FAILED tests/test_mlflow.py::test_autolog_pymc_model[numpyro] - ModuleNotFoundError: No module named 'jax.linear_util'

Some considerations I had myself:

• Each private function takes original_scale, but should we pass posterior_predictive_data directly instead of handling scaling within each function?
• The private plotting functions don't mention 'posterior_predictive' in their names. Is it clear they're always based on posterior predictive data? Would this justify passing it explicitly as an argument?

[wd60622]

Is that a local failure? Or in the CI?

[jsnyde0]

Local failure, seems like the test passed in the CI