rtreeparty

Add horizontal line to terminal barplots in party / partykit trees


I am using the ctree function from partykit.

library(rpart)
library(partykit)

fit <- ctree(Kyphosis ~ Age + Number + Start, data=kyphosis)
plot(fit, terminal_panel=node_barplot)

enter image description here

I want to add an additional horizonal line to each barplot, indicating the average reponse across the dataset, i.e. at 0.79 here.

prop.table(table(kyphosis$Kyphosis))

    absent   present 
 0.7901235 0.2098765

Approach: I started to modify the node_barplot function that is passed to the terminal_panel argument. But the source code is very long and comes with almost no comments. So I tried to go step by step, stripping down the function to its first two lines of code (plus an extra print command). However, if I run it, the object y is NULL and an error is thrown.

node_barplot2 <- function(obj, ...)
{   
  y <- obj$fitted[["(response)"]]   # first lime from node_barplot source
  print(y)
  stopifnot(is.factor(y) || isTRUE(all.equal(round(y), y)) || is.data.frame(y))
}

plot(fit, terminal_panel=node_barplot2)

> Error in round(y) : Non-numeric argument in mathematical function

As its the original code, I do not quite see where I go wrong here and how I could draw the horizontal line. Any ideas?


Solution

  • The partykit distinguishes "panel" functions and "panel-generating" functions:

    To signal that a function is a panel-generating function, it has to have class "grapcon_generator". Hence

    class(node_barplot)
    ## [1] "grapcon_generator"
    

    To add certain graphical elements to the function, I would recommend copying the whole node_barplot source code (including the class assignment at the end) and then add the elements you need, e.g., a horizontal reference line you can draw with grid.lines().