Epidemic renewal models

A renewal model explains daily case counts through a latent infection process: today's infections are the recent infections, weighted by how infectious a case is some days after its own infection, times the reproduction number R(t). The statistical question is the path of R(t); everything between that path and the counts is a fixed mechanical map.

This page writes such a model with @brm, in three steps of growing size:

a reporting-delay model fitted to a linelist, which is where the reporting-delay distribution of the other two models comes from;
one population, where log R(t) is a random walk, plus a forecast and a prior predictive check obtained from the same declaration;
six coupled patches, where infections spread between patches, the patches share one R trend and deviate from it in a spatially correlated way.

Two things are worth seeing here. The statistical parts are formula lines — a reader who knows y ~ 1 + x + (1 | g) can read log_R ~ 1 + rw(time). The mechanical part stays what it is, a function: the renewal recursion is written once as Stan functions with StanBlocks.@deffun and called from the model by name.

All data are simulated from known parameter values, so every figure below shows the truth next to the estimate. The page is built from one checked-in script, research/epi_renewal/renewal.jl: the code blocks are read from it at build time, and the figures are drawn from the summaries it writes.

Formula interfaces for the reproduction number are established practice in R — epidemia models R with a regression formula that accepts random-walk terms, and EpiNow2 and epinowcast expose comparable building blocks. What differs here is the division of labour: the package knows nothing about epidemics, the terms rw and cdar are general-purpose (Formula terms), and the epidemic mechanics are user code.

The process

