Creating a branching process in R

Question

I am trying to simulate a branching process (or Galton-Watson process) with a negative binomial offspring distribution. An example of a branching process is represented below, taken somewhat arbitrarily from Cui, et al.

Briefly, in the context of infectious disease, each branching process starts with one infectious individual (generation 0) and they infect a random number of secondary cases defined by some offspring distribution (I am using a negative binomial offspring with parameters (mean=r0, dispersion=k)). Each case in subsequent generations infects a random number of cases, and so on (i.i.d.).

I imagine the data layout to be a matrix, similar to this this, which represents the first three generations in the figure:

#Example of data layout from first three generations in figure
x<-matrix(c(0,1,1,2,2,2,1,1,2,1,2,3,2,2,1,1,0,3),ncol = 3,nrow = 6)
colnames(x)<-c("Generation","Case Number","Secondary Cases (Z)")

I am unsure how to create a function that teases out the number of individual Z values. I have written a function that simulates a negative binomial branching process, but only sums the total number of cases in each generation (i.e. in the image below, G(0)=1, G(1)=2, G(2)=3, G(3)=4). Furthermore, in this function I have to define the number of generations ("n"), and a more accurate code would allow it to continue naturally until there are 0 cases in a generation.

The below functions provide some benefit to what I am trying to do, but have limited functionality.

#Single branching process
nbbp<-function(n, r0, k){
  # initialize return vector
  Z<-integer(n)
  #Initiate branching process with 1 index case in generation 1
  Z[1]<-1
  for (i in seq_len(n)[-1]){
    if(Z[i-1]>0) {
         x<-rnbinom(Z[i-1], size=k, mu=r0)
         Z[i] <- sum(x)
    }
  }
  return(Z)
}
#Simulate multiple BP 
nbbp.ind<-function(num_sim,n,r0,k){
    x<-replicate(num_sim, nbbp(n,r0,k))
    }

Answer 1

A generalized branching process function that can accept any distribution (distr) with an arbitrary number of parameters can be written as:

bp <- function(distr, gens = 20, init.size = 1,  ...){  
  Z <- list() 
  Z[[1]] <- init.size 
  i <- 1 
  while(sum(Z[[i]]) > 0 && i <= gens) { 
    Z[[i+1]] <- distr(sum(Z[[i]]), ...) 
    i <- i+1 
  } 
  return(Z)
} 

Note gens sets the maximum number of generations to prevent overflow (can be set to Inf if desired), and init.size allows you to change the initial number of events in generation 0.

This returns a list of vectors, where each position position in the list represents a generation, and each position in the vector represents an offspring:

set.seed(12345)

# Negative binomial distribution
bp(distr = rnbinom, mu = 0.9, size = 0.25, gen = Inf, init.size = 1)
# [[1]]
# [1] 1
# 
# [[2]]
# [1] 1
# 
# [[3]]
# [1] 0

# Poisson distiribution
bp(distr = rpois, lambda = 0.9, gen = Inf, init.size = 1)

# [[1]]
# [1] 1
# 
# [[2]]
# [1] 1
# 
# [[3]]
# [1] 4
# 
# [[4]]
# [1] 0 0 1 0
# 
# [[5]]
# [1] 0