ComposableTuringIDModels

Early development

This package is in early development. Expect rough edges and breaking changes as the design settles.

ComposableTuringIDModels builds epidemiological models from small, reusable components — infection processes (each owning its own latent process) and observation models — and turns each one into a Turing / DynamicPPL model through the single generic constructor as_turing_model. Components compose by sampling one another as submodels, so a full model is assembled from parts rather than written by hand.

Getting started

using ComposableTuringIDModels, Distributions

# Compose infections -> observations. The infection model owns its latent
# process (here a random walk in its `Z` slot); the latent is folded in rather
# than threaded as a separate top-level component.
model = IDModel(
    DirectInfections(; Z = RandomWalk(), initialisation_prior = Normal()),
    PoissonError())

# Build a Turing model; `missing` observations simulate from the prior.
turing_model = as_turing_model(model, missing, 20)
draw = rand(turing_model)

The composed model exposes its generated quantities directly:

(; generated_y_t, I_t, Z_t) = turing_model()
length(generated_y_t), length(I_t), length(Z_t)

(20, 20, 20)

See Composable design for how the pieces fit together, the case studies for worked end-to-end examples, and the Public API for the full surface.

Automatic differentiation backend

Hamiltonian samplers such as NUTS need gradients of the log density, computed by automatic differentiation (AD). The recommended default for this package is Mooncake, a reverse-mode backend that handles the DynamicPPL @model + to_submodel composition used throughout well. Select it by passing an adtype to the sampler:

using Turing, Mooncake
using ADTypes: AutoMooncake
chain = sample(model, NUTS(; adtype = AutoMooncake(; config = nothing)), 1_000)

The examples in Composable design and the renewal/delay case studies all use Mooncake. Other backends — ForwardDiff, ReverseDiff, and Enzyme — are supported and exercised by the package's AD test suite; swap them in with AutoForwardDiff(), AutoReverseDiff(), or AutoEnzyme(). ForwardDiff is the right choice for the ODE models in the SIR case study: Mooncake-driven NUTS through an ODE solver is not available yet (see that page).

Adapted from

The modelling code in this package is ported and adapted from the open-source, Apache-2.0 licensed EpiAware package (CDCgov/Rt-without-renewal, ported from the fork seabbs/Rt-without-renewal). ComposableTuringIDModels is a modified, derived work: renamed, re-architected around as_turing_model, and upgraded to the latest Turing. See the NOTICE file for attribution and the list of changes.