R2D2 prior and Gaussian Processes

Modeling
9

Hi again @avehtari. I needed to tidy up some bits of code before sharing but below is my model.

It’s a version of BKMR, which is an MCMC model written in R for flexibly modelling the effects of mixtures of exposures. The model includes an ARD component using inverse length scales to select the most important exposures (with selection currently imposed by a simple cauchy prior on r). The model below works, but for the binomial output needs a large n to detect signal in simulated data. So, I was trying to tune the priors via prior predictive checks and concluded the priors as currently specified are not well thought out. In particular, as the (inverse) length scale priors go to zero, one would need a more tightly defined prior on lambda. So this set me to thinking about R2D2 or similar priors. However, I don’t know how to set R2D2 on a model like this with a binomial outcome and a mix of linear regression and a GP, and also whether lambda should be included in the R2D2 scales or not (it is not included in the R2D2 scales in the brms R2D2 implementation for GP’s, but I don’t know why).

functions {
    // Calculate distance/covariance matrix for exposures - checked vs bkmr code- matches
    matrix cov_Kzr(matrix Z, vector r) {
        // Define variables
        int n_obs = dims(Z)[1];
        int n_exp = dims(Z)[2];
        matrix[n_obs, n_exp] Zr = Z .* rep_matrix(sqrt(r'), n_obs);
        matrix[n_obs, n_obs] K = rep_matrix(0.0, n_obs, n_obs);

        // calculate K
        for (i in 1:n_obs){
            for (j in 1:n_obs){
                if (i < j){
                    for (k in 1:n_exp){
                        K[i, j] += square( Zr[i, k] - Zr[j, k] );
                    }
                    K[j, i] = K[i, j];
                }
            }
        }
        return K;
    }
}

data {
    int n_obs;
    int n_exp;
    int n_cov;
    
    array[n_obs] int y;
    matrix[n_obs, n_exp] Z;
    matrix[n_obs, n_cov] X;
    
    int prior_only;
}

transformed data {
    matrix[n_obs, n_cov] Xc;  // centered version of X without an intercept
    vector[n_cov] means_X;  // column means of X before centering
    for (i in 1:n_cov) {
        means_X[i] = mean(X[, i]);
        Xc[, i] = X[, i] - means_X[i];
    }
}

parameters {
    real<lower=0.0> sigma;
    real<lower=0.0> lambda;
    vector[n_cov] beta;
    real Intcpt;
    
    // Vector of r values for each exposure
    vector<lower=0.0>[n_exp] r;
    
    vector[n_obs] H;
}

model {
    // local params
    matrix[n_obs, n_obs] V;
    V = add_diag( lambda * exp(-cov_Kzr(Z, r)) , 1.0 ) ; // add_diag( , 1.0) as per bkmr code

    // Priors
    sigma ~ inv_gamma(0.001, 0.001);
    lambda ~ inv_gamma(1, 1);
    r ~ cauchy(0, 0.1);
    beta ~ normal(0, 0.5);
    Intcpt ~ normal(0, 0.5);

    // Likelihood
    if (!prior_only) {
        H ~ multi_normal_cholesky(rep_vector(0, n_obs), sigma * cholesky_decompose(V));
        y ~ bernoulli_logit(Intcpt + Xc * beta + H);
    }
}

Any suggestions/insights you might have would be appreciated. Thank you!