Disease transmission at a boarding school

This case study follows Léo Grinsztajn and collaborators' boarding-school influenza case study. The authoritative source notebook fits a susceptible–infectious–removed (SIR) system to daily counts of pupils in bed. The first model below reproduces that Stan program. The other two complete the observation-model variants that the original StanBlocks notebook had only sketched: infections acquired during each day, then reported infections when only a fraction of cases are observed.

All three models use the same latent epidemic. If , , and are the three compartments and is the school population, then

Here is the infection rate, is the removal rate,  , and is the mean infectious period. A negative-binomial observation model allows more day-to-day variation than a Poisson model; phi_inv is given an exponential prior and inverted to obtain Stan's neg_binomial_2 precision .

The three observation models

Infectious prevalence

The official case study treats the number of pupils in bed on day as a noisy measurement of . This is the right interpretation when “in bed” is a snapshot of current infectious prevalence. The posterior-predictive count and pointwise log likelihood are generated automatically from the observed cases ~ neg_binomial_2(...) statement.

Daily incidence

If the recorded count instead means new infections during a day, its latent mean is the susceptible decrement   . The second model changes only that observation statement. It therefore has one fewer observation than the prevalence model: a decrement needs two adjacent solver states. The incidence variants bind that shorter count vector directly, so their generated predictive arrays have the same length and every element is initialized.

Reported daily incidence

The third model estimates a reporting fraction with a beta(1,2) prior and scales incidence by it. This separates infections from observed reports, though the two can be weakly identified without external information about reporting.

StanBlocks implementation

• school_sir_rhs is a typed @deffun: loops and mutation belong in helper functions, while the @slic body remains a flat probabilistic declaration.

• Current Stan uses the variadic ode_rk45 interface. The legacy theta, x_r, and x_i packing arrays from the source Stan program are unnecessary; infection_rate, gamma, and population are ordinary trailing arguments.

• The source parameter beta is named infection_rate, leaving the builtin beta(...) distribution unshadowed when the reporting variant declares p_reported.

• The three @slic definitions deliberately keep the complete model visible. Their only statistical difference is the final observation statement (and p_reported in the third); the shared ODE helper keeps the differential equations themselves in one place.

• The small Julia values establish types and shapes for documentation transpilation. They are not the boarding-school data used for inference; rebind the resulting model to the complete case-study data before fitting.

Use the tabs to inspect the exact Julia source evaluated by this documentation build and each complete generated Stan program. Compare side by side opens the feature-atlas modal.

using StanBlocks

@deffun @stanonly school_sir_rhs(
    t::real, state::vector[n_state],
    infection_rate::real, removal_rate::real, population::real,
)::vector[n_state] = begin
    susceptible = state[1]
    infectious = state[2]

    derivative::vector[n_state]
    derivative[1] = -infection_rate * infectious * susceptible / population
    derivative[2] =
        infection_rate * infectious * susceptible / population -
        removal_rate * infectious
    derivative[3] = removal_rate * infectious
    derivative
end

school_data = (;
    y0=[762.0, 1.0, 0.0],
    t0=0.0,
    times=[1.0, 2.0, 3.0, 4.0],
    population=763.0,
    cases=[3, 8, 15, 20],
    n_days=4,
)
school_incidence_data = Base.merge(
    school_data,
    (; cases=school_data.cases[1:(school_data.n_days - 1)]),
)

school_prevalence = @slic school_data begin
    infection_rate::real ~ normal(2.0, 1.0; lower=0.0)
    gamma::real ~ normal(0.4, 0.5; lower=0.0)
    phi_inv::real ~ exponential(5.0; lower=0.0)

    R0 = infection_rate / gamma
    recovery_time = 1.0 / gamma
    phi = 1.0 / phi_inv
    trajectory = ode_rk45(
        school_sir_rhs, y0, t0, to_array_1d(times),
        infection_rate, gamma, population,
    )
    prevalence = to_vector(trajectory[:, 2])
    susceptible = to_vector(trajectory[:, 1])
    incidence =
        susceptible[1:(n_days - 1)] - susceptible[2:n_days]

    cases ~ neg_binomial_2(prevalence, phi)
end

school_incidence = @slic school_incidence_data begin
    infection_rate::real ~ normal(2.0, 1.0; lower=0.0)
    gamma::real ~ normal(0.4, 0.5; lower=0.0)
    phi_inv::real ~ exponential(5.0; lower=0.0)

    R0 = infection_rate / gamma
    recovery_time = 1.0 / gamma
    phi = 1.0 / phi_inv
    trajectory = ode_rk45(
        school_sir_rhs, y0, t0, to_array_1d(times),
        infection_rate, gamma, population,
    )
    susceptible = to_vector(trajectory[:, 1])
    incidence =
        susceptible[1:(n_days - 1)] - susceptible[2:n_days]

    cases ~ neg_binomial_2(incidence, phi)
end

