A visualization of (discrete time) branching process

branching_process.R

library(dplyr)
library(ggplot2)
###
branching_process <- function(theta, gamma, G){
  gen <- vector("list", G)
  K <- rpois(1, theta)
  id <- 1L:K
  parent <- 1L
  gen[[1]] <- data.frame(id = id,
                         parent = parent,
                         g = 1L)
  parent <- id
  next_id <- K
  for(i in 2L:G){
    K <- rpois(length(parent), gamma)
    sumK <- sum(K)
    if(sumK == 0){ break }
    parent <- rep(parent, K)
    id <- (next_id+1L):(next_id+sumK)
    gen[[i]] <- data.frame(id = id,
                           parent = parent,
                           g = i)
    parent <- id
    next_id <- next_id+sumK
  }
  dplyr::bind_rows(gen)
}

theta <- 5
gamma <- 0.5
G <- 5
set.seed(12345)
df <- branching_process(theta, gamma, G)

ggplot(df)+
  geom_segment(aes(yend=g, xend=id, x = parent, y=g-0.95), 
               arrow = arrow(length = unit(0.2,"cm")))+
  scale_y_reverse()+
  labs(x = "node id", y = "generation")+
  theme_classic()