GP_IV.stan

functions {
  matrix L_cov_exp_quad_ARD(matrix x,
                            real alpha,
                            vector rho,
                            real delta) {
                              int N = rows(x);
                              matrix[N, N] K;
                              real sq_alpha = square(alpha);
                              for (i in 1:(N-1)) {
                                K[i, i] = sq_alpha + delta;
                                for (j in (i + 1):N) {
                                  K[i, j] = sq_alpha
                                  * exp(-0.5 * dot_self((x[i] - x[j])' ./ rho ));
                                  K[j, i] = K[i, j];
                                  }
                              }
                              
                              K[N, N] = sq_alpha + delta;
                              
                              return cholesky_decompose(K);
                              }
}
data {
  int N; // number of observations
  int P; // number of covariates on which we suspect the treatment effect varies
  int P2; // number of instruments
  int P3; // number of covarites 
  matrix[N, P] X; // matrix of covariates on which we suspect the treatment effect varies
  vector[N] treatment;
  matrix[N, P3] Covariates;
  matrix[N, P2] Z; // instruments
  vector[N] Y; // the outcome
}
parameters {
  vector<lower = 0>[P] rho;
  real<lower = 0> alpha;
  vector[N] z;
  
  vector[P3] beta;
  vector[P3] gamma;
  vector[2] intercept;
  vector[P2] instrument_slope;
  cholesky_factor_corr[2] L_Omega;
  vector<lower = 0>[2] scale;
}
transformed parameters {
  vector[N] tau;
  matrix[N, N] L = L_cov_exp_quad_ARD(X,alpha, rho, 1e-9);
  tau =  L* z;
  // first stage
}
model {
  // priors
  rho ~ inv_gamma(1,1);
  alpha ~ normal(0, 2);
  z ~ normal(0, 1);
  beta ~ normal(0, 1);
  gamma ~ normal(0, 1);
  intercept ~ normal(0, 1);
  instrument_slope ~ normal(0, 1);
  L_Omega ~ lkj_corr_cholesky(1);
  scale ~ student_t(5, 0, 2);
  
  // likelihood
  {
    matrix[N, 2] mu;
    matrix[2,2] tauL = diag_matrix(scale) * L_Omega;
    matrix[N, 2] YY = append_col(treatment, Y);
    mu[:,1] = intercept[1] + Z * instrument_slope + Covariates * gamma;
    // second stage
    mu[:,2] = intercept[2] + tau .* treatment + Covariates * beta;
    
    for(n in 1:N) {
      YY[n] ~ multi_normal_cholesky(mu[n], tauL);
    }
  }
}
generated quantities{
  vector[N] hte = multi_normal_cholesky_rng(rep_vector(0.0, N), L);
  real ATE = mean(tau);
}