multinom_reg add factor lvls

patr1ckm · patr1ckm · commit 86330e648235 · 2019-11-01T11:08:18.000-04:00
diff --git a/R/multinom_reg.R b/R/multinom_reg.R
@@ -290,7 +290,7 @@ multi_predict._multnet <-
       pred <-
         tibble(
           .row = rep(1:nrow(new_data), length(penalty)),
-          .pred_class = factor(as.vector(pred)),
+          .pred_class = factor(as.vector(pred), levels = object$lvl),
           penalty = rep(penalty, each = nrow(new_data))
         )
     }
diff --git a/tests/testthat/test_multinom_reg_glmnet.R b/tests/testthat/test_multinom_reg_glmnet.R
@@ -121,7 +121,8 @@ test_that('glmnet probabilities, mulitiple lambda', {
     multi_predict(xy_fit, iris[rows, 1:4], penalty = lams, type = "prob")$.pred
   )
 
-  mult_class <- factor(names(mult_probs)[apply(mult_probs, 1, which.max)])
+  mult_class <- factor(names(mult_probs)[apply(mult_probs, 1, which.max)],
+                       levels = xy_fit$lvl)
   mult_class <- tibble(
     .pred_class = mult_class,
     penalty = rep(lams, each = 3),

Original file line number	Diff line number	Diff line change
`@@ -290,7 +290,7 @@ multi_predict._multnet <-`
`290`	`290`	`pred <-`
`291`	`291`	`tibble(`
`292`	`292`	`.row = rep(1:nrow(new_data), length(penalty)),`
`293`		`- .pred_class = factor(as.vector(pred)),`
	`293`	`+ .pred_class = factor(as.vector(pred), levels = object$lvl),`
`294`	`294`	`penalty = rep(penalty, each = nrow(new_data))`
`295`	`295`	`)`
`296`	`296`	`}`
Original file line number	Diff line number	Diff line change
`@@ -121,7 +121,8 @@ test_that('glmnet probabilities, mulitiple lambda', {`
`121`	`121`	`multi_predict(xy_fit, iris[rows, 1:4], penalty = lams, type = "prob")$.pred`
`122`	`122`	`)`
`123`	`123`
`124`		`- mult_class <- factor(names(mult_probs)[apply(mult_probs, 1, which.max)])`
	`124`	`+ mult_class <- factor(names(mult_probs)[apply(mult_probs, 1, which.max)],`
	`125`	`+ levels = xy_fit$lvl)`
`125`	`126`	`mult_class <- tibble(`
`126`	`127`	`.pred_class = mult_class,`
`127`	`128`	`penalty = rep(lams, each = 3),`