#code for Chapter 3 "Propensity Score Weighting" of book:
#Leite, W. L. (2017). Practical propensity score methods using R. 
#Thousand Oaks, CA: Sage. 
# 
#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


#Estimate the effect of high school student participation in 
#career academies on future income
#using data from the Education Longitudinal Study (ELS)

#load datasets
load(file="Chapter3_ELS_data_imputed_example_career_academy.Rdata")
#------------------------------------

#load required survey library
require(survey)
options(survey.lonely.psu = "adjust")

#normalize the base year student sampling weight so that it sums to the sample size
ELS.data.imputed$bystuwt <- ELS.data.imputed$bystuwt/mean(ELS.data.imputed$bystuwt)
sum(ELS.data.imputed$bystuwt) #not they sum to sample size

#define design object that describes the sample characteristics
    #the variable psu identifies the primary sampling units (cluster ids)
    #the variable STRATA_ID identifies the strata ids
    #the variable bystuwt identifies the base-year sampling weights for
    #respondents of the 2002 and 2004 waves (Base year and 1st follow-up)
surveyDesign <- svydesign(ids=~psu, strata=~STRAT_ID, weights=~bystuwt,
                   data = ELS.data.imputed, nest=T)

#get weighted percentages of treated and untreated cases
#the treatment variable is "treat" - Ever in career academy from BY Student Questionnaire
(treatmentTable <- svymean(~treat, surveyDesign)) #8.8% of cases received treatment

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

#obtain weights for estimating the ATT for propensity scores obtained with logistic regression
ELS.data.imputed$weightATT <- with(ELS.data.imputed,ifelse(treat==1, 
                                                          1, pScores/(1-pScores))) 
with(ELS.data.imputed, by(weightATT,treat,summary))

#obtain weights for estimating the ATE
ELS.data.imputed$weightATE <- with(ELS.data.imputed,ifelse(treat==1, 
                                              1/pScores, 1/(1-pScores)))
                                                   
with(ELS.data.imputed, by(weightATE,treat,summary))

#===============================================================
#create a final weights for estimating the effect of career academy participation
#on income in the second folow up (2006) 


ELS.data.imputed$finalWeight2006 <- with(ELS.data.imputed,bystuwt*weightATT)


#normalize the base-year to second folow-up (2006) final weight
ELS.data.imputed$finalWeight2006 <- ELS.data.imputed$finalWeight2006/mean(ELS.data.imputed$finalWeight2006)


#check distribution of final weights
summary(ELS.data.imputed$finalWeight2006)



#save the dataset with weights
save(ELS.data.imputed, file="Chapter3_ELS_data_imputed_with_weights.Rdata", compress=T)

#==========================================================
#obtain weights for multiple imputed datasets

library(mitools)
#add propensity score mean to imputed datasets
allImputations <- update(allImputations, weightATT = ifelse(treat==1, 
                                                            1, pScores/(1-pScores)))

with(allImputations, by(weightATT,treat,summary))

#save object with multiple imputed datasets
save(allImputations, file="Chapter3_ELS_all_imputed_datasets_with_weights.Rdata", compress=T)