Getting started with composableIDModelR

Source: vignettes/composableIDModelR.Rmd

composableIDModelR builds infectious disease models from interchangeable parts and runs them with ComposableTuringIDModels.jl in Julia. A model has two parts:

  • an infection model, which generates unobserved infections and owns the latent process that drives them (for example the log reproduction number of a renewal model), and
  • an observation model, which links infections to data, possibly through delays or under-ascertainment.

Every part has an R constructor with the same name and arguments as its Julia counterpart, so the ComposableTuringIDModels.jl documentation applies directly.

library(composableIDModelR)
#> 
#> Attaching package: 'composableIDModelR'
#> The following object is masked from 'package:stats':
#> 
#>     Gamma

Setting up Julia

composableIDModelR needs Julia (>= 1.12), the version its pinned dependencies were resolved with, ideally installed with juliaup. The Julia packages are installed and loaded the first time a model is simulated or fitted, which can take several minutes. To follow progress, run the setup yourself:

To run MCMC chains in parallel, set JULIA_NUM_THREADS before Julia starts:

Sys.setenv(JULIA_NUM_THREADS = "4")

Building a model

Distributions come from distspec, so they can be given by their natural parameters or by a mean and standard deviation.

latent <- RandomWalk(
  init = Normal(0, 0.25),
  epsilon_t = HierarchicalNormal(std = HalfNormal(0.05))
)

infection <- Renewal(
  generation_time = Gamma(shape = 6.5, scale = 0.62),
  rt = latent,
  initialisation = Normal(log(10), 0.5)
)

observation <- NegativeBinomialError(cluster_factor = HalfNormal(0.1))

model <- IDModel(infection, observation)
model
#> <julia model spec>
#> IDModel(
#>     Renewal(;
#>         generation_time = Gamma(6.5, 0.62),
#>         rt = RandomWalk(;
#>             init = Normal(0.0, 0.25),
#>             ϵ_t = HierarchicalNormal(; std = HalfNormal(0.05))
#>         ),
#>         initialisation = Normal(2.3025850929940459, 0.5)
#>     ),
#>     NegativeBinomialError(; cluster_factor = HalfNormal(0.1))
#> )

Building a model does not start Julia. The printout is the Julia code the model runs, which you can also get as a string with as_julia(). Arguments left out use the Julia defaults.

Simulating from the prior

simulate() draws parameters from their priors and simulates the latent process, infections and observations. Each quantity is returned as a matrix with one row per simulation.

sims <- simulate(model, nsim = 20, n = 50, seed = 1)
#> Starting Julia ...
names(sims)
#> [1] "generated_y_t" "expected_y_t"  "I_t"           "Z_t"          
#> [5] "Rt"
matplot(t(sims$Rt), type = "l", lty = 1, col = "grey50", log = "y",
        xlab = "Day", ylab = "Rt")

Prior simulations are a quick check that the priors produce plausible epidemics before fitting.

Fitting to data

We fit the model to one of the simulated series. nuts() sets up the sampler; short chains keep this example fast.

y <- sims$generated_y_t[1, ]

fitted <- fit(
  model, y,
  method = nuts(draws = 250, warmup = 250, chains = 2),
  seed = 2
)
fitted
#> <composableIDModelR fit>
#> Model:
#>   IDModel(
#>       Renewal(;
#>           generation_time = Gamma(6.5, 0.62),
#>           rt = RandomWalk(;
#>               init = Normal(0.0, 0.25),
#>               ϵ_t = HierarchicalNormal(; std = HalfNormal(0.05))
#>           ),
#>           initialisation = Normal(2.3025850929940459, 0.5)
#>       ),
#>       NegativeBinomialError(; cluster_factor = HalfNormal(0.1))
#>   )
#> 
#> Data: 50 time points
#> Sampling: 2 chains of 250 draws after 250 warmup
#> 
#> Convergence:
#>   Max Rhat: 1.03 
#>   Min bulk ESS: 254 
#>   Divergent transitions: 1 
#> 
#> Use summary() for parameter estimates and plot() to visualise.

summary() gives posterior summaries of the parameters, and the draws work with the posterior package and anything built on it, such as bayesplot.

