## ----include = FALSE----------------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)

eval_chunks <- requireNamespace("dsFedAte",       quietly = TRUE) &&
               requireNamespace("dsFedAteClient", quietly = TRUE) &&
               requireNamespace("MASS",           quietly = TRUE) &&
               requireNamespace("DSLite",         quietly = TRUE) &&
               requireNamespace("dsBaseClient",   quietly = TRUE)

knitr::opts_hooks$set(eval = function(options) {
  if (isTRUE(options$eval) && !eval_chunks) options$eval <- FALSE
  options
})

## ----setup, eval = FALSE------------------------------------------------------
# library(dsFedAteClient)

## ----echo = FALSE, out.width = "50%", fig.align = "center"--------------------
knitr::include_graphics(ifelse(knitr::is_latex_output(), "figure_federated.pdf", "figure_federated.svg"))

## ----eval = FALSE-------------------------------------------------------------
# install.packages("dsFedAteClient")  # client package (analyst machine)
# install.packages("dsFedAte")        # server package (deployed on each data server)

## ----echo = FALSE, out.width = "35%", fig.align = "center"--------------------
knitr::include_graphics(ifelse(knitr::is_latex_output(), "figure2_small.pdf", "figure2_small.svg"))

## ----eval = TRUE--------------------------------------------------------------
library(dsFedAteClient)
library(dsFedAte)
library(DSLite)
library(MASS)

datashield_env <- new.env(parent = emptyenv())
options(datashield.env = datashield_env)

## ----eval = TRUE--------------------------------------------------------------
K           <- 5
d           <- 10
sigma_noise <- 1.0
alpha_0     <- -2
alpha_1     <- -1.85

beta_1 <- seq(-1.75, 0.5,  by = 0.25)
beta_0 <- seq(-1.8,  0.0,  by = 0.2)

I_d <- diag(d)
J   <- matrix(1, nrow = d, ncol = d)

## ----eval = TRUE--------------------------------------------------------------
set.seed(123)

n_per_study <- 20 * d
p_treatment <- c(0.65, 0.65, 0.65, 0.35, 0.35)

mu_X    <- c(rep(1, d %/% 2), rep(-1, d %/% 2))
Sigma_X <- 0.5 * I_d + 0.5 * J

true_ATE <- (alpha_1 - alpha_0) + sum(mu_X * (beta_1 - beta_0))

data_list <- vector("list", K)
for (k in seq_len(K)) {
  X_k   <- mvrnorm(n = n_per_study, mu = mu_X, Sigma = Sigma_X)
  W_k   <- rbinom(n = n_per_study, size = 1, prob = p_treatment[k])
  Y_0   <- alpha_0 + X_k %*% beta_0 + rnorm(n_per_study, sd = sigma_noise)
  Y_1   <- alpha_1 + X_k %*% beta_1 + rnorm(n_per_study, sd = sigma_noise)
  Y_obs <- W_k * Y_1 + (1 - W_k) * Y_0
  df    <- data.frame(W = W_k, Y = as.vector(Y_obs), X_k)
  colnames(df) <- c("W", "Y", paste0("X", seq_len(d)))
  data_list[[k]] <- df
}
pooled_data <- do.call(rbind, data_list)

## ----eval = TRUE--------------------------------------------------------------
options(
  datashield.privacyControlLevel = "banana",
  nfilter.glm                    = 0.33
)

for (k in seq_len(K)) {
  srv <- DSLite::newDSLiteServer(tables = list(study_data = data_list[[k]]))
  srv$config(DSLite::defaultDSConfiguration(include = c("dsBase", "dsFedAte")))
  srv$aggregateMethod("localAteDS", "localAteDS")
  assign(paste0("dslite.server", k), srv, envir = datashield_env)
}

builder <- DSI::newDSLoginBuilder()
for (k in seq_len(K)) {
  builder$append(
    server = paste0("study", k),
    url    = paste0("dslite.server", k),
    table  = "study_data",
    driver = "DSLiteDriver"
  )
}
conns <- DSI::datashield.login(logins = builder$build(), assign = FALSE)
DSI::datashield.assign.table(conns = conns, table = "study_data", symbol = "base")

## ----eval = TRUE--------------------------------------------------------------
outcome_name    <- "Y"
treatment_name  <- "W"
covariate_names <- paste0("X", seq_len(d))

federated_ate <- ds.EstimateAte(
  datasources     = conns,
  base_name       = "base",
  outcome_name    = outcome_name,
  treatment_name  = treatment_name,
  covariate_names = covariate_names,
  heterogeneous   = FALSE,
  quiet           = TRUE
)

