Super Nerdy Cool

My (technical) adventures through the professional world…
May 29, 2012

Better decision tree graphics for rpart via party and partykit

I’ve been using Graphviz to create better decision tree graphics “by hand” for rpart objects created in R (final tree). I stumbled on this post that shows how one could convert an rpart object to a party project via the as.party function in partykit to utilize the plot functions in party. It looks quite nice.

I might have to do additional hacking as I like to display the node size and percentage of success in every node. For example, in rpart, I do something like

## rpartObj created from rpart
textRpartCustom <- 
{
    nclass <- (ncol(yval) - 1L)/2
    group <- yval[, 1L]
    counts <- yval[, 1L + (1L:nclass)]
    if (!is.null(ylevel)) 
        group <- ylevel[group]
    temp1 <- rpart:::formatg(counts, digits)
    if (nclass > 1) {
        ## temp1 <- apply(matrix(temp1, ncol = nclass), 1, paste, 
        ##     collapse = "/")
      temp1 <- matrix(as.numeric(temp1), ncol=nclass)
      ##temp1 <- paste("p=", round(temp1[, 2] / apply(temp1, 1, sum)100, 1), "%", "; N=", apply(temp1, 1, sum), sep="")
      temp1 <- paste("", round(temp1[, 2] / apply(temp1, 1, sum)100, 1), "%", "; ", apply(temp1, 1, sum), sep="") 
    }
    if (use.n) {
        out <- paste(format(group, justify = "left"), "\n", temp1, 
            sep = "")
    }
    else {
        out <- format(group, justify = "left")
    }
    return(out)
}

rpartObj$functions$text <- textRpartCustom plot(rpartObj) text(rpartObj)

to get these information to display for a classification fit.

Leave a Comment