Add a top-level generic "solver" option with a Frank-Wolfe solver for the SCM weights

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) %>%

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)) {

R/fit_synth.R

@@ -30,21 +30,23 @@ 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]
     ## if no  is supplied, set equal to 1
 
     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))
 }

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

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,22 +226,23 @@ 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))
     }
     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
 

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))
     }

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)
+}