Use .default argument in case_when() when applicable (#164)

EmilHvitfeldt · web-flow · commit 069163851c8a · 2025-11-04T11:30:44.000-08:00
diff --git a/NEWS.md b/NEWS.md
@@ -22,6 +22,8 @@
 
 - Cubist rules will return simplified rules whenever possible to avoid multiplying by 0 and 1. (#152)
 
+- tree based models now uses `.default` argument in produced `case_when()` code when applicable. (#153)
+
 # tidypredict 0.5.1
 
 - Exported a number of internal functions to be used in {orbital} package
diff --git a/R/tree.R b/R/tree.R
@@ -49,16 +49,26 @@
 #' calculating all notes.
 #'
 #' @keywords internal
-generate_case_when_trees <- function(parsedmodel) {
+generate_case_when_trees <- function(parsedmodel, default = TRUE) {
   map(
     parsedmodel$trees,
     generate_case_when_tree,
-    mode = parsedmodel$general$mode
+    mode = parsedmodel$general$mode,
+    default = default
   )
 }
 
-generate_case_when_tree <- function(tree, mode) {
-  expr(case_when(!!!generate_tree_nodes(tree, mode)))
+generate_case_when_tree <- function(tree, mode, default = TRUE) {
+  nodes <- generate_tree_nodes(tree, mode)
+
+  if (default) {
+    default <- nodes[[length(nodes)]]
+    default <- rlang::f_rhs(default)
+    nodes[[length(nodes)]] <- NULL
+    nodes <- c(nodes, .default = default)
+  }
+
+  expr(case_when(!!!nodes))
 }
 
 generate_tree_nodes <- function(tree, mode) {
diff --git a/man/generate_case_when_trees.Rd b/man/generate_case_when_trees.Rd
diff --git a/tests/testthat/_snaps/model-partykit.md b/tests/testthat/_snaps/model-partykit.md
@@ -3,7 +3,7 @@
     Code
       rlang::expr_text(tf)
     Output
-      [1] "case_when(cyl <= 4 ~ 26.6636363636364, cyl <= 6 & cyl > 4 ~ 19.7428571428571, \n    cyl > 6 & cyl > 4 ~ 15.1)"
+      [1] "case_when(cyl <= 4 ~ 26.6636363636364, cyl <= 6 & cyl > 4 ~ 19.7428571428571, \n    .default = 15.1)"
 
 # formulas produces correct predictions
 
diff --git a/tests/testthat/_snaps/model-ranger.md b/tests/testthat/_snaps/model-ranger.md
@@ -3,7 +3,7 @@
     Code
       rlang::expr_text(tf)
     Output
-      [1] "case_when(Petal.Length < 2.6 ~ \"setosa\", Sepal.Length < 6.25 & \n    Petal.Length >= 2.6 ~ \"versicolor\", Sepal.Length >= 6.25 & \n    Petal.Length >= 2.6 ~ \"virginica\") + case_when(Petal.Width < \n    0.75 ~ \"setosa\", Petal.Width < 1.75 & Petal.Width >= 0.75 ~ \n    \"versicolor\", Petal.Width >= 1.75 & Petal.Width >= 0.75 ~ \n    \"virginica\") + case_when(Petal.Length < 2.35 ~ \"setosa\", \n    Petal.Length < 4.75 & Petal.Length >= 2.35 ~ \"versicolor\", \n    Petal.Length >= 4.75 & Petal.Length >= 2.35 ~ \"virginica\")"
+      [1] "case_when(Petal.Length < 2.6 ~ \"setosa\", Sepal.Length < 6.25 & \n    Petal.Length >= 2.6 ~ \"versicolor\", .default = \"virginica\") + \n    case_when(Petal.Width < 0.75 ~ \"setosa\", Petal.Width < 1.75 & \n        Petal.Width >= 0.75 ~ \"versicolor\", .default = \"virginica\") + \n    case_when(Petal.Length < 2.35 ~ \"setosa\", Petal.Length < \n        4.75 & Petal.Length >= 2.35 ~ \"versicolor\", .default = \"virginica\")"
 
 # formulas produces correct predictions
 
diff --git a/tests/testthat/_snaps/model-rf.md b/tests/testthat/_snaps/model-rf.md
@@ -3,7 +3,7 @@
     Code
       rlang::expr_text(tf)
     Output
-      [1] "(case_when(disp < 78.85 & wt < 2.3325 ~ 33.525, disp >= 78.85 & \n    wt < 2.3325 ~ 27.425, disp >= 281 & wt >= 2.3325 ~ 14.1, \n    drat < 3.695 & cyl < 5 & disp < 281 & wt >= 2.3325 ~ 24.4, \n    drat >= 3.695 & cyl < 5 & disp < 281 & wt >= 2.3325 ~ 22.3666666666667, \n    disp < 152.5 & wt < 3.3275 & cyl >= 5 & disp < 281 & wt >= \n        2.3325 ~ 19.7, disp >= 152.5 & wt < 3.3275 & cyl >= 5 & \n        disp < 281 & wt >= 2.3325 ~ 21.1, disp < 196.3 & wt >= \n        3.3275 & cyl >= 5 & disp < 281 & wt >= 2.3325 ~ 18.92, \n    disp >= 196.3 & wt >= 3.3275 & cyl >= 5 & disp < 281 & wt >= \n        2.3325 ~ 18.1) + case_when(hp < 80.5 & disp < 266.9 ~ \n    28.5333333333333, drat < 3.035 & disp >= 266.9 ~ 10.4, carb < \n    2.5 & drat >= 3.035 & disp >= 266.9 ~ 18.7, wt < 1.989 & \n    hp < 118 & hp >= 80.5 & disp < 266.9 ~ 30.4, qsec < 18.6 & \n    hp >= 118 & hp >= 80.5 & disp < 266.9 ~ 19.6, qsec >= 18.6 & \n    hp >= 118 & hp >= 80.5 & disp < 266.9 ~ 17.8, drat >= 3.635 & \n    carb >= 2.5 & drat >= 3.035 & disp >= 266.9 ~ 13.3, hp < \n    96 & wt >= 1.989 & hp < 118 & hp >= 80.5 & disp < 266.9 ~ \n    22.8, wt < 4.5625 & drat < 3.635 & carb >= 2.5 & drat >= \n    3.035 & disp >= 266.9 ~ 15.04, wt >= 4.5625 & drat < 3.635 & \n    carb >= 2.5 & drat >= 3.035 & disp >= 266.9 ~ 14.7, qsec < \n    19.725 & hp >= 96 & wt >= 1.989 & hp < 118 & hp >= 80.5 & \n    disp < 266.9 ~ 21.4, qsec >= 19.725 & hp >= 96 & wt >= 1.989 & \n    hp < 118 & hp >= 80.5 & disp < 266.9 ~ 21.5) + case_when(wt < \n    3.4725 & drat < 3.75 ~ 20.8333333333333, qsec < 16.23 & wt >= \n    3.4725 & drat < 3.75 ~ 13.3, disp < 78.85 & disp < 130.55 & \n    drat >= 3.75 ~ 32.2333333333333, disp >= 78.85 & disp < 130.55 & \n    drat >= 3.75 ~ 27.75, cyl < 5 & disp >= 130.55 & drat >= \n    3.75 ~ 22.8, disp >= 456 & qsec >= 16.23 & wt >= 3.4725 & \n    drat < 3.75 ~ 10.4, qsec < 17.66 & cyl >= 5 & disp >= 130.55 & \n    drat >= 3.75 ~ 21, qsec >= 17.66 & cyl >= 5 & disp >= 130.55 & \n    drat >= 3.75 ~ 19.2, qsec < 17.225 & disp < 456 & qsec >= \n    16.23 & wt >= 3.4725 & drat < 3.75 ~ 19.2, wt < 4.7075 & \n    qsec >= 17.225 & disp < 456 & qsec >= 16.23 & wt >= 3.4725 & \n    drat < 3.75 ~ 16.34, wt >= 4.7075 & qsec >= 17.225 & disp < \n    456 & qsec >= 16.23 & wt >= 3.4725 & drat < 3.75 ~ 14.7))/3L"
+      [1] "(case_when(disp < 78.85 & wt < 2.3325 ~ 33.525, disp >= 78.85 & \n    wt < 2.3325 ~ 27.425, disp >= 281 & wt >= 2.3325 ~ 14.1, \n    drat < 3.695 & cyl < 5 & disp < 281 & wt >= 2.3325 ~ 24.4, \n    drat >= 3.695 & cyl < 5 & disp < 281 & wt >= 2.3325 ~ 22.3666666666667, \n    disp < 152.5 & wt < 3.3275 & cyl >= 5 & disp < 281 & wt >= \n        2.3325 ~ 19.7, disp >= 152.5 & wt < 3.3275 & cyl >= 5 & \n        disp < 281 & wt >= 2.3325 ~ 21.1, disp < 196.3 & wt >= \n        3.3275 & cyl >= 5 & disp < 281 & wt >= 2.3325 ~ 18.92, \n    .default = 18.1) + case_when(hp < 80.5 & disp < 266.9 ~ 28.5333333333333, \n    drat < 3.035 & disp >= 266.9 ~ 10.4, carb < 2.5 & drat >= \n        3.035 & disp >= 266.9 ~ 18.7, wt < 1.989 & hp < 118 & \n        hp >= 80.5 & disp < 266.9 ~ 30.4, qsec < 18.6 & hp >= \n        118 & hp >= 80.5 & disp < 266.9 ~ 19.6, qsec >= 18.6 & \n        hp >= 118 & hp >= 80.5 & disp < 266.9 ~ 17.8, drat >= \n        3.635 & carb >= 2.5 & drat >= 3.035 & disp >= 266.9 ~ \n        13.3, hp < 96 & wt >= 1.989 & hp < 118 & hp >= 80.5 & \n        disp < 266.9 ~ 22.8, wt < 4.5625 & drat < 3.635 & carb >= \n        2.5 & drat >= 3.035 & disp >= 266.9 ~ 15.04, wt >= 4.5625 & \n        drat < 3.635 & carb >= 2.5 & drat >= 3.035 & disp >= \n        266.9 ~ 14.7, qsec < 19.725 & hp >= 96 & wt >= 1.989 & \n        hp < 118 & hp >= 80.5 & disp < 266.9 ~ 21.4, .default = 21.5) + \n    case_when(wt < 3.4725 & drat < 3.75 ~ 20.8333333333333, qsec < \n        16.23 & wt >= 3.4725 & drat < 3.75 ~ 13.3, disp < 78.85 & \n        disp < 130.55 & drat >= 3.75 ~ 32.2333333333333, disp >= \n        78.85 & disp < 130.55 & drat >= 3.75 ~ 27.75, cyl < 5 & \n        disp >= 130.55 & drat >= 3.75 ~ 22.8, disp >= 456 & qsec >= \n        16.23 & wt >= 3.4725 & drat < 3.75 ~ 10.4, qsec < 17.66 & \n        cyl >= 5 & disp >= 130.55 & drat >= 3.75 ~ 21, qsec >= \n        17.66 & cyl >= 5 & disp >= 130.55 & drat >= 3.75 ~ 19.2, \n        qsec < 17.225 & disp < 456 & qsec >= 16.23 & wt >= 3.4725 & \n            drat < 3.75 ~ 19.2, wt < 4.7075 & qsec >= 17.225 & \n            disp < 456 & qsec >= 16.23 & wt >= 3.4725 & drat < \n            3.75 ~ 16.34, .default = 14.7))/3L"
 
 # formulas produces correct predictions
 
diff --git a/tests/testthat/test-tree.R b/tests/testthat/test-tree.R
@@ -18,7 +18,7 @@ test_that("generate_case_when_trees() works", {
   )
 
   expect_identical(
-    generate_case_when_trees(parsedmodel),
+    generate_case_when_trees(parsedmodel, default = FALSE),
     list(
       quote(case_when(disp > 100 ~ 14 + hp * 4 + drat * 2))
     )
@@ -30,7 +30,7 @@ test_that("generate_case_when_trees() works", {
   )
 
   expect_identical(
-    generate_case_when_trees(parsedmodel),
+    generate_case_when_trees(parsedmodel, default = FALSE),
     list(
       quote(case_when(ifelse(disp > 100, 14 + hp * 4 + drat * 2, 0)))
     )
@@ -42,7 +42,7 @@ test_that("generate_case_when_trees() works", {
   )
 
   expect_identical(
-    generate_case_when_trees(parsedmodel),
+    generate_case_when_trees(parsedmodel, default = FALSE),
     list(
       quote(
         case_when(
@@ -63,7 +63,7 @@ test_that("generate_case_when_trees() works", {
   )
 
   expect_identical(
-    generate_case_when_trees(parsedmodel),
+    generate_case_when_trees(parsedmodel, default = FALSE),
     list(
       quote(
         case_when(
@@ -94,18 +94,18 @@ test_that("generate_case_when_tree() works", {
   nodes <- list(node)
 
   expect_identical(
-    generate_case_when_tree(nodes, mode = ""),
+    generate_case_when_tree(nodes, mode = "", default = FALSE),
     quote(case_when(disp > 100 ~ 14 + hp * 4 + drat * 2))
   )
   expect_identical(
-    generate_case_when_tree(nodes, mode = "ifelse"),
+    generate_case_when_tree(nodes, mode = "ifelse", default = FALSE),
     quote(case_when(ifelse(disp > 100, 14 + hp * 4 + drat * 2, 0)))
   )
 
   nodes <- list(node, node)
 
   expect_identical(
-    generate_case_when_tree(nodes, mode = ""),
+    generate_case_when_tree(nodes, mode = "", default = FALSE),
     quote(
       case_when(
         disp > 100 ~ 14 + hp * 4 + drat * 2,
@@ -114,14 +114,45 @@ test_that("generate_case_when_tree() works", {
     )
   )
   expect_identical(
-    generate_case_when_tree(nodes, mode = "ifelse"),
+    generate_case_when_tree(nodes, mode = "ifelse", default = FALSE),
     quote(
       case_when(
         ifelse(disp > 100, 14 + hp * 4 + drat * 2, 0),
         ifelse(disp > 100, 14 + hp * 4 + drat * 2, 0)
       )
     )
   )
+
+  nodes <- list(
+    list(
+      prediction = 25,
+      path = list(list(
+        type = "conditional",
+        col = "cyl",
+        val = 4,
+        op = "less-equal"
+      ))
+    ),
+    list(
+      prediction = 20,
+      path = list(
+        list(type = "conditional", col = "cyl", val = 6, op = "less-equal"),
+        list(type = "conditional", col = "cyl", val = 4, op = "more")
+      )
+    ),
+    list(
+      prediction = 15,
+      path = list(
+        list(type = "conditional", col = "cyl", val = 6, op = "more"),
+        list(type = "conditional", col = "cyl", val = 4, op = "more")
+      )
+    )
+  )
+
+  expect_identical(
+    generate_case_when_tree(nodes, mode = "", default = TRUE),
+    quote(case_when(cyl <= 4 ~ 25, cyl <= 6 & cyl > 4 ~ 20, .default = 15))
+  )
 })
 
 test_that("generate_tree_nodes() works", {
