Composable design

ComposableTuringIDModels treats an epidemiological model as a composition of independent parts. A full model has two top-level parts:

• an infection model generates unobserved infections . It owns a latent model internally — an unobserved process (e.g. a log reproduction number or growth rate) that the infection model maps to ;

• an observation model maps infections to the observed data .

A latent model describes an unobserved process over time. It is no longer a mandatory top-level component: the latent (e.g. ) is not always the estimand, so it is folded into the infection model that consumes it. This gives more flexibility — you choose an infection model and hand it whatever latent process you want to drive it — and decouples the generation interval so that only the Renewal model carries one.

Every part is a plain struct that implements a single method of the generic constructor as_turing_model. There is no deep type hierarchy: a part is identified by the method it implements, not by its place in a tree. This is the central design change from the package this one is adapted from, which used separate generate_latent, generate_latent_infs, and generate_observations functions dispatched over a layered abstract hierarchy.

One constructor, composed by submodels

as_turing_model(component, args...) returns a DynamicPPL.Model. A component that contains another component builds the inner model and samples it as a submodel:

z ~ to_submodel(as_turing_model(inner_model, n), false)

The trailing false disables automatic variable prefixing, so parameter names stay flat. Because every component speaks the same as_turing_model protocol, components nest freely: an AR process can carry a HierarchicalNormal error model, a DiffLatentModel can wrap that AR to produce an ARIMA-style process, and that whole latent process can be folded into a DirectInfections model observed with a NegativeBinomialError.

The latent is folded into the infection model

The latent process is supplied to the infection model rather than to the composer. An infection model takes a latent slot — Z for DirectInfections, rt for ExpGrowthRate and Renewal — and generates that process internally before mapping it to infections. So as_turing_model for an infection model takes only a series length and returns (; I_t, Z_t): the infection path and the internal latent draw, kept accessible as a generated quantity. Only Renewal needs a generation interval, so it alone carries an IDData; the others take a transformation directly.

Swap-in, swap-out

Because the parts share an interface, you change a modelling assumption by swapping one struct for another, leaving the rest untouched:

using ComposableTuringIDModels, Distributions

# An ARIMA-style latent process: a differenced AR.
latent = DiffLatentModel(; model = AR(), init_priors = [Normal(), Normal()])

# Fold the latent into a direct-infections process, then swap the observation
# model without touching the rest.
poisson_model = IDModel(
    DirectInfections(; Z = latent, initialisation_prior = Normal()),
    PoissonError())

negbin_model = IDModel(
    DirectInfections(; Z = latent, initialisation_prior = Normal()),
    NegativeBinomialError())

Inference

A composed model is an ordinary Turing model. Pass observed data instead of missing to condition it, then sample. We set the automatic-differentiation backend explicitly with NUTS(; adtype = ...): Mooncake is the recommended default for this package (see Automatic differentiation backend).

using Turing, Mooncake
using ADTypes: AutoMooncake
y = as_turing_model(poisson_model, fill(missing, 30), 30)().generated_y_t
posterior = as_turing_model(poisson_model, y, 30)
chain = sample(posterior, NUTS(; adtype = AutoMooncake(; config = nothing)), 1_000)

The standard Turing tools — rand for prior draws, fix to pin parameters, condition (or |) to condition on values, and sample for inference — all apply unchanged.