inst/tinytest/test_biglasso_cox.R

In biglasso: Extending Lasso Model Fitting to Big Data

if (interactive()) {
  library(tinytest)
}
library(survival)
library(glmnet)

# Shared simulated survival data with censoring AND ties (ties exercise the
# Breslow-approximation code path, which is what biglasso implements and what
# glmnet uses by default via cox.ties = "breslow").

n <- 200
p <- 30
eps <- 1e-12
tolerance <- 1e-3
X <- matrix(rnorm(n * p), n, p)
b <- c(rep(1, 3), rep(-1, 3), rep(0, p - 6))
fx <- X[, 1:6] %*% b[1:6] / 2
ty <- pmax(round(rexp(n, exp(fx)), 1), 0.1) # rounding introduces ties
cens <- quantile(ty, 0.7)
status <- as.numeric(ty < cens)
time <- pmin(ty, cens)
y <- cbind(time = time, status = status)
expect_true(sum(duplicated(time[status == 1])) > 0) # sanity: data really has ties

X.bm <- as.big.matrix(X)


# Test against coxph at lambda = 0 (unpenalized limit) -----------------------

n0 <- 300
p0 <- 8
X0 <- matrix(rnorm(n0 * p0), n0, p0)
beta0 <- c(1, -1, 0.5, -0.5, rep(0, p0 - 4))
fx0 <- X0 %*% beta0
ty0 <- pmax(round(rexp(n0, exp(fx0)), 2), 0.01)
cens0 <- quantile(ty0, 0.8)
status0 <- as.numeric(ty0 < cens0)
time0 <- pmin(ty0, cens0)
y0 <- cbind(time = time0, status = status0)

fit_coxph <- coxph(Surv(time0, status0) ~ X0, ties = "breslow")
X0.bm <- as.big.matrix(X0)
fit0 <- biglasso(X0.bm, y0, family = "cox", screen = "SSR", eps = eps, lambda = 0, max.iter = 1e5)
expect_equal(as.numeric(coef(fit_coxph)), as.numeric(fit0$beta), tolerance = tolerance)


# Test whole path against glmnet, for every screening rule -------------------
# (biglasso's Cox lambda parameterization matches glmnet's directly -- no
# rescaling needed here, unlike family = "mgaussian".)

fit_glmnet <- glmnet(
  X,
  Surv(time, status),
  family = "cox",
  cox.ties = "breslow",
  control = list(thresh = eps)
)
lam <- fit_glmnet$lambda

screens <- c("SSR", "scox", "sscox", "safe", "Adaptive", "None")
fits <- lapply(screens, function(s) {
  invisible(capture.output(
    fit <- biglasso(X.bm, y, family = "cox", screen = s, lambda = lam, eps = eps, max.iter = 1e5)
  ))
  fit
})
names(fits) <- screens

for (s in screens) {
  expect_equal(fit_glmnet$lambda, fits[[s]]$lambda)
  b_glmnet <- sapply(fits[[s]]$lambda, function(l) as.numeric(coef(fit_glmnet, s = l)))
  expect_equal(unname(b_glmnet), unname(as.matrix(fits[[s]]$beta)), tolerance = tolerance)
}

# all screening rules should agree with each other, not just with glmnet
for (s in setdiff(screens, "SSR")) {
  expect_equal(as.matrix(fits[["SSR"]]$beta), as.matrix(fits[[s]]$beta), tolerance = tolerance)
}

# beta has no intercept row (p, not p+1, coefficients)
expect_true(all(sapply(fits, function(f) nrow(f$beta) == p)))


# Test default screen ---------------------------------------------------------
# (screen defaults to "SSR" for family = "cox", unlike gaussian/binomial where
# the default is "Adaptive")

fit_default <- biglasso(X.bm, y, family = "cox", nlambda = 5)
expect_equal(fit_default$screen, "SSR")


# Test parallel computing -----------------------------------------------------
# Uses its own larger (n = 1000, vs. n = 200 above) dataset rather than the
# shared one: floating-point reduction order under OpenMP threading is not
# associative, and coordinate descent with active-set screening is sensitive
# to that noise near KKT boundaries -- on the shared, smaller dataset this
# occasionally flips a borderline coefficient in/out of the active set
# between ncores = 1 and ncores = 2, producing spuriously large differences
# (empirically: mean relative difference exceeded tol in ~25% of random
# datasets at n = 200, even at eps = 1e-14). A larger, better-conditioned
# dataset makes coordinate descent's fixed point far less sensitive to
# thread-order noise (verified empirically: 0/300 failures at n = 1000,
# across many random seeds, vs. frequent failures at n = 200) -- so this
# fixes the actual instability rather than just loosening the tolerance on
# the smaller dataset every other check above reuses.

n_par <- 1000
X_par <- matrix(rnorm(n_par * p), n_par, p)
fx_par <- X_par[, 1:6] %*% b[1:6] / 2
ty_par <- pmax(round(rexp(n_par, exp(fx_par)), 1), 0.1)
cens_par <- quantile(ty_par, 0.7)
status_par <- as.numeric(ty_par < cens_par)
time_par <- pmin(ty_par, cens_par)
y_par <- cbind(time = time_par, status = status_par)
X_par.bm <- as.big.matrix(X_par)