Let $w_{s}$ be the generation-interval distribution (the probability that a transmission happens $s = 1, \dots, 13$ days after the infector's own infection) and $π_{d}$ the reporting-delay distribution (the probability that an infection is reported $d = 0, \dots, 14$ days later). Both are fixed and enter the model as data. With infections $I_{t}$ , expected reported cases $Y_{t}$ and observed counts $y_{t}$ :

I_{t} = R_{t} \sum_{s = 1}^{13} w_{s} I_{t - s}, Y_{t} = \sum_{d = 0}^{14} π_{d} I_{t - d}, y_{t} \sim NegBin (mean Y_{t}, variance Y_{t} + c^{2} Y_{t}^{2}) .

The recursion needs infections before day 1. They are an exponential history $I_{s} = I_{0} e^{r (s - 1)}$ for $s \leq 0$ , whose growth rate $r$ is the one implied by the first reproduction number through the Euler–Lotka equation $\sum_{s} w_{s} e^{- r s} = 1 / R_{1}$ (two Newton steps from $r = 0$ ). So the infection process has one free scale, $\log I_{0}$ .

The reproduction number is a random walk on the log scale,

\log R_{t} = β_{0} + σ \sum_{u = 2}^{t} z_{u}, z_{u} \sim N (0, 1),

which on the formula surface is the line log_R ~ 1 + rw(time): the intercept is $\log R_{1}$ and rw(time) is the walk.

The mechanics, as Stan functions

The recursion carries state from one day to the next, and with several patches every patch's next day depends on all patches' past. That is a loop, so it is written as one: StanBlocks.@deffun turns the Julia-syntax definitions below into Stan functions that a model can call by name. Sizes such as T and G in the signatures are bound from the arguments. (The wastewater page runs a similar recursion inside a kernel(...) cell, one independent site per cell; here the patches are coupled, so the function takes the whole predictor vector through a top-level assignment.)

julia

    # ── renewal recursion: one population ──
    clamp2(x::real, lo::real, hi::real)::real = x < lo ? lo : (x > hi ? hi : x)
    # growth rate r implied by a reproduction number R (Euler–Lotka, two Newton steps from r = 0)
    growth_rate(R::real, gen_pmf::vector[G])::real = begin
        r = 0.0
        for iter in 1:2
            f = -1.0 / R
            df = 0.0
            for s in 1:G
                f += gen_pmf[s] * exp(-r * s)
                df -= s * gen_pmf[s] * exp(-r * s)
            end
            r = r - f / df
        end
        clamp2(r, -2.0, 2.0)
    end
    # sum_i pmf[i] * x[t - lag_i]; before day 1 the history is x0 * exp(r * (s - 1))
    lagged_sum(x::vector[T], t::int, pmf::vector[G], x0::real, r::real, first_lag::int)::real = begin
        acc = 0.0
        for i in 1:G
            s = t - (first_lag + i - 1)
            acc += pmf[i] * (s >= 1 ? x[s] : x0 * exp(r * (s - 1)))
        end
        acc
    end
    # infections I_t = R_t * sum_s g_s I_{t-s}, then expected reported cases Y_t = sum_d pi_d I_{t-d}
    expected_cases(log_R::vector[T], log_I0::real, gen_pmf::vector[G], delay_pmf::vector[D])::vector[T] = begin
        I0 = exp(log_I0)
        r0 = growth_rate(exp(log_R[1]), gen_pmf)
        I::vector[T]
        for t in 1:T
            I[t] = clamp2(exp(log_R[t]) * lagged_sum(I, t, gen_pmf, I0, r0, 1), 0.0, 1e15)
        end
        Y::vector[T]
        for t in 1:T
            Y[t] = lagged_sum(I, t, delay_pmf, I0, r0, 0)
        end
        Y
    end

The observation family is declared the same way. @lhs @lpxf makes cases ~ nb_cases(Y, cluster, observed) a valid likelihood statement; the _lpmfs variant returns the pointwise log likelihood and _rng the posterior predictive draw, and the generated model uses both. A row with observed == 0 is not scored, but the generator still draws it — that draw is the forecast for that row.

julia

    # ── observation family: negative-binomial counts on the rows marked observed ──
    # cases ~ NegBin(mean Y, variance Y + cluster^2 Y^2); a row with observed == 0 is not scored,
    # and the generator draws every row, so an unobserved row's draw is its forecast
    @lhs @lpxf nb_cases_lpmf(cases::int[N], Y::vector[N], cluster::real, observed::int[N])::real = begin
        lp = 0.0
        for i in 1:N
            if observed[i] == 1
                lp += neg_binomial_2_lpmf(cases[i], Y[i], 1.0 / (cluster * cluster))
            end
        end
        lp
    end
    nb_cases_lpmfs(cases::int[N], Y::vector[N], cluster::real, observed::int[N])::vector[N] = begin
        lp::vector[N]
        for i in 1:N
            lp[i] = observed[i] == 1 ? neg_binomial_2_lpmf(cases[i], Y[i], 1.0 / (cluster * cluster)) : 0.0
        end
        lp
    end
    nb_cases_rng(int[N], Y::vector[N], cluster::real, observed::int[N])::int[N] = begin
        out::int[N]
        for i in 1:N
            out[i] = neg_binomial_2_rng(Y[i], 1.0 / (cluster * cluster))
        end
        out
    end

Step 1: the reporting delay, from a linelist

A linelist records, for each case, the day of the event and the day of the report. Two things bias the delays it shows. Days are intervals, so an event "on day 3" happened somewhere within that day, and the report likewise (double interval censoring). And on the day of the analysis, a recent event is in the data only if its delay was short (right truncation), so during a growing epidemic — when most events are recent — the observed delays are far too short.

The model states the delay as LogNormal and corrects for both: the family censored_delay(mu, sigma, window) scores each delay $d$ with $P (d \leq P + T < d + 1 ∣ P + T < window)$ , where $P$ is uniform within the event day, $T$ is the LogNormal delay, and window is the number of days between the event and the analysis.

julia

    # ── reporting delay: doubly interval-censored, right-truncated LogNormal ──
    # F(x) = P(P + T < x) for P ~ Uniform(0, 1), T ~ LogNormal(mu, sigma), in closed form: G(x) - G(x - 1)
    lnorm_G(a::real, mu::real, sigma::real)::real =
        a <= 0.0 ? 0.0 : a * Phi((log(a) - mu) / sigma) - exp(mu + 0.5 * sigma * sigma) * Phi((log(a) - mu) / sigma - sigma)
    lnorm_F(x::real, mu::real, sigma::real)::real = lnorm_G(x, mu, sigma) - lnorm_G(x - 1.0, mu, sigma)
    # log P(delay = d | delay < window)
    censored_delay_lp1(d::int, mu::real, sigma::real, window::int)::real =
        log(lnorm_F(d + 1.0, mu, sigma) - lnorm_F(d + 0.0, mu, sigma)) - log(lnorm_F(window + 0.0, mu, sigma))
    @lhs @lpxf censored_delay_lpmf(d::int[N], mu::real, sigma::real, window::int[N])::real = begin
        lp = 0.0
        for j in 1:N
            lp += censored_delay_lp1(d[j], mu, sigma, window[j])
        end
        lp
    end
    censored_delay_lpmfs(d::int[N], mu::real, sigma::real, window::int[N])::vector[N] = begin
        lp::vector[N]
        for j in 1:N
            lp[j] = censored_delay_lp1(d[j], mu, sigma, window[j])
        end
        lp
    end
    censored_delay_rng(int[N], mu::real, sigma::real, window::int[N])::int[N] = begin
        out::int[N]
        for j in 1:N
            W = window[j]
            probs::vector[W]
            for k in 1:W
                probs[k] = lnorm_F(k + 0.0, mu, sigma) - lnorm_F(k - 1.0, mu, sigma)
            end
            out[j] = categorical_rng(probs / sum(probs)) - 1
        end
        out
    end
end

# ═══════════════════════════════════════════════════════════════════════════════

Every declaration on this page is rendered in the standard comparison panes: the @brm source, its intermediate representation, the StanBlocks model it lowers to, the generated Stan program, and the Turing backend. These are StanBlocks models. @deffun functions and @lpxf families are Stan code with no Turing counterpart, so the Turing pane reports each model as unsupported and names the statement it cannot lower.

brm-comparison

Reporting delay from a right-truncated linelist

julia

function reporting_delay_model(data = reporting_delay_data())
    @brm data begin
        mu    ~ Normal(1.0, 0.5; lower=0.0, upper=3.0)
        sigma ~ Normal(0.5, 0.25; lower=0.1, upper=2.0)
        delay ~ censored_delay(mu, sigma, window)
    end
end

julia

BRMI:
  mu ~ Normal(1.0, 0.5; lower=0.0, upper=3.0)
  sigma ~ Normal(0.5, 0.25; lower=0.1, upper=2.0)
  window: data (eltype=Int64, n=250)
  delay ~ censored_delay(mu, sigma, window)

julia

SBBRMI with data keys = [:delay, :window]
emitted @slic body:
begin
    mu ~ normal(1.0, 0.5; lower = 0.0, upper = 3.0)
    sigma ~ normal(0.5, 0.25; lower = 0.1, upper = 2.0)
    delay ~ censored_delay(mu, sigma, window)
end

stan

functions {
real censored_delay_lpmf(
    array[] int d,
    real mu,
    real sigma,
    array[] int window
) {
    int N = dims(d)[1];
    if (dims(window)[1] != N) reject("censored_delay_lpmf: dim mismatch — `window` dim 1 (= ", dims(window)[1], ") does not match `N` (= ", N, "), inferred from `d` dim 1. `N` sizes: `d` dim 1 (= ", dims(d)[1], "), `window` dim 1 (= ", dims(window)[1], ").");
    real lp = 0.0;
    for(j in 1:N) {
        lp += censored_delay_lp1(d[j], mu, sigma, window[j]);
    }
    return lp;
}
real censored_delay_lp1(
    int d,
    real mu,
    real sigma,
    int window
) {
    return (
        log((lnorm_F((d + 1.0), mu, sigma) - lnorm_F((d + 0.0), mu, sigma))) -
        log(lnorm_F((window + 0.0), mu, sigma))
    );
}
real lnorm_F(
    real x,
    real mu,
    real sigma
) {
    return (lnorm_G(x, mu, sigma) - lnorm_G((x - 1.0), mu, sigma));
}
real lnorm_G(
    real a,
    real mu,
    real sigma
) {
    return ((a <= 0.0) ? 0.0 : (
        (a * Phi(((log(a) - mu) / sigma))) -
        (exp((mu + (0.5 * sigma * sigma))) * Phi((((log(a) - mu) / sigma) - sigma)))
    ));
}
vector censored_delay_lpmfs(
    array[] int d,
    real mu,
    real sigma,
    array[] int window
) {
    int N = dims(d)[1];
    if (dims(window)[1] != N) reject("censored_delay_lpmfs: dim mismatch — `window` dim 1 (= ", dims(window)[1], ") does not match `N` (= ", N, "), inferred from `d` dim 1. `N` sizes: `d` dim 1 (= ", dims(d)[1], "), `window` dim 1 (= ", dims(window)[1], ").");
    vector[N] lp;
    for(j in 1:N) {
        lp[j] = censored_delay_lp1(d[j], mu, sigma, window[j]);
    }
    return lp;
}
array[] int censored_delay_int_rng(
    int anontok__1,
    real mu,
    real sigma,
    array[] int window
) {
    int N = anontok__1;
    if (dims(window)[1] != N) reject("censored_delay_rng: dim mismatch — `window` dim 1 (= ", dims(window)[1], ") does not match `N` (= ", N, "), inferred from `anontok__1` dim 1. `N` sizes: `anontok__1` dim 1 (= ", anontok__1, "), `window` dim 1 (= ", dims(window)[1], ").");
    array[N] int out;
    for(j in 1:N) {
        int W = window[j];
        vector[W] probs;
        for(k in 1:W) {
            probs[k] = (lnorm_F((k + 0.0), mu, sigma) - lnorm_F((k - 1.0), mu, sigma));
        }
        out[j] = (categorical_rng((probs / sum(probs))) - 1);
    }
    return out;
}
}
data {
    int delay_n;
    array[delay_n] int delay;
    int window_n;
    array[window_n] int window;
}
transformed data {
}
parameters {
    real<lower=0.0, upper=3.0> mu;
    real<lower=0.1, upper=2.0> sigma;
}
transformed parameters {
}
model {
    mu ~ normal(1.0, 0.5);
    sigma ~ normal(0.5, 0.25);
    delay ~ censored_delay(mu, sigma, window);
}
generated quantities {
    vector[window_n] delay_likelihood = censored_delay_lpmfs(delay, mu, sigma, window);
    array[window_n] int delay_gen = censored_delay_int_rng(delay_n, mu, sigma, window);
}

julia

Turing unsupported for this BRM example

Turing backend: observation `delay` calls `censored_delay`, which defines no Julia methods (a Stan `@lpxf` family has no Julia implementation). Each observation row is sampled from the Julia distribution this call returns, so a Stan-only family cannot execute. Use a Distributions.jl callable or fit this model with the Stan backend.

The simulated linelist holds 250 events from a growing epidemic, seen on day 21; the true delay is LogNormal(1.5, 0.5). The points are the delays as they appear in the linelist: events of the last days before the analysis can only be in the data with a delay of zero or one day, so short delays are heavily over-represented and the mean observed delay is 3.7 days against a true 5.5. The dashed line is the truth, and the bands are the posterior of the daily delay distribution (50 %, 80 % and 95 %), which undoes the truncation.

Step 2: one population

One row per day. The data carry the counts, the two fixed distributions gen_pmf and delay_pmf, and the observed mask.

brm-comparison

Renewal model of one population

julia

function renewal_single_model(data = renewal_single_data())
    @brm data begin
        log_I0  ~ Normal(log(50.0), 0.5)                    # log initial infections (scale of the seeded history)
        cluster ~ Normal(0.0, 0.1; lower=0.0)               # overdispersion of the counts
        log_R   ~ 1 + rw(time)                              # log R_t: a random walk over days
        effect(log_R, Intercept) ~ Normal(log(1.3), 0.1)    # ... its first value, log R_1
        sd(:, rw(time)) ~ Normal(0.0, 0.05)                 # ... its daily innovation scale
        Y       = expected_cases(log_R, log_I0, gen_pmf, delay_pmf)
        cases   ~ nb_cases(Y, cluster, observed)
    end
end

julia

BRMI:
  log_I0 ~ Normal(log(50.0), 0.5)
  cluster ~ Normal(0.0, 0.1; lower=0.0)
  time: data (eltype=Float64, n=56)
  log_R ~ 1 + rw(time)
  effect(log_R, Intercept) ~ Normal(log(1.3), 0.1)
  effect(term_sd, rw(time), :) ~ Normal(0.0, 0.05)
  gen_pmf: data (eltype=Float64, n=13)
  delay_pmf: data (eltype=Float64, n=15)
  :Y = expected_cases(log_R, log_I0, gen_pmf, delay_pmf)
  observed: data (eltype=Int64, n=56)
  cases ~ nb_cases(Y, cluster, observed)

julia

SBBRMI with data keys = [:cases, :delay_pmf, :gen_pmf, :observed, :rw_log_R_time_n_steps, :rw_log_R_time_time_idx, :time]
configured submodels:
_sb_rw1_configured_1 = Base.merge(BayesianRegressionModels._sb_rw1, quote
            sigma ~ normal(0.0, 0.05; lower = 0.0)
        end)
emitted @slic body:
begin
    log_I0 ~ normal(3.912023005428146, 0.5)
    cluster ~ normal(0.0, 0.1; lower = 0.0)
    X_log_R = hcat(rep_vector(1.0, num_elements(time)))
    pop_log_R ~ _popefs_normal(; X = X_log_R, beta_loc = [0.26236426446749106], beta_scale = [0.1])
    rw_log_R_time ~ _sb_rw1_configured_1(; n_steps = rw_log_R_time_n_steps, time_idx = rw_log_R_time_time_idx)
    log_R = pop_log_R + rw_log_R_time
    Y = (Main.renewal.expected_cases)(log_R, log_I0, gen_pmf, delay_pmf)
    cases ~ nb_cases(Y, cluster, observed)
end

stan

functions {
matrix hcat(vector x) {
    int n = dims(x)[1];
    return to_matrix(x, n, 1);
}
vector random_walk_rows(
    real sigma,
    vector z,
    array[] int idx
) {
    int N = dims(idx)[1];
    vector[(dims(z)[1] + 1)] x = random_walk_path(sigma, z);
    vector[N] out;
    for(i in 1:N) {
        out[i] = x[idx[i]];
    }
    return out;
}
vector random_walk_path(
    real sigma,
    vector z
) {
    int n = dims(z)[1];
    return append_row(0.0, (sigma * cumulative_sum(z)));
}
vector expected_cases(
    vector log_R,
    real log_I0,
    vector gen_pmf,
    vector delay_pmf
) {
    int T = dims(log_R)[1];
    real I0 = exp(log_I0);
    real r0 = growth_rate(exp(log_R[1]), gen_pmf);
    vector[T] I;
    for(t in 1:T) {
        I[t] = clamp2((exp(log_R[t]) * lagged_sum(I, t, gen_pmf, I0, r0, 1)), 0.0, 1.0e15);
    }
    vector[T] Y;
    for(t in 1:T) {
        Y[t] = lagged_sum(I, t, delay_pmf, I0, r0, 0);
    }
    return Y;
}
real growth_rate(
    real R,
    vector gen_pmf
) {
    int G = dims(gen_pmf)[1];
    real r = 0.0;
    for(iter in 1:2) {
        real f = (-1.0 / R);
        real df = 0.0;
        for(s in 1:G) {
            f += (gen_pmf[s] * exp(((-r) * s)));
            df -= (s * gen_pmf[s] * exp(((-r) * s)));
        }
        r = (r - (f / df));
    }
    return clamp2(r, -2.0, 2.0);
}
real clamp2(real x, real lo, real hi) {
    return ((x < lo) ? lo : ((x > hi) ? hi : x));
}
real lagged_sum(
    vector x,
    int t,
    vector pmf,
    real x0,
    real r,
    int first_lag
) {
    int G = dims(pmf)[1];
    real acc = 0.0;
    for(i in 1:G) {
        int s = (t - ((first_lag + i) - 1));
        acc += (pmf[i] * ((s >= 1) ? x[s] : (x0 * exp((r * (s - 1))))));
    }
    return acc;
}
real nb_cases_lpmf(
    array[] int cases,
    vector Y,
    real cluster,
    array[] int observed
) {
    int N = dims(cases)[1];
    if (dims(Y)[1] != N) reject("nb_cases_lpmf: dim mismatch — `Y` dim 1 (= ", dims(Y)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    if (dims(observed)[1] != N) reject("nb_cases_lpmf: dim mismatch — `observed` dim 1 (= ", dims(observed)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    real lp = 0.0;
    for(i in 1:N) {
        if((observed[i] == 1)) {
            lp += neg_binomial_2_lpmf(cases[i] | Y[i], (1.0 / (cluster * cluster)));
        }
    }
    return lp;
}
vector nb_cases_lpmfs(
    array[] int cases,
    vector Y,
    real cluster,
    array[] int observed
) {
    int N = dims(cases)[1];
    if (dims(Y)[1] != N) reject("nb_cases_lpmfs: dim mismatch — `Y` dim 1 (= ", dims(Y)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    if (dims(observed)[1] != N) reject("nb_cases_lpmfs: dim mismatch — `observed` dim 1 (= ", dims(observed)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    vector[N] lp;
    for(i in 1:N) {
        lp[i] = ((observed[i] == 1) ? neg_binomial_2_lpmf(cases[i] | Y[i], (1.0 / (cluster * cluster))) : 0.0);
    }
    return lp;
}
array[] int nb_cases_int_rng(
    int anontok__1,
    vector Y,
    real cluster,
    array[] int observed
) {
    int N = anontok__1;
    if (dims(Y)[1] != N) reject("nb_cases_rng: dim mismatch — `Y` dim 1 (= ", dims(Y)[1], ") does not match `N` (= ", N, "), inferred from `anontok__1` dim 1. `N` sizes: `anontok__1` dim 1 (= ", anontok__1, "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    if (dims(observed)[1] != N) reject("nb_cases_rng: dim mismatch — `observed` dim 1 (= ", dims(observed)[1], ") does not match `N` (= ", N, "), inferred from `anontok__1` dim 1. `N` sizes: `anontok__1` dim 1 (= ", anontok__1, "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    array[N] int out;
    for(i in 1:N) {
        out[i] = neg_binomial_2_rng(Y[i], (1.0 / (cluster * cluster)));
    }
    return out;
}
}
data {
    int time_n;
    vector[time_n] time;
    int rw_log_R_time_n_steps;
    int rw_log_R_time_time_idx_n;
    array[rw_log_R_time_time_idx_n] int rw_log_R_time_time_idx;
    int gen_pmf_n;
    vector[gen_pmf_n] gen_pmf;
    int delay_pmf_n;
    vector[delay_pmf_n] delay_pmf;
    int cases_n;
    array[cases_n] int cases;
    int observed_n;
    array[observed_n] int observed;
}
transformed data {
    matrix[num_elements(time), 1] X_log_R = hcat(rep_vector(1.0, num_elements(time)));
    int pop_log_R_n_covariates = 1;
}
parameters {
    real log_I0;
    real<lower=0.0> cluster;
    vector[pop_log_R_n_covariates] pop_log_R_beta_pop;
    real<lower=0.0> rw_log_R_time_sigma;
    vector[(rw_log_R_time_n_steps - 1)] rw_log_R_time_z;
}
transformed parameters {
    vector[num_elements(time)] pop_log_R = (X_log_R * pop_log_R_beta_pop);
    vector[rw_log_R_time_time_idx_n] rw_log_R_time = random_walk_rows(rw_log_R_time_sigma, rw_log_R_time_z, rw_log_R_time_time_idx);
    vector[num_elements(time)] log_R = (pop_log_R + rw_log_R_time);
    vector[num_elements(time)] Y = expected_cases(log_R, log_I0, gen_pmf, delay_pmf);
}
model {
    log_I0 ~ normal(3.912023005428146, 0.5);
    cluster ~ normal(0.0, 0.1);
    pop_log_R_beta_pop ~ normal([0.26236426446749106]', [0.1]');
    rw_log_R_time_sigma ~ normal(0.0, 0.05);
    rw_log_R_time_z ~ std_normal();
    cases ~ nb_cases(Y, cluster, observed);
}
generated quantities {
    vector[observed_n] cases_likelihood = nb_cases_lpmfs(cases, Y, cluster, observed);
    array[observed_n] int cases_gen = nb_cases_int_rng(cases_n, Y, cluster, observed);
}

julia

Turing unsupported for this BRM example

Turing backend: assignment `Y` calls `expected_cases`, which defines no Julia methods (a Stan `@deffun` function has no Julia implementation). Row-dependent assignments lower to row-wise Julia comprehensions, so this call cannot execute. Express the computation with row-wise Julia code or fit this model with the Stan backend.

Reading the declaration top to bottom:

log_I0 and cluster are scalar parameters with a prior each.
log_R ~ 1 + rw(time) is a predictor like any other. Its parts are addressed the usual way: effect(log_R, Intercept) is $\log R_{1}$ and sd(:, rw(time)) is the walk's daily scale $σ$ .
Y = expected_cases(...) is a top-level assignment: the whole predictor vector goes into the Stan function, and Y becomes a named quantity of the model.
cases ~ nb_cases(Y, cluster, observed) is the likelihood.

Fifty-six days of counts recover the level and the trend of R(t). The posterior is a smooth path: the counts cannot resolve the day-to-day steps of the true walk, and the walk's scale is only weakly identified (see the table at the end), so single-day excursions of the truth leave the 95 % band. The last days are the least certain, because their infections have mostly not been reported yet.

The posterior predictive bands of the daily counts (cases_gen in the generated model) against the reported counts and the true expectation, on a logarithmic axis:

A prior predictive check from the same declaration

Building the same model with the response column omitted from the data drops every likelihood contribution, so sampling it (fixed_param) draws from the prior, and the forward-simulated cases becomes the prior predictive distribution. By day 56 the prior's 95 % band runs from a handful of daily cases to more than a million; the counts are what pins R(t) down.

A forecast is a mask

Setting observed to zero from day 43 on turns the last two weeks into a forecast. Nothing else changes: the walk rw(time) still spans all 56 days, the innovations of the unobserved days are informed by the prior only, and the recursion carries the uncertainty forward into the counts.

A random walk forecasts "no change": the band fans out around the last estimated level. In this simulation the true R(t) started to fall a few days before day 42 — too late to show in the counts reported by then — and kept falling, so the held-out counts end up at the lower edge of the predictive band.

Step 3: six coupled patches

Six patches of different population sizes sit on a 100 km square. The outbreak starts in the smallest one; the others are seeded with a small fraction of a case and catch the epidemic through mixing.

Infection pressure on patch $g$ is a mixture over all patches $h$ , with gravity weights that grow with both population sizes $N$ and fall with distance $d$ ,

I_{g, t} = R_{g, t} \sum_{h} K_{g h} \sum_{s} w_{s} I_{h, t - s}, K_{g h} \propto {\begin{cases} 1 & g = h \\ \frac{(N_{g} / \bar{N}) (N_{h} / \bar{N})}{d_{g h}^{γ}} & g \neq h \end{cases} (rows sum to one),

and every patch has its own reproduction number: a shared trend plus a weekly deviation,

\log R_{g, t} = β_{0} + W_{t} + δ_{g, w (t)}, δ_{\cdot, w} = ρ δ_{\cdot, w - 1} + σ_{δ} \sqrt{1 - ρ^{2}} L η_{w}, L L^{⊤} = C, C_{g h} = e^{- d_{g h} / 30} .

$W_{t}$ is the same random walk as before. The deviations are damped ( $ρ$ ) and their innovations are correlated between patches: neighbouring patches deviate together.

The mechanics gain a mixing step; the rest is the single-population recursion per patch.

julia

    # ── renewal recursion: coupled patches ──
    # rows of the long (day, patch) frame are ordered by patch, then day: row i = t + (g - 1) * T,
    # so `to_matrix(x, T, P)` puts patch g in column g
    dist_at(dist_flat::vector[PP], g::int, h::int, P::int)::real = dist_flat[g + (h - 1) * P]
    # gravity mixing: K[g, h] is the share of patch g's infection pressure that comes from patch h
    gravity_K(pop::vector[P], dist_flat::vector[PP], gamma::real)::matrix[P, P] = begin
        mean_pop = sum(pop) / P
        K::matrix[P, P]
        for g in 1:P
            rs = 0.0
            for h in 1:P
                K[g, h] = g == h ? 1.0 : (pop[g] / mean_pop) * (pop[h] / mean_pop) / exp(gamma * log(dist_at(dist_flat, g, h, P)))
                rs += K[g, h]
            end
            for h in 1:P
                K[g, h] = K[g, h] / rs
            end
        end
        K
    end
    # I[t, g] = R[t, g] * sum_h K[g, h] * lambda[h, t],  lambda[h, t] = sum_s g_s I[t - s, h]
    patch_infections(R::matrix[T, P], K::matrix[P, P], log_I0::vector[P], gen_pmf::vector[G])::matrix[T, P] = begin
        I0 = exp(log_I0)
        growth::vector[P]
        for h in 1:P
            growth[h] = growth_rate(R[1, h], gen_pmf)
        end
        I::matrix[T, P]
        lambda::vector[P]
        for t in 1:T
            for h in 1:P
                lambda[h] = lagged_sum(col(I, h), t, gen_pmf, I0[h], growth[h], 1)
            end
            for g in 1:P
                pressure = 0.0
                for h in 1:P
                    pressure += K[g, h] * lambda[h]
                end
                I[t, g] = clamp2(R[t, g] * pressure, 0.0, 1e15)
            end
        end
        I
    end
    # the seed predictor is per row and constant within a patch: read each patch's first row
    patch_seeds(log_I0_row::vector[N], pop::vector[P])::vector[P] = begin
        T = N / P
        out::vector[P]
        for g in 1:P
            out[g] = log_I0_row[(g - 1) * T + 1]
        end
        out
    end
    patch_expected_cases(log_R::vector[N], log_I0::vector[P], gamma::real, pop::vector[P], dist_flat::vector[PP],
                         gen_pmf::vector[G], delay_pmf::vector[D])::vector[N] = begin
        T = N / P
        R = to_matrix(exp(log_R), T, P)
        I = patch_infections(R, gravity_K(pop, dist_flat, gamma), log_I0, gen_pmf)
        I0 = exp(log_I0)
        Y::matrix[T, P]
        for g in 1:P
            growth = growth_rate(R[1, g], gen_pmf)
            for t in 1:T
                Y[t, g] = lagged_sum(col(I, g), t, delay_pmf, I0[g], growth, 0)
            end
        end
        to_vector(Y)
    end

The data are one row per (day, patch), 336 rows, ordered by patch and then by day. C is a matrix-valued field; pop, dist_flat and the two distributions are shared vectors.

brm-comparison

Renewal model of six coupled patches

julia

function renewal_patch_model(data = renewal_patch_data())
    @brm data begin
        gamma   ~ Normal(1.5, 0.5; lower=0.0)               # distance decay of the gravity mixing
        cluster ~ Normal(0.0, 0.1; lower=0.0)               # overdispersion of the counts
        log_R   ~ 1 + rw(time) + cdar(week; by=patch, cor=C)   # shared walk + correlated weekly patch deviations
        effect(log_R, Intercept) ~ Normal(log(1.3), 0.1)    # the walk's first value
        sd(:, rw(time)) ~ Normal(0.0, 0.05)                 # the walk's daily innovation scale
        sd(:, cdar(week)) ~ Normal(0.0, 0.2)                # the deviations' stationary scale
        ar(:, cdar(week)) ~ Normal(0.8, 0.1)                # the deviations' week-to-week persistence
        log_I0  ~ 0 + offset(seed_mean) + factor(patch)     # one seed per patch: prior mean + a cell mean
        effect(log_I0, patch) ~ Normal(0.0, 0.5)
        seeds   = patch_seeds(log_I0, pop)
        Y       = patch_expected_cases(log_R, seeds, gamma, pop, dist_flat, gen_pmf, delay_pmf)
        cases   ~ nb_cases(Y, cluster, observed)
    end
end

# ═══════════════════════════════════════════════════════════════════════════════
# 6. Fit (NUTS via WarmupHMC on the BridgeStan log density), constrain, summarise
# ═══════════════════════════════════════════════════════════════════════════════

julia

BRMI:
  gamma ~ Normal(1.5, 0.5; lower=0.0)
  cluster ~ Normal(0.0, 0.1; lower=0.0)
  time: data (eltype=Float64, n=336)
  patch: data (eltype=Int64, n=336)
  C: data (eltype=Float64, n=36)
  week: data (eltype=Int64, n=336)
  log_R ~ 1 + rw(time) + cdar(week; by=patch, cor=C)
  effect(log_R, Intercept) ~ Normal(log(1.3), 0.1)
  effect(term_sd, rw(time), :) ~ Normal(0.0, 0.05)
  effect(term_sd, cdar(week), :) ~ Normal(0.0, 0.2)
  effect(term_ar, cdar(week), :) ~ Normal(0.8, 0.1)
  seed_mean: data (eltype=Float64, n=336)
  log_I0 ~ 0 + offset(seed_mean) + factor(patch)
  effect(log_I0, patch) ~ Normal(0.0, 0.5)
  pop: data (eltype=Float64, n=6)
  :seeds = patch_seeds(log_I0, pop)
  dist_flat: data (eltype=Float64, n=36)
  gen_pmf: data (eltype=Float64, n=13)
  delay_pmf: data (eltype=Float64, n=15)
  :Y = patch_expected_cases(log_R, seeds, gamma, pop, dist_flat, gen_pmf, delay_pmf)
  observed: data (eltype=Int64, n=336)
  cases ~ nb_cases(Y, cluster, observed)

julia

SBBRMI with data keys = [:C, :cases, :cdar_log_R_week_L, :cdar_log_R_week_group_idx, :cdar_log_R_week_n_groups, :cdar_log_R_week_n_steps, :cdar_log_R_week_step_idx, :delay_pmf, :dist_flat, :gen_pmf, :observed, :patch, :patch_idx, :patch_n_levels, :pop, :rw_log_R_time_n_steps, :rw_log_R_time_time_idx, :seed_mean, :time, :week]
configured submodels:
_sb_rw1_configured_1 = Base.merge(BayesianRegressionModels._sb_rw1, quote
            sigma ~ normal(0.0, 0.05; lower = 0.0)
        end)
_sb_cdar_configured_1 = Base.merge(BayesianRegressionModels._sb_cdar, quote
            rho ~ normal(0.8, 0.1; lower = 0.0, upper = 1.0)
        end)
emitted @slic body:
begin
    gamma ~ normal(1.5, 0.5; lower = 0.0)
    cluster ~ normal(0.0, 0.1; lower = 0.0)
    X_log_R = hcat(rep_vector(1.0, num_elements(time)))
    pop_log_R ~ _popefs_normal(; X = X_log_R, beta_loc = [0.26236426446749106], beta_scale = [0.1])
    rw_log_R_time ~ _sb_rw1_configured_1(; n_steps = rw_log_R_time_n_steps, time_idx = rw_log_R_time_time_idx)
    cdar_log_R_week ~ _sb_cdar_configured_1(; n_groups = cdar_log_R_week_n_groups, n_steps = cdar_log_R_week_n_steps, L = cdar_log_R_week_L, group_idx = cdar_log_R_week_group_idx, step_idx = cdar_log_R_week_step_idx)
    log_R = pop_log_R + rw_log_R_time + cdar_log_R_week
    cat_log_I0_patch ~ _sb_cat_cells_normal(; x = patch_idx, n_levels = patch_n_levels, beta_loc = 0.0, beta_scale = 0.5)
    log_I0 = seed_mean + cat_log_I0_patch
    seeds = (Main.renewal.patch_seeds)(log_I0, pop)
    Y = (Main.renewal.patch_expected_cases)(log_R, seeds, gamma, pop, dist_flat, gen_pmf, delay_pmf)
    cases ~ nb_cases(Y, cluster, observed)
end

stan

functions {
matrix hcat(vector x) {
    int n = dims(x)[1];
    return to_matrix(x, n, 1);
}
vector random_walk_rows(
    real sigma,
    vector z,
    array[] int idx
) {
    int N = dims(idx)[1];
    vector[(dims(z)[1] + 1)] x = random_walk_path(sigma, z);
    vector[N] out;
    for(i in 1:N) {
        out[i] = x[idx[i]];
    }
    return out;
}
vector random_walk_path(
    real sigma,
    vector z
) {
    int n = dims(z)[1];
    return append_row(0.0, (sigma * cumulative_sum(z)));
}
vector correlated_damped_walk(
    real sigma,
    real rho,
    vector eta,
    matrix L,
    int W,
    array[] int group_idx,
    array[] int step_idx
) {
    int P = dims(L)[1];
    int N = dims(group_idx)[1];
    if (dims(L)[2] != P) reject("correlated_damped_walk: dim mismatch — `L` dim 2 (= ", dims(L)[2], ") does not match `P` (= ", P, "), inferred from `L` dim 1. `P` sizes: `L` dim 1 (= ", dims(L)[1], "), `L` dim 2 (= ", dims(L)[2], ").");
    if (dims(step_idx)[1] != N) reject("correlated_damped_walk: dim mismatch — `step_idx` dim 1 (= ", dims(step_idx)[1], ") does not match `N` (= ", N, "), inferred from `group_idx` dim 1. `N` sizes: `group_idx` dim 1 (= ", dims(group_idx)[1], "), `step_idx` dim 1 (= ", dims(step_idx)[1], ").");
    matrix[P, W] E = to_matrix(eta, P, W);
    matrix[P, W] delta;
    vector[dims(L)[1]] prev = (sigma * (L * col(E, 1)));
    delta[:, 1] = prev;
    real scale = (sigma * sqrt((1.0 - (rho * rho))));
    for(w in 2:W) {
        prev = ((rho * prev) + (scale * (L * col(E, w))));
        delta[:, w] = prev;
    }
    vector[N] out;
    for(i in 1:N) {
        out[i] = delta[group_idx[i], step_idx[i]];
    }
    return out;
}
vector patch_seeds(
    vector log_I0_row,
    vector pop
) {
    int N = dims(log_I0_row)[1];
    int P = dims(pop)[1];
    int T = (N / P);
    vector[P] out;
    for(g in 1:P) {
        out[g] = log_I0_row[(((g - 1) * T) + 1)];
    }
    return out;
}
vector patch_expected_cases(
    vector log_R,
    vector log_I0,
    real gamma,
    vector pop,
    vector dist_flat,
    vector gen_pmf,
    vector delay_pmf
) {
    int N = dims(log_R)[1];
    int P = dims(log_I0)[1];
    if (dims(pop)[1] != P) reject("patch_expected_cases: dim mismatch — `pop` dim 1 (= ", dims(pop)[1], ") does not match `P` (= ", P, "), inferred from `log_I0` dim 1. `P` sizes: `log_I0` dim 1 (= ", dims(log_I0)[1], "), `pop` dim 1 (= ", dims(pop)[1], ").");
    int T = (N / P);
    matrix[T, P] R = to_matrix(exp(log_R), T, P);
    matrix[T, dims(log_I0)[1]] I = patch_infections(R, gravity_K(pop, dist_flat, gamma), log_I0, gen_pmf);
    vector[dims(log_I0)[1]] I0 = exp(log_I0);
    matrix[T, P] Y;
    for(g in 1:P) {
        real growth = growth_rate(R[1, g], gen_pmf);
        for(t in 1:T) {
            Y[t, g] = lagged_sum(col(I, g), t, delay_pmf, I0[g], growth, 0);
        }
    }
    return to_vector(Y);
}
matrix patch_infections(
    matrix R,
    matrix K,
    vector log_I0,
    vector gen_pmf
) {
    int T = dims(R)[1];
    int P = dims(R)[2];
    if (dims(K)[1] != P) reject("patch_infections: dim mismatch — `K` dim 1 (= ", dims(K)[1], ") does not match `P` (= ", P, "), inferred from `R` dim 2. `P` sizes: `R` dim 2 (= ", dims(R)[2], "), `K` dim 1 (= ", dims(K)[1], "), `K` dim 2 (= ", dims(K)[2], "), `log_I0` dim 1 (= ", dims(log_I0)[1], ").");
    if (dims(K)[2] != P) reject("patch_infections: dim mismatch — `K` dim 2 (= ", dims(K)[2], ") does not match `P` (= ", P, "), inferred from `R` dim 2. `P` sizes: `R` dim 2 (= ", dims(R)[2], "), `K` dim 1 (= ", dims(K)[1], "), `K` dim 2 (= ", dims(K)[2], "), `log_I0` dim 1 (= ", dims(log_I0)[1], ").");
    if (dims(log_I0)[1] != P) reject("patch_infections: dim mismatch — `log_I0` dim 1 (= ", dims(log_I0)[1], ") does not match `P` (= ", P, "), inferred from `R` dim 2. `P` sizes: `R` dim 2 (= ", dims(R)[2], "), `K` dim 1 (= ", dims(K)[1], "), `K` dim 2 (= ", dims(K)[2], "), `log_I0` dim 1 (= ", dims(log_I0)[1], ").");
    vector[dims(log_I0)[1]] I0 = exp(log_I0);
    vector[P] growth;
    for(h in 1:P) {
        growth[h] = growth_rate(R[1, h], gen_pmf);
    }
    matrix[T, P] I;
    vector[P] lambda;
    for(t in 1:T) {
        for(h in 1:P) {
            lambda[h] = lagged_sum(col(I, h), t, gen_pmf, I0[h], growth[h], 1);
        }
        for(g in 1:P) {
            real pressure = 0.0;
            for(h in 1:P) {
                pressure += (K[g, h] * lambda[h]);
            }
            I[t, g] = clamp2((R[t, g] * pressure), 0.0, 1.0e15);
        }
    }
    return I;
}
real growth_rate(
    real R,
    vector gen_pmf
) {
    int G = dims(gen_pmf)[1];
    real r = 0.0;
    for(iter in 1:2) {
        real f = (-1.0 / R);
        real df = 0.0;
        for(s in 1:G) {
            f += (gen_pmf[s] * exp(((-r) * s)));
            df -= (s * gen_pmf[s] * exp(((-r) * s)));
        }
        r = (r - (f / df));
    }
    return clamp2(r, -2.0, 2.0);
}
real clamp2(real x, real lo, real hi) {
    return ((x < lo) ? lo : ((x > hi) ? hi : x));
}
real lagged_sum(
    vector x,
    int t,
    vector pmf,
    real x0,
    real r,
    int first_lag
) {
    int G = dims(pmf)[1];
    real acc = 0.0;
    for(i in 1:G) {
        int s = (t - ((first_lag + i) - 1));
        acc += (pmf[i] * ((s >= 1) ? x[s] : (x0 * exp((r * (s - 1))))));
    }
    return acc;
}
matrix gravity_K(
    vector pop,
    vector dist_flat,
    real gamma
) {
    int P = dims(pop)[1];
    real mean_pop = (sum(pop) / P);
    matrix[P, P] K;
    for(g in 1:P) {
        real rs = 0.0;
        for(h in 1:P) {
            K[g, h] = ((g == h) ? 1.0 : (((pop[g] / mean_pop) * (pop[h] / mean_pop)) / exp((gamma * log(dist_at(dist_flat, g, h, P))))));
            rs += K[g, h];
        }
        for(h in 1:P) {
            K[g, h] = (K[g, h] / rs);
        }
    }
    return K;
}
real dist_at(vector dist_flat, int g, int h, int P) {
    return dist_flat[(g + ((h - 1) * P))];
}
real nb_cases_lpmf(
    array[] int cases,
    vector Y,
    real cluster,
    array[] int observed
) {
    int N = dims(cases)[1];
    if (dims(Y)[1] != N) reject("nb_cases_lpmf: dim mismatch — `Y` dim 1 (= ", dims(Y)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    if (dims(observed)[1] != N) reject("nb_cases_lpmf: dim mismatch — `observed` dim 1 (= ", dims(observed)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    real lp = 0.0;
    for(i in 1:N) {
        if((observed[i] == 1)) {
            lp += neg_binomial_2_lpmf(cases[i] | Y[i], (1.0 / (cluster * cluster)));
        }
    }
    return lp;
}
vector nb_cases_lpmfs(
    array[] int cases,
    vector Y,
    real cluster,
    array[] int observed
) {
    int N = dims(cases)[1];
    if (dims(Y)[1] != N) reject("nb_cases_lpmfs: dim mismatch — `Y` dim 1 (= ", dims(Y)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    if (dims(observed)[1] != N) reject("nb_cases_lpmfs: dim mismatch — `observed` dim 1 (= ", dims(observed)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    vector[N] lp;
    for(i in 1:N) {
        lp[i] = ((observed[i] == 1) ? neg_binomial_2_lpmf(cases[i] | Y[i], (1.0 / (cluster * cluster))) : 0.0);
    }
    return lp;
}
array[] int nb_cases_int_rng(
    int anontok__1,
    vector Y,
    real cluster,
    array[] int observed
) {
    int N = anontok__1;
    if (dims(Y)[1] != N) reject("nb_cases_rng: dim mismatch — `Y` dim 1 (= ", dims(Y)[1], ") does not match `N` (= ", N, "), inferred from `anontok__1` dim 1. `N` sizes: `anontok__1` dim 1 (= ", anontok__1, "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    if (dims(observed)[1] != N) reject("nb_cases_rng: dim mismatch — `observed` dim 1 (= ", dims(observed)[1], ") does not match `N` (= ", N, "), inferred from `anontok__1` dim 1. `N` sizes: `anontok__1` dim 1 (= ", anontok__1, "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    array[N] int out;
    for(i in 1:N) {
        out[i] = neg_binomial_2_rng(Y[i], (1.0 / (cluster * cluster)));
    }
    return out;
}
}
data {
    int time_n;
    vector[time_n] time;
    int rw_log_R_time_n_steps;
    int rw_log_R_time_time_idx_n;
    array[rw_log_R_time_time_idx_n] int rw_log_R_time_time_idx;
    int cdar_log_R_week_n_groups;
    int cdar_log_R_week_n_steps;
    int cdar_log_R_week_step_idx_n;
    int cdar_log_R_week_L_m;
    int cdar_log_R_week_L_n;
    matrix[cdar_log_R_week_L_m, cdar_log_R_week_L_n] cdar_log_R_week_L;
    int cdar_log_R_week_group_idx_n;
    array[cdar_log_R_week_group_idx_n] int cdar_log_R_week_group_idx;
    array[cdar_log_R_week_step_idx_n] int cdar_log_R_week_step_idx;
    int patch_n_levels;
    int patch_idx_n;
    array[patch_idx_n] int patch_idx;
    int seed_mean_n;
    vector[seed_mean_n] seed_mean;
    int pop_n;
    vector[pop_n] pop;
    int dist_flat_n;
    vector[dist_flat_n] dist_flat;
    int gen_pmf_n;
    vector[gen_pmf_n] gen_pmf;
    int delay_pmf_n;
    vector[delay_pmf_n] delay_pmf;
    int cases_n;
    array[cases_n] int cases;
    int observed_n;
    array[observed_n] int observed;
}
transformed data {
    matrix[num_elements(time), 1] X_log_R = hcat(rep_vector(1.0, num_elements(time)));
    int pop_log_R_n_covariates = 1;
}
parameters {
    real<lower=0.0> gamma;
    real<lower=0.0> cluster;
    vector[pop_log_R_n_covariates] pop_log_R_beta_pop;
    real<lower=0.0> rw_log_R_time_sigma;
    vector[(rw_log_R_time_n_steps - 1)] rw_log_R_time_z;
    real<lower=0.0> cdar_log_R_week_sigma;
    real<lower=0.0, upper=1.0> cdar_log_R_week_rho;
    vector[(cdar_log_R_week_n_groups * cdar_log_R_week_n_steps)] cdar_log_R_week_eta;
    vector[patch_n_levels] cat_log_I0_patch_beta;
}
transformed parameters {
    vector[num_elements(time)] pop_log_R = (X_log_R * pop_log_R_beta_pop);
    vector[rw_log_R_time_time_idx_n] rw_log_R_time = random_walk_rows(rw_log_R_time_sigma, rw_log_R_time_z, rw_log_R_time_time_idx);
    vector[cdar_log_R_week_step_idx_n] cdar_log_R_week = correlated_damped_walk(
        cdar_log_R_week_sigma,
        cdar_log_R_week_rho,
        cdar_log_R_week_eta,
        cdar_log_R_week_L,
        cdar_log_R_week_n_steps,
        cdar_log_R_week_group_idx,
        cdar_log_R_week_step_idx
    );
    vector[num_elements(time)] log_R = (pop_log_R + rw_log_R_time + cdar_log_R_week);
    vector[patch_idx_n] cat_log_I0_patch = cat_log_I0_patch_beta[patch_idx];
    vector[seed_mean_n] log_I0 = (seed_mean + cat_log_I0_patch);
    vector[pop_n] seeds = patch_seeds(log_I0, pop);
    vector[num_elements(time)] Y = patch_expected_cases(log_R, seeds, gamma, pop, dist_flat, gen_pmf, delay_pmf);
}
model {
    gamma ~ normal(1.5, 0.5);
    cluster ~ normal(0.0, 0.1);
    pop_log_R_beta_pop ~ normal([0.26236426446749106]', [0.1]');
    rw_log_R_time_sigma ~ normal(0.0, 0.05);
    rw_log_R_time_z ~ std_normal();
    cdar_log_R_week_sigma ~ normal(0.0, 0.2);
    cdar_log_R_week_rho ~ normal(0.8, 0.1);
    cdar_log_R_week_eta ~ std_normal();
    cat_log_I0_patch_beta ~ normal(0.0, 0.5);
    cases ~ nb_cases(Y, cluster, observed);
}
generated quantities {
    vector[observed_n] cases_likelihood = nb_cases_lpmfs(cases, Y, cluster, observed);
    array[observed_n] int cases_gen = nb_cases_int_rng(cases_n, Y, cluster, observed);
}

julia

Turing unsupported for this BRM example

Turing backend: assignment `seeds` calls `patch_seeds`, which defines no Julia methods (a Stan `@deffun` function has no Julia implementation). Row-dependent assignments lower to row-wise Julia comprehensions, so this call cannot execute. Express the computation with row-wise Julia code or fit this model with the Stan backend.

Three formula lines carry the statistical structure:

log_R ~ 1 + rw(time) + cdar(week; by=patch, cor=C) — on a long frame, rw(time) is one walk over the distinct days, shared by all rows of a day, and cdar(week; by=patch, cor=C) gives every patch a damped weekly path whose innovations are correlated through C. sd(:, cdar(week)) addresses $σ_{δ}$ and ar(:, cdar(week)) addresses $ρ$ .
log_I0 ~ 0 + offset(seed_mean) + factor(patch) — a predictor without an intercept codes its first categorical term by cell means (Home): one coefficient per patch, no reference level. With the prior mean as an offset, the line says $\log I_{0, g} = m_{g} + δ_{g}, δ_{g} \sim N (0, 0.5)$ .
gamma is an ordinary scalar parameter; the mixing matrix is computed from it inside the Stan function.

Per-patch reproduction numbers, 56 days of counts in six patches:

Daily counts per patch, on a logarithmic axis. The five patches that start without cases are fitted through the zeros of the early weeks:

The seeds, as the cell means estimate them: posterior median with 50 % and 95 % intervals, the prior mean (grey dot) and the truth (black dot). For the five patches without early cases the counts say little about the seed beyond "very small", so the posterior stays close to the prior; the origin's seed is pinned down by its counts:

The same models through a package

Everything mechanical above is user code, which is the point — but it is also forty lines of Stan functions that the next renewal model would copy. The repository carries them once more as a small example package, examples/EpiRenewal, that sits on top of BayesianRegressionModels the way epidemia sits on rstanarm: the package owns the renewal recursion, the mixing and the delay, and the model keeps one line per mechanical step.

The six-patch model, with using EpiRenewal:

brm-comparison

Six coupled patches with EpiRenewal's operators

julia

function epirenewal_patch_model(data = page.renewal_patch_data())
    @brm data begin
        gamma   ~ Normal(1.5, 0.5; lower=0.0)
        cluster ~ Normal(0.0, 0.1; lower=0.0)
        log_R   ~ 1 + rw(time) + cdar(week; by=patch, cor=C)
        effect(log_R, Intercept) ~ Normal(log(1.3), 0.1)
        sd(:, rw(time)) ~ Normal(0.0, 0.05)
        sd(:, cdar(week)) ~ Normal(0.0, 0.2)
        ar(:, cdar(week)) ~ Normal(0.8, 0.1)
        log_I0  ~ 0 + offset(seed_mean) + factor(patch)
        effect(log_I0, patch) ~ Normal(0.0, 0.5)
        seeds   = per_patch(log_I0, pop)                             # one seed per patch
        mixing  = gravity(pop, dist_flat, gamma)                     # P × P mixing weights
        past    = seeded_history(seeds, exp(log_R), gen_pmf, 14)     # 14 × P
        I       = renewal(gen_pmf, exp(log_R), past, mixing)
        Y       = delay(delay_pmf, I, past)
        cases   ~ nb_cases(Y, cluster, observed)
    end
end

julia

BRMI:
  gamma ~ Normal(1.5, 0.5; lower=0.0)
  cluster ~ Normal(0.0, 0.1; lower=0.0)
  time: data (eltype=Float64, n=336)
  patch: data (eltype=Int64, n=336)
  C: data (eltype=Float64, n=36)
  week: data (eltype=Int64, n=336)
  log_R ~ 1 + rw(time) + cdar(week; by=patch, cor=C)
  effect(log_R, Intercept) ~ Normal(log(1.3), 0.1)
  effect(term_sd, rw(time), :) ~ Normal(0.0, 0.05)
  effect(term_sd, cdar(week), :) ~ Normal(0.0, 0.2)
  effect(term_ar, cdar(week), :) ~ Normal(0.8, 0.1)
  seed_mean: data (eltype=Float64, n=336)
  log_I0 ~ 0 + offset(seed_mean) + factor(patch)
  effect(log_I0, patch) ~ Normal(0.0, 0.5)
  pop: data (eltype=Float64, n=6)
  :seeds = per_patch(log_I0, pop)
  dist_flat: data (eltype=Float64, n=36)
  :mixing = gravity(pop, dist_flat, gamma)
  gen_pmf: data (eltype=Float64, n=13)
  :past = seeded_history(seeds, exp(log_R), gen_pmf, 14)
  :I = renewal(gen_pmf, exp(log_R), past, mixing)
  delay_pmf: data (eltype=Float64, n=15)
  :Y = delay(delay_pmf, I, past)
  observed: data (eltype=Int64, n=336)
  cases ~ nb_cases(Y, cluster, observed)

julia

SBBRMI with data keys = [:C, :cases, :cdar_log_R_week_L, :cdar_log_R_week_group_idx, :cdar_log_R_week_n_groups, :cdar_log_R_week_n_steps, :cdar_log_R_week_step_idx, :delay_pmf, :dist_flat, :gen_pmf, :observed, :patch, :patch_idx, :patch_n_levels, :pop, :rw_log_R_time_n_steps, :rw_log_R_time_time_idx, :seed_mean, :time, :week]
configured submodels:
_sb_rw1_configured_1 = Base.merge(BayesianRegressionModels._sb_rw1, quote
            sigma ~ normal(0.0, 0.05; lower = 0.0)
        end)
_sb_cdar_configured_1 = Base.merge(BayesianRegressionModels._sb_cdar, quote
            rho ~ normal(0.8, 0.1; lower = 0.0, upper = 1.0)
        end)
emitted @slic body:
begin
    gamma ~ normal(1.5, 0.5; lower = 0.0)
    cluster ~ normal(0.0, 0.1; lower = 0.0)
    X_log_R = hcat(rep_vector(1.0, num_elements(time)))
    pop_log_R ~ _popefs_normal(; X = X_log_R, beta_loc = [0.26236426446749106], beta_scale = [0.1])
    rw_log_R_time ~ _sb_rw1_configured_1(; n_steps = rw_log_R_time_n_steps, time_idx = rw_log_R_time_time_idx)
    cdar_log_R_week ~ _sb_cdar_configured_1(; n_groups = cdar_log_R_week_n_groups, n_steps = cdar_log_R_week_n_steps, L = cdar_log_R_week_L, group_idx = cdar_log_R_week_group_idx, step_idx = cdar_log_R_week_step_idx)
    log_R = pop_log_R + rw_log_R_time + cdar_log_R_week
    cat_log_I0_patch ~ _sb_cat_cells_normal(; x = patch_idx, n_levels = patch_n_levels, beta_loc = 0.0, beta_scale = 0.5)
    log_I0 = seed_mean + cat_log_I0_patch
    seeds = (Main.renewal_package.EpiRenewal.per_patch)(log_I0, pop)
    mixing = (Main.renewal_package.EpiRenewal.gravity)(pop, dist_flat, gamma)
    past = (Main.renewal_package.EpiRenewal.seeded_history)(seeds, (exp)(log_R), gen_pmf, 14)
    I = (Main.renewal_package.EpiRenewal.renewal)(gen_pmf, (exp)(log_R), past, mixing)
    Y = (Main.renewal_package.EpiRenewal.delay)(delay_pmf, I, past)
    cases ~ nb_cases(Y, cluster, observed)
end

stan

functions {
matrix hcat(vector x) {
    int n = dims(x)[1];
    return to_matrix(x, n, 1);
}
vector random_walk_rows(
    real sigma,
    vector z,
    array[] int idx
) {
    int N = dims(idx)[1];
    vector[(dims(z)[1] + 1)] x = random_walk_path(sigma, z);
    vector[N] out;
    for(i in 1:N) {
        out[i] = x[idx[i]];
    }
    return out;
}
vector random_walk_path(
    real sigma,
    vector z
) {
    int n = dims(z)[1];
    return append_row(0.0, (sigma * cumulative_sum(z)));
}
vector correlated_damped_walk(
    real sigma,
    real rho,
    vector eta,
    matrix L,
    int W,
    array[] int group_idx,
    array[] int step_idx
) {
    int P = dims(L)[1];
    int N = dims(group_idx)[1];
    if (dims(L)[2] != P) reject("correlated_damped_walk: dim mismatch — `L` dim 2 (= ", dims(L)[2], ") does not match `P` (= ", P, "), inferred from `L` dim 1. `P` sizes: `L` dim 1 (= ", dims(L)[1], "), `L` dim 2 (= ", dims(L)[2], ").");
    if (dims(step_idx)[1] != N) reject("correlated_damped_walk: dim mismatch — `step_idx` dim 1 (= ", dims(step_idx)[1], ") does not match `N` (= ", N, "), inferred from `group_idx` dim 1. `N` sizes: `group_idx` dim 1 (= ", dims(group_idx)[1], "), `step_idx` dim 1 (= ", dims(step_idx)[1], ").");
    matrix[P, W] E = to_matrix(eta, P, W);
    matrix[P, W] delta;
    vector[dims(L)[1]] prev = (sigma * (L * col(E, 1)));
    delta[:, 1] = prev;
    real scale = (sigma * sqrt((1.0 - (rho * rho))));
    for(w in 2:W) {
        prev = ((rho * prev) + (scale * (L * col(E, w))));
        delta[:, w] = prev;
    }
    vector[N] out;
    for(i in 1:N) {
        out[i] = delta[group_idx[i], step_idx[i]];
    }
    return out;
}
vector per_patch(
    vector x,
    vector like
) {
    int N = dims(x)[1];
    int P = dims(like)[1];
    int T = (N / P);
    vector[P] out;
    for(g in 1:P) {
        out[g] = x[(((g - 1) * T) + 1)];
    }
    return out;
}
matrix gravity(
    vector pop,
    vector dist_flat,
    real gamma
) {
    int P = dims(pop)[1];
    real mean_pop = (sum(pop) / P);
    matrix[P, P] K;
    for(g in 1:P) {
        real total = 0.0;
        for(h in 1:P) {
            K[g, h] = ((g == h) ? 1.0 : (((pop[g] / mean_pop) * (pop[h] / mean_pop)) / exp((gamma * log(dist_flat[(g + ((h - 1) * P))])))));
            total += K[g, h];
        }
        for(h in 1:P) {
            K[g, h] = (K[g, h] / total);
        }
    }
    return K;
}
matrix seeded_history(
    vector log_I0,
    vector R,
    vector gen_pmf,
    int H
) {
    int P = dims(log_I0)[1];
    int N = dims(R)[1];
    int T = (N / P);
    matrix[H, P] past;
    for(g in 1:P) {
        real r = growth_rate(R[(((g - 1) * T) + 1)], gen_pmf);
        for(h in 1:H) {
            past[h, g] = (exp(log_I0[g]) * exp((r * ((h - H) - 1))));
        }
    }
    return past;
}
real growth_rate(
    real R,
    vector gen_pmf
) {
    int G = dims(gen_pmf)[1];
    real r = 0.0;
    for(iter in 1:2) {
        real f = (-1.0 / R);
        real df = 0.0;
        for(s in 1:G) {
            f += (gen_pmf[s] * exp(((-r) * s)));
            df -= (s * gen_pmf[s] * exp(((-r) * s)));
        }
        r = (r - (f / df));
    }
    return clamp_between(r, -2.0, 2.0);
}
real clamp_between(real x, real lo, real hi) {
    return ((x < lo) ? lo : ((x > hi) ? hi : x));
}
vector renewal(
    vector gen_pmf,
    vector R,
    matrix past,
    matrix K
) {
    int G = dims(gen_pmf)[1];
    int N = dims(R)[1];
    int P = dims(past)[2];
    if (dims(K)[1] != P) reject("renewal: dim mismatch — `K` dim 1 (= ", dims(K)[1], ") does not match `P` (= ", P, "), inferred from `past` dim 2. `P` sizes: `past` dim 2 (= ", dims(past)[2], "), `K` dim 1 (= ", dims(K)[1], "), `K` dim 2 (= ", dims(K)[2], ").");
    if (dims(K)[2] != P) reject("renewal: dim mismatch — `K` dim 2 (= ", dims(K)[2], ") does not match `P` (= ", P, "), inferred from `past` dim 2. `P` sizes: `past` dim 2 (= ", dims(past)[2], "), `K` dim 1 (= ", dims(K)[1], "), `K` dim 2 (= ", dims(K)[2], ").");
    int T = (N / P);
    matrix[T, P] Rm = to_matrix(R, T, P);
    matrix[T, P] I;
    vector[P] force;
    for(t in 1:T) {
        for(h in 1:P) {
            force[h] = 0.0;
            for(s in 1:G) {
                force[h] += (gen_pmf[s] * at_day(col(I, h), col(past, h), (t - s)));
            }
        }
        for(g in 1:P) {
            real pressure = 0.0;
            for(h in 1:P) {
                pressure += (K[g, h] * force[h]);
            }
            I[t, g] = clamp_between((Rm[t, g] * pressure), 0.0, 1.0e15);
        }
    }
    return to_vector(I);
}
real at_day(
    vector x,
    vector past,
    int s
) {
    int H = dims(past)[1];
    return ((s >= 1) ? x[s] : (((s + H) >= 1) ? past[(s + H)] : 0.0));
}
vector delay(
    vector pmf,
    vector I,
    matrix past
) {
    int D = dims(pmf)[1];
    int N = dims(I)[1];
    int P = dims(past)[2];
    int T = (N / P);
    matrix[T, P] Im = to_matrix(I, T, P);
    matrix[T, P] Y;
    for(g in 1:P) {
        for(t in 1:T) {
            real acc = 0.0;
            for(d in 0:(D - 1)) {
                acc += (pmf[(d + 1)] * at_day(col(Im, g), col(past, g), (t - d)));
            }
            Y[t, g] = acc;
        }
    }
    return to_vector(Y);
}
real nb_cases_lpmf(
    array[] int cases,
    vector Y,
    real cluster,
    array[] int observed
) {
    int N = dims(cases)[1];
    if (dims(Y)[1] != N) reject("nb_cases_lpmf: dim mismatch — `Y` dim 1 (= ", dims(Y)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    if (dims(observed)[1] != N) reject("nb_cases_lpmf: dim mismatch — `observed` dim 1 (= ", dims(observed)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    real lp = 0.0;
    for(i in 1:N) {
        if((observed[i] == 1)) {
            lp += neg_binomial_2_lpmf(cases[i] | Y[i], (1.0 / (cluster * cluster)));
        }
    }
    return lp;
}
vector nb_cases_lpmfs(
    array[] int cases,
    vector Y,
    real cluster,
    array[] int observed
) {
    int N = dims(cases)[1];
    if (dims(Y)[1] != N) reject("nb_cases_lpmfs: dim mismatch — `Y` dim 1 (= ", dims(Y)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    if (dims(observed)[1] != N) reject("nb_cases_lpmfs: dim mismatch — `observed` dim 1 (= ", dims(observed)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    vector[N] lp;
    for(i in 1:N) {
        lp[i] = ((observed[i] == 1) ? neg_binomial_2_lpmf(cases[i] | Y[i], (1.0 / (cluster * cluster))) : 0.0);
    }
    return lp;
}
array[] int nb_cases_int_rng(
    int anontok__1,
    vector Y,
    real cluster,
    array[] int observed
) {
    int N = anontok__1;
    if (dims(Y)[1] != N) reject("nb_cases_rng: dim mismatch — `Y` dim 1 (= ", dims(Y)[1], ") does not match `N` (= ", N, "), inferred from `anontok__1` dim 1. `N` sizes: `anontok__1` dim 1 (= ", anontok__1, "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    if (dims(observed)[1] != N) reject("nb_cases_rng: dim mismatch — `observed` dim 1 (= ", dims(observed)[1], ") does not match `N` (= ", N, "), inferred from `anontok__1` dim 1. `N` sizes: `anontok__1` dim 1 (= ", anontok__1, "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    array[N] int out;
    for(i in 1:N) {
        out[i] = neg_binomial_2_rng(Y[i], (1.0 / (cluster * cluster)));
    }
    return out;
}
}
data {
    int time_n;
    vector[time_n] time;
    int rw_log_R_time_n_steps;
    int rw_log_R_time_time_idx_n;
    array[rw_log_R_time_time_idx_n] int rw_log_R_time_time_idx;
    int cdar_log_R_week_n_groups;
    int cdar_log_R_week_n_steps;
    int cdar_log_R_week_step_idx_n;
    int cdar_log_R_week_L_m;
    int cdar_log_R_week_L_n;
    matrix[cdar_log_R_week_L_m, cdar_log_R_week_L_n] cdar_log_R_week_L;
    int cdar_log_R_week_group_idx_n;
    array[cdar_log_R_week_group_idx_n] int cdar_log_R_week_group_idx;
    array[cdar_log_R_week_step_idx_n] int cdar_log_R_week_step_idx;
    int patch_n_levels;
    int patch_idx_n;
    array[patch_idx_n] int patch_idx;
    int seed_mean_n;
    vector[seed_mean_n] seed_mean;
    int pop_n;
    vector[pop_n] pop;
    int dist_flat_n;
    vector[dist_flat_n] dist_flat;
    int gen_pmf_n;
    vector[gen_pmf_n] gen_pmf;
    int delay_pmf_n;
    vector[delay_pmf_n] delay_pmf;
    int cases_n;
    array[cases_n] int cases;
    int observed_n;
    array[observed_n] int observed;
}
transformed data {
    matrix[num_elements(time), 1] X_log_R = hcat(rep_vector(1.0, num_elements(time)));
    int pop_log_R_n_covariates = 1;
}
parameters {
    real<lower=0.0> gamma;
    real<lower=0.0> cluster;
    vector[pop_log_R_n_covariates] pop_log_R_beta_pop;
    real<lower=0.0> rw_log_R_time_sigma;
    vector[(rw_log_R_time_n_steps - 1)] rw_log_R_time_z;
    real<lower=0.0> cdar_log_R_week_sigma;
    real<lower=0.0, upper=1.0> cdar_log_R_week_rho;
    vector[(cdar_log_R_week_n_groups * cdar_log_R_week_n_steps)] cdar_log_R_week_eta;
    vector[patch_n_levels] cat_log_I0_patch_beta;
}
transformed parameters {
    vector[num_elements(time)] pop_log_R = (X_log_R * pop_log_R_beta_pop);
    vector[rw_log_R_time_time_idx_n] rw_log_R_time = random_walk_rows(rw_log_R_time_sigma, rw_log_R_time_z, rw_log_R_time_time_idx);
    vector[cdar_log_R_week_step_idx_n] cdar_log_R_week = correlated_damped_walk(
        cdar_log_R_week_sigma,
        cdar_log_R_week_rho,
        cdar_log_R_week_eta,
        cdar_log_R_week_L,
        cdar_log_R_week_n_steps,
        cdar_log_R_week_group_idx,
        cdar_log_R_week_step_idx
    );
    vector[num_elements(time)] log_R = (pop_log_R + rw_log_R_time + cdar_log_R_week);
    vector[patch_idx_n] cat_log_I0_patch = cat_log_I0_patch_beta[patch_idx];
    vector[seed_mean_n] log_I0 = (seed_mean + cat_log_I0_patch);
    vector[pop_n] seeds = per_patch(log_I0, pop);
    matrix[pop_n, pop_n] mixing = gravity(pop, dist_flat, gamma);
    matrix[14, pop_n] past = seeded_history(seeds, exp(log_R), gen_pmf, 14);
    vector[num_elements(time)] I = renewal(gen_pmf, exp(log_R), past, mixing);
    vector[num_elements(time)] Y = delay(delay_pmf, I, past);
}
model {
    gamma ~ normal(1.5, 0.5);
    cluster ~ normal(0.0, 0.1);
    pop_log_R_beta_pop ~ normal([0.26236426446749106]', [0.1]');
    rw_log_R_time_sigma ~ normal(0.0, 0.05);
    rw_log_R_time_z ~ std_normal();
    cdar_log_R_week_sigma ~ normal(0.0, 0.2);
    cdar_log_R_week_rho ~ normal(0.8, 0.1);
    cdar_log_R_week_eta ~ std_normal();
    cat_log_I0_patch_beta ~ normal(0.0, 0.5);
    cases ~ nb_cases(Y, cluster, observed);
}
generated quantities {
    vector[observed_n] cases_likelihood = nb_cases_lpmfs(cases, Y, cluster, observed);
    array[observed_n] int cases_gen = nb_cases_int_rng(cases_n, Y, cluster, observed);
}

julia

Turing unsupported for this BRM example

Turing backend: assignment `seeds` calls `per_patch`, which defines no Julia methods (a Stan `@deffun` function has no Julia implementation). Row-dependent assignments lower to row-wise Julia comprehensions, so this call cannot execute. Express the computation with row-wise Julia code or fit this model with the Stan backend.

seeded_history, renewal and delay are ordinary Stan functions that the package declares with @deffun; past, I and Y are named quantities of the model. This is the same posterior as the hand-written version: the package's test requires the same dimension and the same log density and gradient at arbitrary points for every model on this page, so the fits and figures above are this model's fits and figures.

Each operator takes its kernel first, as a data vector or as a function. A function can be written as a do block, and its body may read sampled parameters (Formula terms) — which a data vector cannot. So the reporting delay of step 1 no longer has to be a fixed input: here it is estimated inside the renewal model, from the counts.

brm-comparison

Reporting delay estimated inside the renewal model

julia

function epirenewal_estimated_delay_model(data = page.renewal_single_data())
    @brm data begin
        mu_delay    ~ Normal(1.5, 0.2)
        sigma_delay ~ Normal(0.5, 0.1; lower=0.0)
        log_I0  ~ Normal(log(50.0), 0.5)
        cluster ~ Normal(0.0, 0.1; lower=0.0)
        log_R   ~ 1 + rw(time)
        effect(log_R, Intercept) ~ Normal(log(1.3), 0.1)
        sd(:, rw(time)) ~ Normal(0.0, 0.05)
        past    = seeded_history(log_I0, exp(log_R), gen_pmf, 14)
        I       = renewal(gen_pmf, exp(log_R), past)
        Y       = delay(I, past, 15) do d                            # d = 0 … 14
            censored_lognormal_mass(d, mu_delay, sigma_delay, 15)
        end
        cases   ~ nb_cases(Y, cluster, observed)
    end
end

julia

BRMI:
  mu_delay ~ Normal(1.5, 0.2)
  sigma_delay ~ Normal(0.5, 0.1; lower=0.0)
  log_I0 ~ Normal(log(50.0), 0.5)
  cluster ~ Normal(0.0, 0.1; lower=0.0)
  time: data (eltype=Float64, n=56)
  log_R ~ 1 + rw(time)
  effect(log_R, Intercept) ~ Normal(log(1.3), 0.1)
  effect(term_sd, rw(time), :) ~ Normal(0.0, 0.05)
  gen_pmf: data (eltype=Float64, n=13)
  :past = seeded_history(log_I0, exp(log_R), gen_pmf, 14)
  :I = renewal(gen_pmf, exp(log_R), past)
  :Y = delay((d,)->begin
        #= brm-docs-example.jl:13 =#
        censored_lognormal_mass(d, mu_delay, sigma_delay, 15)
    end, I, past, 15)
  observed: data (eltype=Int64, n=56)
  cases ~ nb_cases(Y, cluster, observed)

julia

SBBRMI with data keys = [:cases, :gen_pmf, :observed, :rw_log_R_time_n_steps, :rw_log_R_time_time_idx, :time]
configured submodels:
_sb_rw1_configured_1 = Base.merge(BayesianRegressionModels._sb_rw1, quote
            sigma ~ normal(0.0, 0.05; lower = 0.0)
        end)
emitted @slic body:
begin
    mu_delay ~ normal(1.5, 0.2)
    sigma_delay ~ normal(0.5, 0.1; lower = 0.0)
    log_I0 ~ normal(3.912023005428146, 0.5)
    cluster ~ normal(0.0, 0.1; lower = 0.0)
    X_log_R = hcat(rep_vector(1.0, num_elements(time)))
    pop_log_R ~ _popefs_normal(; X = X_log_R, beta_loc = [0.26236426446749106], beta_scale = [0.1])
    rw_log_R_time ~ _sb_rw1_configured_1(; n_steps = rw_log_R_time_n_steps, time_idx = rw_log_R_time_time_idx)
    log_R = pop_log_R + rw_log_R_time
    past = (Main.renewal_package.EpiRenewal.seeded_history)(log_I0, (exp)(log_R), gen_pmf, 14)
    I = (Main.renewal_package.EpiRenewal.renewal)(gen_pmf, (exp)(log_R), past)
    Y = (Main.renewal_package.EpiRenewal.delay)(((d,)->begin
                    #= brm-docs-example.jl:13 =#
                    censored_lognormal_mass(d, mu_delay, sigma_delay, 15)
                end), I, past, 15)
    cases ~ nb_cases(Y, cluster, observed)
end

stan

functions {
matrix hcat(vector x) {
    int n = dims(x)[1];
    return to_matrix(x, n, 1);
}
vector random_walk_rows(
    real sigma,
    vector z,
    array[] int idx
) {
    int N = dims(idx)[1];
    vector[(dims(z)[1] + 1)] x = random_walk_path(sigma, z);
    vector[N] out;
    for(i in 1:N) {
        out[i] = x[idx[i]];
    }
    return out;
}
vector random_walk_path(
    real sigma,
    vector z
) {
    int n = dims(z)[1];
    return append_row(0.0, (sigma * cumulative_sum(z)));
}
vector seeded_history(
    real log_I0,
    vector R,
    vector gen_pmf,
    int H
) {
    real r = growth_rate(R[1], gen_pmf);
    vector[H] past;
    for(h in 1:H) {
        past[h] = (exp(log_I0) * exp((r * ((h - H) - 1))));
    }
    return past;
}
real growth_rate(
    real R,
    vector gen_pmf
) {
    int G = dims(gen_pmf)[1];
    real r = 0.0;
    for(iter in 1:2) {
        real f = (-1.0 / R);
        real df = 0.0;
        for(s in 1:G) {
            f += (gen_pmf[s] * exp(((-r) * s)));
            df -= (s * gen_pmf[s] * exp(((-r) * s)));
        }
        r = (r - (f / df));
    }
    return clamp_between(r, -2.0, 2.0);
}
real clamp_between(real x, real lo, real hi) {
    return ((x < lo) ? lo : ((x > hi) ? hi : x));
}
vector renewal(
    vector gen_pmf,
    vector R,
    vector past
) {
    int G = dims(gen_pmf)[1];
    int T = dims(R)[1];
    vector[T] I;
    for(t in 1:T) {
        real force = 0.0;
        for(s in 1:G) {
            force += (gen_pmf[s] * at_day(I, past, (t - s)));
        }
        I[t] = clamp_between((R[t] * force), 0.0, 1.0e15);
    }
    return I;
}
real at_day(
    vector x,
    vector past,
    int s
) {
    int H = dims(past)[1];
    return ((s >= 1) ? x[s] : (((s + H) >= 1) ? past[(s + H)] : 0.0));
}
vector delay_closure_1(
    real mu_delay,
    real sigma_delay,
    vector I,
    vector past,
    int D
) {
    int T = dims(I)[1];
    vector[T] Y;
    for(t in 1:T) {
        real acc = 0.0;
        for(d in 0:(D - 1)) {
            acc += (censored_lognormal_mass(d, mu_delay, sigma_delay, 15) * at_day(I, past, (t - d)));
        }
        Y[t] = acc;
    }
    return Y;
}
real censored_lognormal_mass(
    int d,
    real mu,
    real sigma,
    int D
) {
    return ((lnorm_F((d + 1.0), mu, sigma) - lnorm_F((d + 0.0), mu, sigma)) / lnorm_F((D + 0.0), mu, sigma));
}
real lnorm_F(
    real x,
    real mu,
    real sigma
) {
    return (lnorm_G(x, mu, sigma) - lnorm_G((x - 1.0), mu, sigma));
}
real lnorm_G(
    real a,
    real mu,
    real sigma
) {
    return ((a <= 0.0) ? 0.0 : (
        (a * Phi(((log(a) - mu) / sigma))) -
        (exp((mu + (0.5 * sigma * sigma))) * Phi((((log(a) - mu) / sigma) - sigma)))
    ));
}
real nb_cases_lpmf(
    array[] int cases,
    vector Y,
    real cluster,
    array[] int observed
) {
    int N = dims(cases)[1];
    if (dims(Y)[1] != N) reject("nb_cases_lpmf: dim mismatch — `Y` dim 1 (= ", dims(Y)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    if (dims(observed)[1] != N) reject("nb_cases_lpmf: dim mismatch — `observed` dim 1 (= ", dims(observed)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    real lp = 0.0;
    for(i in 1:N) {
        if((observed[i] == 1)) {
            lp += neg_binomial_2_lpmf(cases[i] | Y[i], (1.0 / (cluster * cluster)));
        }
    }
    return lp;
}
vector nb_cases_lpmfs(
    array[] int cases,
    vector Y,
    real cluster,
    array[] int observed
) {
    int N = dims(cases)[1];
    if (dims(Y)[1] != N) reject("nb_cases_lpmfs: dim mismatch — `Y` dim 1 (= ", dims(Y)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    if (dims(observed)[1] != N) reject("nb_cases_lpmfs: dim mismatch — `observed` dim 1 (= ", dims(observed)[1], ") does not match `N` (= ", N, "), inferred from `cases` dim 1. `N` sizes: `cases` dim 1 (= ", dims(cases)[1], "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    vector[N] lp;
    for(i in 1:N) {
        lp[i] = ((observed[i] == 1) ? neg_binomial_2_lpmf(cases[i] | Y[i], (1.0 / (cluster * cluster))) : 0.0);
    }
    return lp;
}
array[] int nb_cases_int_rng(
    int anontok__1,
    vector Y,
    real cluster,
    array[] int observed
) {
    int N = anontok__1;
    if (dims(Y)[1] != N) reject("nb_cases_rng: dim mismatch — `Y` dim 1 (= ", dims(Y)[1], ") does not match `N` (= ", N, "), inferred from `anontok__1` dim 1. `N` sizes: `anontok__1` dim 1 (= ", anontok__1, "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    if (dims(observed)[1] != N) reject("nb_cases_rng: dim mismatch — `observed` dim 1 (= ", dims(observed)[1], ") does not match `N` (= ", N, "), inferred from `anontok__1` dim 1. `N` sizes: `anontok__1` dim 1 (= ", anontok__1, "), `Y` dim 1 (= ", dims(Y)[1], "), `observed` dim 1 (= ", dims(observed)[1], ").");
    array[N] int out;
    for(i in 1:N) {
        out[i] = neg_binomial_2_rng(Y[i], (1.0 / (cluster * cluster)));
    }
    return out;
}
}
data {
    int time_n;
    vector[time_n] time;
    int rw_log_R_time_n_steps;
    int rw_log_R_time_time_idx_n;
    array[rw_log_R_time_time_idx_n] int rw_log_R_time_time_idx;
    int gen_pmf_n;
    vector[gen_pmf_n] gen_pmf;
    int cases_n;
    array[cases_n] int cases;
    int observed_n;
    array[observed_n] int observed;
}
transformed data {
    matrix[num_elements(time), 1] X_log_R = hcat(rep_vector(1.0, num_elements(time)));
    int pop_log_R_n_covariates = 1;
}
parameters {
    real mu_delay;
    real<lower=0.0> sigma_delay;
    real log_I0;
    real<lower=0.0> cluster;
    vector[pop_log_R_n_covariates] pop_log_R_beta_pop;
    real<lower=0.0> rw_log_R_time_sigma;
    vector[(rw_log_R_time_n_steps - 1)] rw_log_R_time_z;
}
transformed parameters {
    vector[num_elements(time)] pop_log_R = (X_log_R * pop_log_R_beta_pop);
    vector[rw_log_R_time_time_idx_n] rw_log_R_time = random_walk_rows(rw_log_R_time_sigma, rw_log_R_time_z, rw_log_R_time_time_idx);
    vector[num_elements(time)] log_R = (pop_log_R + rw_log_R_time);
    vector[14] past = seeded_history(log_I0, exp(log_R), gen_pmf, 14);
    vector[num_elements(time)] I = renewal(gen_pmf, exp(log_R), past);
    vector[num_elements(time)] Y = delay_closure_1(mu_delay, sigma_delay, I, past, 15);
}
model {
    mu_delay ~ normal(1.5, 0.2);
    sigma_delay ~ normal(0.5, 0.1);
    log_I0 ~ normal(3.912023005428146, 0.5);
    cluster ~ normal(0.0, 0.1);
    pop_log_R_beta_pop ~ normal([0.26236426446749106]', [0.1]');
    rw_log_R_time_sigma ~ normal(0.0, 0.05);
    rw_log_R_time_z ~ std_normal();
    cases ~ nb_cases(Y, cluster, observed);
}
generated quantities {
    vector[observed_n] cases_likelihood = nb_cases_lpmfs(cases, Y, cluster, observed);
    array[observed_n] int cases_gen = nb_cases_int_rng(cases_n, Y, cluster, observed);
}

julia

Turing unsupported for this BRM example

Turing backend: assignment `I` calls `renewal`, which defines no Julia methods (a Stan `@deffun` function has no Julia implementation). Row-dependent assignments lower to row-wise Julia comprehensions, so this call cannot execute. Express the computation with row-wise Julia code or fit this model with the Stan backend.

The operators read plain vectors, so they rely on the frame's layout — one row per (patch, day), ordered by patch and then by day, every patch covering the same days. The package's epi_frame(rows; time, by) sorts a table into that order and refuses one that cannot satisfy it.

Parameters and sampler diagnostics

Posterior medians and 95 % intervals next to the values the data were simulated from. The table is generated from the summaries of the fits shown above.

Model	Parameter	Truth	Median	95 % interval	Truth inside
reporting delay	`mu`	1.5	1.54	[1.44, 1.65]	yes
reporting delay	`sigma`	0.5	0.49	[0.427, 0.559]	yes
one population	`log R_1 (intercept of log_R)`	0.262	0.248	[0.141, 0.384]	yes
one population	`walk scale sd(:, rw(time))`	0.05	0.038	[0.0176, 0.0809]	yes
one population	`log_I0`	3.91	3.97	[3.84, 4.11]	yes
one population	`cluster`	0.1	0.115	[0.0833, 0.154]	yes
one population, days 1-42	`log R_1 (intercept of log_R)`	0.262	0.22	[0.139, 0.338]	yes
one population, days 1-42	`walk scale sd(:, rw(time))`	0.05	0.0203	[0.000985, 0.0661]	yes
one population, days 1-42	`log_I0`	3.91	3.97	[3.86, 4.08]	yes
one population, days 1-42	`cluster`	0.1	0.108	[0.0693, 0.156]	yes
six patches	`log R_1 (intercept of log_R)`	0.182	0.313	[0.148, 0.472]	yes
six patches	`walk scale sd(:, rw(time))`	0.03	0.0328	[0.0112, 0.0709]	yes
six patches	`deviation scale sd(:, cdar(week))`	0.15	0.139	[0.0889, 0.225]	yes
six patches	`persistence ar(:, cdar(week))`	0.8	0.757	[0.606, 0.892]	yes
six patches	`gamma`	1.5	1.52	[1.33, 1.7]	yes
six patches	`cluster`	0.1	0.0933	[0.0785, 0.111]	yes

How often the truth lies inside the pointwise 95 % band, per figure:

Quantity	Truth inside the 95 % band
`R(t)`, one population, 56 days	44 of 56
`R(t)`, fitted on days 1–42: fitted days	39 of 42
`R(t)`, fitted on days 1–42: held-out days	9 of 14
`R(g, t)`, six patches, 336 patch-days	334 of 336
seeds, six patches	5 of 6

Each model was sampled with one chain of NUTS (WarmupHMC.jl) on the log density that BridgeStan compiles from the generated Stan program. The one-population posterior couples the walk scale with 55 innovations that the counts determine tightly, which calls for a small step size; those two fits therefore run at a target acceptance of 0.995. The settings and diagnostics of every fit:

Model	Parameters	Draws	Target acceptance	Divergences	Min. ESS	Max. R-hat	Seconds
reporting delay	2	1000	0.9	0	395	1.003	14
one population	59	2000	0.995	1	157	1.005	59
one population, prior only	59	1000	0.8	0	492	1.015	1
one population, days 1-42	59	2000	0.995	0	296	1.009	40
six patches	115	2000	0.9	2	89	1.006	841

The route from a declaration to draws is short. mod=@__MODULE__ tells the backend where the @deffun functions live:

julia

function build(brmi; held_out=())
    sb = SBBRMI(brmi; mod=@__MODULE__, held_out)
    code = BayesianRegressionModels.stan_code(sb)
    stanc = StanBlocks.stanc_check(code; warn_pedantic=false)
    stanc.ok || error("stanc rejected the model:\n" * stanc.output)
    (; sb, code, problem=StanBlocks.stan_instantiate(sb.model))
end
function fit(name, problem; n_draws=1000, seed=1, target_acceptance_rate=0.9, tolerate_gq_failures=false)
    seconds = @elapsed f = WarmupHMC.adaptive_warmup_mcmc(Xoshiro(seed), problem; n_draws, target_acceptance_rate, progress=nothing)
    q = convert(Matrix{Float64}, f.posterior_position)
    ess, rhat = MCMCDiagnosticTools.ess_rhat(reshape(permutedims(q), size(q, 2), 1, size(q, 1)))
    names = BridgeStan.param_names(problem.model; include_tp=true, include_gq=true)
    rng = BridgeStan.StanRNG(problem.model, seed)
    cons = Matrix{Float64}(undef, length(names), size(q, 2)); failed = 0
    for j in 1:size(q, 2)
        try
            cons[:, j] = BridgeStan.param_constrain(problem.model, q[:, j]; include_tp=true, include_gq=true, rng)
        catch
            tolerate_gq_failures || rethrow()               # a prior draw can overflow the count generator
            failed += 1; cons[:, j] .= NaN
        end
    end
    diag = (; name, seed, draws=size(q, 2), seconds, divergences=f.n_divergent_samples, dim=size(q, 1),
              min_ess=minimum(ess), max_rhat=maximum(rhat), target_acceptance_rate, gq_failed=failed)
    println(@sprintf("%s: %d draws in %.0fs, divergent=%d, dim=%d, min ESS=%.0f, max Rhat=%.3f%s", name, diag.draws, seconds,
                     diag.divergences, diag.dim, diag.min_ess, diag.max_rhat, failed > 0 ? " ($failed draws without generated quantities)" : ""))
    flush(stdout)
    (; names, cons, diag)
end

Scope

The generation-interval and reporting-delay distributions are fixed inputs of the fitted renewal models. Step 1 estimates a delay distribution, but its uncertainty is not propagated into steps 2 and 3. The model that estimates the delay inside the renewal model is shown and checked (finite density and gradient, and the same expected cases as the fixed-delay model at the true delay parameters), not fitted here.
The correlation matrix C of cdar is data, not a sampled covariance; its length scale (30 km) is fixed.
The data are simulated from the model family that is fitted, so the figures show that the declarations recover their own parameters — not how the model behaves under misspecification, reporting artefacts or day-of-week effects.
The models use the StanBlocks backend. The Turing backend refuses top-level assignments that call @deffun functions, and custom @lpxf families, naming the statement.

Run it

After bootstrapping the repository's test environment, the first command fits all models and writes the summaries; the second draws the figures and needs an environment that provides AlgebraOfVega.jl and the vl-convert command-line tool.

julia --startup-file=no --project=test research/epi_renewal/renewal.jl
julia --startup-file=no --project=<plot env> research/epi_renewal/renewal_figures.jl

Epidemic renewal models ​

The process ​

The mechanics, as Stan functions ​

Step 1: the reporting delay, from a linelist ​

Step 2: one population ​

A prior predictive check from the same declaration ​

A forecast is a mask ​

Step 3: six coupled patches ​

The same models through a package ​

Parameters and sampler diagnostics ​

Scope ​

Run it ​

Epidemic renewal models

The process

The mechanics, as Stan functions

Step 1: the reporting delay, from a linelist

Step 2: one population

A prior predictive check from the same declaration

A forecast is a mask

Step 3: six coupled patches

The same models through a package

Parameters and sampler diagnostics

Scope

Run it