tmle_itt_analysis.R

##################################
# ITT analysis with TMLE         #
##################################

library(purrr)
library(origami)
library(weights)
library(dplyr)
library(reporttools)
library(ggplot2)
library(reshape2)
library(ranger)
library(xgboost)
library(glmnet)
library(tictoc)
library(grid)
library(sl3)
options(sl3.verbose = FALSE)
library(car)
library(latex2exp)
library(WeightIt)
library(cobalt)

# Setup parallel processing
library(parallel)
library(doParallel)
library(foreach)

cores <- parallel::detectCores()
print(paste0("number of cores used: ", cores))

cl <- parallel::makeCluster(cores, outfile="")

doParallel::registerDoParallel(cl) # register cluster

# load utils
source('./src/tmle_calculation.R')
source('./src/tmleContrast.R')
source('./src/misc_fns.R')

# command line args
args <- commandArgs(trailingOnly = TRUE) # command line arguments
outcome <- as.character(args[1]) # "combined" (death or diabetes), "diabetes", or "death"
outcome.type <-  as.character(args[2]) # binomial" 
condition <-  as.character(args[3]) # "none", "schizophrenia", or "mdd"
weights.loc <- as.character(args[4])
use.SL <- as.logical(args[5])  # When TRUE, use Super Learner for initial Y model and treatment model estimation; if FALSE, use GLM
use.simulated <- as.logical(args[6])  # When TRUE, use simulated data; if FALSE, use real data.

scale.continuous <- TRUE # standardize continuous covariates

# output directory
output_dir <- './outputs/'
simulation_version <- ifelse(weights.loc=='none', paste0(format(Sys.time(), "%Y%m%d"),"/"), weights.loc)
if(!dir.exists(output_dir)){
  print(paste0('create folder for outputs at: ', output_dir))
  
  dir.create(output_dir)
}
output_dir = paste0(output_dir, simulation_version)
if(!dir.exists(output_dir)){
  print(paste0('create folder for outputs at: ', output_dir))
  dir.create(output_dir)
}

# load data
if(use.simulated){
  load("simdata_from_basevars.RData")
  cms.data <- simdata_from_basevars
}else{
  load("/data/MedicaidAP_associate/poulos/fup3yr_episode_months_deid_fixmonthout.RData") # run on Argos
  paste0("Original data dimension: ", dim(fup3yr_episode_months_deid))
  
  # keep all the patients who died before 3 years were done or had complete 3-year follow-up
  cms.data <- fup3yr_episode_months_deid[((as.Date('1/1/2011',format='%m/%d/%Y')- fup3yr_episode_months_deid$srvc_dt)>=0) & ((!is.na(fup3yr_episode_months_deid$dod_both_sources) & as.numeric(fup3yr_episode_months_deid$dod_both_sources - fup3yr_episode_months_deid$fup3yr_end_date)==0 &  as.numeric(fup3yr_episode_months_deid$index_plus_3yrs_date -fup3yr_episode_months_deid$fup3yr_end_date) >= 1) | (as.numeric(fup3yr_episode_months_deid$index_plus_3yrs_date -fup3yr_episode_months_deid$fup3yr_end_date) == 1)),]
  paste0("Data dimension after exclusion: ", dim(cms.data))
}

cms.data <- cms.data[cms.data$month_number==0,] # keep index month

# create outcome

Y.combined <- ifelse((!is.na(cms.data$days_to_death) & cms.data$days_to_death <=1095) | (!is.na(cms.data$days_to_diabetes) & cms.data$days_to_diabetes <= 1095), 1, 0)
Y.death <- ifelse((!is.na(cms.data$days_to_death) & cms.data$days_to_death <=1095), 1, 0)
Y.diabetes <- ifelse((!is.na(cms.data$days_to_diabetes) & cms.data$days_to_diabetes <= 1095), 1, 0)

if(outcome=='combined'){ # outcome is death or diabetes within 36 months
  Y <- Y.combined
}

if(outcome=='death'){ # outcome is death within 36 months
  Y <- Y.death
}

if(outcome=='diabetes'){ # outcome is diabetes within 36 months
  Y <- Y.diabetes
}

# treatment

