feat(state-space): exact discrete NB-branching latent with particle Gibbs (issue #48)

[seabbs-bot]

Note

Draft — agent-generated. Not for merge without human review.

Candidate redesign for issue #48: a state-space architecture in which the latent daily infection counts follow the exact discrete-time NB-overdispersed branching renewal step, with particle Gibbs as the inference route for the integer latent path and Mooncake-differentiable continuous nuisances for NUTS inside Gibbs.
Sibling to the LNA route in #101 and the deterministic renewal route in #103.

Architecture

• Latent state: a daily grid t = 1, …, n. Day 1 seeded from a small-cluster prior. Each subsequent day I_t | I_{<t}, R_t ~ NegativeBinomial(mean = R_t · Σ_s I_{t-s} g_s, dispersion = φ), the exact NB-overdispersed branching renewal step. φ → ∞ recovers the Poisson branching limit.
• Continuous nuisances: weekly piecewise-linear log-R_t with non-centred Gaussian-walk knots; Gamma generation interval (mean and SD priors); Gamma incubation, onset-to-death, onset-to-report and onset-to-detection delays (each with priors on parameters); CFR; surveillance dispersion (NB 1/√k prior); pooled DRC/Uganda ascertainment (non-centred, same as baseline); traveller volume; seed size; offspring dispersion φ.
• Onset staging: infections are convolved with the incubation PMF to get the daily onset incidence O_t. The four observation streams (deaths, reports, exports, deaths-among-exports) are built by convolving O_t with onset-to-X delays once and shared across streams.
• Observation likelihoods: Poisson for the small exports / export-deaths streams, NegBinomial(mean, k) for the large DRC streams, conditioned on cumulative totals to the cut-off. Any stream can be missing so the composer doubles as a prior-predictive generator. Genetic TMRCA enters as the same censored upper-tail soft lower bound the baseline uses.

Design rules

• Every prior is a submodel injected via a keyword argument. The composer accepts rt, gi, incubation, onset_to_death, onset_to_report, onset_to_detect, cfr, dispersion, ascertainment, traveller, seed, offspring, all overridable.
• All delays carry uncertainty, including the generation interval (Gamma mean and SD priors). No fixed delay constant in the model body.
• Observation distributions are swappable through the dispersion submodel.
• No inline using / import in src/state_space.jl; all imports stay on the module page.
• Non-centred parameterisation for the weekly R_t walk and the pooled ascertainment offsets, mirroring the baseline.
• The R_t walk uses a typed cumulative-sum allocation so Libtask's IR transformation handles it under PG; the abstract Vector{Real} allocation pattern that the renewal track uses is rejected by the PG transform.
• Delay discretisation by hand-rolled trapezoidal bins with density-at-zero forced to zero, matching the renewal track's AD-safe trick (CensoredDistributions' analytical primary-censored Gamma CDF is not Mooncake-differentiable; double-interval-censored PMF approximation bias is documented in the proposal).

Inference and runtime

The natural sampler is Gibbs(:I => PG(N), :nuisances => NUTS()) — particle Gibbs over the integer latent path, NUTS over the continuous nuisances with Mooncake gradients.

Demonstrated:

• Compiles end-to-end.
• Prior-predictive draw gives finite, integer latent path and finite expected counts on all four streams.
• Continuous-only Mooncake gradient (latent path conditioned) is finite and non-zero on every dimension.
• A PG(50) smoke test runs to completion in ~80s on the n = 60 day grid.

Honest runtime estimate for a full fit: roughly an order of magnitude or more above the LNA route, because per-Gibbs-sweep cost scales as O(N · n) likelihood evaluations against a single evaluation per NUTS leapfrog on the LNA. For n ≈ 130 d, N ≈ 200 particles, expect ~5–25 h per chain for 1000 post-warmup samples versus ~15–25 min for the LNA route.

PF-vs-LNA verdict

What the exact discrete latent buys in principle: heavier-than-Gaussian early-phase tails, integer realism at small counts, and one extra mass term in the posterior for minimum outbreak age consistent with an early-cluster fluctuation.

What it does not buy on the available data: the four primary observation streams are all cumulative aggregates, which by CLT are nearly Gaussian regardless of the latent's discreteness. The earliest-onset bound and the TMRCA bound only depend on T, not on path discreteness, so the LNA captures both faithfully. The PG path-mixing rate is also weakest at the start of the path — exactly where the architecture is meant to be most informative — because the observational anchor is weakest there.

Recommendation: do not pursue this as the default architecture for the current data. Keep #101 (LNA) as the production route while streams remain aggregate. The state-space route becomes worthwhile when a future data drop adds time-resolved early-phase records (e.g. dated first-onset trail, or sequenced cases with sampling dates that anchor the early branching) — at which point the path's discreteness starts to matter.

Identifiability risks

• Offspring φ, process noise, and surveillance dispersion k all add to observed-total spread and are weakly separated under cumulative-only data; priors carry the inference.
• Seed size I_1 and R_1 near-degenerate, same as T / r in the baseline.
• Particle ancestry collapse at the start of long latent paths under mild observation noise.

Files

• src/state_space.jl, test/test_state_space.jl, scripts/state_space.jl, docs/src/proposals/state-space-particle.md.
• docs/make.jl wires the proposal under a new "Redesign proposals" page group.

This was opened by a bot. Please ping @seabbs for any questions.

[github-actions]

📖 Documentation preview is ready!

View the docs for this PR at: http://epiforecasts.io/BVDOutbreakSize/previews/PR116/

This preview will be updated automatically when you push new commits.