ComposableTuringIDModels.jl

A toolkit for composable probabilistic infectious disease modelling in Julia.

Why ComposableTuringIDModels?

• Composable models: Assemble a model from interchangeable infection and observation parts — each infection model owning its own latent process — instead of writing one monolithic model.

• Swap a part to change an assumption: Change the latent process, infection process, or observation model on its own to test how each assumption shapes your conclusions.

• One interface: Every part becomes a Turing / DynamicPPL model through the single as_turing_model constructor, so parts nest freely as submodels. The full Turing inference toolbox (NUTS, prior simulation) applies.

• Simulate and infer: Generate synthetic data from any model, then run Bayesian inference on the same model with real data.

• A library of parts: Random walks, AR/MA/ARIMA latent processes, renewal and exponential-growth infection models, ODE (SIR/SEIR) processes, Poisson and negative-binomial observations, reporting delays, ascertainment, and aggregation — all interchangeable.

Getting started

You assemble a model from parts, and ComposableTuringIDModels turns the assembly into a single Turing model you can simulate from and fit. Each part is itself a model, joined through the generic as_turing_model constructor.

Compose a model: an ARIMA-style latent process (a differenced AR) folded into a direct-infections process, observed with Poisson error.

using ComposableTuringIDModels, Distributions, Turing
model = IDModel(
    DirectInfections(;
        Z = DiffLatentModel(; model = AR(), init = [Normal(), Normal()]),
        initialisation = Normal()),
    PoissonError())

IDModel
├─ infection: DirectInfections
│  └─ Z: DiffLatentModel
│     └─ model: AR
│        └─ ϵ_t: HierarchicalNormal
└─ observation: PoissonError

Build a Turing model; missing observations simulate from the prior.

n = 30
prior_model = as_turing_model(model, missing, n)

DynamicPPL.Model{typeof(ComposableTuringIDModels._as_turing_model_idmodel), (:model, :y_t, :n), (), (:y_t,), Tuple{IDModel{DirectInfections{DiffLatentModel{AR{Distributions.Truncated{Distributions.Normal{Float64}, Distributions.Continuous, Float64, Float64, Float64}, Distributions.Normal{Float64}, Int64, HierarchicalNormal{Float64, Distributions.Truncated{Distributions.Normal{Float64}, Distributions.Continuous, Float64, Float64, Float64}, Bool}, typeof(identity)}, Vector{Distributions.Normal{Float64}}}, typeof(exp), Distributions.Normal{Float64}}, PoissonError}, Missing, Int64}, Tuple{}, DynamicPPL.DefaultContext, false}(ComposableTuringIDModels._as_turing_model_idmodel, (model = IDModel, y_t = missing, n = 30), NamedTuple(), DynamicPPL.DefaultContext())

Sample from the prior and inspect the generated quantities.

draw = rand(prior_model)
(; generated_y_t, I_t, Z_t) = prior_model()

(generated_y_t = Union{Missing, Int64}[2, 2, 2, 1, 0, 0, 0, 0, 1, 0  …  0, 0, 0, 0, 0, 0, 0, 0, 0, 0], expected_y_t = [3.3831271516466574, 3.044906224486632, 1.9888167594063646, 1.2906128911118944, 0.8375481080596769, 0.5462780912891562, 0.3569046385115774, 0.23350911640815994, 0.15262979280814262, 0.09976049334460445  …  0.0010256871941593398, 0.0006842237410396727, 0.00045598498378983556, 0.0003031939933640396, 0.00020241781969428755, 0.0001345503311412326, 8.934675774518192e-5, 5.9352088838894625e-5, 3.949981861573408e-5, 2.6280604020382287e-5], I_t = [3.3831271516466574, 3.044906224486632, 1.9888167594063646, 1.2906128911118944, 0.8375481080596769, 0.5462780912891562, 0.3569046385115774, 0.23350911640815994, 0.15262979280814262, 0.09976049334460445  …  0.0010256871941593398, 0.0006842237410396727, 0.00045598498378983556, 0.0003031939933640396, 0.00020241781969428755, 0.0001345503311412326, 8.934675774518192e-5, 5.9352088838894625e-5, 3.949981861573408e-5, 2.6280604020382287e-5], Z_t = [0.5579507401465119, 0.4526203697152922, 0.026690134216360684, -0.40573251948231415, -0.8381263088895851, -1.2654768426098513, -1.6911363863165563, -2.115383894618318, -2.540589678950194, -2.9658327667240143  …  -7.543242192121835, -7.948075321462567, -8.353900413693781, -8.761987449613143, -9.166026316874836, -9.574421954108926, -9.983835338560691, -10.392872976680566, -10.800064212490703, -11.207529114995411])