summary(fitted)
#> # A tibble: 53 × 10
#>    variable     mean  median     sd    mad      q5    q95  rhat ess_bulk
#>    <chr>       <dbl>   <dbl>  <dbl>  <dbl>   <dbl>  <dbl> <dbl>    <dbl>
#>  1 rw_init  -0.0353  -0.0315 0.130  0.128  -0.257  0.177  0.999     273.
#>  2 std       0.0592   0.0555 0.0215 0.0217  0.0295 0.0986 1.01      289.
#>  3 ϵ_t[1]   -0.0530  -0.0584 0.971  0.905  -1.65   1.54   1.01      565.
#>  4 ϵ_t[2]   -0.00416 -0.0275 0.993  0.985  -1.60   1.68   1.00      357.
#>  5 ϵ_t[3]   -0.101   -0.0899 0.937  0.871  -1.73   1.40   1.01      725.
#>  6 ϵ_t[4]   -0.0787  -0.0409 0.985  0.998  -1.65   1.60   0.999     667.
#>  7 ϵ_t[5]    0.182    0.151  0.976  1.00   -1.28   1.79   1.01      545.
#>  8 ϵ_t[6]   -0.0388  -0.0377 0.961  1.03   -1.59   1.53   1.00      491.
#>  9 ϵ_t[7]   -0.0212  -0.0730 1.01   0.970  -1.66   1.78   1.01      407.
#> 10 ϵ_t[8]    0.143    0.172  0.949  0.933  -1.37   1.72   1.01      532.
#> # ℹ 43 more rows
#> # ℹ 1 more variable: ess_tail <dbl>

Generated quantities (the latent process Z_t, infections I_t, expected and predicted observations, and Rt for renewal models) are matrices with one row per posterior draw:

str(fitted$generated)
#> List of 6
#>  $ generated_y_t: num [1:500, 1:50] 17 17 17 17 17 17 17 17 17 17 ...
#>  $ expected_y_t : num [1:500, 1:50] 11 11.2 11.1 15 13.9 ...
#>  $ I_t          : num [1:500, 1:50] 11 11.2 11.1 15 13.9 ...
#>  $ Z_t          : num [1:500, 1:50] 0.1297 0.2278 0.1527 0.1978 0.0429 ...
#>  $ predicted_y_t: num [1:500, 1:50] 9 5 8 14 10 14 12 14 11 17 ...
#>  $ Rt           : num [1:500, 1:50] 1.14 1.26 1.16 1.22 1.04 ...

Plotting and forecasting 

plot(fitted, type = "Rt")

type = "cases" compares posterior predictions with the data, and horizon adds a forecast, which continues the latent process beyond the data.

plot(fitted, type = "cases", horizon = 14)

The forecast draws themselves come from predict():

forecasts <- predict(fitted, horizon = 14)
dim(forecasts)
#> [1] 500  64

Swapping components

Changing an assumption means replacing one part. Here observations follow a delay from infection and are Poisson distributed, and the reproduction number follows an AR(1) process:

IDModel(
  Renewal(
    generation_time = Gamma(shape = 6.5, scale = 0.62),
    rt = AR(damp = truncated(Normal(0.9, 0.05), 0, 1)),
    initialisation = Normal(log(10), 0.5)
  ),
  LatentDelay(PoissonError(), delay = LogNormal(1.6, 0.42))
)
#> <julia model spec>
#> IDModel(
#>     Renewal(;
#>         generation_time = Gamma(6.5, 0.62),
#>         rt = AR(; damp = truncated(Normal(0.9, 0.05), 0.0, 1.0)),
#>         initialisation = Normal(2.3025850929940459, 0.5)
#>     ),
#>     LatentDelay(PoissonError(), LogNormal(1.6, 0.42))
#> )

Components without an R wrapper

julia_spec() constructs any Julia component by name, and julia() inserts Julia code verbatim, for example a function argument. Positional R arguments become positional Julia arguments and named ones become keyword arguments. .role says what the component is, so that it is only accepted where it belongs: "distribution", "latent", "infection", "observation" or "model".

IDModel(
  DirectInfections(
    Z = julia_spec("HilbertSpaceGP", .role = "latent"),
    initialisation = Normal(log(100), 1)
  ),
  LatentDelay(PoissonError(), delay = julia_spec("Frechet", 3, 2,
                                                 .role = "distribution"))
)
#> <julia model spec>
#> IDModel(
#>     DirectInfections(;
#>         Z = HilbertSpaceGP(),
#>         initialisation = Normal(4.6051701859880918, 1.0)
#>     ),
#>     LatentDelay(PoissonError(), Frechet(3.0, 2.0))
#> )

The Mishra et al. (2020) case study applies these steps to COVID-19 data from South Korea.