---
title: "Federated ATE for Randomised Controlled Trials"
output: rmarkdown::html_vignette
bibliography: biblio.bib
vignette: >
  %\VignetteIndexEntry{Federated ATE for Randomised Controlled Trials}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r, 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
})
```

```{r setup, eval = FALSE}
library(dsFedAteClient)
```

Randomised Controlled Trials (RCTs) are the gold standard for estimating causal treatment
effects, but individual trials are often limited in size or patient diversity, making
multi-site analyses highly desirable. The obstacle is that individual-level data are
typically siloed across institutions by privacy regulations and data governance policies,
making naive data centralisation infeasible. **DataSHIELD** addresses this through
federated analysis: each site returns only aggregate statistics so that raw records never
leave their originating server. This vignette walks through four worked examples of
federated Average Treatment Effect (ATE) estimation using the estimator of @khellaf2025,
which fits arm-specific outcome models across sites and aggregates local plug-in G-formula
estimates into a global ATE.

```{r, echo = FALSE, out.width = "50%", fig.align = "center"}
knitr::include_graphics(ifelse(knitr::is_latex_output(), "figure_federated.pdf", "figure_federated.svg"))
```

**Organisation.** The Estimator section introduces the notation and the three-step
estimation procedure. Four examples then illustrate the method under progressively
richer conditions: Example 1 covers the baseline homogeneous covariate distributions setting and verifies that
the federated estimator matches the centralised one; Example 2 shows how DataSHIELD's
`nfilter.glm` disclosure threshold governs whether estimation is feasible under extreme
treatment imbalance; Example 3 extends
the analysis to sites with heterogeneous covariate distributions; Example 4 addresses
site-specific intercept shifts that would otherwise confound the unadjusted estimator.

## Installation

```{r, eval = FALSE}
install.packages("dsFedAteClient")  # client package (analyst machine)
install.packages("dsFedAte")        # server package (deployed on each data server)
```

## Estimator

Consider $n$ individuals distributed across $K$ sites, with site $k$ holding $n_k$
observations ($\sum_k n_k = n$). For individual $i$ in site $k$, let $X_{ki} \in \mathbb{R}^d$
denote baseline covariates, $W_{ki} \in \{0,1\}$ the treatment indicator, and $Y_{ki} \in \mathbb{R}$
the observed outcome, related to the potential outcomes [@rubin1974] by
$Y_{ki} = W_{ki}\,Y_{ki}^{(1)} + (1-W_{ki})\,Y_{ki}^{(0)}$.

The quantity of interest is the **Average Treatment Effect**
$$\tau^\star = \mathbb{E}\!\left[Y^{(1)} - Y^{(0)}\right].$$

```{r, echo = FALSE, out.width = "35%", fig.align = "center"}
knitr::include_graphics(ifelse(knitr::is_latex_output(), "figure2_small.pdf", "figure2_small.svg"))
```

Under an outcome model as given in the DAG above, for each arm $w \in \{0,1\}$, an arm-specific linear model is fitted by **federated
linear regression** — each site contributes its local cross-products of the design
matrix and the solution is identical to the centralised one:
$$(\hat\alpha_w,\,\hat\beta_w) = \underset{\alpha,\,\beta}{\arg\min}
  \sum_{k=1}^K \sum_{i:\,W_{ki}=w}
  \bigl(Y_{ki} - \alpha - \beta^\top X_{ki}\bigr)^2.$$
No individual-level records are shared at any point.
Each site $k$ then applies the shared coefficients via the **plug-in G-formula**:
$$\hat\tau_k = \frac{1}{n_k}\sum_{i=1}^{n_k}
  \Bigl[
    \bigl(\hat\alpha_1 + \hat\beta_1^\top X_{ki}\bigr) -
    \bigl(\hat\alpha_0 + \hat\beta_0^\top X_{ki}\bigr)
  \Bigr],$$
and the global estimate is the sample-size-weighted average
$$\hat\tau = \sum_{k=1}^K \frac{n_k}{n}\,\hat\tau_k.$$

> **Remark.** In Example 4 we consider an outcome model as in the DAG above but with an
> additional arrow from the site indicator $H$ to $Y$, known as *study effects* in the
> literature. In that setting the arm-specific regressions are also fitted over $H$ — by
> augmenting the design matrix with site-membership dummies — in order to absorb the
> site-level intercept shifts and obtain an unbiased estimate of $\tau^\star$. See
> Example 4 for details.

## Local simulation with DSLite

All examples use **DSLite** to emulate a federated DataSHIELD infrastructure locally.
In a real deployment each site would run a dedicated institutional server.

```{r, eval = TRUE}
library(dsFedAteClient)
library(dsFedAte)
library(DSLite)
library(MASS)

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

