Chain statistics, likelihood, and fitting

Compute chain size and length statistics, evaluate likelihoods, and fit offspring distributions — analytically where possible, simulation-based otherwise.

Chain statistics

using EpiBranch
using Distributions
using DataFrames
using StableRNGs

model = BranchingProcess(Poisson(0.9), Exponential(5.0))

rng = StableRNG(42)
state = simulate(model; n_initial = 20, max_cases = 10_000, rng = rng)

cs = chain_statistics(state)
first(cs, 10)

10×3 DataFrame

Row	chain_id	size	length
	Int64	Int64	Int64
1	1	1	0
2	2	6	4
3	3	2	1
4	4	1	0
5	5	1	0
6	6	47	19
7	7	1	0
8	8	14	6
9	9	2	1
10	10	1	0

println("Mean size: $(round(mean(cs.size), digits=2)), Max: $(maximum(cs.size))")
println("Mean length: $(round(mean(cs.length), digits=2))")

Mean size: 6.75, Max: 47
Mean length: 2.6

Analytical chain size distributions

For certain offspring distributions, the chain size distribution has a closed form:

# Poisson → Borel
d_borel = chain_size_distribution(Poisson(0.8))
println("Borel(0.8) mean: $(round(mean(d_borel), digits=2))")

# NegBin → Lagrange inversion formula
d_nb = chain_size_distribution(NegBin(0.8, 0.5))
println("NegBin chain size P(1): $(round(pdf(d_nb, 1), digits=4))")

Borel(0.8) mean: 5.0
NegBin chain size P(1): 0.6202

Likelihood

Given observed chain sizes, evaluate the log-likelihood using loglikelihood with data wrapped in ChainSizes:

data = ChainSizes([1, 1, 2, 1, 3, 1, 1, 5, 1, 2])

# Compare offspring parameters
for λ in [0.3, 0.5, 0.7, 0.9]
    ll = loglikelihood(data, Poisson(λ))
    println("Poisson($λ): LL = $(round(ll, digits=2))")
end

Poisson(0.3): LL = -12.98
Poisson(0.5): LL = -12.49
Poisson(0.7): LL = -13.4
Poisson(0.9): LL = -14.99

With Negative Binomial offspring:

ll = loglikelihood(data, NegBin(0.9, 0.5))
println("NegBin(R=0.9, k=0.5): LL = $(round(ll, digits=2))")

NegBin(R=0.9, k=0.5): LL = -14.07

Imperfect observation

Account for incomplete case ascertainment by attaching a PerCaseObservation to the process with a ModelSpec: pass observation = … to the wrapper. Each case is detected independently with the given probability:

data = ChainSizes([1, 1, 2, 1, 3, 1, 1, 5, 1, 2])
full = ModelSpec(BranchingProcess(Poisson(0.9)); observation = PerCaseObservation(detection_prob = 1.0))
partial = ModelSpec(BranchingProcess(Poisson(0.9)); observation = PerCaseObservation(detection_prob = 0.7))
println("Full observation:  $(round(loglikelihood(data, full), digits=2))")
println("70% observation:   $(round(loglikelihood(data, partial), digits=2))")

Full observation:  -14.99
70% observation:   -16.16

Offspring counts

If you have direct observations of secondary case counts (who infected whom), wrap them in OffspringCounts:

offspring_data = OffspringCounts([0, 1, 2, 0, 3, 1, 0, 2, 5, 0])
ll = loglikelihood(offspring_data, Poisson(1.4))
println("Offspring LL: $(round(ll, digits=2))")

Offspring LL: -17.25

Chain lengths

length_data = ChainLengths([0, 1, 0, 2, 1, 0, 0, 3, 0, 1])
ll = loglikelihood(length_data, Poisson(0.5))
println("Chain length LL: $(round(ll, digits=2))")

Chain length LL: -12.62

Simulation-based likelihood

For models with interventions, use the simulation-based likelihood by passing a ModelSpec instead of a distribution. Because the interventions and attributes are composed onto the process by the spec, the likelihood is evaluated under exactly the process that produced the data, with nothing to pass twice:

iso = Isolation(onset_to_isolation_delay = Exponential(2.0))

model = ModelSpec(BranchingProcess(Poisson(2.0), Exponential(5.0));
    interventions = [iso],
    attributes = clinical_presentation(incubation_period = LogNormal(1.5, 0.5)),
)

ll = loglikelihood(ChainSizes([1, 1, 2, 1, 3, 1, 1, 5, 1, 2]), model;
    n_sim = 1000,
    rng = StableRNG(42),
)
println("LL with interventions: $(round(ll, digits=2))")

LL with interventions: -48.05

Forward simulation and likelihood evaluation share the same code path. Likelihoods can therefore be evaluated under interventions — something not possible when simulation and inference are in separate packages — and the same model drives both simulate(model) and loglikelihood.

Fitting

The loglikelihood interface is the backend for both maximum-likelihood and Bayesian fitting. For an MLE, maximise loglikelihood over the parameter — with Optim.jl, Turing's maximum_likelihood, or, for a single parameter, a simple grid:

rng = StableRNG(42)
truth = BranchingProcess(Poisson(0.5), Exponential(5.0))
states = simulate(truth, 500; rng = rng)
sizes = Int[]
for s in states
    cs = chain_statistics(s)
    append!(sizes, cs.size)
end
data = ChainSizes(sizes)

Rgrid = 0.05:0.01:0.95
R̂ = Rgrid[argmax([loglikelihood(data, Poisson(R)) for R in Rgrid])]
println("Poisson MLE from chain sizes: R = $(round(R̂, digits=2))")

┌ Warning: Assignment to `cs` in soft scope is ambiguous because a global variable by the same name exists: `cs` will be treated as a new local. Disambiguate by using `local cs` to suppress this warning or `global cs` to assign to the existing global variable.
└ @ chains.md:142
Poisson MLE from chain sizes: R = 0.48

Bayesian inference with Turing.jl

Call chain_size_distribution (or chain_length_distribution, offspring_distribution) on the model and the data sits directly on the right-hand side of ~:

using Turing

@model function chain_model(data)
    R ~ LogNormal(-0.5, 1.0)
    data ~ chain_size_distribution(BranchingProcess(Poisson(R)))
end

For a bare offspring law these return the analytical distribution (Borel, GammaBorel, …) where one exists; when the spec composes interventions onto the process (or you pass seeds, prob_concluded, n_sim, …) they return a wrapper that routes through the simulation-based loglikelihood. Either way the spec carries its own interventions and attributes, so the ~ line stays clean.

data ~ chain_size_distribution(BranchingProcess(Poisson(R));
    seeds = seeds, prob_concluded = prob_concluded)