## ----eval = TRUE--------------------------------------------------------------
cat(sprintf("ATE: %.4f  SE: %.4f  95%% CI: [%.4f, %.4f]  True ATE: %.4f\n",
            federated_ate$ate, sqrt(federated_ate$variance),
            federated_ate$ci_lower, federated_ate$ci_upper, true_ATE))

## ----eval = TRUE--------------------------------------------------------------
pooled_treated <- pooled_data[pooled_data$W == 1, ]
pooled_control <- pooled_data[pooled_data$W == 0, ]
formula        <- paste0(outcome_name, " ~ ", paste(covariate_names, collapse = " + "))
m1             <- glm(formula, family = gaussian, data = pooled_treated)
m0             <- glm(formula, family = gaussian, data = pooled_control)
x_bar          <- colMeans(as.matrix(pooled_data[, covariate_names]))
pooled_ate     <- as.numeric(
  (m1$coefficients["(Intercept)"] - m0$coefficients["(Intercept)"]) +
  sum(x_bar * (m1$coefficients[covariate_names] - m0$coefficients[covariate_names]))
)
isTRUE(all.equal(pooled_ate, federated_ate$ate))

## ----eval = TRUE--------------------------------------------------------------
DSI::datashield.logout(conns)

## ----eval = TRUE--------------------------------------------------------------
p_treatment_2 <- c(0.9, 0.9, 0.9, 0.1, 0.1)

data_list_2 <- vector("list", K)
for (k in seq_len(K)) {
  X_k   <- mvrnorm(n = n_per_study, mu = mu_X, Sigma = Sigma_X)
  W_k   <- rbinom(n = n_per_study, size = 1, prob = p_treatment_2[k])
  Y_0   <- alpha_0 + X_k %*% beta_0 + rnorm(n_per_study, sd = sigma_noise)
  Y_1   <- alpha_1 + X_k %*% beta_1 + rnorm(n_per_study, sd = sigma_noise)
  Y_obs <- W_k * Y_1 + (1 - W_k) * Y_0
  df    <- data.frame(W = W_k, Y = as.vector(Y_obs), X_k)
  colnames(df) <- c("W", "Y", paste0("X", seq_len(d)))
  data_list_2[[k]] <- df
}

## ----eval = TRUE--------------------------------------------------------------
options(datashield.privacyControlLevel = "banana", nfilter.glm = 0.33)

for (k in seq_len(K)) {
  srv <- DSLite::newDSLiteServer(tables = list(study_data = data_list_2[[k]]))
  srv$config(DSLite::defaultDSConfiguration(include = c("dsBase", "dsFedAte")))
  srv$aggregateMethod("localAteDS", "localAteDS")
  assign(paste0("dslite.server2_", k), srv, envir = datashield_env)
}

builder2 <- DSI::newDSLoginBuilder()
for (k in seq_len(K)) {
  builder2$append(
    server = paste0("study", k),
    url    = paste0("dslite.server2_", k),
    table  = "study_data",
    driver = "DSLiteDriver"
  )
}
conns_2 <- DSI::datashield.login(logins = builder2$build(), assign = FALSE)
DSI::datashield.assign.table(conns = conns_2, table = "study_data", symbol = "base")

## ----eval = TRUE, error = TRUE------------------------------------------------
try({
ds.EstimateAte(
  datasources     = conns_2,
  base_name       = "base",
  outcome_name    = outcome_name,
  treatment_name  = treatment_name,
  covariate_names = covariate_names,
  heterogeneous   = FALSE
)
})

## ----eval = TRUE--------------------------------------------------------------
options(datashield.privacyControlLevel = "banana", nfilter.glm = 1.0)

for (k in seq_len(K)) {
  srv <- DSLite::newDSLiteServer(tables = list(study_data = data_list_2[[k]]))
  srv$config(DSLite::defaultDSConfiguration(include = c("dsBase", "dsFedAte")))
  srv$aggregateMethod("localAteDS", "localAteDS")
  assign(paste0("dslite.server2_", k), srv, envir = datashield_env)
}
conns_2 <- DSI::datashield.login(logins = builder2$build(), assign = FALSE)
DSI::datashield.assign.table(conns = conns_2, table = "study_data", symbol = "base")

## ----eval = TRUE--------------------------------------------------------------
federated_ate_2 <- ds.EstimateAte(
  datasources     = conns_2,
  base_name       = "base",
  outcome_name    = outcome_name,
  treatment_name  = treatment_name,
  covariate_names = covariate_names,
  heterogeneous   = FALSE,
  quiet           = TRUE
)

