#code for Chapter 6 "Propensity Score Methods for Multiple Treatments" of book:
#Leite, W. L. (2017). Practical propensity score methods using R. 
#Thousand Oaks, CA: Sage. 
# 
#PART 3 - OBTAIN IPTW AND CHECK COVARIATE BALANCE WITH TWANG

#this is the code that was used to generate the example results in the book
#As the R software and the R packages used in this example are updated frequently
#some incompatibilities between the current code and new R versions or package versions
#may appear
#Any updates to the code will be posted at:
# http://www.practicalpropensityscore.com

#this code estimates  the effects of assigning mentors 
#of different areas to new teachers on the probability 
#that they will continue in the teaching profession 
#in the following year
#using data from the 1999--2000 School and Staffing Survey (SASS) 
#and 2000--2001 Teacher Follow-up Survey (TFS)

#load data with generalized propensity scores
load("chapter6_generalized_propensity_score_estimation_results.Rdata")


#------------------------------

#calculate the IPTW from the propensity scores obtained with generalized boosting,
#and multiply by the sampling weights
imputedData$finalWeightATE2 <- get.weights(boost.ps, stop.method = "es.max",   estimand = "ATE",
                                 withSampW = TRUE)
#normalize weights
imputedData$finalWeightATE2 <- with(imputedData,finalWeightATE2 /
                                    mean(finalWeightATE2 ))
#check distribution of weights
with(imputedData, by(finalWeightATE2,Treat,summary))


#=============================================
#evaluate covariate balance with the twang function using the 
#generalized propensity scores obtained with generalized boosted regression

#visual inspection of covariate balance
require(twang)
tiff("Chapter6_figure6-4.tif", res=600, compression = "lzw", height=6, width=15, units="in")
plot(boost.ps, plots = 3, color=F, pairwiseMax = F)
dev.off()

#request a balance table
balance.boost.ps <- bal.table(boost.ps, digits = 2)
summary.balance.boost <- aggregate(std.eff.sz~tmt1+tmt2+stop.method, data=balance.boost.ps, FUN=max)
write.csv(summary.balance.boost, file="summary_balance_with_generalized_boosted_regression.csv")

#identify covariates with standardized mean difference above 0.25, which will be 
#used in covariance adjustment
(balance.above.0.25.gbm <- subset(balance.boost.ps, (std.eff.sz > 0.25 & stop.method == "es.max" )))
unbalanced.covariates2 <- as.character(balance.above.0.25.gbm$var)
unbalanced.covariates2 <- strsplit(unbalanced.covariates2, ":")
unbalanced.covariates2 <- unique(sapply(unbalanced.covariates2, "[", 1))
print(unbalanced.covariates2)