Last active
October 16, 2023 09:37
-
-
Save khakieconomics/7f62e5e38b0e1baed7590bc858171832 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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); | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment