Refactor formatting functions for glmnet methods for multi_predict() (#930)

* refactor `format_glmnet_multi_linear_reg()`

to match the pattern of `format_glmnet_multi_logistic_reg()`

* refactor `format_glmnet_multi_multinom_reg()`

to more closely match the other formatting function for glmnet multi_predict() methods

* fix column name

* required for `list_rbind()`

Author: hfrick

DESCRIPTION

@@ -30,7 +30,7 @@ Imports:
     magrittr,
     pillar,
     prettyunits,
-    purrr,
+    purrr (>= 1.0.0),
     rlang (>= 0.3.1),
     stats,
     tibble (>= 2.1.1),

R/glmnet-engines.R

@@ -223,8 +223,8 @@ multi_predict_glmnet <- function(object,
     "multinom_reg" = format_glmnet_multi_multinom_reg(pred,
                                                       penalty = penalty,
                                                       type = type,
-                                                      n_rows = nrow(new_data),
-                                                      lvl = object$lvl)
+                                                      lvl = object$lvl,
+                                                      n_obs = nrow(new_data))
   )
 
   res
@@ -248,26 +248,28 @@ multi_predict._multnet <- multi_predict_glmnet
 multi_predict._glmnetfit <- multi_predict_glmnet
 
 format_glmnet_multi_linear_reg <- function(pred, penalty) {
-  param_key <- tibble(group = colnames(pred), penalty = penalty)
+  penalty_key <- tibble(s = colnames(pred), penalty = penalty)
+
   pred <- as_tibble(pred)
-  pred$.row <- 1:nrow(pred)
-  pred <- gather(pred, group, .pred, -.row)
+  pred$.row <- seq_len(nrow(pred))
+  pred <- tidyr::pivot_longer(pred, -.row, names_to = "s", values_to = ".pred")
+
   if (utils::packageVersion("dplyr") >= "1.0.99.9000") {
-    pred <- full_join(param_key, pred, by = "group", multiple = "all")
+    pred <- dplyr::full_join(penalty_key, pred, by = "s", multiple = "all")
   } else {
-    pred <- full_join(param_key, pred, by = "group")
+    pred <- dplyr::full_join(penalty_key, pred, by = "s")
   }
-  pred$group <- NULL
-  pred <- arrange(pred, .row, penalty)
-  .row <- pred$.row
-  pred$.row <- NULL
-  pred <- split(pred, .row)
-  names(pred) <- NULL
-  tibble(.pred = pred)
+
+  pred <- pred %>%
+    dplyr::select(-s) %>%
+    dplyr::arrange(penalty) %>%
+    tidyr::nest(.by = .row, .key = ".pred") %>%
+    dplyr::select(-.row)
+
+  pred
 }
 
 format_glmnet_multi_logistic_reg <- function(pred, penalty, type, lvl) {
-
   type <- rlang::arg_match(type, c("class", "prob"))
 
   penalty_key <- tibble(s = colnames(pred), penalty = penalty)
@@ -303,36 +305,46 @@ format_glmnet_multi_logistic_reg <- function(pred, penalty, type, lvl) {
   pred
 }
 
-format_glmnet_multi_multinom_reg <- function(pred, penalty, type, n_rows, lvl) {
-  format_probs <- function(x) {
-    x <- as_tibble(x)
-    names(x) <- paste0(".pred_", names(x))
-    nms <- names(x)
-    x$.row <- 1:nrow(x)
-    x[, c(".row", nms)]
-  }
+format_glmnet_multi_multinom_reg <- function(pred, penalty, type, lvl, n_obs) {
+  type <- rlang::arg_match(type, c("class", "prob"))
 
-  if (type == "prob") {
-    pred <- apply(pred, 3, format_probs)
-    names(pred) <- NULL
-    pred <- map_dfr(pred, function(x) x)
-    pred$penalty <- rep(penalty, each = n_rows)
-    pred <- dplyr::relocate(pred, penalty)
-  } else {
-    pred <-
-      tibble(
-        .row = rep(1:n_rows, length(penalty)),
-        penalty = rep(penalty, each = n_rows),
-        .pred_class = factor(as.vector(pred), levels = lvl)
-      )
-  }
+  pred <- switch(
+    type,
+    prob = format_glmnet_multinom_prob(pred, penalty, lvl, n_obs),
+    class = format_glmnet_multinom_class(pred, penalty, lvl, n_obs)
+  )
+
+  pred <- pred %>%
+    dplyr::arrange(.row, penalty) %>%
+    tidyr::nest(.by = .row, .key = ".pred") %>%
+    dplyr::select(-.row)
+
+  pred
+}
+
+format_glmnet_multinom_prob <- function(pred, penalty, lvl, n_obs) {
+  # pred is an array with
+  # dim 1 = observations
+  # dim 2 = levels of the response
+  # dim 3 = penalty values
+  apply(pred, 3, as_tibble) %>%
+    purrr::list_rbind() %>%
+    rlang::set_names(paste0(".pred_", lvl)) %>%
+    dplyr::mutate(
+      .row = rep(seq_len(n_obs), times = length(penalty)),
+      penalty = rep(penalty, each = n_obs)
+    ) %>%
+    dplyr::relocate(penalty)
+}
 
-  pred <- arrange(pred, .row, penalty)
-  .row <- pred$.row
-  pred$.row <- NULL
-  pred <- split(pred, .row)
-  names(pred) <- NULL
-  tibble(.pred = pred)
+format_glmnet_multinom_class <- function(pred, penalty, lvl, n_obs) {
+  # pred is a matrix n_obs x n_penalty
+  # unless n_obs == 1, then it's a vector of length n_penalty
+  tibble(
+    .row = rep(seq_len(n_obs), times = length(penalty)),
+    penalty = rep(penalty, each = n_obs),
+    .pred_class = factor(as.vector(pred), levels = lvl)
+  )
 }
 
 # -------------------------------------------------------------------------