Disease Transmission

analysis.py

import arviz as az
import numpy as np
import xarray as xr

# load data from data.json
import json
with open("data.json") as f:
    data = json.load(f)

observed_cases = xr.DataArray(name="cases",
                              data=data["cases"], 
                              coords={"day": data["ts"]})

# arviz massaging 
draws_az = draws.to_arviz()

pred_cases = draws_az.posterior.pred_cases\
                     .assign_coords(day = ("pred_cases_dim_0", data["ts"]))\
                     .swap_dims(pred_cases_dim_0="day")
draws_az.add_groups(observed_data=observed_cases, 
                    posterior_predictive=pred_cases)

# plot (can also try plot_ts, or plot_ppc)
az.plot_lm(idata=draws_az, y="cases", y_hat="pred_cases",
           backend_kwargs={"figsize": (11,11)})

analysis.R

# posterior predictive check using the pred_cases generated quantity

install.packages("bayesplot")
library(posterior)
library(ggplot2)

# load from data
d <- jsonlite::read_json('./data.json')
cases <- unlist(d$cases)
n_days <- d$n_days
ts <- unlist(d$ts)

# Extract posterior predictive checks
pred_cases <- as_draws_matrix(as_draws_rvars(draws)$pred_cases)

bayesplot::ppc_ribbon(y = cases, yrep = pred_cases,
                      x = ts, y_draw = "point") +
  theme_bw() +
  ylab("cases") + xlab("days")

data.json

{
  "_source": "This data comes from the R package 'outbreaks'",
  "_description": "Influenza in a boarding school in England, 1978",
  "N": 763,
  "cases": [
    3,
    8,
    26,
    76,
    225,
    298,
    258,
    233,
    189,
    128,
    68,
    29,
    14,
    4
  ],
  "n_days": 14,
  "ts": [
    1,
    2,
    3,
    4,
    5,
    6,
    7,
    8,
    9,
    10,
    11,
    12,
    13,
    14
  ],
  "t0": 0,
  "y0": [
    762,
    1,
    0
  ]
}

data.py

# The source for this data seems to live _only_
# in the R package 'outbreaks'
# This code downloads the .RData file and reads
# it using the rdata package

# We will need to install some extra packages to do this --
# we can do this from inside Python by using micropip
import micropip

# download data from R package 'outbreaks' source
await micropip.install("lzma")
await micropip.install("requests")
await micropip.install('pyodide-http')

import pyodide_http
pyodide_http.patch_all() 

import requests

resp = requests.request(
    "GET",
    "https://raw.githubusercontent.com/reconverse/outbreaks/refs/heads/master/data/influenza_england_1978_school.RData",
)
with open("influenza.RData", "wb") as fd:
    for chunk in resp.iter_content(chunk_size=128):
        fd.write(chunk)

# read .RData
await micropip.install("rdata")
import rdata

data_from_R = rdata.read_rda("influenza.RData")["influenza_england_1978_school"]

print(data_from_R.head())

# prepare data

N = 763  # given in documentation
cases = data_from_R.in_bed
n_days = len(data_from_R)
ts = list(range(1, n_days + 1))
t0 = 0
# started with a single infection
y0 = [N - 1, 1, 0]

data = {
    "N": N,
    "cases": cases,
    "n_days": n_days,
    "ts": ts,
    "t0": t0,
    "y0": y0,
}

data.R

install.packages("outbreaks")

library(outbreaks)
print(head(influenza_england_1978_school))

N <- 763 # given in documentation

cases <- influenza_england_1978_school$in_bed
n_days <- length(cases)
ts <- 1:n_days
t0 <- 0 

# started with a single infection
y0 <- c(N-1, 1, 0)

data <- list(
    "_source"="This data comes from the R package 'outbreaks'",
    "_description"="Influenza in a boarding school in England, 1978",
    N=N,
    cases=cases,
    n_days=n_days,
    ts=ts,
    t0=t0,
    y0=y0
)

main.stan

// the "susceptible-infected-recovered" model
functions {
  vector sir(real t, vector y, real beta, real gamma, int N) {
    real S = y[1];
    real I = y[2];
    real R = y[3];

    real dS_dt = -beta * I * S / N;
    real dI_dt = beta * I * S / N - gamma * I;
    real dR_dt = gamma * I;

    return [dS_dt, dI_dt, dR_dt]';
  }
}
data {
  int<lower=1> n_days;
  vector[3] y0;
  real t0;
  array[n_days] real ts;
  int N;
  array[n_days] int cases; // how many people are sick on day n
}
parameters {
  real<lower=0> gamma;
  real<lower=0> beta;
  real<lower=0> phi_inv;
}
transformed parameters {
  real phi = 1. / phi_inv;
  array[n_days] vector[3] y = ode_rk45(sir, y0, t0, ts, beta, gamma, N);
}
model {
  //priors
  beta ~ normal(2, 1);
  gamma ~ normal(0.4, 0.5);
  phi_inv ~ exponential(5);

  cases ~ poisson(y[ : , 2]);
  // try it: a overdispersed likelihood instead
  // cases ~ neg_binomial_2(y[,2], phi);
}
generated quantities {
  // R0 is the expected number of new infections caused by a single infected individual
  real R0 = beta / gamma;
  real recovery_time = 1 / gamma;
  array[n_days] real pred_cases = poisson_rng(y[ : , 2]);
  // array[n_days] real pred_cases = neg_binomial_2_rng(y[,2], phi);
}

sampling_opts.json

{"num_chains":4,"num_warmup":100,"num_samples":100,"init_radius":0.2}