A <- factor(cms.data$drug_group) # categorical
obs.treatment <- dummify(A) # dummy matrix

# covariates
L.unscaled <- cbind(dummify(factor(cms.data$state_character),show.na = FALSE)[,-c(6)],
                    dummify(factor(cms.data$race_defined_character),show.na = FALSE)[,-c(3)], 
                    dummify(factor(cms.data$smi_condition))[,-c(1)], # omit bipolar
                    "payer_index_mdcr"=dummify(factor(cms.data$payer_index))[,2],
                    "year"=cms.data$year, 
                    "female"=cms.data$female, 
                    "preperiod_ever_psych"=cms.data$preperiod_ever_psych,
                    "preperiod_ever_metabolic"=cms.data$preperiod_ever_metabolic,
                    "preperiod_ever_other"=cms.data$preperiod_ever_other,
                    "preperiod_drug_use_days"=cms.data$preperiod_drug_use_days,
                    "preperiod_ever_mt_gluc_or_lip"=cms.data$preperiod_ever_mt_gluc_or_lip,
                    "preperiod_ever_rx_antidiab"=cms.data$preperiod_ever_rx_antidiab,
                    "preperiod_ever_rx_cm_nondiab"=cms.data$preperiod_ever_rx_cm_nondiab,
                    "preperiod_ever_rx_other"=cms.data$preperiod_ever_rx_other,
                    "calculated_age"=cms.data$calculated_age,
                    "preperiod_er_mhsa"=cms.data$preperiod_er_mhsa,
                    "preperiod_er_nonmhsa"=cms.data$preperiod_er_nonmhsa,
                    "preperiod_er_injury"=cms.data$preperiod_er_injury,
                    "preperiod_cond_mhsa"=cms.data$preperiod_cond_mhsa,
                    "preperiod_cond_nonmhsa"=cms.data$preperiod_cond_nonmhsa,
                    "preperiod_cond_injury"=cms.data$preperiod_cond_injury,
                    "preperiod_los_mhsa"=cms.data$preperiod_los_mhsa,
                    "preperiod_los_nonmhsa"=cms.data$preperiod_los_nonmhsa,
                    "preperiod_los_injury"=cms.data$preperiod_los_injury)

rm(cms.data)

colnames(L.unscaled)[which(colnames(L.unscaled)=="West Virginia")] <- "West_Virginia"

L <- L.unscaled

if(scale.continuous){
  continuous.vars <- c("calculated_age","preperiod_drug_use_days","preperiod_er_mhsa","preperiod_er_nonmhsa","preperiod_er_injury","preperiod_cond_mhsa",
                       "preperiod_cond_nonmhsa","preperiod_cond_injury","preperiod_los_mhsa","preperiod_los_nonmhsa","preperiod_los_injury")
  
  L[,continuous.vars] <- scale(L[,continuous.vars]) # scale continuous vars
}

if(use.SL==FALSE){ # exclude multicolinear variables for GLM (VIF)
  glm.exlude <- c("California")
  L <- L[,!colnames(L)%in%glm.exlude]
}

## create condition dummy

if(condition=="schizophrenia"){
  x <- ifelse(L[,'schiz']==1,1,0)
} else if(condition=="mdd"){
  x <- ifelse(L[,'mdd']==1,1,0)
} else if(condition=="bipolar"){
  x <- ifelse(L[,'mdd']==0 & L[,'schiz']==0,1,0)
}else if(condition=="black"){
  x <- ifelse(L[,'black']==1,1,0)
}else if(condition=="latino"){
  x <- ifelse(L[,'latino']==1,1,0)
}else if(condition=="white"){
  x <- ifelse(L[,'white']==1,1,0)
}else if(condition=="other"){
  x <- ifelse(L[,'white']==0 & L[,'black']==0 & L[,'latino']==0,1,0)
} else if(condition=="no_drug"){
  x <- ifelse(L.unscaled[,'preperiod_drug_use_days']==0,1,0)
}else if(condition=="young"){
  x <- ifelse(L.unscaled[,"calculated_age"]<45,1,0)
}else if(condition=="no_drug_young"){
  x <- ifelse(L.unscaled[,"calculated_age"]<45 & L.unscaled[,'preperiod_drug_use_days']==0,1,0)
}