## ----eval = TRUE--------------------------------------------------------------
cat(sprintf("ATE: %.4f  SE: %.4f  95%% CI: [%.4f, %.4f]  True ATE: %.4f\n",
            federated_ate_2$ate, sqrt(federated_ate_2$variance),
            federated_ate_2$ci_lower, federated_ate_2$ci_upper, true_ATE))

## ----eval = TRUE--------------------------------------------------------------
DSI::datashield.logout(conns_2)

## ----eval = TRUE--------------------------------------------------------------
set.seed(123)

n_list          <- c(600, 120, 120, 120, 120)
p_treatment_het <- c(0.65, 0.65, 0.65, 0.35, 0.35)

mu_list <- list(
  c(rep( 1.0, d / 2), rep(-1.0, d / 2)),
  c(rep(-1.0, d / 2), rep( 1.0, d / 2)),
  rep(0, d),
  c(rep( 0.5, d / 2), rep(-1.0, d / 2)),
  c(rep( 1.2, d / 2), rep( 0.8, d / 2))
)

Sigma_1    <- 0.5 * I_d + 0.5 * J
Sigma_list <- list(
  Sigma_1,
  20 * Sigma_1 + 0.5 * I_d - 0.5 * J,
  0.02 * Sigma_1 + 0.7  * I_d,
  1.0  * Sigma_1 + 0.5  * I_d - 0.15 * J,
  1.5  * Sigma_1 + 0.5  * I_d - 0.15 * J
)

tau_k_vec    <- sapply(mu_list, function(mu_k) (alpha_1 - alpha_0) + sum(mu_k * (beta_1 - beta_0)))
true_ATE_het <- sum(n_list * tau_k_vec) / sum(n_list)

data_list_het <- vector("list", K)
for (k in seq_len(K)) {
  n_k   <- n_list[k]
  X_k   <- mvrnorm(n = n_k, mu = mu_list[[k]], Sigma = Sigma_list[[k]])
  W_k   <- rbinom(n = n_k, size = 1, prob = p_treatment_het[k])
  Y_0   <- alpha_0 + X_k %*% beta_0 + rnorm(n_k, sd = sigma_noise)
  Y_1   <- alpha_1 + X_k %*% beta_1 + rnorm(n_k, sd = sigma_noise)
  Y_obs <- W_k * Y_1 + (1 - W_k) * Y_0
  df    <- data.frame(W = W_k, Y = as.vector(Y_obs), X_k)
  colnames(df) <- c("W", "Y", paste0("X", seq_len(d)))
  data_list_het[[k]] <- df
}
pooled_data_het <- do.call(rbind, data_list_het)

## ----eval = TRUE--------------------------------------------------------------
options(datashield.privacyControlLevel = "banana", nfilter.glm = 0.33)

for (k in seq_len(K)) {
  srv <- DSLite::newDSLiteServer(tables = list(study_data = data_list_het[[k]]))
  srv$config(DSLite::defaultDSConfiguration(include = c("dsBase", "dsFedAte")))
  srv$aggregateMethod("localAteDS", "localAteDS")
  assign(paste0("dslite.server_het", k), srv, envir = datashield_env)
}

builder_het <- DSI::newDSLoginBuilder()
for (k in seq_len(K)) {
  builder_het$append(
    server = paste0("study", k),
    url    = paste0("dslite.server_het", k),
    table  = "study_data",
    driver = "DSLiteDriver"
  )
}
conns_het <- DSI::datashield.login(logins = builder_het$build(), assign = FALSE)
DSI::datashield.assign.table(conns = conns_het, table = "study_data", symbol = "base")

## ----eval = TRUE--------------------------------------------------------------
federated_ate_het <- ds.EstimateAte(
  datasources     = conns_het,
  base_name       = "base",
  outcome_name    = outcome_name,
  treatment_name  = treatment_name,
  covariate_names = covariate_names,
  heterogeneous   = TRUE,
  quiet           = TRUE
)

## ----eval = TRUE--------------------------------------------------------------
cat(sprintf("ATE: %.4f  SE: %.4f  95%% CI: [%.4f, %.4f]  True ATE: %.4f\n",
            federated_ate_het$ate, sqrt(federated_ate_het$variance),
            federated_ate_het$ci_lower, federated_ate_het$ci_upper, true_ATE_het))

## ----eval = TRUE--------------------------------------------------------------
x_bar_het      <- colMeans(as.matrix(pooled_data_het[, covariate_names]))
pooled_trt_het <- pooled_data_het[pooled_data_het$W == 1, ]
pooled_ctl_het <- pooled_data_het[pooled_data_het$W == 0, ]
m1_het         <- glm(formula, family = gaussian, data = pooled_trt_het)
m0_het         <- glm(formula, family = gaussian, data = pooled_ctl_het)
pooled_ate_het <- as.numeric(
  (m1_het$coefficients["(Intercept)"] - m0_het$coefficients["(Intercept)"]) +
  sum(x_bar_het * (m1_het$coefficients[covariate_names] - m0_het$coefficients[covariate_names]))
)
isTRUE(all.equal(pooled_ate_het, federated_ate_het$ate))

