tidy implementation; tests, printing; etc.

Author: CoryMcCartan

R/ei_ridge.R

@@ -91,8 +91,9 @@
 #'   if it is contained within \eqn{[0, 1]}, for instance, then the bounds will
 #'   be `c(0, 1)`. The default `bounds = FALSE` uses an unbounded outcome.
 #' @param sum_one If `TRUE`, the outcome variables are constrained to sum to one.
-#'   Can only apply when `bounds` are enforced and there are more than one
-#'   outcome variables.
+#'   Can only apply when `bounds` are enforced and there is more than one
+#'   outcome variable. The default `NULL` infers `sum_one = TRUE` when the bounds
+#'   are `c(0, 1)` the outcome variables sum to 1.
 #' @param scale If `TRUE`, scale covariates `z` to have unit variance.
 #' @param vcov If `TRUE`, calculate and return the covariance matrix of the
 #'    estimated coefficients. Ignored when `bounds` are provided.
@@ -120,14 +121,14 @@
 #' min(fitted(ei_ridge(spec)))
 #' min(fitted(ei_ridge(spec, bounds = 0:1)))
 #' @export
-ei_ridge <- function(x, ..., weights, bounds = FALSE, sum_one = FALSE, penalty = NULL, scale = TRUE, vcov = TRUE) {
+ei_ridge <- function(x, ..., weights, bounds = FALSE, sum_one = NULL, penalty = NULL, scale = TRUE, vcov = TRUE) {
     UseMethod("ei_ridge")
 }
 
 
 #' @export
 #' @rdname ei_ridge
-ei_ridge.formula <- function(formula, data, weights, bounds=FALSE, sum_one = FALSE,
+ei_ridge.formula <- function(formula, data, weights, bounds=FALSE, sum_one = NULL,
                              penalty=NULL, scale=TRUE, vcov=TRUE, ...) {
     forms = ei_forms(formula)
     form_preds = terms(rlang::new_formula(lhs=NULL, rhs=forms$predictors))
@@ -154,7 +155,7 @@ ei_ridge.formula <- function(formula, data, weights, bounds=FALSE, sum_one = FAL
 
 #' @export
 #' @rdname ei_ridge
-ei_ridge.ei_spec <- function(x, weights, bounds=FALSE, sum_one = FALSE, penalty=NULL,
+ei_ridge.ei_spec <- function(x, weights, bounds=FALSE, sum_one = NULL, penalty=NULL,
                              scale=TRUE, vcov=TRUE, ...) {
     spec = x
     validate_ei_spec(spec)
@@ -184,7 +185,7 @@ ei_ridge.ei_spec <- function(x, weights, bounds=FALSE, sum_one = FALSE, penalty=
 
 #' @export
 #' @rdname ei_ridge
-ei_ridge.data.frame <- function(x, y, z, weights, bounds=FALSE, sum_one = FALSE, penalty=NULL,
+ei_ridge.data.frame <- function(x, y, z, weights, bounds=FALSE, sum_one = NULL, penalty=NULL,
                                 scale=TRUE, vcov=TRUE, ...) {
     if (length(both <- intersect(colnames(x), colnames(z))) > 0) {
         cli_abort(c("Predictors and covariates must be distinct",
@@ -213,7 +214,7 @@ ei_ridge.data.frame <- function(x, y, z, weights, bounds=FALSE, sum_one = FALSE,
 
 #' @export
 #' @rdname ei_ridge
-ei_ridge.matrix <- function(x, y, z, weights, bounds=FALSE, sum_one = FALSE, penalty=NULL,
+ei_ridge.matrix <- function(x, y, z, weights, bounds=FALSE, sum_one = NULL, penalty=NULL,
                             scale=TRUE, vcov=TRUE, ...) {
     ei_ridge.data.frame(x, y, z, weights, penalty, sum_one, bounds, scale, vcov, ...)
 }
@@ -275,6 +276,9 @@ ei_ridge_bridge <- function(processed, vcov, ...) {
     if (ncol(z) == 0) {
         bp$penalty = 0
     }
+    if (is.null(bp$sum_one) && all(bp$bounds == c(0, 1))) {
+        bp$sum_one = isTRUE(all.equal(rowSums(y), rep(1, nrow(y))))
+    }
 
     fit <- ei_ridge_impl(x, y, z, weights, bp$bounds, bp$sum_one, bp$penalty, vcov)
 
@@ -315,7 +319,7 @@ ei_ridge_bridge <- function(processed, vcov, ...) {
 #' @rdname ei-impl
 #' @export
 ei_ridge_impl <- function(x, y, z, weights=rep(1, nrow(x)),
-                          bounds=c(-Inf, Inf), sum_one=FALSE, penalty=NULL, vcov=TRUE) {
+                          bounds=c(-Inf, Inf), sum_one=NULL, penalty=NULL, vcov=TRUE) {
     int_scale = if (!is.null(penalty) && penalty == 0) 1 + 1e2*sqrt(penalty) else 1e4
     xz = row_kronecker(x, z, int_scale)
     sqrt_w = sqrt(weights / mean(weights))
@@ -336,6 +340,7 @@ ei_ridge_impl <- function(x, y, z, weights=rep(1, nrow(x)),
         if (is.null(penalty)) {
             penalty = ridge_auto(udv, y, sqrt_w, FALSE)$penalty
         }
+
         ridge_bounds(xz, z, y, weights, bounds, sum_one, penalty)
     }
 
@@ -415,7 +420,9 @@ print.ei_ridge <- function(x, ...) {
              nrow(x$fitted), " observations")
     bounds = x$blueprint$bounds
     if (any(is.finite(bounds))) {
-        cat_line("With outcome bounded in (", bounds[1], ", ", bounds[2], ")")
+        sumt1 = if (isTRUE(x$blueprint$sum_one)) " and constrained to sum to 1" else ""
+        pl = if (ncol(m$y) > 1) "s" else ""
+        cat_line("With outcome", pl, " bounded in (", bounds[1], ", ", bounds[2], ")", sumt1)
     }
     cat_line("Fit with penalty = ", signif(x$penalty))
 }

R/rr_impl.R

@@ -122,10 +122,19 @@ ridge_bounds <- function(xz, z, y, weights, bounds, sum_one=FALSE, penalty=0) {
         cli_abort("{.fn ridge_bounds} requires at least one finite bound.")
     }
 
+    fit_err = \(e) {
+        cli_abort(c(
+            "Constrained ridge regression failed with inconsistent constraints.",
+            ">" = "Try setting {.arg sum_one=FALSE} or relaxing the bounds."
+        ), call = NULL)
+    }
     if (isFALSE(sum_one)) {
         coefs = matrix(nrow = nrow(dvecs), ncol = ncol(dvecs))
         for (i in seq_len(n_y)) {
-            fit = quadprog::solve.QP.compact(R, dvecs[, i], Amat, Aind, bvec, factorized = TRUE)
+            fit = tryCatch(
+                quadprog::solve.QP.compact(R, dvecs[, i], Amat, Aind, bvec, factorized = TRUE),
+                error = fit_err
+            )
             coefs[, i] = fit$solution
         }
     } else {
@@ -149,14 +158,17 @@ ridge_bounds <- function(xz, z, y, weights, bounds, sum_one=FALSE, penalty=0) {
         }
         bvec_y = c(rep(1, n * n_x), rep(1, n_y) %x% bvec)
 
-        fit = quadprog::solve.QP.compact(
-            R_y,
-            c(dvecs),
-            Amat_y,
-            Aind_y,
-            bvec_y,
-            meq = n * n_x,
-            factorized = TRUE
+        do_fit = function(eq) {
+            quadprog::solve.QP.compact(R_y, c(dvecs), Amat_y, Aind_y, bvec_y, meq = eq, factorized = TRUE)
+        }
+
+        # relax to inequality constraint if sum-to-one fails
+        fit <- tryCatch(
+            do_fit(n * n_x),
+            error = \(e_outer) {
+                cli_warn("Relaxing sum-to-one constraint to inequality to achieve feasible solution.", call=NULL)
+                tryCatch(do_fit(0), error = fit_err)
+            }
         )
         coefs = matrix(fit$solution, nrow = nrow(dvecs), ncol = ncol(dvecs))
     }

man/ei-impl.Rd

@@ -92,17 +92,27 @@ test_that("leave-one-out shortcut is correct for Riesz regression", {
     expect_equal(fit_naive$loo, loo_act, tolerance = 0.2)
 })
 
-test_that("ridge bounds work", {
+test_that("ridge constraints work", {
     d = elec_1968
-    form = pres_rep_nix ~ vap_white + vap_black + vap_other | state + pop_city +
+    form = pres_dem_hum + pres_rep_nix + pres_ind_wal + pres_abs ~ vap_white  |
         pop_urban + pop_rural + farm + educ_elem + educ_hsch + educ_coll +
         inc_00_03k + inc_03_08k + inc_08_25k + inc_25_99k + log(pop) + pres_turn
 
     m = ei_ridge(form, data=elec_1968)
-    m01 = ei_ridge(form, data=elec_1968, bounds=0:1)
+    m01 = ei_ridge(form, data=elec_1968, bounds=0:1, sum_one=FALSE)
+    m01s = ei_ridge(form, data=elec_1968, bounds=c(0, 1), sum_one=TRUE)
+    m01def = ei_ridge(form, data=elec_1968, bounds=NULL, sum_one=NULL)
 
     expect_true(min(fitted(m)) < 0)
-    expect_true(min(fitted(m01)) > 0)
+    expect_true(min(fitted(m01)) > -.Machine$double.eps)
     expect_true(all(ei_est(m01, data=elec_1968, total=pres_total)$estimate > 0))
     expect_true(all(ei_est(m01, data=elec_1968, total=pres_total)$estimate < 1))
+
+    expect_true(min(fitted(m01s)) > -.Machine$double.eps)
+    expect_true(all(ei_est(m01s, data=elec_1968, total=pres_total)$estimate > 0))
+    expect_true(all(ei_est(m01s, data=elec_1968, total=pres_total)$estimate < 1))
+
+    tots = rowSums(as.matrix(ei_est(m01s, data=elec_1968, total=pres_total)))
+    expect_true(all.equal(tots, c(vap_white=1, .other=1)))
+    expect_identical(m01def, m01s) # check defaults infer correctly
 })