## set parameters

J <- 6
n <- length(Y) # sample size

gbound <- c(0.001,0.999) # define bounds to be used for the propensity score

ybound <- c(0.0001, 0.9999) # define bounds to be used for the Y predictions

n.folds <- 5 # number of folds for SL

target.gwt <- TRUE # When TRUE, moves propensity weights from denominator of clever covariate to regression weight when fitting updated model for Y

#create contrast matrix (1 is treated, -1s are reference)
loc <- t(combn(J,2))
colnames(loc) <- c("reference","treated")

C <- matrix(0, nrow(loc), J)
for(i in 1:nrow(C)){
  C[i,loc[i,1]] = -1
  C[i,loc[i,2]] = 1
}
colnames(C) <-1:J

C <- C[which(C[,1]==-1),] # reference

m <- nrow(C) # number of comparisons

## Manual TMLE (ours)

# stack learners into a model
learner_stack_A <- make_learner_stack(list("Lrnr_xgboost",nrounds=20, objective="multi:softprob", eval_metric="mlogloss",num_class=J), list("Lrnr_ranger",num.trees=100),list("Lrnr_ranger",num.trees=500), list("Lrnr_glmnet",nfolds = n.folds,alpha = 1, family = "multinomial"), list("Lrnr_glmnet",nfolds = n.folds,alpha = 0.25, family = "multinomial"), list("Lrnr_glmnet",nfolds = n.folds,alpha = 0.5, family = "multinomial"), list("Lrnr_glmnet",nfolds = n.folds,alpha = 0.75, family = "multinomial"))  
learner_stack_A_bin <- make_learner_stack(list("Lrnr_xgboost",nrounds=20, objective = "reg:logistic"), list("Lrnr_ranger",num.trees=100),list("Lrnr_ranger",num.trees=500), list("Lrnr_glmnet",nfolds = n.folds,alpha = 1, family = "binomial"), list("Lrnr_glmnet",nfolds = n.folds,alpha = 0.25, family = "binomial"), list("Lrnr_glmnet",nfolds = n.folds,alpha = 0.5, family = "binomial"), list("Lrnr_glmnet",nfolds = n.folds,alpha = 0.75, family = "binomial"))  

# metalearner defaults 
if(outcome.type=="continuous"){
  metalearner_Y <- make_learner(Lrnr_solnp,learner_function=metalearner_linear,eval_function=loss_squared_error) # nonlinear optimization via augmented Lagrange
  learner_stack_Y <- make_learner_stack(list("Lrnr_xgboost",nrounds=20, objective = "reg:squarederror"), list("Lrnr_ranger",num.trees=100), list("Lrnr_ranger",num.trees=500), list("Lrnr_glmnet",nfolds = n.folds,alpha = 1, family = "gaussian"), list("Lrnr_glmnet",nfolds = n.folds,alpha = 0.25, family = "gaussian"), list("Lrnr_glmnet",nfolds = n.folds,alpha = 0.5, family = "gaussian"), list("Lrnr_glmnet",nfolds = n.folds,alpha = 0.75, family = "gaussian")) 
}else if (outcome.type =="binomial"){
  metalearner_Y <- make_learner(Lrnr_solnp,learner_function=metalearner_logistic_binomial,eval_function=loss_loglik_binomial)
  learner_stack_Y <- make_learner_stack(list("Lrnr_xgboost",nrounds=20, objective = "reg:logistic"), list("Lrnr_ranger",num.trees=100), list("Lrnr_ranger",num.trees=500), list("Lrnr_glmnet",nfolds = n.folds,alpha = 1, family = "binomial"), list("Lrnr_glmnet",nfolds = n.folds,alpha = 0.25, family = "binomial"), list("Lrnr_glmnet",nfolds = n.folds,alpha = 0.5, family = "binomial"), list("Lrnr_glmnet",nfolds = n.folds,alpha = 0.75, family = "binomial")) 
}
metalearner_A <- make_learner(Lrnr_solnp,learner_function=metalearner_linear_multinomial, eval_function=loss_loglik_multinomial)
metalearner_A_bin <- make_learner(Lrnr_solnp,learner_function=metalearner_logistic_binomial,eval_function=loss_loglik_binomial)

