Standardized output types for multi_predict and predict in multinom_reg

NEWS.md

@@ -8,14 +8,16 @@ parsnip model object, and is printed when the model object is printed.
 * Some default parameter ranges were updated for SVM, KNN, and MARS models. 
 
 ## Fixes
+* [A bug](https://github.com/tidymodels/parsnip/issues/222) was fixed standardizing
+the output column types of `multi_predict` and `predict` for `multinom_reg`.
 
 * [A bug](https://github.com/tidymodels/parsnip/issues/208) was fixed related to using data descriptors and `fit_xy()`. 
 
 * A bug was fixed related to the column names generated by `multi_predict()`. The top-level tibble will always have a column named `.pred` and this list column contains tibbles across sub-models. The column names for these sub-model tibbles will have names consistent with `predict()` (which was previously incorrect). See [43c15db](https://github.com/tidymodels/parsnip/commit/43c15db377ea9ef27483ff209f6bd0e98cb830d2).
 
 * The model `udpate()` methods gained a `parameters` argument for cases when the parameters are contained in a tibble or list. 
 
-# [A bug](https://github.com/tidymodels/parsnip/issues/174) was fixed standardizing the column names of `nnet` class probability predictions.
+* [A bug](https://github.com/tidymodels/parsnip/issues/174) was fixed standardizing the column names of `nnet` class probability predictions.
 
 
 # parsnip 0.0.3.1

R/multinom_reg.R

@@ -297,7 +297,7 @@ multi_predict._multnet <-
       pred <-
         tibble(
           .row = rep(1:nrow(new_data), length(penalty)),
-          .pred_class = as.vector(pred),
+          .pred_class = factor(as.vector(pred), levels = object$lvl),
           penalty = rep(penalty, each = nrow(new_data))
         )
     }

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 <- 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),
@@ -149,3 +150,12 @@ test_that('glmnet probabilities, mulitiple lambda', {
   )
 
 })
+
+test_that("class predictions are factors with all levels", {
+  basic <- multinom_reg() %>% set_engine("glmnet") %>% fit(Species ~ ., data = iris)
+  nd <- iris[iris$Species == "setosa", ]
+  yhat <- predict(basic, new_data = nd, penalty = .1)
+  yhat_multi <- multi_predict(basic, new_data =  nd, penalty = .1)$.pred
+  expect_is(yhat_multi[[1]]$.pred_class, "factor")
+  expect_equal(levels(yhat_multi[[1]]$.pred_class), levels(iris$Species))
+})