calculate a decision tree with data.tree, and plot it with DiagrammeR / mermaid

decisionMermaid.R

### This demo calculates and plots a simple decision tree
### It requires the yaml and the ape packages to be installed

library(data.tree)
library(yaml)


#load from file

fileName <- 'jennylind.yaml'
l <- yaml.load_file(fileName)
jl <- as.Node(l)
print(jl, "type", "payoff", "p")

#calculate decision tree

payoff <- function(x) {
  if (x$type == 'chance') x$payoff <- Aggregate(x, function(node) node$payoff * node$p, sum)
  else if (x$type == 'decision') x$payoff <- Aggregate(x, "payoff", max)
}

jl$Do(payoff, traversal = "post-order", filterFun = isNotLeaf)

decision <- function(x) {
  po <- sapply(x$children, function(child) child$payoff)
  x$decision <- names(po[po == x$payoff])
}

jl$Do(decision, filterFun = function(x) x$type == 'decision')


#plot the same with DiagrammeR / mermaid
library("DiagrammeR")

jl$Set(id = letters[1:(jl$totalCount)])

FromLabel <- function(node) {
  if(node$parent$isRoot) return (ToLabel(node$parent))
  return (as.character(node$parent$id))
}

EdgeLabel <- function(node) {
  if (node$type == "decision") {
    return ('')
  } else if (node$type == "chance") {
    lbl <- node$name
  } else if (node$type == "terminal") {
    lbl <- paste0(node$name,": ", node$p)
  }
  lbl <- paste0(" --> |", lbl, "|")
  return (lbl)
}

FormatPayoff <- function(payoff) {
  paste0("$", payoff/1000, "k")
}

ToLabel <- function(node) {
  if (node$type == "decision") {
    lbl <- paste0("[", FormatPayoff(node$payoff), "]")
  } else if (node$type == "chance") {
    lbl <- paste0("((", FormatPayoff(node$payoff), "))")
  } else if (node$type == "terminal") {
    lbl <- paste0("[", FormatPayoff(node$payoff), "]")
  }
  lbl <- paste0(" ", node$id, lbl)
  return (lbl)
}


format <- paste0(
  "classDef default fill:none, bg:none, stroke-width:0px;
  classDef decision fill:#9f6,stroke:#333,stroke-width:1px;
  classDef chance fill:red,stroke:#333,stroke-width:1px;
  class ", paste(jl$Get("id", filterFun = function(x) x$type == "decision"), collapse = ","), " decision;
  class ", paste(jl$Get("id", filterFun = function(x) x$type == "chance"), collapse = ","), " chance;")
format <- ""
t <- Traverse(jl, traversal = "level", filterFun = isNotRoot)
df <- data.frame(from = Get(t, FromLabel), edge = Get(t, EdgeLabel), to = Get(t, ToLabel))

diagram <- paste("graph LR", 
                 paste( paste0(df$from, 
                               df$edge, 
                               df$to), 
                        collapse = "\n"), 
                 format, sep = "\n")

DiagrammeR(diagram)

jennylind.yaml

name: Jenny Lind
type: decision
Sign with Movie Company:
  type: chance
  Small Box Office:
    type: terminal
    p: 0.3
    payoff: 200000
  Medium Box Office:
    type: terminal
    p: 0.6
    payoff: 1000000
  Large Box Office:
    type: terminal
    p: 0.1
    payoff: 3000000
Sign with TV Network:
  type: chance
  Small Box Office:
    type: terminal
    p: 0.3
    payoff: 900000
  Medium Box Office:
    type: terminal
    p: 0.6
    payoff: 900000
  Large Box Office:
    type: terminal
    p: 0.1
    payoff: 900000