tic(paste0("Initialize manual TMLE","outcome = ", outcome, "; outcome.type = ", outcome.type, "condition = ", condition, "use.SL =",use.SL, "use.simulated =",use.simulated))

## fit initial outcome model
if(use.SL){
  
  # define task and candidate learners
  initial_model_for_Y_task <- make_sl3_Task(data.frame(Y,L,obs.treatment), covariates = c(colnames(L),colnames(obs.treatment)), outcome = "Y", outcome_type="binomial", 
                                            folds = origami::make_folds(data.frame(Y,L,obs.treatment), fold_fun = folds_vfold, V = n.folds))
  
  # Super Learner algorithm
  
  initial_model_for_Y_sl <- make_learner(Lrnr_sl, # cross-validates base models
                                         learners = learner_stack_Y,
                                         metalearner = metalearner_Y,
                                         keep_extra=TRUE)
  if(weights.loc!='none'){
    initial_model_for_Y_sl_fit <-  readRDS(paste0(output_dir, 
                                                  "tmle_", outcome,"_",outcome.type, "_", "none", "_", use.SL, "_", use.simulated,"_",
                                                  "initial_model_for_Y_sl_fit.rds")) # same model used for all conditions
  }else{
    initial_model_for_Y_sl_fit <- initial_model_for_Y_sl$train(initial_model_for_Y_task)
    saveRDS(initial_model_for_Y_sl_fit,paste0(output_dir,"tmle_",outcome,"_",outcome.type, "_", condition, "_", use.SL, "_", use.simulated,"_", "initial_model_for_Y_sl_fit.rds")) # save model
  }
  initial_model_for_Y_preds <- initial_model_for_Y_sl_fit$predict(initial_model_for_Y_task) # predicted probs.
  
  if(outcome=="combined" & condition=="none" & outcome.type =="binomial"){
    Y.risks <- initial_model_for_Y_sl_fit$cv_risk(loss_loglik_binomial)
    print("Outcome model - SL CV risk")
    print(Y.risks)
  }
  
  for(j in 1:J){
    newdata <- data.frame(Y,L,matrix(0, nrow=nrow(obs.treatment), ncol=ncol(obs.treatment), dimnames = dimnames(obs.treatment)))
    newdata[,which(colnames(newdata)==colnames(obs.treatment)[j])] <- 1
    assign(paste0("Q",j), initial_model_for_Y_sl_fit$predict(sl3_Task$new(newdata, covariates = c(colnames(L),colnames(obs.treatment)), outcome = "Y", outcome_type="binomial")))
    rm(newdata)
  }
} else{
  if(outcome.type=="continuous"){
    initial_model_for_Y <- glm(Y~.,data=data.frame(cbind(Y,obs.treatment[,-2],L)), family=gaussian(), control = glm.control(maxit = 100)) 
    initial_model_for_Y_preds <- predict(initial_model_for_Y)
  }else{
    initial_model_for_Y <- glm(Y~.,data=data.frame(cbind(Y,obs.treatment[,-2],L)), family=binomial(), control = glm.control(maxit = 100)) 
    initial_model_for_Y_preds <- predict(initial_model_for_Y,type="response") # predicted probs.
  }
  
  if(condition=="none"){
    saveRDS(initial_model_for_Y,paste0(output_dir,"tmle_",outcome,"_",outcome.type, "_", condition, "_", use.SL, "_", use.simulated,"_", "initial_model_for_Y.rds")) # save model
    
    print("VIFs for binomial logistic regression - outcome model")
    print(vif(initial_model_for_Y))
  }
  
  for(j in 1:J){
    newdata <- data.frame(cbind(Y,L,matrix(0, nrow=nrow(obs.treatment), ncol=ncol(obs.treatment), dimnames = dimnames(obs.treatment))))
    newdata[,which(colnames(newdata)==colnames(obs.treatment)[j])] <- 1
    assign(paste0("Q",j), predict(initial_model_for_Y,newdata=data.frame(newdata), type="response"))
    rm(newdata)
  }
}