## ----eval = TRUE--------------------------------------------------------------
DSI::datashield.logout(conns_het)

## ----eval = TRUE--------------------------------------------------------------
set.seed(123)

n_list        <- c(600, 300, 300, 300, 300)
h_k           <- c(1, 2, -1, 30, 2)
p_treatment_3 <- c(0.65, 0.65, 0.65, 0.35, 0.35)
sigma_0       <- 0.5
sigma_1       <- 3.0

tau_k_vec <- sapply(mu_list, function(mu_k) (alpha_1 - alpha_0) + sum(mu_k * (beta_1 - beta_0)))
true_ATE  <- sum(n_list * tau_k_vec) / sum(n_list)

data_list_3 <- vector("list", K)
for (k in seq_len(K)) {
  n_k   <- n_list[k]
  X_k   <- mvrnorm(n = n_k, mu = mu_list[[k]], Sigma = Sigma_list[[k]])
  W_k   <- rbinom(n = n_k, size = 1, prob = p_treatment_3[k])
  Y_0   <- alpha_0 + h_k[k] + X_k %*% beta_0 + rnorm(n_k, sd = sigma_0)
  Y_1   <- alpha_1 + h_k[k] + X_k %*% beta_1 + rnorm(n_k, sd = sigma_1)
  Y_obs <- W_k * Y_1 + (1 - W_k) * Y_0
  df    <- data.frame(W = W_k, Y = as.vector(Y_obs), X_k)
  colnames(df) <- c("W", "Y", paste0("X", seq_len(d)))
  data_list_3[[k]] <- df
}

## ----eval = TRUE--------------------------------------------------------------
options(datashield.privacyControlLevel = "banana", nfilter.glm = 0.33)

for (k in seq_len(K)) {
  srv <- DSLite::newDSLiteServer(tables = list(study_data = data_list_3[[k]]))
  srv$config(DSLite::defaultDSConfiguration(include = c("dsBase", "dsFedAte")))
  srv$aggregateMethod("localAteDS", "localAteDS")
  srv$assignMethod("addStudyIdDS", "addStudyIdDS")
  assign(paste0("dslite.server3_", k), srv, envir = datashield_env)
}

builder3 <- DSI::newDSLoginBuilder()
for (k in seq_len(K)) {
  builder3$append(
    server = paste0("study", k),
    url    = paste0("dslite.server3_", k),
    table  = "study_data",
    driver = "DSLiteDriver"
  )
}
conns_3 <- DSI::datashield.login(logins = builder3$build(), assign = FALSE)
DSI::datashield.assign.table(conns = conns_3, table = "study_data", symbol = "base")

## ----eval = TRUE--------------------------------------------------------------
federated_ate_unadj <- ds.EstimateAte(
  datasources     = conns_3,
  base_name       = "base",
  outcome_name    = outcome_name,
  treatment_name  = treatment_name,
  covariate_names = covariate_names,
  study_effects   = FALSE,
  heterogeneous   = TRUE,
  quiet           = TRUE
)

## ----eval = TRUE--------------------------------------------------------------
cat(sprintf("Unadjusted ATE: %.4f   bias: %.4f   (true ATE: %.4f)\n",
            federated_ate_unadj$ate,
            federated_ate_unadj$ate - true_ATE,
            true_ATE))

## ----eval = TRUE--------------------------------------------------------------
federated_ate_adj <- ds.EstimateAte(
  datasources          = conns_3,
  base_name            = "base",
  outcome_name         = outcome_name,
  treatment_name       = treatment_name,
  covariate_names      = covariate_names,
  study_effects        = TRUE,
  heterogeneous        = TRUE,
  heteroscedastic_arms = TRUE,
  quiet                = TRUE
)

## ----eval = TRUE--------------------------------------------------------------
cat(sprintf("Adjusted ATE: %.4f  bias: %.4f  SE: %.4f  95%% CI: [%.4f, %.4f]  True ATE: %.4f\n",
            federated_ate_adj$ate, federated_ate_adj$ate - true_ATE,
            sqrt(federated_ate_adj$variance),
            federated_ate_adj$ci_lower, federated_ate_adj$ci_upper, true_ATE))

## ----eval = TRUE--------------------------------------------------------------
DSI::datashield.logout(conns_3)