feat: add option to predict probabilities

R/LearnerClassifSpatial.R

@@ -6,14 +6,15 @@ LearnerClassifSpatial = R6::R6Class("LearnerClassifSpatial",
     initialize = function(learner) {
       self$learner = assert_learner(learner)
       super$initialize(
-        id = "classif.ranger",
+        id = learner$id,
         param_set = learner$param_set,
         predict_types = learner$predict_types,
         feature_types = learner$feature_types,
         properties = union(learner$properties, "missings"),
         packages = learner$packages,
         man = "mlr3learners::mlr_learners_classif.spatial"
       )
+      self$predict_type = learner$predict_type
     },
 
     predict = function(task, row_ids = NULL) {
@@ -25,6 +26,13 @@ LearnerClassifSpatial = R6::R6Class("LearnerClassifSpatial",
       pred$data$row_ids = seq_len(nrow(data))
       pred$data$response = response
       pred$data$truth = rep(NaN, nrow(data))
+      if (self$learner$predict_type == "prob") {
+        prob = matrix(NaN, nrow = nrow(data), ncol = 2)
+        prob[ids, 1] = pred$data$prob[, 1]
+        prob[ids, 2] = pred$data$prob[, 2]
+        attributes(prob) = attributes(pred$prob)
+        pred$data$prob = prob
+      }
       pred
     }
   )

R/predict_spatial.R

@@ -15,6 +15,9 @@
 #'   For vector data only `"sf"` is supported.
 #' @param filename (`character(1)`)\cr
 #'   Path where the spatial object should be written to.
+#' @param predict_type (`character(1)`)\cr
+#'  Type of prediction to return.
+#'  Accepted values are `"response"` and `"prob"`.
 #'
 #' @return Spatial object of class given in argument `format`.
 #' @examples
@@ -31,10 +34,11 @@
 #' # predict land cover classes
 #' pred = predict_spatial(stack, learner, chunksize = 1L)
 #' @export
-predict_spatial = function(newdata, learner, chunksize = 200L, format = "terra", filename = NULL) {
+predict_spatial = function(newdata, learner, chunksize = 200L, format = "terra", filename = NULL, predict_type = "response") {
   task = as_task_unsupervised(newdata)
   assert_multi_class(task$backend, c("DataBackendRaster", "DataBackendVector"))
   assert_learner(learner)
+  assert_choice(predict_type, c("response", "prob"))
 
   if (test_class(task$backend, "DataBackendRaster")) {
     assert_number(chunksize)
@@ -63,7 +67,8 @@ predict_spatial = function(newdata, learner, chunksize = 200L, format = "terra",
 
       stack = task$backend$stack
       pred = learner$predict(task, row_ids = cells_seq:((cells_seq + cells_to_read - 1)))
-      terra::writeValues(x = target_raster, v = pred$response,
+      vals = if (predict_type == "prob") pred$prob[, learner$learner$state$train_task$positive] else pred$response
+      terra::writeValues(x = target_raster, v = vals,
         start = terra::rowFromCell(stack, cells_seq), # start row number
         nrows = terra::rowFromCell(stack, cells_to_read)) # how many rows
       lg$info("Chunk %i of %i finished", n, length(bs$cells_seq))
@@ -72,7 +77,7 @@ predict_spatial = function(newdata, learner, chunksize = 200L, format = "terra",
     terra::writeStop(target_raster)
     lg$info("Finished raster prediction in %i seconds", as.integer(proc.time()[3] - start_time))
 
-    if (learner$task_type == "classif") {
+    if (learner$task_type == "classif" && predict_type == "response") {
       levels = learner$learner$state$train_task$levels()[[learner$learner$state$train_task$target_names]]
       value = data.table(ID = seq_along(levels), categories = levels)
       target_raster = terra::categories(target_raster, value = value, index = 2)