## Shared parameters

All four examples share the model structure parameters below.
`sigma_noise` is the common noise standard deviation used in Examples 1–3;
Example 4 uses arm-specific values defined in its own data section.

```{r, 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)
```

## Example 1 — Homogeneous sites

Five sites, balanced sample sizes, same covariate distribution everywhere.
The federated estimator matches the centralised one exactly.

### Data

```{r, 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)
```

### Servers

```{r, 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")
```

### Estimate

```{r, 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
)
```

```{r, 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))
```

The federated ATE equals the centralised estimator to machine precision:

```{r, 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))
```

```{r, eval = TRUE}
DSI::datashield.logout(conns)
```

## Example 2 — Disclosure thresholds

Same model, but extreme treatment probabilities (0.9 / 0.1) produce very small arm sizes
at some sites. With $n_k = 200$ and $p = 0.1$, the expected minority arm has only 20
observations. The default `nfilter.glm = 0.33` blocks estimation because the GLM has
$d + 1 = 11$ parameters and $11/20 = 0.55 > 0.33$.

DataSHIELD enforces two distinct mechanisms: the *privacy control level* (`"banana"`,
`"permissive"`, etc.) gates which functions may be called at all, while `nfilter.*`
thresholds gate what values those functions will return. Both `localAteDS` and
`addStudyIdDS` are permitted under `"banana"` (the default), so the privacy level is not
the binding constraint. Raising `nfilter.glm` alone — while keeping `"banana"` — is
sufficient to unblock estimation.

### Data

```{r, 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
}
```

### Default settings — blocked

```{r, 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")
```

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

### Raised GLM threshold — succeeds

```{r, 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")
```

```{r, 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
)
```

```{r, 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))
```

The standard error is larger than in Example 1: extreme treatment imbalance reduces the
effective sample size for both arms. The privacy control level (`"banana"`) is unchanged
throughout — only the GLM disclosure threshold was relaxed.

```{r, eval = TRUE}
DSI::datashield.logout(conns_2)
```

## Example 3 — Heterogeneous covariate distributions

Each site now recruits from a different patient population: covariates are drawn as
$X_{ki} \sim \mathcal{N}(\mu_k, \Sigma_k)$ with site-specific mean $\mu_k \in \mathbb{R}^d$
and covariance $\Sigma_k \in \mathbb{R}^{d \times d}$.
Pass `heterogeneous = TRUE` to use the OLS-covariance-based variance formula, which is
valid under distributional shift.

### Data

```{r, 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)
```

### Servers

```{r, 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")
```

### Estimate

```{r, 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
)
```

```{r, 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))
```

```{r, 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))
```

```{r, eval = TRUE}
DSI::datashield.logout(conns_het)
```

## Example 4 — Study effects

Sites differ in baseline outcome levels: a site-specific scalar $h_k \in \mathbb{R}$
shifts both potential outcomes at site $k$,
$$Y_{ki}^{(w)} = \alpha_w + h_k + \beta_w^\top X_{ki} + \varepsilon_{ki}^{(w)}.$$
Because $h_k$ appears in both arms it cancels in $Y^{(1)} - Y^{(0)}$ and does not affect
the ATE. However, when treatment probabilities are heterogeneous across sites, $h_k$
confounds the unadjusted estimator.
The fix is to augment $X_{ki}$ with $K-1$ study-membership dummies
$H_{j,ki} = \mathbf{1}\{k = j\}$, $j = 2, \ldots, K$, and estimate the federated GLMs
on the extended design matrix, absorbing the shifts.
Pass `study_effects = TRUE` to do this automatically.

### Data

```{r, 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
}
```

### Servers

```{r, 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")
```

### Unadjusted — biased

```{r, 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
)
```

```{r, 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))
```

### Adjusted — unbiased

`heteroscedastic_arms = TRUE` estimates arm-specific noise variances $\hat\sigma_0^2$ and
$\hat\sigma_1^2$ separately, giving a correctly specified standard error when
$\sigma_0 \neq \sigma_1$.

```{r, 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
)
```

```{r, 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))
```

The adjusted estimate recovers the true ATE and the CI covers it.

```{r, eval = TRUE}
DSI::datashield.logout(conns_3)
```

## References
