A visualization of multi-type discrete-time branching process

multitype_branching_process.R

library(ggplot2)
library(dplyr)
multitype_branching_process <- function(M, G, initial_type = 1L) {
  # M[i, j] = type i の個体1個体あたりの
  #          type j の平均子孫数
  K <- nrow(M)
  if (ncol(M) != K) {
    stop("M must be a square matrix.")
  }
  if (any(M < 0)) {
    stop("All elements of M must be non-negative.")
  }
  if (initial_type < 1 || initial_type > K) {
    stop("Invalid initial_type.")
  }
  
  type_names <- rownames(M)
  if (is.null(type_names)) {
    type_names <- paste0("type", seq_len(K))
  }
  
  gen <- vector("list", G)
  
  # 個体ID
  next_id <- 0L
  
  # 第1世代
  next_id <- next_id + 1L
  
  gen[[1]] <- data.frame(
    id = next_id,
    parent = 0L,
    g = 1L,
    type = initial_type
  )
  
  # 現在の世代の個体
  parent_id <- next_id
  parent_type <- initial_type
  
  # 第2世代以降
  for (g in 2L:G) {
    
    child_list <- vector("list", length(parent_id))
    
    for (k in seq_along(parent_id)) {
      
      i <- parent_type[k]
      
      # 親タイプ i から各タイプの子供数
      n_child <- rpois(K, M[i, ])
      
      if (sum(n_child) > 0) {
        
        child_type <- rep(seq_len(K), n_child)
        child_id <- (next_id + 1L):(next_id + length(child_type))
        
        child_list[[k]] <- data.frame(
          id = child_id,
          parent = parent_id[k],
          g = g,
          type = child_type
        )
        
        next_id <- next_id + length(child_type)
      }
    }
    
    children <- dplyr::bind_rows(child_list)
    
    if (nrow(children) == 0) {
      break
    }
    
    gen[[g]] <- children
    
    parent_id <- children$id
    parent_type <- children$type
  }
  
  result <- dplyr::bind_rows(gen)
  
  result$type_name <- type_names[result$type]
  
  result
}

M <- matrix(
  c(
    1.2, 0.3, 0.1,
    0.2, 0.8, 0.4,
    0.1, 0.2, 1.0
  ),
  nrow = 3,
  byrow = TRUE,
  dimnames = list(
    c("A", "B", "C"),
    c("A", "B", "C")
  )
)

set.seed(123)
tree <- multitype_branching_process(
  M = M,
  G = 5,
  initial_type = 1
)

ggplot(tree)+
  geom_point(aes(x=id, y=g, colour=factor(type), shape=factor(type)), size=3)+
  geom_segment(aes(yend=g, xend=id, x = parent, y=g-0.95, colour=factor(type)), 
               arrow = arrow(length = unit(0.2,"cm")))+
  scale_y_reverse()+
  scale_shape_manual(values = c(1,2,5))+
  labs(x = "node id", y = "generation", colour="type", shape="type")+
  theme_classic()

ggsave("tree.png", width = 7, height = 7)