plot_cap() can have different model parameters as output

[tomicapretto]

The function plot_cap() gains a new argument, target. This determines what parameter of the distribution is plotted as a function of the covariates. This is useful in distributional models. See the example

import bambi as bmb
import numpy as np
import pandas as pd

from bambi.plots import plot_cap

rng = np.random.default_rng(121195)
N = 200
a, b = 0.5, 1.1
x = rng.uniform(-1.5, 1.5, N)
shape = np.exp(0.3 + x * 0.5 + rng.normal(scale=0.1, size=N))
y = rng.gamma(shape, np.exp(a + b * x) / shape, N)
data_gamma = pd.DataFrame({"x": x, "y": y})

formula = bmb.Formula("y ~ x", "alpha ~ x")
model = bmb.Model(formula, data_gamma, family="gamma")
idata = model.fit(random_seed=1234)

The default behavior still shows the mean parameter.

plot_cap(model, idata, "x")

But we can have other targets as well

plot_cap(model, idata, "x", target="alpha")

Closes #622

[codecov-commenter]

Codecov Report

Merging #627 (2534768) into main (6fec04f) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #627   +/-   ##
=======================================
  Coverage   86.63%   86.63%           
=======================================
  Files          37       37           
  Lines        2200     2200           
=======================================
  Hits         1906     1906           
  Misses        294      294           

Impacted Files	Coverage Δ
bambi/plots/plot_cap.py	90.82% <100.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more