encapsulate glmnet formatting code (#867)

* encapsulate formatting code

* Apply suggestions from code review

Co-authored-by: Simon P. Couch <simonpatrickcouch@gmail.com>

---------

Co-authored-by: Simon P. Couch <simonpatrickcouch@gmail.com>

Author: hfrick

R/linear_reg.R

@@ -241,20 +241,25 @@ multi_predict._elnet <-
 
     pred <- predict._elnet(object, new_data = new_data, type = "raw",
                            opts = dots, penalty = penalty, multi = TRUE)
-    param_key <- tibble(group = colnames(pred), penalty = penalty)
-    pred <- as_tibble(pred)
-    pred$.row <- 1:nrow(pred)
-    pred <- gather(pred, group, .pred, -.row)
-    if (utils::packageVersion("dplyr") >= "1.0.99.9000") {
-      pred <- full_join(param_key, pred, by = "group", multiple = "all")
-    } else {
-      pred <- full_join(param_key, pred, by = "group")
-    }
-    pred$group <- NULL
-    pred <- arrange(pred, .row, penalty)
-    .row <- pred$.row
-    pred$.row <- NULL
-    pred <- split(pred, .row)
-    names(pred) <- NULL
-    tibble(.pred = pred)
+
+    format_glmnet_multi_linear_reg(pred, penalty = penalty)
   }
+
+format_glmnet_multi_linear_reg <- function(pred, penalty) {
+  param_key <- tibble(group = colnames(pred), penalty = penalty)
+  pred <- as_tibble(pred)
+  pred$.row <- 1:nrow(pred)
+  pred <- gather(pred, group, .pred, -.row)
+  if (utils::packageVersion("dplyr") >= "1.0.99.9000") {
+    pred <- full_join(param_key, pred, by = "group", multiple = "all")
+  } else {
+    pred <- full_join(param_key, pred, by = "group")
+  }
+  pred$group <- NULL
+  pred <- arrange(pred, .row, penalty)
+  .row <- pred$.row
+  pred$.row <- NULL
+  pred <- split(pred, .row)
+  names(pred) <- NULL
+  tibble(.pred = pred)
+}

R/logistic_reg.R

@@ -285,34 +285,41 @@ multi_predict._lognet <-
     pred <- predict._lognet(object, new_data = new_data, type = "raw",
                             opts = dots, penalty = penalty, multi = TRUE)
 
-    param_key <- tibble(group = colnames(pred), penalty = penalty)
-    pred <- as_tibble(pred)
-    pred$.row <- 1:nrow(pred)
-    pred <- gather(pred, group, .pred_class, -.row)
-    if (dots$type == "class") {
-      pred[[".pred_class"]] <- factor(pred[[".pred_class"]], levels = object$lvl)
-    } else {
-      if (dots$type == "response") {
-        pred[[".pred2"]] <- 1 - pred[[".pred_class"]]
-        names(pred) <- c(".row", "group", paste0(".pred_", rev(object$lvl)))
-        pred <- pred[, c(".row", "group", paste0(".pred_", object$lvl))]
-      }
-    }
-    if (utils::packageVersion("dplyr") >= "1.0.99.9000") {
-      pred <- full_join(param_key, pred, by = "group", multiple = "all")
-    } else {
-      pred <- full_join(param_key, pred, by = "group")
-    }
-    pred$group <- NULL
-    pred <- arrange(pred, .row, penalty)
-    .row <- pred$.row
-    pred$.row <- NULL
-    pred <- split(pred, .row)
-    names(pred) <- NULL
-    tibble(.pred = pred)
+    format_glmnet_multi_logistic_reg(
+      pred,
+      penalty,
+      type = dots$type,
+      lvl = object$lvl
+    )
   }
 
-
+format_glmnet_multi_logistic_reg <- function(pred, penalty, type, lvl) {
+  param_key <- tibble(group = colnames(pred), penalty = penalty)
+  pred <- as_tibble(pred)
+  pred$.row <- 1:nrow(pred)
+  pred <- gather(pred, group, .pred_class, -.row)
+  if (type == "class") {
+    pred[[".pred_class"]] <- factor(pred[[".pred_class"]], levels = lvl)
+  } else {
+    if (type == "response") {
+      pred[[".pred2"]] <- 1 - pred[[".pred_class"]]
+      names(pred) <- c(".row", "group", paste0(".pred_", rev(lvl)))
+      pred <- pred[, c(".row", "group", paste0(".pred_", lvl))]
+    }
+  }
+  if (utils::packageVersion("dplyr") >= "1.0.99.9000") {
+    pred <- full_join(param_key, pred, by = "group", multiple = "all")
+  } else {
+    pred <- full_join(param_key, pred, by = "group")
+  }
+  pred$group <- NULL
+  pred <- arrange(pred, .row, penalty)
+  .row <- pred$.row
+  pred$.row <- NULL
+  pred <- split(pred, .row)
+  names(pred) <- NULL
+  tibble(.pred = pred)
+}
 
 
 

R/multinom_reg.R

@@ -224,35 +224,13 @@ multi_predict._multnet <-
     pred <- predict._multnet(object, new_data = new_data, type = "raw",
                              opts = dots, penalty = penalty, multi = TRUE)
 
-    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)]
-    }
-
-    if (type == "prob") {
-      pred <- apply(pred, 3, format_probs)
-      names(pred) <- NULL
-      pred <- map_dfr(pred, function(x) x)
-      pred$penalty <- rep(penalty, each = nrow(new_data))
-      pred <- dplyr::relocate(pred, penalty)
-    } else {
-      pred <-
-        tibble(
-          .row = rep(1:nrow(new_data), length(penalty)),
-          penalty = rep(penalty, each = nrow(new_data)),
-          .pred_class = factor(as.vector(pred), levels = object$lvl)
-        )
-    }
-
-    pred <- arrange(pred, .row, penalty)
-    .row <- pred$.row
-    pred$.row <- NULL
-    pred <- split(pred, .row)
-    names(pred) <- NULL
-    tibble(.pred = pred)
+    format_glmnet_multi_multinom_reg(
+      pred,
+      penalty = penalty,
+      type = type,
+      n_rows = nrow(new_data),
+      lvl = object$lvl
+    )
   }
 
 #' @export
@@ -273,3 +251,35 @@ predict_raw._multnet <- function(object, new_data, opts = list(), ...) {
   opts$s <- object$spec$args$penalty
   predict_raw.model_fit(object, new_data = new_data, opts = opts, ...)
 }
+
+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)]
+  }
+
+  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 <- arrange(pred, .row, penalty)
+  .row <- pred$.row
+  pred$.row <- NULL
+  pred <- split(pred, .row)
+  names(pred) <- NULL
+  tibble(.pred = pred)
+}