############################################################################### ## swcrtcov.R ## ------------------------------------------------------------------------- ## Covariate-adjusted average treatment effect (ATE) estimation for ## cross-sectional STEPPED-WEDGE cluster-randomized trials (SW-CRTs), ## with two variance estimators: ## ## (1) a design-based (DB) "plug-in" cluster covariance, and ## (2) a cluster-robust sandwich (CRSE, type "CR0"). ## ## The package implements the estimators used in the accompanying simulation ## study. For each of several weighting schemes (estimands) it fits: ## ## * unadjusted (saturated period x treatment model) ## * ANCOVA I (additive covariate main effects) ## * ANCOVA II (period x covariate interactions) ## * ANCOVA III (additive period effects + treatment x covariate) ## * ANCOVA IV (period x covariate AND period x covariate x trt) ## ## and returns, for each, the weighted ATE point estimate together with the ## design-based and cluster-robust standard errors. ## ## Design assumptions (standard "balanced" stepped wedge): ## * J periods, indexed 1..J. Period 1 (all control) and period J ## (all treated) carry no treatment contrast and are dropped from the ## analysis; analysis periods are 2..(J-1). ## * Equal allocation: I/(J-1) clusters adopt treatment at each of the ## J-1 adoption times {2,...,J}. Adoption is monotone (once treated, ## always treated). ## * Each cluster's adoption time is recorded in a `schedule` column. ## ## No external package is required. If `clubSandwich` is installed it can be ## used for the CR0 standard error (set crse_backend = "clubSandwich"); the ## built-in base-R backend reproduces clubSandwich's CR0 to machine precision. ## ## Dependencies: base R only (stats). ############################################################################### ## =========================================================================== ## 1. Treatment-schedule helper ## =========================================================================== #' Monotone treatment indicator for a single adoption time. #' #' Returns a length-J 0/1 vector giving the treatment status across periods #' 1..J for a cluster that first adopts treatment in period `j`. #' #' @param j integer adoption time (1..J); j = 1 means treated in all periods. #' @param J integer total number of periods. #' @return integer vector of length J. swcrt_treat <- function(j, J) { if (j == 1) { out <- rep(1, J) } else if (j > J) { out <- rep(0, J) } else { out <- c(rep(0, j - 1), rep(1, J - j + 1)) } out } ## =========================================================================== ## 2. Data generator (for tutorials / simulation studies) ## =========================================================================== #' Simulate one individual-level data set from a balanced stepped-wedge CRT. #' #' Reproduces the data-generating process used in the simulation study: #' cluster-period sizes grow with period, two "structured" covariates (x1, x2) #' are centred within period, three i.i.d. normal covariates (x3, x4, x5) are #' added, and treated/control potential outcomes are generated under a known #' individual-level treatment effect. Three individual weight columns are #' returned, corresponding to three estimands: #' #' w1 : equal individual weights -> individual-average ATE #' w2 : inverse period-total-size -> period-average (cell) ATE #' w3 : inverse cluster-period size -> cluster-period-average ATE #' #' The returned data frame is already restricted to the analysis periods #' 2..(J-1) and contains period dummy columns period_2..period_{J-1}. #' #' @param I number of clusters (must be a multiple of J-1). #' @param J number of periods including the all-control and all-treated periods. #' @param rho0 control-arm intra-cluster correlation (cluster random-effect var). #' @param seed optional integer RNG seed. #' @return a data.frame with individual-level rows. swcrt_simulate <- function(I = 18, J = 7, rho0 = 0.1, seed = NULL) { if (!is.null(seed)) set.seed(seed) if (I %% (J - 1) != 0) stop("I must be a multiple of (J-1) for a balanced stepped wedge.") ## ---- cluster-period level attributes ---- cluster <- rep(1:I, each = J) period <- rep(1:J, I) size <- sapply(period, function(x) round(runif(1, min = 10 + 2.5 * x^2, max = 90 + 2.5 * x^2), 0)) adoption_time <- sample(rep(2:J, each = I / (J - 1)), I, replace = FALSE) schedule <- rep(adoption_time, each = J) treatment <- unlist(lapply(adoption_time, function(x) swcrt_treat(x, J))) ## ---- expand to individual level ---- cluster_ind <- rep(cluster, size) period_ind <- rep(period, size) individual_ind <- unlist(lapply(size, function(x) seq_len(x))) schedule_ind <- rep(schedule, size) treatment_ind <- rep(treatment, size) n_ind <- length(cluster_ind) ## ---- period indicator matrix ---- time_ind <- matrix(0, nrow = n_ind, ncol = J) for (j in 1:J) time_ind[, j] <- as.numeric(period_ind == j) colnames(time_ind) <- paste("period", 1:J, sep = "_") ## ---- covariates (x1, x2 centred within period) ---- x_ij1 <- rbinom(n = I * J, size = 1, prob = 0.5) x_ij1_ind <- rep(x_ij1, size) x_j1_mean <- tapply(x_ij1, period, mean) x_ij1_ind_c <- x_ij1_ind - rowSums(time_ind * (rep(1, n_ind) %o% x_j1_mean)) x_ijk2 <- cluster_ind / I + runif(n_ind, min = -1, max = 1) x_j2_mean <- tapply(x_ijk2, period_ind, mean) x_ijk2_c <- x_ijk2 - rowSums(time_ind * (rep(1, n_ind) %o% x_j2_mean)) x_ijk3 <- rnorm(n_ind) x_ijk4 <- rnorm(n_ind) x_ijk5 <- rnorm(n_ind) ## ---- potential outcomes ---- Y_ijk0 <- period_ind / J + x_ij1_ind + x_ijk2_c^2 + x_ijk3 + x_ijk4 + x_ijk5 + rep(rep(rnorm(I, sd = sqrt(rho0)), each = J), size) + rnorm(n_ind, sd = sqrt(1 - rho0)) Y_ijk1 <- 2 * rep(size, size) * I / mean(tapply(size, period, sum)[-1]) + 0.5 * x_ij1_ind + x_ijk2_c^3 Y_ijk <- treatment_ind * Y_ijk1 + (1 - treatment_ind) * Y_ijk0 ## ---- individual weights (three estimands) ---- w1 <- rep(1, n_ind) w2 <- rowSums(time_ind * (rep(1, n_ind) %o% (1 / tapply(size, period, sum)))) w3 <- 1 / rep(size, size) ## ---- assemble ---- data_out <- data.frame(time_ind) data_out$cluster <- cluster_ind data_out$period <- period_ind data_out$individual <- individual_ind data_out$schedule <- schedule_ind data_out$treatment <- treatment_ind data_out$x1 <- x_ij1_ind_c data_out$x2 <- x_ijk2_c data_out$x3 <- x_ijk3 data_out$x4 <- x_ijk4 data_out$x5 <- x_ijk5 data_out$y1 <- Y_ijk1 data_out$y0 <- Y_ijk0 data_out$y_obs <- Y_ijk data_out$w1 <- w1 data_out$w2 <- w2 data_out$w3 <- w3 ## restrict to analysis periods 2..(J-1) and drop the two boundary dummies keep_rows <- data_out$period %in% (2:(J - 1)) drop_cols <- c("period_1", paste0("period_", J)) data_out <- data_out[keep_rows, !(names(data_out) %in% drop_cols)] rownames(data_out) <- NULL data_out } ## =========================================================================== ## 3. Internal helpers: design inference & coefficient identification ## =========================================================================== ## Infer (I, J) and the analysis-period dummy names from an analysis data frame. .swcrt_design <- function(data, period_var, cluster_var, J = NULL) { periods <- sort(unique(data[[period_var]])) if (is.null(J)) J <- max(periods) + 1L # analysis periods are 2..J-1 I <- length(unique(data[[cluster_var]])) period_dummies <- paste("period", periods, sep = "_") list(I = I, J = J, periods = periods, period_dummies = period_dummies) } ## Ensure the period dummy columns exist; create them from `period` if absent. .swcrt_add_period_dummies <- function(data, periods, period_var) { for (p in periods) { nm <- paste("period", p, sep = "_") if (is.null(data[[nm]])) data[[nm]] <- as.numeric(data[[period_var]] == p) } data } ## TRUE iff a coefficient name is exactly one period dummy (a period main effect) .is_period_main <- function(nm, period_dummies) { nm %in% period_dummies } ## TRUE iff a coefficient name is the 2-way interaction {period dummy, treatment} .is_period_trt <- function(nm, period_dummies, treatment_var) { parts <- strsplit(nm, ":", fixed = TRUE)[[1]] length(parts) == 2L && treatment_var %in% parts && any(parts %in% period_dummies) } ## Return the period-ascending order index for a set of period:treatment names. .period_order <- function(names_vec, period_dummies) { ## numeric period suffix embedded in each name suff <- sapply(names_vec, function(nm) { parts <- strsplit(nm, ":", fixed = TRUE)[[1]] pd <- parts[parts %in% period_dummies] as.integer(sub("period_", "", pd)) }) order(suff) } ## =========================================================================== ## 4. Design-based plug-in cluster covariance ## =========================================================================== #' Design-based plug-in covariance of the period-specific treatment effects. #' #' Computes the (J-2) x (J-2) covariance matrix of the vector of #' period-specific weighted treatment-effect estimates, using the #' randomization/design-based "plug-in" estimator. This is the matrix that is #' sandwiched by the estimand weights `varpi` to obtain the design-based #' variance of the overall ATE. #' #' @param data analysis data frame (periods 2..J-1). #' @param residual numeric vector: the partial residual r_x for each row #' (raw residual plus the fitted period & period x treatment part). #' @param weights numeric vector of individual weights (the estimand weights). #' @param I number of clusters. #' @param J number of periods (analysis periods are 2..J-1). #' @param period_var,treatment_var,cluster_var,schedule_var column names. #' @return a (J-2) x (J-2) numeric covariance matrix (period-ascending order). swcrt_plugin_cov <- function(data, residual, weights, I, J, period_var = "period", treatment_var = "treatment", cluster_var = "cluster", schedule_var = "schedule") { d <- data.frame( period = data[[period_var]], cluster = data[[cluster_var]], schedule = data[[schedule_var]], treatment = data[[treatment_var]], res_x = as.numeric(residual), weights = as.numeric(weights) ) Aperiods <- 2:(J - 1) # analysis periods P <- J - 2L # number of analysis periods ## balanced-design constants I_a <- I / (J - 1) # clusters per adoption group I_tr_seq <- cumsum(rep(I_a, P)) # treated count per period I_ct_seq <- I - I_tr_seq # control count per period ## period-specific weighted mean residual among treated / control u1 <- sapply(Aperiods, function(x) weighted.mean(d$res_x[d$treatment == 1 & d$period == x], d$weights[d$treatment == 1 & d$period == x])) u0 <- sapply(Aperiods, function(x) weighted.mean(d$res_x[d$treatment == 0 & d$period == x], d$weights[d$treatment == 0 & d$period == x])) ## average cluster weight per period within each arm wj1 <- tapply(d$weights[d$treatment == 1], d$period[d$treatment == 1], sum) / I_tr_seq wj0 <- tapply(d$weights[d$treatment == 0], d$period[d$treatment == 0], sum) / I_ct_seq ## accumulate the per-adoption-group cluster covariance cluster_cov_ar <- array(0, dim = c(J - 1, P, P)) for (a in 2:J) { da <- d[d$schedule == a, ] ## (period, cluster)-level weighted-mean residual and weight total key <- interaction(da$period, da$cluster, drop = TRUE, sep = "\r") w_sum <- tapply(da$weights, key, sum) rx_wmean <- tapply(seq_along(da$weights), key, function(ix) weighted.mean(da$res_x[ix], da$weights[ix])) kp <- strsplit(names(w_sum), "\r", fixed = TRUE) kper <- as.integer(vapply(kp, `[`, "", 1)) kcl <- as.integer(vapply(kp, `[`, "", 2)) agg <- data.frame(period = kper, cluster = kcl, mean_res_x = as.numeric(rx_wmean), w_ij = as.numeric(w_sum)) ## expected residual mean for this adoption group, by period treated_at <- a <= Aperiods # logical length P uj_mean_a <- u1 * treated_at + u0 * (!treated_at) cl_ids <- unique(agg$cluster) acc <- matrix(0, P, P) for (cc in cl_ids) { sub <- agg[agg$cluster == cc, ] sub <- sub[order(sub$period), ] # period-ascending vec <- sub$mean_res_x - uj_mean_a wd <- (sub$w_ij * I_a) / (wj1 * I_tr_seq) * treated_at - (sub$w_ij * I_a) / (wj0 * I_ct_seq) * (!treated_at) v <- as.numeric(vec * wd) # vec %*% diag(wd) acc <- acc + (v %o% v) } cluster_cov_ar[a - 1, , ] <- (1 / I_a) * (1 / (I_a - 1)) * acc } cov_mat <- apply(cluster_cov_ar, c(2, 3), sum) cov_mat } ## =========================================================================== ## 5. Cluster-robust (CR0) sandwich covariance -- base-R implementation ## =========================================================================== #' CR0 cluster-robust covariance for a (weighted) lm. #' #' Reproduces clubSandwich::vcovCR(fit, cluster, type = "CR0") for linear #' models, including the weighted-least-squares case: #' V = (X'WX)^{-1} [ sum_g (X_g' W_g e_g)(X_g' W_g e_g)' ] (X'WX)^{-1}. #' #' @param fit an lm object. #' @param cluster vector of cluster ids (length = number of model rows). #' @return full p x p covariance matrix with dimnames = coefficient names. swcrt_vcovCR0 <- function(fit, cluster) { X <- model.matrix(fit) e <- as.numeric(residuals(fit)) w <- fit$weights if (is.null(w)) w <- rep(1, length(e)) ## drop aliased (NA) coefficients to keep X full column rank b <- coef(fit) keep <- !is.na(b) X <- X[, keep, drop = FALSE] XtWX <- crossprod(X, X * w) # X' W X bread <- solve(XtWX) score <- X * (w * e) # row i: w_i e_i x_i cl <- as.character(cluster) meat <- crossprod(rowsum(score, cl)) # sum_g s_g s_g' V <- bread %*% meat %*% bread dimnames(V) <- list(colnames(X), colnames(X)) V } ## =========================================================================== ## 6. Formula construction for the five estimators ## =========================================================================== ## Build the model formula (as a string) for a given adjustment method. ## method %in% c("unadjusted","ANCOVA1","ANCOVA2","ANCOVA3","ANCOVA4") .swcrt_formula <- function(method, outcome, period_dummies, treatment_var, covariates) { pd <- period_dummies trt <- treatment_var cv <- covariates ## core terms shared by every model: period_main <- pd # period main effects period_trt <- paste(pd, trt, sep = ":") # period x treatment terms <- switch( method, "unadjusted" = c(period_main, period_trt), "ANCOVA1" = c(cv, period_main, period_trt), "ANCOVA2" = c(period_main, as.vector(outer(pd, cv, paste, sep = ":")), period_trt), "ANCOVA3" = c(cv, period_main, paste(trt, cv, sep = ":"), period_trt), "ANCOVA4" = c(period_main, as.vector(outer(pd, cv, paste, sep = ":")), as.vector(outer(outer(pd, cv, paste, sep = ":"), trt, paste, sep = ":")), period_trt), stop("Unknown method: ", method) ) paste(outcome, "~", paste(terms, collapse = " + "), "- 1") } ## =========================================================================== ## 7. Single-estimator fit: point estimate + DB SE + CRSE ## =========================================================================== #' Fit one covariate-adjustment estimator for one estimand (weight column). #' #' @param data analysis data frame. #' @param outcome name of the outcome column. #' @param weight name of the weight column (defines the estimand). #' @param method one of "unadjusted","ANCOVA1","ANCOVA2","ANCOVA3","ANCOVA4". #' @param covariates character vector of covariate column names. #' @param period_var,treatment_var,cluster_var,schedule_var column names. #' @param J optional override for number of periods. #' @param crse logical: also compute the cluster-robust SE? #' @param crse_backend "base" (default) or "clubSandwich". #' @return a list with estimate, se_db, se_crse, and supporting pieces. swcrt_fit <- function(data, outcome = "y_obs", weight = "w1", method = "unadjusted", covariates = c("x1", "x2", "x3", "x4", "x5"), period_var = "period", treatment_var = "treatment", cluster_var = "cluster", schedule_var = "schedule", J = NULL, crse = TRUE, crse_backend = "base") { des <- .swcrt_design(data, period_var, cluster_var, J = J) J <- des$J; I <- des$I data <- .swcrt_add_period_dummies(data, des$periods, period_var) pd <- des$period_dummies ## --- fit the weighted linear model --- if (method == "unadjusted") covariates <- character(0) fml <- .swcrt_formula(method, outcome, pd, treatment_var, covariates) fit <- lm(as.formula(fml), data = data, weights = data[[weight]]) b <- coef(fit) nm <- names(b) ## --- identify period-main and period:treatment coefficients by NAME --- main_idx <- which(vapply(nm, .is_period_main, logical(1), period_dummies = pd)) trt_idx <- which(vapply(nm, .is_period_trt, logical(1), period_dummies = pd, treatment_var = treatment_var)) ## order both period-ascending so they align with each other and with varpi main_idx <- main_idx[.period_order(nm[main_idx], pd)] trt_idx <- trt_idx [.period_order(nm[trt_idx], pd)] beta_trt <- b[trt_idx] beta_main <- b[main_idx] ## --- estimand weights varpi (period totals / grand total) --- w_vec <- data[[weight]] per <- data[[period_var]] per_tot <- tapply(w_vec, per, sum) per_tot <- per_tot[order(as.integer(names(per_tot)))] varpi <- as.numeric(per_tot / sum(w_vec)) ## --- point estimate: weighted average of period-specific effects --- estimate <- as.numeric(sum(varpi * beta_trt)) ## --- partial residual r_x = e + [period | trt*period] %*% [main | trt] --- Pdum <- as.matrix(data[, pd, drop = FALSE]) Pdum <- Pdum[, order(as.integer(sub("period_", "", colnames(Pdum)))), drop = FALSE] Xpt <- cbind(Pdum, data[[treatment_var]] * Pdum) res_x <- as.numeric(residuals(fit)) + as.numeric(Xpt %*% c(beta_main, beta_trt)) ## --- design-based SE --- covDB <- swcrt_plugin_cov(data, res_x, w_vec, I, J, period_var, treatment_var, cluster_var, schedule_var) se_db <- sqrt(as.numeric(varpi %*% covDB %*% varpi)) ## --- cluster-robust SE --- se_crse <- NA_real_ if (crse) { if (crse_backend == "clubSandwich" && requireNamespace("clubSandwich", quietly = TRUE)) { Vcr <- clubSandwich::vcovCR(fit, cluster = data[[cluster_var]], type = "CR0") Vcr <- as.matrix(Vcr) } else { Vcr <- swcrt_vcovCR0(fit, cluster = data[[cluster_var]]) } sub <- Vcr[nm[trt_idx], nm[trt_idx], drop = FALSE] se_crse <- sqrt(as.numeric(varpi %*% sub %*% varpi)) } list(estimate = estimate, se_db = se_db, se_crse = se_crse, method = method, weight = weight, beta_period_trt = beta_trt, varpi = varpi, cov_db = covDB, fit = fit, formula = fml) } ## =========================================================================== ## 8. Convenience wrapper: all methods x all estimands ## =========================================================================== #' Estimate the ATE under every method/estimand combination. #' #' @param data analysis data frame. #' @param outcome outcome column name. #' @param weights named character vector of weight columns; names label the #' estimands (default w1/w2/w3). #' @param methods character vector of methods to run. #' @param covariates covariate column names. #' @param ... passed to swcrt_fit (period_var, crse_backend, J, ...). #' @return a tidy data.frame, one row per (estimand, method). swcrt_ate <- function(data, outcome = "y_obs", weights = c(w1 = "w1", w2 = "w2", w3 = "w3"), methods = c("unadjusted", "ANCOVA1", "ANCOVA2", "ANCOVA3", "ANCOVA4"), covariates = c("x1", "x2", "x3", "x4", "x5"), ...) { if (is.null(names(weights))) names(weights) <- weights rows <- list() for (wlab in names(weights)) { for (m in methods) { r <- swcrt_fit(data, outcome = outcome, weight = weights[[wlab]], method = m, covariates = covariates, ...) rows[[length(rows) + 1]] <- data.frame( estimand = wlab, method = m, estimate = r$estimate, se_db = r$se_db, se_crse = r$se_crse, stringsAsFactors = FALSE) } } out <- do.call(rbind, rows) rownames(out) <- NULL out } ## =========================================================================== ## 9. True (oracle) estimand value, from potential outcomes ## =========================================================================== #' Compute the true weighted ATE from known potential outcomes (for sims). #' #' @param data data frame containing potential outcome columns y1, y0. #' @param weight name of the weight column. #' @return numeric scalar: sum((y1-y0)*w)/sum(w). swcrt_true_ate <- function(data, weight = "w1", y1 = "y1", y0 = "y0") { w <- data[[weight]] sum((data[[y1]] - data[[y0]]) * w) / sum(w) } ## =========================================================================== ## 10. Monte-Carlo driver (bias / RMSE / SE / coverage) ## =========================================================================== #' Run a Monte-Carlo simulation reproducing the study's summary metrics. #' #' For each replication a data set is generated with `gen_fun`, the true #' estimand is computed from the potential outcomes, and every method/estimand #' combination is estimated. The function returns the per-replication #' estimates/SEs and a summary table of bias, RMSE, average SEs and 95% #' confidence-interval coverage (for both the DB and CRSE intervals). #' #' @param n_rep number of replications. #' @param gen_fun function() returning one simulated data frame (with y0/y1). #' @param weights,methods,covariates as in swcrt_ate. #' @param seed optional RNG seed set once before the loop. #' @param progress print progress every 10 reps? #' @param ... passed to swcrt_fit. #' @return list(estimates = , summary = ). swcrt_run_simulation <- function(n_rep = 200, gen_fun = function() swcrt_simulate(18, 7), weights = c(w1 = "w1", w2 = "w2", w3 = "w3"), methods = c("unadjusted", "ANCOVA1", "ANCOVA2", "ANCOVA3", "ANCOVA4"), covariates = c("x1", "x2", "x3", "x4", "x5"), seed = 1, progress = TRUE, ...) { if (!is.null(seed)) set.seed(seed) if (is.null(names(weights))) names(weights) <- weights combos <- expand.grid(estimand = names(weights), method = methods, stringsAsFactors = FALSE) K <- nrow(combos) truth <- matrix(NA_real_, n_rep, length(weights), dimnames = list(NULL, names(weights))) est <- se_db <- se_cr <- matrix(NA_real_, n_rep, K) for (ii in seq_len(n_rep)) { d <- gen_fun() for (wl in names(weights)) truth[ii, wl] <- swcrt_true_ate(d, weights[[wl]]) res <- swcrt_ate(d, weights = weights, methods = methods, covariates = covariates, ...) ## align res rows to combos order key_res <- paste(res$estimand, res$method) key_cmb <- paste(combos$estimand, combos$method) idx <- match(key_cmb, key_res) est[ii, ] <- res$estimate[idx] se_db[ii, ] <- res$se_db[idx] se_cr[ii, ] <- res$se_crse[idx] if (progress && ii %% 10 == 0) { cat("\r", paste("repetition", ii, "of", n_rep, "completed")) flush.console() } } if (progress) cat("\n") ## ---- summary metrics ---- summ <- data.frame(combos, true = NA, mean_est = NA, rel_bias = NA, rmse = NA, emp_sd = NA, mean_se_db = NA, mean_se_crse = NA, cover_db = NA, cover_crse = NA) for (k in seq_len(K)) { wl <- combos$estimand[k] th <- truth[, wl] e <- est[, k] summ$true[k] <- mean(th) summ$mean_est[k] <- mean(e) summ$rel_bias[k] <- mean((e - th) / th) summ$rmse[k] <- sqrt(mean((e - th)^2)) summ$emp_sd[k] <- sd(e) summ$mean_se_db[k] <- mean(se_db[, k]) summ$mean_se_crse[k] <- mean(se_cr[, k]) summ$cover_db[k] <- mean((e - 1.96 * se_db[, k] <= th) & (e + 1.96 * se_db[, k] > th)) summ$cover_crse[k] <- mean((e - 1.96 * se_cr[, k] <= th) & (e + 1.96 * se_cr[, k] > th)) } colnames(est) <- paste(combos$estimand, combos$method, "est", sep = ".") colnames(se_db) <- paste(combos$estimand, combos$method, "seDB", sep = ".") colnames(se_cr) <- paste(combos$estimand, combos$method, "seCR", sep = ".") estimates <- data.frame(rep = seq_len(n_rep), truth, est, se_db, se_cr, check.names = FALSE) list(estimates = estimates, summary = summ) } ############################################################################### ## End of swcrtcov.R ###############################################################################