How does one do predictive inference with the sum of gaussian processes?

Modeling
2

I am going to answer my own question as I discover the solution from reading another thread and the manual:

//https://mc-stan.org/docs/stan-users-guide/gaussian-processes.html#predictive-inference-with-a-gaussian-process
//https://likelyllc.com/2022/08/03/additive-gaussian-process-time-series-regression-in-stan/

functions {
   
  vector gp_pred_rng(array[] vector x2,
                     vector y, 
                     array[] vector x1,
                     real m, 
                     real l,
                     real sig,
                     real sigma,
                     real delta) {
    int N1 = rows(y);
    int N2 = size(x2);
    vector[N2] f2;
    {
      matrix[N1, N1] K = add_diag(gp_exp_quad_cov(x1, m, l) +  gp_dot_prod_cov(x1, sig), square(sigma));
      matrix[N1, N1] L_K = cholesky_decompose(K);
      vector[N1] L_K_div_y = mdivide_left_tri_low(L_K, y);
      vector[N1] R_K_div_y = mdivide_right_tri_low(L_K_div_y', L_K)';
      matrix[N1, N2] k_x1_x2 = gp_exp_quad_cov(x1, x2, m, l) + gp_dot_prod_cov(x1, x2, sig);
      vector[N2] f2_mu = (k_x1_x2' * R_K_div_y);
      matrix[N1, N2] v_pred = mdivide_left_tri_low(L_K, k_x1_x2);
      matrix[N2, N2] cov_f2 = gp_exp_quad_cov(x2, m, l) + gp_dot_prod_cov(x2, sig) - v_pred' * v_pred;
      matrix[N2, N2] diag_delta = diag_matrix(rep_vector(delta, N2));
      f2 = multi_normal_rng(f2_mu, cov_f2 + diag_delta);
    }
    return f2;
  }
}
data {
  int<lower=1> N;
  int<lower=1> D;
  array[N] vector[D] x;
  vector[N] y;

  int<lower=1> N_pred;
  array[N_pred] vector[D] x_pred;
}
transformed data {
  real delta = 1e-9;
  vector[N] mu = rep_vector(0, N);
}
parameters {
  real<lower=0> l;
  real<lower=0> m;
  real<lower=0> sig0;

  real<lower=0> sigma;
}
transformed parameters {
  matrix[N, N] L_K;
  {
    matrix[N, N] K = gp_exp_quad_cov(x, m, l) + gp_dot_prod_cov(x, sig0);
    
    K = add_diag(K, square(sigma));
    L_K = cholesky_decompose(K);
  }
}
model {
  m ~ normal(0, 1);
  l ~ normal(0, 1);
  sig0 ~ normal(0, 1);

  sigma ~ normal(0, 1);
  
  y ~ multi_normal_cholesky(mu, L_K);
}
generated quantities {
  
   vector[N_pred] f_pred;
  vector[N_pred] y_pred;

  f_pred = gp_pred_rng(x_pred, y, x, m, l, sig0, sigma, delta);
  for (n2 in 1:N_pred) {
    y_pred[n2] = normal_rng(f_pred[n2], sigma);
  }
  
}

NOTE: A bug in the code was discovered on 2024-12-11 and has been corrected.

2 Likes