Issue61

R/orsf_R6.R

@@ -467,6 +467,7 @@ ObliqueForest <- R6::R6Class(
    self$verbose_progress <- verbose_progress
    self$pred_aggregate   <- pred_aggregate
 
+
    out <- try(
     expr = {
 
@@ -4262,6 +4263,8 @@ ObliqueForestClassification <- R6::R6Class(
 
    out <- do.call(orsf_cpp, args = cpp_args)$pred_new
 
+   out <- private$clean_pred_new(out)
+
    if(self$pred_type == 'prob' && self$pred_aggregate){
 
     colnames(out) <- self$class_levels
@@ -4285,6 +4288,12 @@ ObliqueForestClassification <- R6::R6Class(
 
    out
 
+  },
+
+  clean_pred_new_internal = function(preds){
+
+   preds
+
   }
 
  )
@@ -4510,10 +4519,18 @@ ObliqueForestRegression <- R6::R6Class(
 
    out <- do.call(orsf_cpp, args = cpp_args)$pred_new
 
+   out <- private$clean_pred_new(out)
+
    if(simplify) dim(out) <- NULL
 
    out
 
+  },
+
+  clean_pred_new_internal = function(preds){
+
+   preds
+
   }
 
  )

tests/testthat/test-na_action.R

@@ -111,3 +111,49 @@ test_that(
 
  })
 
+test_that(
+ desc = "na action of pass works with new preds",
+ code = {
+
+  mtcars_orsf <- mtcars
+  mtcars_orsf$vs <- factor(mtcars_orsf$vs)
+  mtcars_na <- mtcars_orsf
+
+  set_to_miss <- c(1, 4, 18)
+  mtcars_na$cyl[set_to_miss] <- NA
+
+  aorsf_regr_fit <- orsf(
+   data = mtcars_orsf,
+   formula = mpg ~ .,
+   n_tree = n_tree_test,
+   tree_seeds = seeds_standard
+  )
+
+  aorsf_regr_pred <- predict(
+   aorsf_regr_fit,
+   new_data = mtcars_na,
+   na_action = 'pass'
+  )
+
+  expect_equal(which(is.na(aorsf_regr_pred)), set_to_miss)
+
+  aorsf_clsf_fit <- aorsf::orsf(
+   data = mtcars_orsf,
+   formula = vs ~ .,
+   n_tree = n_tree_test,
+   tree_seeds = seeds_standard
+  )
+
+  aorsf_clsf_pred <- predict(
+   aorsf_clsf_fit,
+   new_data = mtcars_na,
+   na_action = 'pass'
+  )
+
+  expect_equal(which(is.na(aorsf_clsf_pred[,1])), set_to_miss)
+  expect_equal(which(is.na(aorsf_clsf_pred[,2])), set_to_miss)
+
+  # this test is already done for survival in test-orsf_predict
+
+ }
+)