added fix for probability type predictions for two class GAM models

Author: oj713

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"))