From 6f00332cfbd6a59f8540493f544e74387957451c Mon Sep 17 00:00:00 2001
From: Marvin Wright <github@wrig.de>
Date: Mon, 21 Oct 2019 12:55:16 +0200
Subject: [PATCH] add test for sparse survival data

---
 tests/testthat/test_sparse.R | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/tests/testthat/test_sparse.R b/tests/testthat/test_sparse.R
index 00b870269..7c54290ce 100644
--- a/tests/testthat/test_sparse.R
+++ b/tests/testthat/test_sparse.R
@@ -16,6 +16,11 @@ dat <- data.frame(y = y, x)
 dat_matrix <- data.matrix(dat)
 dat_sparse <- Matrix(dat_matrix, sparse = TRUE)
 
+# Survival sparse data
+dat_survival <- data.frame(x, time = round(runif(n, 0, 10)), status = rbinom(n, 1, .7))
+dat_survival_matrix <- data.matrix(dat_survival)
+dat_survival_sparse <- Matrix(dat_survival_matrix, sparse = TRUE)
+
 test_that("Same result with sparse data for iris classification", {
   set.seed(56)
   rf1 <- ranger(data = iris_sparse, dependent.variable.name = "Species", classification = TRUE, num.trees = 5)
@@ -86,6 +91,29 @@ test_that("Same result with sparse data for 0/1 probability prediction", {
   expect_equal(pred1, pred2)
 })
 
+test_that("Same result with sparse data for survival", {
+  set.seed(56)
+  rf1 <- ranger(data = dat_survival_sparse, dependent.variable.name = "time", status.variable.name = "status", num.trees = 5)
+  
+  set.seed(56)
+  rf2 <- ranger(data = dat_survival, dependent.variable.name = "time", status.variable.name = "status", num.trees = 5)
+  
+  expect_equal(rf1$prediction.error, rf2$prediction.error)
+  
+  pred1 <- rf1$survival[!is.na(rf1$survival)]
+  pred2 <- rf2$survival[!is.na(rf2$survival)]
+  expect_equal(pred1, pred2)
+})
+
+test_that("Survival prediction is the same with or without outcome in prediction data", {
+  rf <- ranger(data = dat_survival_sparse, dependent.variable.name = "time", status.variable.name = "status", num.trees = 5)
+
+  pred1 <- predict(rf, dat_survival_sparse)$survival  
+  pred2 <- predict(rf, dat_survival_sparse[, c(-6, -7)])$survival  
+  
+  expect_equal(pred1, pred2)
+})
+
 test_that("Prediction is the same if training or testing data is sparse", {
   idx <- sample(nrow(iris), 2/3*nrow(iris))
   train <- iris[idx, ]