Qs <- t(do.call(mapply, c(FUN = cbind, mget(paste0("Q", 1:J)))))
colnames(Qs) <- paste0("Q", 1:J)

QAW <- data.frame(apply(cbind(QA=initial_model_for_Y_preds,Qs), 2, boundProbs, bounds=ybound)) # bound predictions

##  fit initial treatment model

if(use.SL){
  
  # multinomial
  
  initial_model_for_A_task <- make_sl3_Task(data.frame(A,L), covariates = colnames(L), outcome = "A", outcome_type="categorical", 
                                            folds = origami::make_folds(data.frame(A,L), fold_fun = folds_vfold, V = n.folds))
  
  initial_model_for_A_sl <- make_learner(Lrnr_sl, # cross-validates base models
                                         learners = learner_stack_A,
                                         metalearner = metalearner_A,
                                         keep_extra=TRUE)
  if(weights.loc!='none'){
    initial_model_for_A_sl_fit <-  readRDS(paste0(output_dir, # use combined outcome model
                                                  "tmle_", "combined","_",outcome.type, "_", "none", "_", use.SL, "_", use.simulated,"_",
                                                  "initial_model_for_A_sl_fit.rds")) # treatment model used in combined outcome analysis
  }else{
    initial_model_for_A_sl_fit <- initial_model_for_A_sl$train(initial_model_for_A_task)
    saveRDS(initial_model_for_A_sl_fit,paste0(output_dir,"tmle_",outcome,"_",outcome.type, "_", condition, "_", use.SL, "_", use.simulated, "_", "initial_model_for_A_sl_fit.rds")) # save model
  }
  
  g_preds  <- initial_model_for_A_sl_fit$predict(initial_model_for_A_task)  # estimated propensity scores 
  
  g_preds  <-  data.frame(matrix(unlist(lapply(g_preds, unlist)), nrow=length(lapply(g_preds, unlist)), byrow=TRUE)) 
  
  colnames(g_preds) <- colnames(obs.treatment)
  
  if(outcome=="combined" & condition=="none" & use.SL){
    A.risks <- initial_model_for_A_sl_fit$cv_risk(loss_loglik_multinomial)
    print("Treatment model - multinomial - SL CV risk")
    print(A.risks)
  }
  
  # binomial 
  
  tic(paste0("TMLE multiple binary treatment","outcome = ", outcome, "; outcome.type = ", outcome.type, "condition = ", condition))
  
  initial_model_for_A_task_bin <- lapply(1:J, function(j) make_sl3_Task(data.frame("A"=obs.treatment[,j],L), covariates = colnames(L), outcome = "A", outcome_type="binomial",
                                                                        folds = origami::make_folds(data.frame(A,L), fold_fun = folds_vfold, V = n.folds)))
  
  initial_model_for_A_sl_bin <- make_learner(Lrnr_sl, # cross-validates base models
                                             learners = learner_stack_A_bin,
                                             metalearner = metalearner_A_bin,
                                             keep_extra=TRUE)
  
  if(weights.loc!='none'){ 
    initial_model_for_A_sl_fit_bin <-  readRDS(paste0(output_dir, 
                                                      "tmle_", "combined","_",outcome.type, "_", "none", "_", use.SL, "_", use.simulated,"_",
                                                      "initial_model_for_A_sl_fit_bin.rds")) # treatment model used in combined outcome analysis
  }else{
    initial_model_for_A_sl_fit_bin <- lapply(1:J, function(j) initial_model_for_A_sl_bin$train(initial_model_for_A_task_bin[[j]]))
    saveRDS(initial_model_for_A_sl_fit_bin,paste0(output_dir,"tmle_",outcome,"_",outcome.type, "_", condition, "_", use.SL, "_", use.simulated,"_","initial_model_for_A_sl_fit_bin.rds")) # save model
  }
  
  g_preds_bin  <- sapply(1:J, function(j) initial_model_for_A_sl_fit_bin[[j]]$predict(initial_model_for_A_task_bin[[j]]))  # estimated propensity scores 
  
  colnames(g_preds_bin) <- colnames(obs.treatment)
  
  if(outcome=="combined" & condition=="none" & use.SL){
    A.risks.bin <- lapply(1:J, function(j) initial_model_for_A_sl_fit_bin[[j]]$cv_risk(loss_loglik_binomial))
    print("Treatment model - binomial - SL CV risk")
    print(A.risks.bin)
  }
} else{
  
  # multinomial logistic regression
  initial_model_for_A <- nnet::multinom(formula=A ~ ., data=data.frame(cbind(A, L)), maxit = 500, trace = FALSE, model = TRUE) 
  g_preds <- fitted(initial_model_for_A)
  colnames(g_preds) <- colnames(obs.treatment)
  
  if(condition=="none"){
    saveRDS(initial_model_for_A,paste0(output_dir,"tmle_",outcome,"_",outcome.type, "_", condition, "_", use.SL, "_", use.simulated,"_","initial_model_for_A.rds")) # save model
  }
  
  # binomial logistic regression
  initial_model_for_A_bin <- lapply(1:J, function(j) glm(formula=treatment ~ ., data=data.frame(cbind("treatment"=obs.treatment[,j], L)), family="binomial"))
  g_preds_bin <- sapply(1:J, function(j) predict(initial_model_for_A_bin[[j]], type="response"))
  colnames(g_preds_bin) <- colnames(obs.treatment)
  
  if(condition=="none"){
    saveRDS(initial_model_for_A_bin,paste0(output_dir,"tmle_",outcome,"_",outcome.type, "_", condition, "_", use.SL, "_", use.simulated,"_","initial_model_for_A_bin.rds")) # save model
  }
  
  if(condition=="none"){
    print("VIFs for binomial logistic regression - treatment model")
    lapply(1:J, function(j) print(vif(initial_model_for_A_bin[[j]])))
  }
}

