Extending EpiBranch

Sometimes the outbreak you want to model has a feature the built-in models and interventions lack: a reproduction number that falls once a gathering ban starts, a control measure the package does not include, a ward where patients can only infect the people in the next beds. This guide shows how to add such a feature in a few lines of Julia, in your own script, without editing the package.

If what you need is already built in (isolation, contact tracing, ring or mass vaccination), start with Isolation and contact tracing or Vaccination instead: using the built-in measures needs nothing beyond keyword arguments.

The recipes below assume a little Julia: writing a function, and the (rng, ind) -> ... form of an anonymous function. Julia for R users covers what you need.

Which tool do I need?

What you want to model	What you write	Where
R that changes over time, between people or once a policy starts; a cap on how many people one case can infect	a function in place of the offspring distribution	Change how many people each case infects
A generation time that depends on the case, such as one linked to its incubation period	a function returning a distribution	Generation time that depends on the case
Population characteristics the built-in builders do not set (risk group, region, a household's reporting rate)	a function of the random number generator and the individual	Population characteristics of your own
A rule of your own for when a simulation stops	a small type and one method	Stopping rules
A control measure that blocks some transmissions (a border closure, prophylaxis)	a type and a `competing_risk` method	A custom intervention: closing a border
Treatment that ends infections early, vaccination aimed at a group, a measure limited by a start date or by capacity	an intervention type with further methods	Writing an intervention
A rule of your own for which cases have their contacts traced	a small type and one `is_eligible` method	Who triggers contact tracing
A clinical event of your own (testing, treatment, loss to follow-up)	a clinical transition type	Custom clinical transitions
A latent period, several routes of transmission (community, household, funeral), seasonal transmission	infectiousness windows and routes	New transmission structures
A contact structure the package lacks (a ward, a school, a network you build), or a closed population with structured mixing	a transmission model type	New transmission structures
A new way cases are observed, or new data to fit	an observation model or a likelihood method	New transmission structures
The per-person values the package reserves, which interventions can be fitted exactly, numerical caveats	nothing; look things up	Extension reference

Words used in these pages

Writing a new piece of a model means using a few programming words. Each has a plain meaning here.

Word	Meaning in these pages
type (`struct`)	A named record holding the parameters of something you add, such as a border closure's leakage. `struct BorderClosure <: AbstractIntervention ... end` declares one; `<: AbstractIntervention` tells the package it is an intervention and where to use it.
method	One version of a package function, written for your type. `EpiBranch.competing_risk(bc::BorderClosure, ...) = ...` tells the package's `competing_risk` what a border closure does. Writing methods like this is how you extend EpiBranch: you never edit its source.
hook	A package function the simulation calls at a fixed step: when a person is created, before a case's contacts are drawn, when a transmission is about to happen. You write a method of it for your intervention. A hook you leave out leaves the simulation unchanged.
risk (`Risk`)	Something that can stop one transmission from a given time on, with a given probability. Isolation on day 5 is a risk with probability 1 from day 5.
state	Each simulated person (`ind`, an `Individual`) has a dictionary `ind.state` of named values, such as `ind.state[:age]` or `ind.state[:isolated]`. The whole outbreak so far is a `SimulationState`, usually called `state`.
model (`ModelSpec`)	A transmission model together with its natural history (`progression`), population characteristics (`attributes`), interventions and observation. This is what you pass to `simulate` and `loglikelihood`.
transmission model	Who can infect whom, and when: a branching process, a network, a household model.
generation-based and continuous-time models	Every built-in model is simulated in continuous time, one case at a time in order of infection time. A branching process creates each case's contacts as it goes; the network, household, homogeneous and structured-mixing (`MixingProcess`) models simulate a fixed population, in which everyone exists from the start and the simulation works out when each person is infected (by the Sellke construction). A model written against `generate_offspring` or `contacts_of` is simulated one generation at a time, and some hooks are called by only one kind.
contact interval (`kernel` in code)	The time from the start of a case's infectiousness to a contact that would infect if nothing intervened.
`!` at the end of a name	The function changes its argument in place, as `abort_infection!(ind, t)` changes `ind`.

Change how many people each case infects

Gathering limits, event-size caps and a reproduction number that changes over time all change how many people a case infects. Pass a function in place of the offspring distribution to BranchingProcess. The function receives the random number generator rng (pass it to every rand call so runs are reproducible) and the infector ind, and returns the number of secondary cases.

Here no case can infect more than five people, as under a limit on gathering size. NegBin(2.5, 0.16) has mean R = 2.5 and dispersion k = 0.16, and the generation time Exponential(5.0) has a mean of 5 days:

julia

using EpiBranch
using Distributions
using StableRNGs

capped_offspring(rng, ind) = min(rand(rng, NegBin(2.5, 0.16)), 5)

uncapped = simulate(BranchingProcess(NegBin(2.5, 0.16), Exponential(5.0)), 500;
    max_cases = 500, rng = StableRNG(1))
capped = simulate(BranchingProcess(capped_offspring, Exponential(5.0)), 500;
    max_cases = 500, rng = StableRNG(1))
(uncapped = containment_probability(uncapped),
 capped = containment_probability(capped))

(uncapped = 0.816, capped = 0.884)

The containment probability is the proportion of the 500 simulated outbreaks that ended before reaching 500 cases. Removing the largest superspreading events lets more outbreaks die out.

A function with a third argument, (rng, ind, state), can also read the outbreak so far. Here the cap only applies once there have been 20 cases, as a policy brought in partway through an outbreak would. c ? a : b is Julia for R's if (c) a else b:

julia

function policy_offspring(rng, ind, state)
    n = rand(rng, NegBin(2.5, 0.16))
    return state.cumulative_cases >= 20 ? min(n, 5) : n
end
model_policy = BranchingProcess(policy_offspring, Exponential(5.0))

BranchingProcess(offspring=Function, generation_time=Distributions.Exponential{Float64}, population_size=unlimited)

A reproduction number that changes over time reads the infector's infection time, ind.infection_time, in days. Here R falls from 3 to 1 over the first 50 days:

julia

r_at_time(t) = max(1.0, 3.0 - 2.0 * t / 50.0)
time_varying = (rng, ind) -> rand(rng, Poisson(r_at_time(ind.infection_time)))
model_rt = BranchingProcess(time_varying, Exponential(5.0))

BranchingProcess(offspring=Function, generation_time=Distributions.Exponential{Float64}, population_size=unlimited)

Use ind.infection_time when R follows each infector's own infection time, and state.max_infection_time (with the three-argument form) when R follows the outbreak's own clock.

R can also depend on who the infector is. Below it depends on a risk group set by a population characteristic (see Population characteristics of your own for how :risk_group is set):

julia

function risk_group!(rng, ind)
    if rand(rng) < 0.2
        ind.state[:risk_group] = :high
    else
        ind.state[:risk_group] = :low
    end
    return nothing
end

function risk_offspring(rng, ind)
    R = ind.state[:risk_group] == :high ? 4.0 : 1.5
    return rand(rng, Poisson(R))
end

model_risk = ModelSpec(BranchingProcess(risk_offspring, Exponential(5.0));
    attributes = risk_group!)
results = simulate(model_risk, 200; max_cases = 500, rng = StableRNG(42))
containment_probability(results)

0.32

A fifth of cases are high-risk, giving an average R of 0.2 × 4 + 0.8 × 1.5 = 2, and most outbreaks grow past 500 cases.

Or on the infector's generation, here for transmission that wanes as the outbreak goes on:

julia

function waning_offspring(rng, ind)
    R = 3.0 * exp(-0.1 * ind.generation)
    return rand(rng, Poisson(R))
end

model_waning = BranchingProcess(waning_offspring, Exponential(5.0))
results = simulate(model_waning, 200; max_cases = 500, rng = StableRNG(42))
containment_probability(results)

0.09

R drops below 1 after about 11 generations, by which time most outbreaks have already passed 500 cases.

Closed-form results need a distribution

A function works for simulation only. Extinction probability, chain-size distributions and the other analytical functions need an offspring distribution they can work with; see Offspring distributions of your own for writing one.

Generation time that depends on the case

generation_time can be one Distribution for everyone, or a function of the infector that returns a distribution. The function can read anything stored on the infector. A common use links the generation time to the case's own incubation period, read with incubation_period. LogNormal(1.5, 0.5) takes the mean and standard deviation of the log incubation period, in days, and Gamma(2.0, θ) has shape 2 and scale θ:

julia

gt = ind -> Gamma(2.0, incubation_period(ind) / 2)
linked = ModelSpec(BranchingProcess(NegBin(2.5, 0.16), gt);
    attributes = clinical_presentation(incubation_period = LogNormal(1.5, 0.5)))

state = simulate(linked; max_cases = 500, rng = StableRNG(42))
state.cumulative_cases

Each case's incubation period is drawn once and its generation time is built from it. The two are therefore correlated: a case with a late onset also tends to transmit late. incubation_linked_generation_time is a ready-made version, the skew-normal model of Hellewell et al. (2020).

The generation time can depend on any value a population characteristic has stored, not only the incubation period. Here one per-person draw sets both the onset time and the mean generation time, in days:

julia

function host!(rng, ind)
    scale = 3.0 + rand(rng)
    ind.state[:gt_scale] = scale
    ind.state[:onset_time] = ind.infection_time + scale
    return nothing
end

scaled = ModelSpec(
    BranchingProcess(Poisson(2.0), ind -> Exponential(ind.state[:gt_scale]));
    attributes = host!)

state = simulate(scaled; max_cases = 500, rng = StableRNG(42))
state.cumulative_cases

Use this whenever generation time and onset should come from one per-person draw instead of two independent ones.

Population characteristics of your own

The attributes argument of a model is a function (rng, ind) -> ... that sets values on each person when they are created, before any intervention acts. The built-in builders clinical_presentation, demographics and transmission_traits return such functions. For any other value, write your own, as risk_group! above does.

Characteristics can feed each other. Here transmission_traits sets each person's susceptibility (a number between 0 and 1) from the risk group:

julia

attrs = [
    risk_group!,
    transmission_traits(
        susceptibility = (rng, ind) -> ind.state[:risk_group] == :high ? 0.8 : 0.3,
    ),
]

model = ModelSpec(BranchingProcess(NegBin(2.5, 0.16), Exponential(5.0)); attributes = attrs)
runs = simulate(model, 200; max_cases = 100, rng = StableRNG(42))
people = reduce(vcat, [s.individuals for s in runs])
high_share(group) = count(ind -> ind.state[:risk_group] == :high, group) / length(group)
(all_contacts = high_share(people), infected = high_share(filter(is_infected, people)))

(all_contacts = 0.17401285583103765, infected = 0.3252375923970433)

The output keeps contacts who were exposed but not infected. About a fifth of all contacts are high-risk, but because they are more susceptible they make up a larger share of those infected (0.2 × 0.8 / (0.2 × 0.8 + 0.8 × 0.3) = 0.4, ignoring the index cases). The list is applied in order, so a later entry can read what an earlier one set. Built-in builders and your own functions mix freely:

julia

combined = [
    clinical_presentation(incubation_period = LogNormal(1.5, 0.5)),
    demographics(age_distribution = Normal(40, 15)),
    risk_group!,
    transmission_traits(
        susceptibility = (rng, ind) -> ind.state[:risk_group] == :high ? 0.8 : 0.3,
    ),
]

model_combined = ModelSpec(BranchingProcess(NegBin(2.5, 0.16), Exponential(5.0));
    attributes = combined)
state = simulate(model_combined; max_cases = 100, rng = StableRNG(42))
ind = state.individuals[1]
(age = ind.state[:age], sex = ind.state[:sex], risk = ind.state[:risk_group])

(age = 60, sex = :male, risk = :high)

group_attribute gives every member of a household, community or other group the same value. It draws once for the first member of each group and keeps the value for the run. Here each household has its own reporting probability, drawn from Beta(6, 4) (mean 0.6):

julia

reporting_attributes = [groups(50; key = :household),
    group_attribute(:reporting_probability; value = Beta(6, 4),
        group_key = :household)]
reporting_model = ModelSpec(BranchingProcess(Poisson(0.5), Exponential(5.0));
    attributes = reporting_attributes,
    progression = [Reporting(
        delay = Dirac(0.0),
        probability = (rng, ind) -> ind.state[:reporting_probability],
        from = :infection)])
reporting_state = simulate(reporting_model; n_initial = 10, rng = StableRNG(42))

SimulationState(cases=18, individuals=18, gen=2, extinct)

Set the group label before the shared value, as the order above does. Reusing these builders in further simulations draws fresh values, including when simulations run in parallel.

Rules with parameters of their own

The probabilities and delays of clinical transitions, the parameters of population characteristics and the conditions of interventions accept a type with parameters wherever they accept a function, called with the same arguments. A reporting rule can then keep its age threshold as a parameter. (rule::AgeDetection)(rng, ind) = ... makes an AgeDetection usable as a function:

julia

struct AgeDetection
    minimum_age::Float64
end
(rule::AgeDetection)(rng, ind) = ind.state[:age] >= rule.minimum_age ? 1.0 : 0.0
age_reporting = Reporting(delay = Dirac(0.0), probability = AgeDetection(50.0))

Reporting{Distributions.Dirac{Float64}, Main.AgeDetection, Symbol}(Distributions.Dirac{Float64}(value=0.0), Main.AgeDetection(50.0), :onset_time)

A reference time given to Reporting as a function (from) takes only ind, and a condition for Scheduled takes the simulation state. Numbers and distributions keep their usual meaning wherever those are accepted.

Stopping rules

A simulation stops at the first step where any of its stopping rules says so. The built-in rules Extinction, MaxCases, MaxGenerations and MaxTime cover most needs, and the max_cases, max_generations and max_time keywords build them (see SimOpts). A rule of your own is the smallest type you will write: a struct subtyping AbstractStoppingRule, and a should_stop method that returns true when the simulation should end.

Here a rule stops once any chain of transmission reaches a given number of generations, whatever the case count:

julia

struct MaxChainLength <: AbstractStoppingRule
    n::Int
end
EpiBranch.should_stop(r::MaxChainLength, state::SimulationState) =
    maximum(ind.generation for ind in state.individuals; init = 0) >= r.n

chain_model = BranchingProcess(NegBin(2.5, 0.16), Exponential(5.0))
chain_state = simulate(
    chain_model; stopping_rules = [Extinction(), MaxChainLength(5)], rng = StableRNG(3)
)
maximum(ind.generation for ind in chain_state.individuals; init = 0)

The printed value is the last generation reached: 5 if the rule stopped the run, fewer if transmission died out first. init = 0 gives the maximum a value when there are no individuals.

Extinction is added automatically unless your stopping_rules already include it. A simulation therefore still ends when transmission dies out: MaxChainLength alone would never stop a chain that goes extinct before 5 generations.

Write the method as EpiBranch.should_stop(...), with the EpiBranch. prefix (or after import EpiBranch: should_stop). Without it, Julia creates a new function of your own called should_stop that the simulation never calls. Type the state argument as state::SimulationState, as above, so your method does not clash with the package's default one.

Stopping rules on continuous-time models

The network, household, homogeneous and mixing models run a continuous-time race rather than stepping by generation. A rule that should end such a run at a time or a generation depth defines EpiBranch.time_bound(rule) or EpiBranch.generation_bound(rule), as MaxTime and MaxGenerations do; the race applies the bound directly, dropping a case's own proposals past a generation bound rather than waiting to be asked, since cases there settle in time order rather than generation order. If one of these bounds is all the rule tests, it also declares EpiBranch.honoured_without_should_stop(rule) = true; otherwise the race polls should_stop after each settled case to apply the rest of it, as it does for MaxCases and any rule of your own with no bound at all. The max_cases and max_generations keyword defaults have no effect on these models, so an existing run over a large population is not capped unless you ask for it with an explicit value or a stopping_rules entry.

A custom intervention: closing a border

An intervention that stops some transmissions needs one type and one method of competing_risk. The simulation asks every intervention, for each pair of infector and contact, whether anything blocks that transmission. Whatever blocks it first wins (survival analysts call these competing risks).

Here two regions close the border between them on day 10. After that, a transmission between people in different regions goes ahead only with a small probability, the leakage.

julia

using EpiBranch
using Distributions
using StableRNGs

struct BorderClosure <: AbstractIntervention
    start_time::Float64   # day the border closes
    leakage::Float64      # proportion of cross-border transmissions still happening
end

function EpiBranch.competing_risk(bc::BorderClosure, parent, contact, state)
    if parent.state[:region] == contact.state[:region]
        return nothing    # same region: the closure does not apply
    end
    return Risk(event_time = bc.start_time, block_probability = 1.0 - bc.leakage)
end

parent is the infector and contact the person who would be infected. Returning nothing means the intervention does not block this transmission. Risk(event_time = 10.0, block_probability = 0.95) blocks a transmission happening on or after day 10 with probability 0.95; transmissions before the closure are unaffected.

Each person needs a region. In this example people are equally likely to live in either, and about half of all contacts cross the border:

julia

region! = (rng, ind) -> (ind.state[:region] = rand(rng, (:north, :south)))

process = BranchingProcess(NegBin(1.6, 0.5), Gamma(2.0, 2.5))
open_border = ModelSpec(process; attributes = region!)
closed_border = ModelSpec(process; attributes = region!,
    interventions = [BorderClosure(10.0, 0.05)])

open_runs = simulate(open_border, 500; max_cases = 500, rng = StableRNG(1))
closed_runs = simulate(closed_border, 500; max_cases = 500, rng = StableRNG(1))
(open = containment_probability(open_runs),
 closed = containment_probability(closed_runs))

(open = 0.76, closed = 1.0)

With R = 1.6 and the border open, a sizeable share of outbreaks grow past 500 cases. After the closure, cross-border transmission falls by 95%. The effective reproduction number is then about 1.6 × (0.5 + 0.5 × 0.05) ≈ 0.84 and nearly every outbreak dies out.

Your intervention combines with the built-in ones: pass interventions = [Isolation(...), BorderClosure(10.0, 0.05)] and both apply. The analytical functions, such as extinction_probability, stop with an error on closed_border, because the closure has no exact effect on the offspring distribution for them to use; see All hooks for the rare intervention that has one.

Check that your intervention acts

Building a ModelSpec warns when your intervention has no method of any hook, when a method takes a number of arguments no call uses, or when it was written without the EpiBranch. prefix. It cannot tell an intervention that leaves out a hook on purpose from one missing a hook it needs, so compare a small simulation with and without your intervention, as above, before relying on it.

Where next

Writing an intervention: every hook an intervention can use, the order in which they are called, ending an infection early, scheduling, capacity limits, custom vaccinations.
New transmission structures: latent periods and infectiousness windows, several routes of transmission, seasonal transmission, new transmission models and observation models.
Extension reference: reserved per-person values, which interventions the likelihoods can fit exactly, and details of the continuous-time models.

Extending EpiBranch ​

Which tool do I need? ​

Words used in these pages ​

Change how many people each case infects ​

Generation time that depends on the case ​

Population characteristics of your own ​

Sharing a value within groups ​

Rules with parameters of their own ​

Stopping rules ​

A custom intervention: closing a border ​

Where next ​

Extending EpiBranch

Which tool do I need?

Words used in these pages

Change how many people each case infects

Generation time that depends on the case

Population characteristics of your own

Sharing a value within groups

Rules with parameters of their own

Stopping rules

A custom intervention: closing a border

Where next