R/glmnet.R

In offsetreg: An Extension of 'Tidymodels' Supporting Offset Terms

#' Fit Penalized Generalized Linear Models with an Offset
#'
#' This function is a wrapper around [glmnet::glmnet()] that uses a column from
#' `x` as an offset.
#'
#' @details
#' Outside of the `tidymodels` ecosystem, `glmnet_offset()` has no advantages
#' over [glmnet::glmnet()] since that function allows for offsets to be
#' specified in its `offset` argument.
#'
#' Within `tidymodels`, `glmnet_offset()` provides an advantage because it will
#' ensure that offsets are included in the data whenever resamples are created.
#'
#' The `x`, `y`, `family`, `lambda`, `alpha` and `weights` arguments have the
#' same meanings as [glmnet::glmnet()]. See that function's documentation for
#' full details.
#'
#' @param x Input matrix
#' @param y Response variable
#' @param lambda A numeric vector of regularization penalty values
#' @param alpha A number between zero and one denoting the proportion of L1
#' (lasso) versus L2 (ridge) regularization.
#' - `alpha = 1`: Pure lasso model
#' - `alpha = 0`: Pure ridge model
#' @param ... Additional named arguments passed to [glmnet::glmnet()].
#' @inheritParams glm_offset
#'
#' @returns A `glmnet` object. See [glmnet::glmnet()] for full details.
#'
#' @examples
#' if (interactive()) {
#'   us_deaths$off <- log(us_deaths$population)
#'   x <- model.matrix(~ age_group + gender + off, us_deaths)[, -1]
#'   glmnet_offset(x, us_deaths$deaths, family = "poisson", offset_col = "off")
#' }
#'
#' @seealso [glmnet::glmnet()]
#' @export
glmnet_offset <- function(
  x,
  y,
  family,
  offset_col = "offset",
  weights = NULL,
  lambda = NULL,
  alpha = 1,
  ...
) {
  rlang::check_installed("glmnet")

  if (!offset_col %in% colnames(x)) {
    cli::cli_abort("A column named `{offset_col}` must be present.")
  }

  offsets <- x[, offset_col]
  x <- x[, !colnames(x) %in% offset_col, drop = FALSE]

  glmnet::glmnet(
    x,
    y,
    family = family,
    weights,
    offset = offsets,
    lambda = lambda,
    alpha = alpha,
    ...
  )
}

# code from the parsnip package
set_glmnet_penalty_path <- function(x) {
  if (any(names(x$eng_args) == "path_values")) {
    # Since we decouple the parsnip `penalty` argument from being the same
    # as the glmnet `lambda` value, `path_values` allows users to set the
    # path differently from the default that glmnet uses. See
    # https://github.com/tidymodels/parsnip/issues/431
    x$method$fit$args$lambda <- x$eng_args$path_values
    x$eng_args$path_values <- NULL
    x$method$fit$args$path_values <- NULL
  } else {
    # See discussion in https://github.com/tidymodels/parsnip/issues/195
    x$method$fit$args$lambda <- NULL
  }
  x
}