Merge pull request #709 from oj713/main

Fixing inconsistency for probability predictions for two class GAM models

Author: EmilHvitfeldt

NEWS.md

@@ -6,6 +6,10 @@
 
 * Exported `xgb_predict()` which wraps xgboost's `predict()` method for use with parsnip extension packages (#688).
 
+## Bug fixes
+
+*  An inconsistency for probability type predictions for two-class GAM models was fixed (#708)
+
 # parsnip 0.2.1
 
 * Fixed a major bug in spark models induced in the previous version (#671).

R/gen_additive_mod_data.R

@@ -166,6 +166,9 @@ set_pred(
   value  = list(
     pre = NULL,
     post = function(x, object) {
+      if (is.array(x)) {
+        x <- as.vector(x)
+      }
       x <- tibble(v1 = 1 - x, v2 = x)
       colnames(x) <- object$lvl
       x

tests/testthat/test_gen_additive_model.R

@@ -83,6 +83,8 @@ test_that('classification', {
   mgcv_pred <- predict(mgcv_mod, head(two_class_dat), type = "response")
   expect_equal(names(f_pred), c(".pred_Class1", ".pred_Class2"))
   expect_equal(f_pred[[".pred_Class2"]], mgcv_pred, ignore_attr = TRUE)
+  expect_equal(class(f_pred[[".pred_Class1"]]), "numeric")
+  expect_equal(class(f_pred[[".pred_Class2"]]), "numeric")
 
   f_cls <- predict(f_res, head(two_class_dat), type = "class")
   expect_true(all(f_cls$.pred_class[mgcv_pred < 0.5] == "Class1"))