simulate branching process trees with unobserved nodes/ancestors

bp_trees.r

library('tidyverse')
library('bpmodels')
library('epicontacts')

obs_tree <- function(obs_prob=0.5, stop_threshold=100, ...) {
    ## generate chains
    chains <- chain_sim(infinite=stop_threshold, tree=T, ...)

    observed <- chains %>%
        mutate(observed=factor(rbinom(n=n(), size=1, prob=obs_prob),
                               levels=c(0, 1),
                               labels=c("unobserved", "observed"))) %>%
        select(n, id, observed)

    observed_ancestor <- observed %>%
        select(n, ancestor=id, observed_ancestor=observed)

    ochains <- chains %>%
        left_join(observed, by=c("n", "id")) %>%
        left_join(observed_ancestor, by=c("n", "ancestor")) %>%
        replace_na(list(observed_ancestor="unobserved"))

    return(ochains)
}

## number of outbreaks to simulate
n <- 100
## observation probability
obs_prob <- 0.7
## R
R <- 1.3
## overdispersion
k <- 0.5
## maximum outbreak size
max_size <- 200

## negative binomial offspring
trees <- obs_tree(n=n, obs_prob=obs_prob, stop_threshold=max_size, mu=R, offspring=rnbinom, size=1/k)

## filter outbreaks that didn't go extinct
big_trees <- trees %>%
    group_by(n) %>%
    mutate(max_id=max(id)) %>%
    ungroup %>%
    filter(max_id >= max_size)

## proportion observed and with ancestor observed
big_trees %>%
    filter(observed=="observed") %>%
    group_by(n) %>%
    summarise(prop_anc_ob=mean(observed_ancestor=="observed")) %>%
    ungroup %>%
    summarise(mean=mean(prop_anc_ob)) %>%
    .$mean

## plot a tree
rand_n <- sample(big_trees$n, 1)
rand_tree <- big_trees %>%
    filter(n==rand_n)

linelist <- rand_tree %>%
    mutate(observed=if_else(is.na(ancestor), "index", as.character(observed))) %>%
    select(id, observed)
contacts <- rand_tree %>%
    filter(!is.na(ancestor)) %>%
    select(from=ancestor, to=id)
ec <- make_epicontacts(linelist, contacts)

plot(ec, node_color="observed")