school_reported_incidence = @slic school_incidence_data begin
    infection_rate::real ~ normal(2.0, 1.0; lower=0.0)
    gamma::real ~ normal(0.4, 0.5; lower=0.0)
    phi_inv::real ~ exponential(5.0; lower=0.0)
    p_reported ~ beta(1.0, 2.0)

    R0 = infection_rate / gamma
    recovery_time = 1.0 / gamma
    phi = 1.0 / phi_inv
    trajectory = ode_rk45(
        school_sir_rhs, y0, t0, to_array_1d(times),
        infection_rate, gamma, population,
    )
    susceptible = to_vector(trajectory[:, 1])
    incidence =
        susceptible[1:(n_days - 1)] - susceptible[2:n_days]

    cases ~ neg_binomial_2(incidence * p_reported, phi)
end

school_models = (;
    infectious_prevalence=school_prevalence,
    daily_incidence=school_incidence,
    reported_daily_incidence=school_reported_incidence,
)

infectious_prevalence

functions {
vector school_sir_rhs(
    real t,
    vector state,
    real infection_rate,
    real removal_rate,
    real population
) {
    int n_state = dims(state)[1];
    real susceptible = state[1];
    real infectious = state[2];
    vector[n_state] derivative;
    derivative[1] = (((-infection_rate) * infectious * susceptible) / population);
    derivative[2] = (((infection_rate * infectious * susceptible) / population) - (removal_rate * infectious));
    derivative[3] = (removal_rate * infectious);
    return derivative;
}
vector neg_binomial_2_lpmfs(
    array[] int obs,
    vector mu,
    real phi
) {
    return jbroadcasted_neg_binomial_2_lpmfs(obs, mu, phi);
}
vector jbroadcasted_neg_binomial_2_lpmfs(
    array[] int x1,
    vector x2,
    real x3
) {
    int n = dims(x1)[1];
    vector[n] rv;
    for(i in 1:n) {
        rv[i] = neg_binomial_2_lpmfs(broadcasted_getindex(x1, i), broadcasted_getindex(x2, i), x3);
    }
    return rv;
}
real neg_binomial_2_lpmfs(
    int args1,
    real args2,
    real args3
) {
    return neg_binomial_2_lpmf(args1 | args2, args3);
}
int broadcasted_getindex(array[] int x, int i) {
    return x[i];
}
real broadcasted_getindex(vector x, int i) {
    return x[i];
}
array[] int neg_binomial_2_int_rng(
    int anontok__1,
    vector a,
    real b
) {
    int n = anontok__1;
    return neg_binomial_2_rng(a, b);
}
}
data {
    int times_n;
    int y0_n;
    vector[y0_n] y0;
    real t0;
    vector[times_n] times;
    real population;
    int n_days;
    int cases_n;
    array[cases_n] int cases;
}
transformed data {
}
parameters {
    real<lower=0.0> infection_rate;
    real<lower=0.0> gamma;
    real<lower=0.0> phi_inv;
}
transformed parameters {
    real phi = (1.0 / phi_inv);
    array[times_n] vector[y0_n] trajectory = ode_rk45(school_sir_rhs, y0, t0, to_array_1d(times), infection_rate, gamma, population);
    vector[times_n] prevalence = to_vector(trajectory[:, 2]);
}
model {
    infection_rate ~ normal(2.0, 1.0);
    gamma ~ normal(0.4, 0.5);
    phi_inv ~ exponential(5.0);
    cases ~ neg_binomial_2(prevalence, phi);
}
generated quantities {
    real R0 = (infection_rate / gamma);
    real recovery_time = (1.0 / gamma);
    vector[times_n] susceptible = to_vector(trajectory[:, 1]);
    vector[(1 + ((n_days - 1) - 1))] incidence = (susceptible[1:(n_days - 1)] - susceptible[2:n_days]);
    vector[cases_n] cases_likelihood = neg_binomial_2_lpmfs(cases, prevalence, phi);
    array[cases_n] int cases_gen = neg_binomial_2_int_rng(cases_n, prevalence, phi);
}

daily_incidence

