tests for oobag predict()

R/orsf_R6.R

@@ -496,7 +496,8 @@ ObliqueForest <- R6::R6Class(
    # object is restored to its original state.
    if(is_error(out)) stop(out, call. = FALSE)
 
-   return(out)
+   # NaNs may occur with oobag = TRUE and small n_tree
+   coerce_nans(out, to = NA_real_)
 
 
   },
@@ -3940,6 +3941,7 @@ ObliqueForestSurvival <- R6::R6Class(
     # put the oob predictions into the same order as the training data.
     unsorted <- collapse::radixorder(private$data_row_sort)
     out_values <- out_values[unsorted, , drop = FALSE]
+
    }
 
    private$clean_pred_new(out_values)
@@ -4177,6 +4179,14 @@ ObliqueForestClassification <- R6::R6Class(
 
   },
 
+  clean_pred_oobag_internal = function(){
+
+   if(self$pred_type %in% c("prob") && is.matrix(self$pred_oobag)){
+    colnames(self$pred_oobag) <- self$class_levels
+   }
+
+  },
+
   predict_internal = function(simplify, oobag){
 
    # resize y to have the right number of columns

tests/testthat/test-orsf_predict.R

@@ -641,6 +641,28 @@ test_that(
  }
 )
 
+test_that(
+ desc = "oobag predictions match predict with oobag = TRUE",
+ code = {
+
+  expect_equal(
+   fit_standard_pbc$fast$pred_oobag,
+   predict(fit_standard_pbc$fast, oobag = TRUE)
+  )
+
+  expect_equal(
+   fit_standard_penguin_bills$net$pred_oobag,
+   predict(fit_standard_penguin_bills$net, oobag = TRUE)
+  )
+
+  expect_equal(
+   fit_standard_penguin_species$custom$pred_oobag,
+   predict(fit_standard_penguin_species$custom, oobag = TRUE)
+  )
+
+ }
+)
+
 new_data <- pbc_test[1:10, ]
 
 test_that(