diff --git a/R/augsynth.R b/R/augsynth.R index 1f4671c..600d38b 100644 --- a/R/augsynth.R +++ b/R/augsynth.R @@ -20,6 +20,17 @@ #' @param scm Whether the SCM weighting function is used. If FALSE, then package will fit the outcome model, but not calculate new donor weights to match pre-treatment covariates. Instead, each donor unit will be equally weighted. If TRUE, weights on donor pool will be calculated. #' @param fixedeff Whether to include a unit fixed effect, default F #' @param cov_agg Covariate aggregation functions, if NULL then use mean with NAs omitted +#' @param solver Solver for the synthetic control weights: "osqp" (the +#' default, solves the quadratic program exactly after forming the +#' n0 x n0 donor Gram matrix) or "frank_wolfe" (Frank-Wolfe with +#' exact line search — the synthdid algorithm — to identify the +#' active donor set, then an exact QP restricted to that support +#' with a KKT re-admission screen; the full n0 x n0 Gram matrix is +#' never formed, so memory stays O(n0 t0) — useful for very large +#' donor pools). May also be a function (X1, X0, V) -> weights to +#' plug in a custom solver. Applies wherever SCM weights are solved, +#' including the ridge lambda cross-validation refits, and is +#' remembered by inference/permutation refits. #' @return augsynth object that contains: #' \itemize{ #' \item{"weights"}{Ridge ASCM weights} @@ -35,6 +46,7 @@ single_augsynth <- function(form, unit, time, t_int, data, scm=T, fixedeff = FALSE, cov_agg=NULL, + solver = "osqp", ...) { call_name <- match.call() @@ -63,9 +75,10 @@ single_augsynth <- function(form, unit, time, t_int, data, Z <- NULL } - # fit augmented SCM + # fit augmented SCM; `solver` travels through `...` so it lands in + # extra_args and is replayed by inference/permutation/cv refits augsynth <- fit_augsynth_internal(wide, synth_data, Z, progfunc, - scm, fixedeff, ...) + scm, fixedeff, solver = solver, ...) # add some extra data augsynth$data$time <- data %>% distinct(!!time) %>% diff --git a/R/augsynth_pre.R b/R/augsynth_pre.R index 8593d2a..67ffeea 100644 --- a/R/augsynth_pre.R +++ b/R/augsynth_pre.R @@ -25,6 +25,7 @@ #' \item{"scm"}{Whether the SCM weighting function is used} #' \item{"fixedeff"}{Whether to include a unit fixed effect, default is FALSE } #' \item{"cov_agg"}{Covariate aggregation functions, if NULL then use mean with NAs omitted} +#' \item{"solver"}{Solver for the SCM weights: "osqp" (default, exact QP on the full n0 x n0 donor Gram matrix) or "frank_wolfe" (Frank-Wolfe support identification plus an exact QP on that support; never forms the full Gram matrix, useful for very large donor pools), or a custom function (X1, X0, V) -> weights} #' } #' \item Multi period (staggered) augsynth #' \itemize{ @@ -79,6 +80,9 @@ augsynth <- function(form, unit, time, data, t_int=NULL, ...) { if("progfunc" %in% names(list(...))) { warning("`progfunc` is not an argument for multisynth, so it is ignored") } + if("solver" %in% names(list(...))) { + warning("`solver` is not an argument for multisynth, so it is ignored") + } return(multisynth(form, !!enquo(unit), !!enquo(time), data, ...)) } else { if (is.null(t_int)) { diff --git a/R/fit_synth.R b/R/fit_synth.R index 17c006c..5cec4cd 100644 --- a/R/fit_synth.R +++ b/R/fit_synth.R @@ -30,13 +30,15 @@ make_V_matrix <- function(t0, V) { #' Fit synthetic controls on outcomes after formatting data #' @param synth_data Panel data in format of Synth::dataprep #' @param V Matrix to scale the obejctive by +#' @param solver Solver for the synth QP: "osqp" (default), "frank_wolfe", +#' or a function (X1, X0, V) -> weights; see R/solvers.R #' @noRd #' @return \itemize{ #' \item{"weights"}{Synth weights} #' \item{"l2_imbalance"}{Imbalance in pre-period outcomes, measured by the L2 norm} #' \item{"scaled_l2_imbalance"}{L2 imbalance scaled by L2 imbalance of uniform weights} #' } -fit_synth_formatted <- function(synth_data, V = NULL) { +fit_synth_formatted <- function(synth_data, V = NULL, solver = "osqp") { t0 <- dim(synth_data$Z0)[1] @@ -44,7 +46,7 @@ fit_synth_formatted <- function(synth_data, V = NULL) { V <- make_V_matrix(t0, V) - weights <- synth_qp(synth_data$X1, t(synth_data$X0), V) + weights <- synth_qp(synth_data$X1, t(synth_data$X0), V, solver = solver) l2_imbalance <- sqrt(sum((synth_data$Z0 %*% weights - synth_data$Z1)^2)) ## primal objective value scaled by least squares difference for mean @@ -57,27 +59,14 @@ fit_synth_formatted <- function(synth_data, V = NULL) { scaled_l2_imbalance=scaled_l2_imbalance)) } -#' Solve the synth QP directly +#' Solve the synth QP with the requested solver #' @param X1 Target vector #' @param X0 Matrix of control outcomes #' @param V Scaling matrix +#' @param solver Solver name or function; resolved by get_synth_solver() #' @noRd -synth_qp <- function(X1, X0, V) { +synth_qp <- function(X1, X0, V, solver = "osqp") { - Pmat <- X0 %*% V %*% t(X0) - qvec <- - t(X1) %*% V %*% t(X0) - - n0 <- nrow(X0) - A <- rbind(rep(1, n0), diag(n0)) - l <- c(1, numeric(n0)) - u <- c(1, rep(1, n0)) - - settings = osqp::osqpSettings(verbose = FALSE, - eps_rel = 1e-8, - eps_abs = 1e-8) - sol <- osqp::solve_osqp(P = Pmat, q = qvec, - A = A, l = l, u = u, - pars = settings) - - return(sol$x) + solve_fn <- get_synth_solver(solver) + return(solve_fn(X1, X0, V)) } diff --git a/R/highdim.R b/R/highdim.R index 1d663d9..5ec464f 100644 --- a/R/highdim.R +++ b/R/highdim.R @@ -82,7 +82,7 @@ fit_augsyn_formatted <- function(wide_data, synth_data, #' } fit_augsyn <- function(wide_data, synth_data, progfunc=c("EN", "RF", "GSYN", "MCP","CITS", "CausalImpact", "seq2seq"), - scm=T, ...) { + scm=T, solver = "osqp", ...) { ## prognostic score and weight functions to use progfunc = tolower(progfunc) if(progfunc == "en") { @@ -104,7 +104,7 @@ fit_augsyn <- function(wide_data, synth_data, } if(scm) { - weightf <- fit_synth_formatted + weightf <- function(sd) fit_synth_formatted(sd, solver = solver) } else { ## still fit synth even if none ## TODO: This is a dumb wasteful hack diff --git a/R/ridge.R b/R/ridge.R index a797174..0d126c2 100644 --- a/R/ridge.R +++ b/R/ridge.R @@ -37,7 +37,8 @@ fit_ridgeaug_formatted <- function(wide_data, synth_data, lambda_max = NULL, holdout_length = 1, min_1se = T, V = NULL, - residualize = FALSE, ...) { + residualize = FALSE, solver = "osqp", + ...) { extra_params = list(...) if (length(extra_params) > 0) { warning("Unused parameters in using ridge augmented weights: ", paste(names(extra_params), collapse = ", ")) @@ -122,7 +123,7 @@ fit_ridgeaug_formatted <- function(wide_data, synth_data, lambda, ridge, scm, lambda_min_ratio, n_lambda, lambda_max, - holdout_length, min_1se) + holdout_length, min_1se, solver = solver) weights <- out$weights synw <- out$synw @@ -225,14 +226,14 @@ fit_ridgeaug_inner <- function(X_c, X_1, trt, synth_data, lambda, ridge, scm, lambda_min_ratio, n_lambda, lambda_max, - holdout_length, min_1se) { + holdout_length, min_1se, solver = "osqp") { lambda_errors <- NULL lambda_errors_se <- NULL lambdas <- NULL ## if SCM fit scm if(scm) { - syn <- fit_synth_formatted(synth_data)$weights + syn <- fit_synth_formatted(synth_data, solver = solver)$weights } else { ## else use uniform weights syn <- rep(1 / sum(trt == 0), sum(trt == 0)) @@ -240,7 +241,8 @@ fit_ridgeaug_inner <- function(X_c, X_1, trt, synth_data, if(ridge) { if(is.null(lambda)) { cv_out <- cv_lambda(X_c, X_1, synth_data, trt, holdout_length, scm, - lambda_max, lambda_min_ratio, n_lambda, min_1se) + lambda_max, lambda_min_ratio, n_lambda, min_1se, + solver = solver) lambda <- cv_out$lambda lambda_errors <- cv_out$lambda_errors @@ -322,7 +324,8 @@ choose_lambda <- function(lambdas, lambda_errors, lambda_errors_se, min_1se) { #' \item{"lambda_errors_se"}{"The SE of the MSE associated with each lambda term} #' } cv_lambda <- function(X_c, X_1, synth_data, trt, holdout_length, scm, - lambda_max, lambda_min_ratio, n_lambda, min_1se) { + lambda_max, lambda_min_ratio, n_lambda, min_1se, + solver = "osqp") { if(is.null(lambda_max)) { lambda_max <- get_lambda_max(X_c) } @@ -331,7 +334,8 @@ cv_lambda <- function(X_c, X_1, synth_data, trt, holdout_length, scm, lambda_out <- get_lambda_errors(lambdas, X_c, X_1, synth_data, trt, - holdout_length, scm) + holdout_length, scm, + solver = solver) lambda_errors <- lambda_out$lambda_errors lambda_errors_se <- lambda_out$lambda_errors_se diff --git a/R/ridge_lambda.R b/R/ridge_lambda.R index c17b9a4..6cbc1e9 100644 --- a/R/ridge_lambda.R +++ b/R/ridge_lambda.R @@ -14,7 +14,7 @@ #' @param scm Include SCM or not #' @noRd #' @return List of lambda errors for each corresponding lambda in the lambdas parameter. -get_lambda_errors <- function(lambdas, X_c, X_t, synth_data, trt, holdout_length=1, scm=T) { +get_lambda_errors <- function(lambdas, X_c, X_t, synth_data, trt, holdout_length=1, scm=T, solver = "osqp") { # vector that stores the sum MSE across all CV sets errors <- matrix(0, nrow = ncol(X_c) - holdout_length, ncol = length(lambdas)) lambda_errors = numeric(length(lambdas)) @@ -32,7 +32,7 @@ get_lambda_errors <- function(lambdas, X_c, X_t, synth_data, trt, holdout_length new_synth_data$X0 <- t(X_0) if(scm) { - syn <- fit_synth_formatted(new_synth_data)$weights + syn <- fit_synth_formatted(new_synth_data, solver = solver)$weights } else { syn <- rep(1/sum(trt==0), sum(trt==0)) } diff --git a/R/solvers.R b/R/solvers.R new file mode 100644 index 0000000..ff394c0 --- /dev/null +++ b/R/solvers.R @@ -0,0 +1,146 @@ +################################################################################ +## Solvers for the synthetic control QP +## min_w (X0' w - X1)' V (X0' w - X1) s.t. w >= 0, sum(w) = 1 +## Each solver is a function with signature (X1, X0, V) -> numeric weights, +## where X1 is the t0 x 1 target, X0 is the n0 x t0 matrix of control +## outcomes, and V is a t0 x t0 scaling matrix. To add a new solver, add it +## to the switch in get_synth_solver() or pass the function itself as the +## `solver` argument. +################################################################################ + +#' Resolve a synth solver from a name or a function +#' @param solver "osqp" (default), "frank_wolfe", or a function with +#' signature (X1, X0, V) returning a weight vector +#' @noRd +get_synth_solver <- function(solver = "osqp") { + if (is.function(solver)) { + return(solver) + } + switch(match.arg(solver, c("osqp", "frank_wolfe")), + osqp = synth_qp_osqp, + frank_wolfe = synth_qp_frank_wolfe) +} + +#' Solve the synth QP with OSQP (forms the n0 x n0 Gram matrix) +#' @param X1 Target vector +#' @param X0 Matrix of control outcomes (n0 x t0) +#' @param V Scaling matrix +#' @noRd +synth_qp_osqp <- function(X1, X0, V) { + + Pmat <- X0 %*% V %*% t(X0) + qvec <- - t(X1) %*% V %*% t(X0) + + n0 <- nrow(X0) + A <- rbind(rep(1, n0), diag(n0)) + l <- c(1, numeric(n0)) + u <- c(1, rep(1, n0)) + + settings = osqp::osqpSettings(verbose = FALSE, + eps_rel = 1e-8, + eps_abs = 1e-8) + sol <- osqp::solve_osqp(P = Pmat, q = qvec, + A = A, l = l, u = u, + pars = settings) + + return(sol$x) +} + +#' Solve the synth QP via Frank-Wolfe support identification +#' +#' Runs Frank-Wolfe with exact line search (the algorithm of the synthdid +#' package: Arkhangelsky, Athey, Hirshberg, Imbens & Wager 2021, dual +#' BSD-3/GPL >= 2) to identify the sparse active donor set, then solves the +#' QP exactly on that support (OSQP on a support-sized problem), and +#' verifies KKT optimality with a full gradient screen, re-admitting any +#' donors the support restriction missed. The n0 x n0 donor Gram matrix of +#' the full QP is never formed: Frank-Wolfe costs O(n0 t0) per iteration +#' and the polished QP is at most (5 t0) x (5 t0), so memory stays +#' O(n0 t0) -- the constraint that binds with very large donor pools. +#' +#' @param X1 Target vector +#' @param X0 Matrix of control outcomes (n0 x t0) +#' @param V Scaling matrix +#' @param max_iter Frank-Wolfe iteration cap +#' @param support_size Donors kept for the QP polish (top Frank-Wolfe +#' weights); default min(n0, 5 t0) +#' @noRd +synth_qp_frank_wolfe <- function(X1, X0, V, max_iter = 2000L, + support_size = NULL) { + + ## fold V into the design: A = V^{1/2} X0' (t0 x n0), b = V^{1/2} X1 + if (all(V == diag(diag(V), nrow(V)))) { + s <- sqrt(diag(V)) + A <- t(X0) * s + b <- as.vector(X1) * s + } else { + E <- eigen(V, symmetric = TRUE) + R <- diag(sqrt(pmax(E$values, 0)), nrow(V)) %*% t(E$vectors) + A <- R %*% t(X0) + b <- as.vector(R %*% X1) + } + n0 <- ncol(A) + t0 <- nrow(A) + + ## Frank-Wolfe with exact line search from uniform weights + w <- rep(1 / n0, n0) + Aw <- as.vector(A %*% w) + val_old <- Inf + for (it in seq_len(max_iter)) { + grad <- crossprod(A, Aw - b) # half-gradient, O(n0 t0) + i <- which.min(grad) + dA <- A[, i] - Aw + denom <- sum(dA^2) + if (denom <= 0) break + step <- max(0, min(1, -sum((Aw - b) * dA) / denom)) + if (step <= 1e-12) break + w <- (1 - step) * w + w[i] <- w[i] + step + Aw <- (1 - step) * Aw + step * A[, i] + val <- sum((Aw - b)^2) + if (val_old - val < 1e-9 * val) break # relative-decrease stop + val_old <- val + } + + ## exact QP restricted to a donor support (OSQP on a small problem) + qp_support <- function(As, bs) { + ns <- ncol(As) + sol <- osqp::solve_osqp( + P = crossprod(As), q = -crossprod(As, bs), + A = rbind(rep(1, ns), diag(ns)), + l = c(1, numeric(ns)), u = c(1, rep(1, ns)), + pars = osqp::osqpSettings(verbose = FALSE, + eps_rel = 1e-8, eps_abs = 1e-8)) + ws <- pmax(sol$x, 0) + ws / sum(ws) + } + + k <- if (is.null(support_size)) min(n0, 5L * t0) + else min(n0, support_size) + ## FW starts uniform and never zeroes untouched donors: its support is + ## the set pushed above the shrunken-uniform baseline + support <- which(w > min(w) * (1 + 1e-9)) + if (!length(support)) support <- seq_len(min(n0, t0)) + if (length(support) > k) + support <- support[order(w[support], decreasing = TRUE)[seq_len(k)]] + + ## polish, then KKT gradient screening: at the optimum every donor + ## outside the active set has (half-)gradient >= the common active-set + ## level mu; one O(n0 t0) pass finds donors FW missed, add them and + ## re-polish until no violations (typically 1-2 rounds) + for (round in seq_len(10L)) { + w <- numeric(n0) + w[support] <- qp_support(A[, support, drop = FALSE], b) + g <- as.vector(crossprod(A, A %*% w - b)) + active <- support[w[support] > 1e-12] + mu <- stats::median(g[active]) + tol_g <- 1e-8 * max(abs(g)) + violated <- setdiff(which(g < mu - tol_g), support) + if (!length(violated)) break + violated <- violated[order(g[violated])[seq_len(min(length(violated), + k))]] + support <- c(support, violated) + } + + return(w) +} diff --git a/man/augsynth.Rd b/man/augsynth.Rd index 722fa2a..42985ed 100644 --- a/man/augsynth.Rd +++ b/man/augsynth.Rd @@ -26,6 +26,7 @@ only)} \item{"scm"}{Whether the SCM weighting function is used} \item{"fixedeff"}{Whether to include a unit fixed effect, default is FALSE } \item{"cov_agg"}{Covariate aggregation functions, if NULL then use mean with NAs omitted} + \item{"solver"}{Solver for the SCM weights: "osqp" (default, exact QP on the full n0 x n0 donor Gram matrix) or "frank_wolfe" (Frank-Wolfe support identification plus an exact QP on that support; never forms the full Gram matrix, useful for very large donor pools), or a custom function (X1, X0, V) -> weights} } \item Multi period (staggered) augsynth \itemize{ diff --git a/man/single_augsynth.Rd b/man/single_augsynth.Rd index aec95fc..18e3806 100644 --- a/man/single_augsynth.Rd +++ b/man/single_augsynth.Rd @@ -14,6 +14,7 @@ single_augsynth( scm = T, fixedeff = FALSE, cov_agg = NULL, + solver = "osqp", ... ) } @@ -42,6 +43,18 @@ causalimpact=Bayesian structural time series with CausalImpact} \item{cov_agg}{Covariate aggregation functions, if NULL then use mean with NAs omitted} +\item{solver}{Solver for the synthetic control weights: "osqp" (the +default, solves the quadratic program exactly after forming the +n0 x n0 donor Gram matrix) or "frank_wolfe" (Frank-Wolfe with +exact line search — the synthdid algorithm — to identify the +active donor set, then an exact QP restricted to that support +with a KKT re-admission screen; the full n0 x n0 Gram matrix is +never formed, so memory stays O(n0 t0) — useful for very large +donor pools). May also be a function (X1, X0, V) -> weights to +plug in a custom solver. Applies wherever SCM weights are solved, +including the ridge lambda cross-validation refits, and is +remembered by inference/permutation refits.} + \item{...}{optional arguments for outcome model} } \value{ diff --git a/tests/testthat/test_solvers.R b/tests/testthat/test_solvers.R new file mode 100644 index 0000000..7c02bf9 --- /dev/null +++ b/tests/testthat/test_solvers.R @@ -0,0 +1,122 @@ +context("Synth QP solvers agree (osqp vs frank_wolfe)") + +library(tidyverse) + +data(basque, package = "Synth") +basque <- basque %>% mutate(trt = case_when(year < 1975 ~ 0, + regionno != 17 ~ 0, + regionno == 17 ~ 1)) %>% + filter(regionno != 1) + + +test_that("solvers agree on the QP objective directly", { + + set.seed(42) + n0 <- 40 + t0 <- 15 + X0 <- matrix(rnorm(n0 * t0), n0, t0) + X1 <- matrix(rnorm(t0), t0, 1) + V <- diag(t0) + + w_qp <- augsynth:::synth_qp_osqp(X1, X0, V) + w_fw <- augsynth:::synth_qp_frank_wolfe(X1, X0, V) + + objective <- function(w) sum((t(X0) %*% w - X1)^2) + + ## frank_wolfe respects the simplex and reaches the osqp objective + expect_equal(sum(w_fw), 1, tolerance = 1e-10) + expect_true(all(w_fw >= -1e-10)) + expect_lt(objective(w_fw), objective(w_qp) * (1 + 1e-4) + 1e-10) +}) + + +test_that("frank_wolfe matches osqp without ridge augmentation", { + + syn_qp <- single_augsynth(gdpcap ~ trt, regionno, year, basque, + progfunc = "None", scm = TRUE, t_int = 1975) + syn_fw <- single_augsynth(gdpcap ~ trt, regionno, year, basque, + progfunc = "None", scm = TRUE, t_int = 1975, + solver = "frank_wolfe") + + expect_equal(sum(syn_fw$weights), 1, tolerance = 1e-8) + expect_true(all(syn_fw$weights >= -1e-8)) + + ## same imbalance (the objective) and the same ATT path + expect_equal(syn_fw$l2_imbalance, syn_qp$l2_imbalance, tolerance = 1e-4) + expect_lt(max(abs(predict(syn_fw, att = TRUE) - + predict(syn_qp, att = TRUE))), 1e-2) +}) + + +test_that("frank_wolfe matches osqp with ridge augmentation (fixed lambda)", { + + syn_qp <- single_augsynth(gdpcap ~ trt, regionno, year, basque, + progfunc = "Ridge", scm = TRUE, t_int = 1975, + lambda = 8) + syn_fw <- single_augsynth(gdpcap ~ trt, regionno, year, basque, + progfunc = "Ridge", scm = TRUE, t_int = 1975, + lambda = 8, solver = "frank_wolfe") + + expect_equal(syn_fw$lambda, syn_qp$lambda) + expect_lt(max(abs(syn_fw$weights - syn_qp$weights)), 1e-2) + expect_lt(max(abs(predict(syn_fw, att = TRUE) - + predict(syn_qp, att = TRUE))), 1e-2) +}) + + +test_that("frank_wolfe matches osqp with ridge augmentation (CV lambda)", { + + syn_qp <- single_augsynth(gdpcap ~ trt, regionno, year, basque, + progfunc = "Ridge", scm = TRUE, t_int = 1975) + syn_fw <- single_augsynth(gdpcap ~ trt, regionno, year, basque, + progfunc = "Ridge", scm = TRUE, t_int = 1975, + solver = "frank_wolfe") + + ## the CV refits run through the requested solver; the selected lambdas + ## (and so the estimates) should agree on this problem + expect_equal(syn_fw$lambda, syn_qp$lambda, tolerance = 1e-8) + expect_lt(max(abs(predict(syn_fw, att = TRUE) - + predict(syn_qp, att = TRUE))), 1e-2) +}) + + +test_that("frank_wolfe matches osqp with a fixed effect", { + + syn_qp <- single_augsynth(gdpcap ~ trt, regionno, year, basque, + progfunc = "Ridge", scm = TRUE, t_int = 1975, + fixedeff = TRUE, lambda = 8) + syn_fw <- single_augsynth(gdpcap ~ trt, regionno, year, basque, + progfunc = "Ridge", scm = TRUE, t_int = 1975, + fixedeff = TRUE, lambda = 8, + solver = "frank_wolfe") + + expect_lt(max(abs(predict(syn_fw, att = TRUE) - + predict(syn_qp, att = TRUE))), 1e-2) +}) + + +test_that("a custom solver function can be plugged in", { + + ## passing the osqp solver as a bare function must reproduce the default + syn_default <- single_augsynth(gdpcap ~ trt, regionno, year, basque, + progfunc = "None", scm = TRUE, t_int = 1975) + syn_custom <- single_augsynth(gdpcap ~ trt, regionno, year, basque, + progfunc = "None", scm = TRUE, t_int = 1975, + solver = function(X1, X0, V) { + augsynth:::synth_qp_osqp(X1, X0, V) + }) + + expect_equal(syn_custom$weights, syn_default$weights, tolerance = 1e-10) +}) + + +test_that("the solver is remembered by inference refits", { + + syn_fw <- single_augsynth(gdpcap ~ trt, regionno, year, basque, + progfunc = "Ridge", scm = TRUE, t_int = 1975, + lambda = 8, solver = "frank_wolfe") + expect_equal(syn_fw$extra_args$solver, "frank_wolfe") + + ## summary (jackknife+ / conformal machinery) replays extra_args + expect_error(summary(syn_fw, inf_type = "jackknife+"), NA) +})