# propensity scores

g_preds_bounded <- boundProbs(g_preds, gbound) # (trimmed) estimated propensity scores
g_preds_bin_bounded <- boundProbs(g_preds_bin, gbound) 

# Update equations, calculate ATE and variance
TMLE <- getTMLE(QAW, g_preds_bounded, obs.treatment, Y, outcome.type, iptw=FALSE, aiptw = FALSE, target.gwt)
TMLE_bin <- getTMLE(QAW, g_preds_bin_bounded, obs.treatment, Y, outcome.type, iptw=FALSE, aiptw = FALSE, target.gwt)

IPTW <- getTMLE(QAW=NULL, g_preds_bounded, obs.treatment, Y, outcome.type, iptw=TRUE, aiptw = FALSE, target.gwt)
IPTW_bin <- getTMLE(QAW=NULL, g_preds_bin_bounded, obs.treatment, Y, outcome.type, iptw=TRUE, aiptw = FALSE, target.gwt)

AIPTW <- getTMLE(QAW, g_preds_bounded, obs.treatment, Y, outcome.type, iptw=FALSE, aiptw = TRUE, target.gwt)
AIPTW_bin <- getTMLE(QAW, g_preds_bin_bounded, obs.treatment, Y, outcome.type, iptw=FALSE, aiptw = TRUE, target.gwt)

# calculate contrasts
if(condition=="none"){
  tmle_contrasts <- tmle_contrast(Qstar=TMLE, C, g_preds_bounded, obs.treatment, Y, outcome.type, alpha=0.05, gcomp = FALSE, iptw = FALSE, aiptw = FALSE, multi.adjust=FALSE)
}else{
  tmle_contrasts <- tmle_contrast(Qstar=TMLE, C, g_preds_bounded, obs.treatment, Y, outcome.type, alpha=0.05, gcomp = FALSE, iptw = FALSE, aiptw = FALSE, multi.adjust=FALSE, x)
}

if(condition=="none"){
  tmle_contrasts_bin <- tmle_contrast(Qstar=TMLE_bin, C, g_preds_bin_bounded, obs.treatment, Y, outcome.type, alpha=0.05, gcomp = FALSE, iptw = FALSE, aiptw = FALSE, multi.adjust=FALSE)
}else{
  tmle_contrasts_bin <- tmle_contrast(Qstar=TMLE_bin, C, g_preds_bin_bounded, obs.treatment, Y, outcome.type, alpha=0.05, gcomp = FALSE, iptw = FALSE, aiptw = FALSE, multi.adjust=FALSE, x)
}