functions {
vector school_sir_rhs(
    real t,
    vector state,
    real infection_rate,
    real removal_rate,
    real population
) {
    int n_state = dims(state)[1];
    real susceptible = state[1];
    real infectious = state[2];
    vector[n_state] derivative;
    derivative[1] = (((-infection_rate) * infectious * susceptible) / population);
    derivative[2] = (((infection_rate * infectious * susceptible) / population) - (removal_rate * infectious));
    derivative[3] = (removal_rate * infectious);
    return derivative;
}
vector neg_binomial_2_lpmfs(
    array[] int obs,
    vector mu,
    real phi
) {
    return jbroadcasted_neg_binomial_2_lpmfs(obs, mu, phi);
}
vector jbroadcasted_neg_binomial_2_lpmfs(
    array[] int x1,
    vector x2,
    real x3
) {
    int n = dims(x1)[1];
    vector[n] rv;
    for(i in 1:n) {
        rv[i] = neg_binomial_2_lpmfs(broadcasted_getindex(x1, i), broadcasted_getindex(x2, i), x3);
    }
    return rv;
}
real neg_binomial_2_lpmfs(
    int args1,
    real args2,
    real args3
) {
    return neg_binomial_2_lpmf(args1 | args2, args3);
}
int broadcasted_getindex(array[] int x, int i) {
    return x[i];
}
real broadcasted_getindex(vector x, int i) {
    return x[i];
}
array[] int neg_binomial_2_int_rng(
    int anontok__1,
    vector a,
    real b
) {
    int n = anontok__1;
    return neg_binomial_2_rng(a, b);
}
}
data {
    int times_n;
    int y0_n;
    vector[y0_n] y0;
    real t0;
    vector[times_n] times;
    real population;
    int n_days;
    int cases_n;
    array[cases_n] int cases;
}
transformed data {
}
parameters {
    real<lower=0.0> infection_rate;
    real<lower=0.0> gamma;
    real<lower=0.0> phi_inv;
}
transformed parameters {
    real phi = (1.0 / phi_inv);
    array[times_n] vector[y0_n] trajectory = ode_rk45(school_sir_rhs, y0, t0, to_array_1d(times), infection_rate, gamma, population);
    vector[times_n] susceptible = to_vector(trajectory[:, 1]);
    vector[(1 + ((n_days - 1) - 1))] incidence = (susceptible[1:(n_days - 1)] - susceptible[2:n_days]);
}
model {
    infection_rate ~ normal(2.0, 1.0);
    gamma ~ normal(0.4, 0.5);
    phi_inv ~ exponential(5.0);
    cases ~ neg_binomial_2(incidence, phi);
}
generated quantities {
    real R0 = (infection_rate / gamma);
    real recovery_time = (1.0 / gamma);
    vector[cases_n] cases_likelihood = neg_binomial_2_lpmfs(cases, incidence, phi);
    array[cases_n] int cases_gen = neg_binomial_2_int_rng(cases_n, incidence, phi);
}

reported_daily_incidence

functions {
vector school_sir_rhs(
    real t,
    vector state,
    real infection_rate,
    real removal_rate,
    real population
) {
    int n_state = dims(state)[1];
    real susceptible = state[1];
    real infectious = state[2];
    vector[n_state] derivative;
    derivative[1] = (((-infection_rate) * infectious * susceptible) / population);
    derivative[2] = (((infection_rate * infectious * susceptible) / population) - (removal_rate * infectious));
    derivative[3] = (removal_rate * infectious);
    return derivative;
}
vector neg_binomial_2_lpmfs(
    array[] int obs,
    vector mu,
    real phi
) {
    return jbroadcasted_neg_binomial_2_lpmfs(obs, mu, phi);
}
vector jbroadcasted_neg_binomial_2_lpmfs(
    array[] int x1,
    vector x2,
    real x3
) {
    int n = dims(x1)[1];
    vector[n] rv;
    for(i in 1:n) {
        rv[i] = neg_binomial_2_lpmfs(broadcasted_getindex(x1, i), broadcasted_getindex(x2, i), x3);
    }
    return rv;
}
real neg_binomial_2_lpmfs(
    int args1,
    real args2,
    real args3
) {
    return neg_binomial_2_lpmf(args1 | args2, args3);
}
int broadcasted_getindex(array[] int x, int i) {
    return x[i];
}
real broadcasted_getindex(vector x, int i) {
    return x[i];
}
array[] int neg_binomial_2_int_rng(
    int anontok__1,
    vector a,
    real b
) {
    int n = anontok__1;
    return neg_binomial_2_rng(a, b);
}
}
data {
    int times_n;
    int y0_n;
    vector[y0_n] y0;
    real t0;
    vector[times_n] times;
    real population;
    int n_days;
    int cases_n;
    array[cases_n] int cases;
}
transformed data {
}
parameters {
    real<lower=0.0> infection_rate;
    real<lower=0.0> gamma;
    real<lower=0.0> phi_inv;
    real<lower=0, upper=1> p_reported;
}
transformed parameters {
    real phi = (1.0 / phi_inv);
    array[times_n] vector[y0_n] trajectory = ode_rk45(school_sir_rhs, y0, t0, to_array_1d(times), infection_rate, gamma, population);
    vector[times_n] susceptible = to_vector(trajectory[:, 1]);
    vector[(1 + ((n_days - 1) - 1))] incidence = (susceptible[1:(n_days - 1)] - susceptible[2:n_days]);
}
model {
    infection_rate ~ normal(2.0, 1.0);
    gamma ~ normal(0.4, 0.5);
    phi_inv ~ exponential(5.0);
    p_reported ~ beta(1.0, 2.0);
    cases ~ neg_binomial_2((incidence * p_reported), phi);
}
generated quantities {
    real R0 = (infection_rate / gamma);
    real recovery_time = (1.0 / gamma);
    vector[cases_n] cases_likelihood = neg_binomial_2_lpmfs(cases, (incidence * p_reported), phi);
    array[cases_n] int cases_gen = neg_binomial_2_int_rng(cases_n, (incidence * p_reported), phi);
}

The original article continues through fitting, diagnostics, prior-predictive simulation, and posterior-predictive checks. This page reproduces the model definitions rather than claiming those numerical analyses for the tiny shape data above.