R/plot.sparsegl.R

In hierNest: Penalized Regression with Hierarchical Nested Parameterization Structure

plot.sparsegl <- function(x,
                          y_axis = c("coef", "group"),
                          x_axis = c("lambda", "penalty"),
                          add_legend = n_legend_values < 20,
                          ...) {

  rlang::check_dots_empty()
  y_axis <- match.arg(y_axis)
  x_axis <- match.arg(x_axis)

  xb <- x$beta
  nonzeros <- sort(unique(xb@i)) + 1
  if (length(nonzeros) == 0) {
    cli::cli_abort(
      c("No nonzero betas / groups are available to plot.",
        i = "All coefficient estimates are exactly 0.")
    )
  }

  xb <- xb[nonzeros, , drop = FALSE]
  g <- x$group[nonzeros]
  uni_group <- unique(g)
  sgnorm <- apply(xb, 2, sp_group_norm, gr = g, asparse = x$asparse)

  if (y_axis == "group") {
    xb <- apply(xb, 2, grouped_sp_norm, gr = g, asparse = x$asparse)
    rownames(xb) <- paste0("group", uni_group)
  }

  n_legend_values <- nrow(xb)
  df <- as.data.frame(t(as.matrix(xb)))

  df$lambda <- x$lambda
  df$penalty <- sgnorm / max(sgnorm)
  df <- df %>%
    tidyr::pivot_longer(!c(.data$lambda, .data$penalty), names_to = y_axis)

  plot_layer <- ggplot2::ggplot(
    df, ggplot2::aes(
      x = !!rlang::sym(x_axis), y = .data$value, color = !!rlang::sym(y_axis)
    )) + ggplot2::geom_hline(yintercept = 0)

  if (x_axis == "penalty") {
    xlab_layer <- ggplot2::xlab("penalty / max (penalty)")
  }
  if (x_axis == "lambda") {
    xlab_layer <- ggplot2::xlab("lambda") + ggplot2::scale_x_log10()
  }
  if (y_axis == "group") {
    plot_layer <- plot_layer + ggplot2::ylab("group norm") +
      ggplot2::geom_line(ggplot2::aes(group = .data$group))
  }
  if (y_axis == "coef") {
    plot_layer <- plot_layer + ggplot2::geom_line() +
      ggplot2::ylab("coefficients")
  }

  legend_layer <- ggplot2::scale_color_viridis_d(name = "")

  theme_layer <- ggplot2::theme_bw()
  if (!add_legend)
    theme_layer <- theme_layer + ggplot2::theme(legend.position = "none")

  p <- plot_layer + xlab_layer + legend_layer + theme_layer
  return(p)
}