gcomp_contrasts <- tmle_contrast(Qstar=NULL, C, g_preds_bounded=NULL, obs.treatment, Y, outcome.type, alpha=0.05, gcomp = TRUE, iptw = FALSE, aiptw = FALSE, multi.adjust=FALSE, QAW=QAW)

iptw_contrasts <- tmle_contrast(Qstar=IPTW, C, g_preds_bounded, obs.treatment, Y, outcome.type, alpha=0.05, gcomp = FALSE, iptw = TRUE, aiptw = FALSE, multi.adjust=FALSE)
iptw_contrasts_bin <- tmle_contrast(Qstar=IPTW_bin, C, g_preds_bin_bounded, obs.treatment, Y, outcome.type, alpha=0.05, gcomp = FALSE, iptw = TRUE, aiptw = FALSE, multi.adjust=FALSE)

aiptw_contrasts <- tmle_contrast(Qstar=AIPTW, C, g_preds_bounded, obs.treatment, Y, outcome.type, alpha=0.05, gcomp = FALSE, iptw = FALSE, aiptw = TRUE, multi.adjust=FALSE)
aiptw_contrasts_bin <- tmle_contrast(Qstar=AIPTW_bin, C, g_preds_bin_bounded, obs.treatment, Y, outcome.type, alpha=0.05, gcomp = FALSE, iptw = FALSE, aiptw = TRUE, multi.adjust=FALSE)

# store TMLE diagnostics (estimated initial Y, treatment probs.)

yinitial_tmle  <- QAW[-1] # estimated initial Y, treatment probs.

Ahat_tmle  <- g_preds
Ahat_tmle_bin  <- g_preds_bin

yhat_tmle <- colMeans(tmle_contrasts$Qstar)
yhat_tmle_bin <- colMeans(tmle_contrasts_bin$Qstar)

ess_ate_tmle <- (colSums(obs.treatment*(1/g_preds))**2)/(colSums(obs.treatment*(1/g_preds)**2)) 
ess_ate_tmle_bin <- (colSums(obs.treatment*(1/g_preds_bin))**2)/(colSums(obs.treatment*(1/g_preds_bin)**2)) 

# point and variance estimates for ATE

ATE_tmle  <- tmle_contrasts$taus$ATE# point and variance estimates for ATE (multinomial)
ATE_CI_tmle_upper <- sapply(tmle_contrasts$CI,"[[",1)[2,]
ATE_CI_tmle_lower <- sapply(tmle_contrasts$CI,"[[",1)[1,]

if(condition!="none"){
  CATE_tmle <- tmle_contrasts$taus$CATE # point and variance estimates for CATE (binomial)
  CATE_CI_tmle_upper <- sapply(tmle_contrasts$CI,"[[",2)[2,]
  CATE_CI_tmle_lower <- sapply(tmle_contrasts$CI,"[[",2)[1,] 
}else if (condition=="none"){
  CATE_tmle  <- NULL
  CATE_CI_tmle_upper <- NULL
  CATE_CI_tmle_lower <- NULL
}

ATE_tmle_bin  <- tmle_contrasts_bin$taus$ATE # point and variance estimates for ATE (binomial)
ATE_CI_tmle_bin_upper  <- sapply(tmle_contrasts_bin$CI,"[[",1)[2,]
ATE_CI_tmle_bin_lower <- sapply(tmle_contrasts_bin$CI,"[[",1)[1,]

if(condition!="none"){
  CATE_tmle_bin  <- tmle_contrasts_bin$taus$CATE # point and variance estimates for CATE (binomial)
  CATE_CI_tmle_bin_upper  <- sapply(tmle_contrasts_bin$CI,"[[",2)[2,]
  CATE_CI_tmle_bin_lower <- sapply(tmle_contrasts_bin$CI,"[[",2)[1,] 
}else if (condition=="none"){
  CATE_tmle_bin  <- NULL
  CATE_CI_tmle_bin_upper  <- NULL
  CATE_CI_tmle_bin_lower <- NULL
}

ATE_gcomp  <- gcomp_contrasts$taus$ATE
ATE_CI_gcomp_upper <- sapply(gcomp_contrasts$CI,"[[",1)[2,]
ATE_CI_gcomp_lower <- sapply(gcomp_contrasts$CI,"[[",1)[1,]

