Implementation of the R2D2 prior in bambi

[yannmclatchie]

This is a feature suggestion based on the growing interest in priors based on prior $R^2$ in literature (Zhang et al., 2022; Yanchenko et al., 2021; Aguilar and Bürkner, 2022; Yanchenko et al., 2023).

In particular, for a linear regression, the R2D2 model takes the form

$$ \begin{aligned} y_i &\sim \text{normal}(\beta_0 + \sum_{k=1}^{p}x_{k,i}\beta_k,\sigma^2) \nonumber \\ \beta_0 &\sim \pi(\beta_0) \nonumber \\ \sigma &\sim \pi(\sigma) \nonumber \\ \beta_k &\sim \text{normal}(0, \sigma^2 \tau^2 \phi_k) \nonumber \\ R^2 &\sim \text{beta}(\mu_{R^2}, \varphi_{R^2}) \nonumber \\ \phi &\sim \text{Dirichlet}(\xi,\dotsc,\xi) \nonumber \\ \tau^2 &= \frac{R^2}{1-R^2}. \nonumber \end{aligned} $$

Note by this notation we use the mean and a pseudo-precision to parametrise the beta distribution for $R^2$ as opposed to the conventional shape parameters $a > 0$ and $b > 0$. The relationship between the two parametrisations is as follows:

$$ \begin{aligned} \mu_{R^2} &= \frac{a}{a+b}, \nonumber\\ \varphi_{R^2} &= a + b. \nonumber \end{aligned} $$

Yanchenko et al. (2021) have shown how this can be extended to generalised linear models, as may be of interest to bambi.

Initial discussions with @tomicapretto reveal that the primary blocker to this is that bambi does not currently allow different model terms to share prior parameters (in our case, the same $\sigma$ features in both $\beta_k$ and $y$).

[SuryaMudimi]

I looked at blog post from Maxim Kochurov https://ferrine.github.io/posts/2023/Feb/01/linear-regression-r2m2d2/. @tomicapretto Any thoughts on implemetation in bambi? How would the user API look like ?

[tomicapretto]

I'm not sure yet. Before that, we need to allow Bambi to share priors between terms.