missings.stan

data {
	int N;
	vector[N] read; // Data vectors
	vector[N] parents;
	vector[N] ses;
	vector[N] iq;
	vector[N] treat;

	int read_nMiss; // N missing from variables; excluded have no missings
	int ses_nMiss;
	int iq_nMiss;

	int read_index[read_nMiss];
	int ses_index[ses_nMiss];
	int iq_index[iq_nMiss];
	
}

parameters {
	// Structural
	vector[4] beta;
	real intercept;
	real<lower=0> sigma;

	// MVN covariates; for imputation
	cholesky_factor_cov[3] Sigma;
	vector[3] Mu;

	// Missings
	vector[read_nMiss] read_imp;
	vector[ses_nMiss] ses_imp;
	vector[iq_nMiss] iq_imp;
}

transformed parameters {
	// Recombine data
	matrix[N,5] Data; // 1: read, 2: parents, 3: ses, 4: iq, 5: treat
	Data[,2] = parents;
	Data[,3] = ses; Data[ses_index,3] = ses_imp;
	Data[,4] = iq; Data[iq_index,4] = iq_imp;
	Data[,5] = treat;
	Data[,1] = read; Data[read_index,1] = read_imp;
}

model {
	

	// MVN imputation
	for(n in 1:N){
		Data[n,2:4] ~ multi_normal_cholesky(Mu,Sigma);
	}

	// Priors
	beta ~ normal(0,2);
	intercept ~ normal(0,2);
	sigma ~ cauchy(0,1);
	Mu[3] ~ normal(100,10);
	Mu[1:2] ~ normal(0,2);
	//// Sigma given uniform prior implicitly

	Data[,1] ~ normal(Data[,2:5]*beta + intercept, sigma);
}

generated quantities {
	cov_matrix[3] Omega = Sigma*(Sigma');
	real<lower=0,upper=1> R2;
	{
		vector[N] muHat = Data[,2:5]*beta + intercept;
		R2 = (variance(muHat))/(variance(muHat) + variance(Data[,1]-muHat));
	}
}