diff --git a/R/orsf_R6.R b/R/orsf_R6.R index e7352e68..0f9bf4b1 100644 --- a/R/orsf_R6.R +++ b/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 + } ) diff --git a/tests/testthat/test-na_action.R b/tests/testthat/test-na_action.R index 3049fb2e..f3d637ad 100644 --- a/tests/testthat/test-na_action.R +++ b/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 + + } +)