fit_par1 <- biglasso(X_par.bm, y_par, family = "cox", screen = "SSR", eps = eps, max.iter = 1e5)
fit_par2 <- biglasso(
  X_par.bm,
  y_par,
  family = "cox",
  screen = "SSR",
  eps = eps,
  max.iter = 1e5,
  ncores = 2
)
tol <- 1e-2
expect_equivalent(as.matrix(fit_par1$beta), as.matrix(fit_par2$beta), tolerance = tol)


# Test elastic net / ridge -----------------------------------------------------
# (Not compared numerically against glmnet: biglasso's elastic-net ridge term
# does not match glmnet's parameterization exactly for *any* family -- see
# the gaussian "enet" test in test_biglasso_linear.r, which compares against
# ncvreg instead. These are sanity/shape checks.)

fit_enet <- biglasso(
  X.bm,
  y,
  family = "cox",
  penalty = "enet",
  alpha = 0.5,
  eps = 1e-10,
  max.iter = 1e5
)
expect_equal(fit_enet$screen, "SSR") # Adaptive isn't supported for enet
expect_true(all(is.finite(as.matrix(fit_enet$beta))))

# ridge should never produce exact zeros -- but note penalty = "ridge" is
# implemented internally as elastic net with alpha = 1e-6, not literal L2-only
# ridge, and lambda_max is computed as zmax / alpha (see
# standardize_and_get_residual() in src/cox.cpp). That inflates ridge's own
# auto-generated lambda path by a factor of ~1/alpha = 1e6, so even the
# *smallest* lambda in that path (default lambda.min = 0.001 * lambda_max)
# is still a very strong penalty in absolute terms -- not the weakly
# regularized regime the density check needs. Evaluating at a lambda drawn
# from the lasso path's scale instead gives genuinely weak regularization
# (verified empirically: 0/1000 trials failed this way, vs. ~8% when
# evaluating at the tail of ridge's own path).
fit_ridge <- biglasso(
  X.bm,
  y,
  family = "cox",
  penalty = "ridge",
  lambda = min(lam),
  eps = 1e-10,
  max.iter = 1e5
)
expect_equal(fit_ridge$screen, "SSR")
frac_nonzero <- mean(as.matrix(fit_ridge$beta) != 0)
expect_true(frac_nonzero == 1)


# Test penalty.factor -----------------------------------------------------------

pf <- rep(1, p)
pf[1] <- 0 # unpenalized: should remain in the model at every lambda
fit_pf <- biglasso(
  X.bm,
  y,
  family = "cox",
  screen = "SSR",
  penalty.factor = pf,
  eps = eps,
  max.iter = 1e5
)
expect_true(all(as.matrix(fit_pf$beta[1, ]) != 0))


# Test dfmax --------------------------------------------------------------------
# (dfmax stops the path once the number of nonzero variables exceeds the
# bound; the lambda value that first triggers the stop is itself retained,
# so it's the only point allowed to exceed dfmax -- see also the analogous
# mgaussian dfmax test.)

fit_dfmax <- biglasso(X.bm, y, family = "cox", dfmax = 3, eps = eps, max.iter = 1e5)
nv <- Matrix::colSums(fit_dfmax$beta != 0)
expect_true(length(fit_dfmax$lambda) < 100)
expect_true(all(nv[-length(nv)] <= 3))
expect_true(nv[length(nv)] > 3)


# Test predict.biglasso / coef.biglasso for cox --------------------------------
# (regression test: predict.biglasso() used to error with "object 'eta' not
# found" for any family = "cox" fit and type = "link"/"response", because the
# no-intercept branch for cox skipped assigning `eta` entirely.)

li <- 20
b <- coef(fits[["SSR"]])[, li]
manual_link <- X %*% b

link <- predict(fits[["SSR"]], X.bm, type = "link")
expect_equal(as.numeric(link[, li]), as.numeric(manual_link), tolerance = 1e-6)

resp <- predict(fits[["SSR"]], X.bm, type = "response")
expect_equal(resp, link) # cox has no separate response-scale transform

coefs <- predict(fits[["SSR"]], X.bm, type = "coefficients")
expect_equal(as.matrix(coefs), as.matrix(coef(fits[["SSR"]])))
expect_equal(nrow(coefs), p) # no intercept row for cox

nv2 <- predict(fits[["SSR"]], X.bm, type = "nvars")
expect_equal(as.numeric(nv2), as.numeric(Matrix::colSums(fits[["SSR"]]$beta != 0)))

expect_error(predict(fits[["SSR"]], X.bm, type = "class")) # class is binomial-only


# Test plot.biglasso for cox ---------------------------------------------------
# smoke test only: plot.biglasso() strips the intercept row based on whether
# length(penalty.factor) == nrow(beta), which happens to already be correct
# for cox (no intercept row is ever added), but this had never been exercised.

tmp_pdf <- tempfile(fileext = ".pdf")
pdf(tmp_pdf)
plot(fits[["SSR"]])
dev.off()
unlink(tmp_pdf)
expect_true(TRUE) # reaching this point means plot.biglasso() didn't error


# Test cv.biglasso() rejects family = "cox" ------------------------------------

expect_error(cv.biglasso(X.bm, y, family = "cox"))


# Test invalid screen argument errors ------------------------------------------

expect_error(biglasso(X.bm, y, family = "cox", screen = "bogus"))