heckman_correction.stan

functions {
  vector inverse_mills(vector z) {
    vector[rows(z)] out;
    
    for(i in 1:rows(z)) {
       out[i] = exp(normal_lpdf(z[i] | 0, 1)) / (Phi(z[i]));
    }
    return(out);
  }
}

data {
  int N; // number of observations
  int P; // number of covariates
  int P2; // number of instruments
  vector[N] Y; // income or log of income
  vector<lower = 0, upper = 1>[N] participation; // participation in the workforce
  matrix[N, P] X; // covariates
  matrix[N, P2] Z_raw; // instruments
  int estimate_model;
}
transformed data {
  matrix[N, P + P2] Z = append_col(X, Z_raw);
}
parameters {
  real alpha;
  real alpha_1;
  vector[P] beta;
  vector[P + P2] gamma;
  real<lower = 0> sigma_u;
  real<lower = -1, upper = 1> rho;
}
transformed parameters {
  vector[N] p = Phi(alpha_1 + Z * gamma);
}
model {
  // priors
  alpha ~ student_t(3, 0, 2);
  alpha_1 ~ student_t(3, 0, 2);
  beta ~ normal(0, 1);
  gamma ~ normal(0, 1);
  sigma_u ~ inv_gamma(1.5, 2);
  rho ~ normal(0, .5);
  
  
  // log likelihood for selection model
  target += participation' * log(p) + (1 - participation)' * log(1 - p);
    
  // log likelihood for outcome model
  for(n in 1:N) {
    if(participation[n] == 1.0) {
      Y[n] ~ normal(alpha + X[n] * beta + sigma_u * rho * inverse_mills(rep_vector(p[n],1) ), sigma_u);
    }
  }
}