library(MASS)

set.seed(4338)

#Expit transform
expit<-function(x){return(1/(1+exp(-x)))}

n<-2000
p<-10

x.index<-1:p
al.select<-be.select<-rep(0,p)
al.select[c(1:5)]<-1
be.select[c(3:7)]<-1

#Confounders
confound.list<-x.index[al.select == 1 & be.select==1]

#Instruments
intrument.list<-x.index[al.select ==1 & be.select==0]

#Predictors of outcome only
outcomeonly.list<-x.index[al.select ==0  & be.select==1]

#Spurious variables
spurious.list<-x.index[al.select ==0  & be.select==0]

#Predictors of outcome
outcome.list<-sort(c(confound.list,outcomeonly.list))

#Predictors of outcome
treatment.list<-x.index[al.select==1]

################################################

set.seed(437973)
#Generate alpha parameters (for treatment allocation)
al<-rnorm(p+1,0,2.0)

#Generate beta parameters (for outcome regression)
be<-rnorm(p+1,0,2.0)

theta<-2.0
sig.Y<-1

#Covariance matrix of predictors
Sig<-rWishart(1,p+1,diag(p)/(p+1))[,,1]

#Predictors
X<-mvrnorm(n,mu=rep(0,p),Sigma=Sig)

#Data simulation
Xmat<-cbind(rep(1,n),X)
eta.Z<-Xmat %*% (c(1,al.select)*al)
p.Z<-expit(eta.Z)
Z<-rbinom(n,1,p.Z)
eta.Y<-Z*theta+Xmat %*% (c(1,be.select)*be)
	
Y<-eta.Y+rnorm(n)*sig.Y
boxplot(Y~Z)

#Direct comparison
mean(Y[Z==1])-mean(Y[Z==0])


#############################################################
#Stratification by the propensity score: all predictors

p1<-glm(Z~X,family=binomial)
ps.vals<-fitted(p1)
	
nstrata<-8
n.in.stratum<-n/nstrata
nreq<-n.in.stratum*0.1
ps.q<-c(0,quantile(ps.vals,prob=c(1:(nstrata-1))/nstrata),1)
ps.gp<-as.numeric(cut(ps.vals,ps.q))
Zbygp<-table(Z,ps.gp)
comparable.list<-c(1:nstrata)[Zbygp[1,] > nreq & Zbygp[2,] > nreq]
n.common<-length(comparable.list)*n.in.stratum

effect<-0
for(i in 1:length(comparable.list)){
	igp<-comparable.list[i]
	Ysub<-Y[ps.gp==igp]	
	Zsub<-Z[ps.gp==igp]
	effect<-effect+(mean(Ysub[Zsub==1])-mean(Ysub[Zsub==0]))
}
effect/length(comparable.list)

par(mfrow=c(2,5),mar=c(2.25, 1.25, 2.25, 1.25))
for(j in 1:10){
	boxplot(X[,j]~Z,col=c('white','red'))
}

par(mfrow=c(2,5),mar=c(2.25, 1.25, 2.25, 1.25))
for(j in 1:10){
	boxplot(X[ps.gp==3,j]~Z[ps.gp==3],col=c('white','red'))
}

par(mfrow=c(2,5),mar=c(2.25, 1.25, 2.25, 1.25))
for(j in 1:10){
	boxplot(X[ps.gp==5,j]~Z[ps.gp==5],col=c('white','red'))
}

#############################################################

#Stratification by the propensity score: using predictors of treatment only

p1<-glm(Z~X[,treatment.list],family=binomial)
ps.vals<-fitted(p1)

nstrata<-8
n.in.stratum<-n/nstrata
nreq<-n.in.stratum*0.1
ps.q<-c(0,quantile(ps.vals,prob=c(1:(nstrata-1))/nstrata),1)
ps.gp<-as.numeric(cut(ps.vals,ps.q))
Zbygp<-table(Z,ps.gp)
comparable.list<-c(1:nstrata)[Zbygp[1,] > nreq & Zbygp[2,] > nreq]
n.common<-length(comparable.list)*n.in.stratum

effect<-0
for(i in 1:length(comparable.list)){
	igp<-comparable.list[i]
	Ysub<-Y[ps.gp==igp]	
	Zsub<-Z[ps.gp==igp]
	effect<-effect+(mean(Ysub[Zsub==1])-mean(Ysub[Zsub==0]))
}
effect/length(comparable.list)

par(mfrow=c(2,5),mar=c(2.25, 1.25, 2.25, 1.25))
for(j in 1:10){
	boxplot(X[,j]~Z,col=c('white','red'))
}

par(mfrow=c(2,5),mar=c(2.25, 1.25, 2.25, 1.25))
for(j in 1:10){
	boxplot(X[ps.gp==3,j]~Z[ps.gp==3],col=c('white','red'))
}

par(mfrow=c(2,5),mar=c(2.25, 1.25, 2.25, 1.25))
for(j in 1:10){
	boxplot(X[ps.gp==5,j]~Z[ps.gp==5],col=c('white','red'))
}

#############################################################

#Stratification by the propensity score: using confounders only

p1<-glm(Z~X[,confound.list],family=binomial)
ps.vals<-fitted(p1)
	
nstrata<-8
n.in.stratum<-n/nstrata
nreq<-n.in.stratum*0.1
ps.q<-c(0,quantile(ps.vals,prob=c(1:(nstrata-1))/nstrata),1)
ps.gp<-as.numeric(cut(ps.vals,ps.q))
Zbygp<-table(Z,ps.gp)
comparable.list<-c(1:nstrata)[Zbygp[1,] > nreq & Zbygp[2,] > nreq]

n.common<-length(comparable.list)*n.in.stratum

effect<-0
for(i in 1:length(comparable.list)){
	igp<-comparable.list[i]
	Ysub<-Y[ps.gp==igp]	
	Zsub<-Z[ps.gp==igp]
	effect<-effect+(mean(Ysub[Zsub==1])-mean(Ysub[Zsub==0]))
}
effect/length(comparable.list)

par(mfrow=c(2,5),mar=c(2.25, 1.25, 2.25, 1.25))
for(j in 1:10){
	boxplot(X[,j]~Z,col=c('white','red'))
}

par(mfrow=c(2,5),mar=c(2.25, 1.25, 2.25, 1.25))
for(j in 1:10){
	boxplot(X[ps.gp==3,j]~Z[ps.gp==3],col=c('white','red'))
}

par(mfrow=c(2,5),mar=c(2.25, 1.25, 2.25, 1.25))
for(j in 1:10){
	boxplot(X[ps.gp==5,j]~Z[ps.gp==5],col=c('white','red'))
}