Condition on data and run inference.

using StatsBase
posterior_model = as_turing_model(model, generated_y_t, n)
chain = sample(posterior_model, NUTS(), 1_000)
summarystats(chain)

╭─FlexiSummary (9 statistics) ─────────────────────────────────────────────────╮
│   iter    collapsed                                                          │
│   chain   collapsed                                                          │
│ ↓ stat  = [mean, std, mcse, ess_bulk, ess_tail, rhat, q5, q50, q95]          │
│                                                                              │
│ Parameters (34) ── AbstractPPL.VarName                                       │
│  Float64  latent_init[1], latent_init[2], ar_init, damp_AR, ρ, std, ϵ_t[1],  │
│           ϵ_t[2], ϵ_t[3], ϵ_t[4], ϵ_t[5], ϵ_t[6], ϵ_t[7], ϵ_t[8], ϵ_t[9],    │
│           ϵ_t[10], ϵ_t[11], ϵ_t[12], ϵ_t[13], ϵ_t[14], ϵ_t[15], ϵ_t[16],     │
│           ϵ_t[17], ϵ_t[18], ϵ_t[19], ϵ_t[20], ϵ_t[21], ϵ_t[22], ϵ_t[23],     │
│           ϵ_t[24], ϵ_t[25], ϵ_t[26], ϵ_t[27], init_incidence                 │
│                                                                              │
│ Extras (14)                                                                  │
│  Float64  n_steps, is_accept, acceptance_rate, log_density,                  │
│           hamiltonian_energy, hamiltonian_energy_error,                      │
│           max_hamiltonian_energy_error, tree_depth, numerical_error,         │
│           step_size, nom_step_size, logprior, loglikelihood, logjoint        │
│                                                                              │
│ Summary                                                                      │
│          param     mean     std    mcse   ess_bulk  ess_tail    rhat  …      │
│   latent_init…   0.2143  0.5837  0.0252   537.3283  398.5904  1.0001  …      │
│   latent_init…  -0.1853  0.7318  0.0311   552.3372  483.8228  1.0004  …      │
│        ar_init  -0.4855  0.5688  0.0255   518.4010  395.5476  1.0008  …      │
│        damp_AR   0.0410  0.0312  0.0012   502.6466  341.2595  1.0007  …      │
│              ρ   0.0410  0.0312  0.0012   502.6466  341.2595  1.0007  …      │
│            std   0.0747  0.0578  0.0022   564.5858  392.3295  1.0062  …      │
│         ϵ_t[1]  -0.0938  1.0227  0.0330   967.4025  477.3459  1.0027  …      │
│         ϵ_t[2]  -0.0572  0.9533  0.0337   814.0101  592.9734  0.9996  …      │
│         ϵ_t[3]  -0.0096  1.0091  0.0375   751.6606  332.2904  1.0017  …      │
│         ϵ_t[4]  -0.0208  0.9410  0.0358   680.4574  304.2566  1.0016  …      │
│         ϵ_t[5]  -0.0030  0.9753  0.0316   953.0170  653.7758  1.0014  …      │
│         ϵ_t[6]  -0.0782  1.0480  0.0318  1083.8486  563.0251  1.0007  …      │
│         ϵ_t[7]  -0.0896  0.9728  0.0344   797.7954  317.8816  1.0043  …      │
│         ϵ_t[8]  -0.1613  0.8985  0.0275  1042.8182  601.4149  1.0019  …      │
│         ϵ_t[9]  -0.0657  0.9838  0.0312   980.1175  443.8196  1.0156  …      │
│        ϵ_t[10]  -0.0318  0.9098  0.0301   904.0984  550.7320  1.0028  …      │
│        ϵ_t[11]  -0.0011  1.0410  0.0349   920.0724  532.2438  1.0362  …      │
│        ϵ_t[12]   0.0270  0.9361  0.0424   492.0724  201.0408  1.0007  …      │
│        ϵ_t[13]  -0.0276  1.0284  0.0350   864.8820  564.1153  1.0007  …      │
│        ϵ_t[14]  -0.0335  0.9583  0.0341   796.4005  652.0758  1.0041  …      │
│        ϵ_t[15]  -0.1180  0.9670  0.0327   845.1258  575.3035  1.0096  …      │
│        ϵ_t[16]  -0.0443  1.0466  0.0367   801.0280  422.2519  1.0012  …      │
│        ϵ_t[17]  -0.0244  1.0639  0.0364   852.2540  526.1526  1.0011  …      │
│        ϵ_t[18]   0.0187  0.9407  0.0289  1052.6816  693.2795  0.9995  …      │
│        ϵ_t[19]  -0.0137  1.0133  0.0301  1131.4570  665.8263  1.0006  …      │
│        ϵ_t[20]   0.0002  0.9645  0.0286  1146.0815  486.3383  1.0017  …      │
│        ϵ_t[21]   0.0330  0.9442  0.0413   556.0160  290.6345  0.9992  …      │
│        ϵ_t[22]   0.0805  1.1013  0.0336  1079.0167  519.4991  0.9992  …      │
│        ϵ_t[23]   0.0277  1.0332  0.0372   755.2522  554.9285  1.0085  …      │
│        ϵ_t[24]   0.0037  0.9618  0.0305   999.4109  495.6441  1.0128  …      │
│        ϵ_t[25]  -0.0166  1.0269  0.0346   909.4720  364.4678  1.0008  …      │
│        ϵ_t[26]   0.0202  1.0687  0.0354   921.2297  501.3851  1.0028  …      │
│        ϵ_t[27]   0.0230  1.0331  0.0303  1172.1007  583.2867  0.9991  …      │
│   init_incide…   0.3395  0.7167  0.0300   582.2983  528.8758  1.0048  …      │
╰──────────────────────────────────────────────────────────────────────────────╯