ATE_iptw  <- iptw_contrasts$taus$ATE
ATE_CI_iptw_upper <- sapply(iptw_contrasts$CI,"[[",1)[2,]
ATE_CI_iptw_lower <- sapply(iptw_contrasts$CI,"[[",1)[1,]

ATE_iptw_bin  <- iptw_contrasts_bin$taus$ATE
ATE_CI_iptw_bin_upper  <- sapply(iptw_contrasts_bin$CI,"[[",1)[2,]
ATE_CI_iptw_bin_lower <- sapply(iptw_contrasts_bin$CI,"[[",1)[1,]

ATE_aiptw  <- aiptw_contrasts$taus$ATE
ATE_CI_aiptw_upper <- sapply(aiptw_contrasts$CI,"[[",1)[2,]
ATE_CI_aiptw_lower <- sapply(aiptw_contrasts$CI,"[[",1)[1,]

ATE_aiptw_bin  <- aiptw_contrasts_bin$taus$ATE 
ATE_CI_aiptw_bin_upper  <- sapply(aiptw_contrasts_bin$CI,"[[",1)[2,]
ATE_CI_aiptw_bin_lower <- sapply(aiptw_contrasts_bin$CI,"[[",1)[1,]

print(toc())

# save results
saveRDS(list("yinitial_tmle"=yinitial_tmle,"Ahat_tmle"=Ahat_tmle,"yhat_tmle"= yhat_tmle, "ess_ate_tmle"=ess_ate_tmle,
             "Ahat_tmle_bin"=Ahat_tmle_bin,"yhat_tmle_bin"= yhat_tmle_bin, "ess_ate_tmle_bin"=ess_ate_tmle_bin,
             "ATE_tmle"= ATE_tmle,"ATE_CI_tmle_upper"=ATE_CI_tmle_upper, "ATE_CI_tmle_lower"=ATE_CI_tmle_lower,
             "CATE_tmle"= CATE_tmle,"CATE_CI_tmle_upper"=CATE_CI_tmle_upper, "CATE_CI_tmle_lower"=CATE_CI_tmle_lower,
             "ATE_tmle_bin"= ATE_tmle_bin,"ATE_CI_tmle_bin_upper"=ATE_CI_tmle_bin_upper, "ATE_CI_tmle_bin_lower"=ATE_CI_tmle_bin_lower,
             "CATE_tmle_bin"= CATE_tmle_bin,"CATE_CI_tmle_bin_upper"=CATE_CI_tmle_bin_upper, "CATE_CI_tmle_bin_lower"=CATE_CI_tmle_bin_lower,
             "ATE_gcomp"= ATE_gcomp,"ATE_CI_gcomp_upper"=ATE_CI_gcomp_upper, "ATE_CI_gcomp_lower"=ATE_CI_gcomp_lower,
             "ATE_iptw"= ATE_iptw,"ATE_CI_iptw_upper"=ATE_CI_iptw_upper, "ATE_CI_iptw_lower"=ATE_CI_iptw_lower,
             "ATE_iptw_bin"= ATE_iptw_bin,"ATE_CI_iptw_bin_upper"=ATE_CI_iptw_bin_upper, "ATE_CI_iptw_bin_lower"=ATE_CI_iptw_bin_lower,
             "ATE_aiptw"= ATE_aiptw,"ATE_CI_aiptw_upper"=ATE_CI_aiptw_upper, "ATE_CI_aiptw_lower"=ATE_CI_aiptw_lower,
             "ATE_aiptw_bin"= ATE_aiptw_bin,"ATE_CI_aiptw_bin_upper"=ATE_CI_aiptw_bin_upper, "ATE_CI_aiptw_bin_lower"=ATE_CI_aiptw_bin_lower),
        paste0(output_dir,"tmle_",outcome,"_",outcome.type, "_", condition, "_", use.SL, "_", use.simulated,"_","itt_analysis.rds"))

## Summary stats and results plots
source('./tmle_itt_analysis_eda.R')

## Censoring EDA
source('./tmle_itt_analysis_censoring.R')