Visualization of continuous-time branching processes

branching_process.R

library(dplyr)
library(ggplot2)

simulate_branching <- function(mu,
                               lambda,
                               Tmax) {
  
  N0 <- rpois(1, Tmax*mu)
  t <- sort(runif(N0, 0, Tmax))
  individuals <- data.frame(
    id = 1L:N0,
    parent = 0L,
    birth_time = t,
    generation = 1L
  )
  
  next_id <- N0 + 1L
  i <- 1L
  while (i <= nrow(individuals)) {
    birth <- individuals$birth_time[i]
    t <- birth
    repeat {
      t <- t + rexp(1, rate = lambda)
      if (t > Tmax){
        break
      }
      individuals <- rbind(
        individuals,
        data.frame(
          id = next_id,
          parent = individuals$id[i],
          birth_time = t,
          generation = individuals$generation[i] + 1L
        )
      )
      next_id <- next_id + 1L
    }
    i <- i + 1L
  }
  
  individuals <- individuals[order(individuals$birth_time), ]
  return(individuals)
}

set.seed(0805)
tree <- simulate_branching(
  mu = 1,
  lambda = 0.6,
  Tmax = 5
)
tree

segments <- dplyr::filter(tree, parent > 0) |> 
  rowwise() |> 
  mutate(
    x    = tree$birth_time[tree$id == parent],
    y    = tree$generation[tree$id == parent],
    xend = birth_time,
    yend = generation - 0.05
  )

ggplot(tree) +
  geom_segment(
    data = segments,
    aes(x = x, y = y, xend = xend, yend = yend),
    arrow = arrow(length = unit(0.02, "npc")),
    colour="darkgrey") +
  geom_point(aes(birth_time, generation), size=3, shape=1) +
  geom_rug(aes(x=birth_time))+
  scale_y_reverse() +
    labs(x = "time", y="generation") +
    theme_classic()