Swap a part

Because every part is interchangeable, the same swap-in/swap-out pattern applies throughout. Replace PoissonError() with NegativeBinomialError(), wrap the observation in a LatentDelay, or change the latent process — without touching the rest of the model.

Same infection process, two observation assumptions.

latent = DiffLatentModel(; model = AR(), init = [Normal(), Normal()])
infections = DirectInfections(; Z = latent, initialisation = Normal())

poisson_model = IDModel(infections, PoissonError())
negbin_model = IDModel(infections, NegativeBinomialError())

IDModel
├─ infection: DirectInfections
│  └─ Z: DiffLatentModel
│     └─ model: AR
│        └─ ϵ_t: HierarchicalNormal
└─ observation: NegativeBinomialError

That is the point: you compare modelling assumptions by swapping parts, not by rewriting models.

Related packages

• ModifiedDistributions.jl modifies the behaviour of the underlying Distributions.jl distributions.

• ReparameterisedDistributions.jl switches distributions between parameterisations, so a prior can be stated in whichever one is natural.

• LoweredDistributions.jl lowers a distribution onto a backend-agnostic dynamical-systems representation, so a delay distribution and a compartmental model become two views of one thing.

• Turing.jl is the probabilistic programming backend every model here compiles down to.

Where to learn more

• New here? Start with the Overview and the Composable design page.

• Want worked examples? See the tutorials, which fit complete models to real surveillance data.

• Want the full interface? Browse the Public API.

• Want to see the code or report a problem? Check the GitHub repository.

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: it has been renamed, re-architected around the generic as_turing_model constructor, and upgraded to build against the latest Turing.jl. See the NOTICE file for full attribution and a summary of the changes, and LICENSE for the Apache-2.0 terms.

Part of the EpiAware ecosystem

ComposableTuringIDModels is part of EpiAware, a set of composable tools for infectious disease modelling. See the other packages in the ecosystem.

Contributing

We welcome contributions and new contributors! Please open an issue or pull request on GitHub. This package follows ColPrac and is formatted with Runic.

How to cite

If you use ComposableTuringIDModels in your work, please cite it. Citation metadata lives in CITATION.cff, which GitHub renders as a "Cite this repository" button on the repository page.

Code of conduct

Please note that the ComposableTuringIDModels project is released with a Contributor Code of Conduct. By contributing, you agree to abide by its terms.

License

Apache License 2.0. See LICENSE and NOTICE.