add filter method for PredictionClust (#16)

* add filter * fix unit test * update required packages Co-authored-by: Michel Lang <[email protected]>
mlr-org · Feb 14, 2022 · de94e10 · de94e10
1 parent 2e78d94
commit de94e10
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 2 deletions.
diff --git a/DESCRIPTION b/DESCRIPTION
@@ -18,14 +18,14 @@ URL: https://mlr3cluster.mlr-org.com,
 BugReports: https://github.com/mlr-org/mlr3cluster/issues
 Depends:
     R (>= 3.1.0),
-    mlr3 (>= 0.10.0)
+    mlr3 (>= 0.13.0)
 Imports:
     backports (>= 1.1.10),
     checkmate,
     clue,
     clusterCrit,
     data.table,
-    mlr3misc (>= 0.4.0),
+    mlr3misc (>= 0.9.4),
     paradox,
     R6,
     stats

diff --git a/NAMESPACE b/NAMESPACE
@@ -10,6 +10,7 @@ S3method(as_task_clust,data.frame)
 S3method(as_task_clust,formula)
 S3method(c,PredictionDataClust)
 S3method(check_prediction_data,PredictionDataClust)
+S3method(filter_prediction_data,PredictionDataClust)
 S3method(is_missing_prediction_data,PredictionDataClust)
 export(LearnerClust)
 export(LearnerClustAP)

diff --git a/R/PredictionDataClust.R b/R/PredictionDataClust.R
@@ -74,3 +74,19 @@ c.PredictionDataClust = function(..., keep_duplicates = TRUE) {
 
   set_class(result, "PredictionDataClust")
 }
+
+#' @export
+filter_prediction_data.PredictionDataClust = function(pdata, row_ids) {
+  keep = pdata$row_ids %in% row_ids
+  pdata$row_ids = pdata$row_ids[keep]
+
+  if (!is.null(pdata$partition)) {
+    pdata$partition = pdata$partition[keep]
+  }
+
+  if (!is.null(pdata$prob)) {
+    pdata$prob = pdata$prob[keep,, drop = FALSE]
+  }
+
+  pdata
+}
diff --git a/tests/testthat/test_PredictionClust.R b/tests/testthat/test_PredictionClust.R
@@ -17,3 +17,15 @@ test_that("Internally constructed Prediction", {
   expect_prediction(p)
   expect_prediction_clust(p)
 })
+
+test_that("filter works", {
+  task = tsk("usarrests")
+  lrn = mlr_learners$get("clust.featureless")
+  lrn$param_set$values = list(num_clusters = 1L)
+  p = lrn$train(task)$predict(task)
+  pdata = p$data
+
+  pdata = filter_prediction_data(pdata, row_ids = 1:3)
+  expect_set_equal(pdata$row_ids, 1:3)
+  expect_integer(pdata$partition, len = 3)
+})