StanBlocks.jl

Appearance

EpiSewer — a composable renewal library ​

R EpiSewer packs a large model space into one config-driven Stan program: several reproduction-number processes, several observation families, and seeding / shedding / limit-of-detection / digital-PCR / outlier options, each selected at runtime by an integer flag. EpiSewer.jl describes the same model space the opposite way — as composable Turing components, with no monolith.

StanBlocks realizes both the same way, and it is the way the two designs already point at: a small library of pieces assembled at construction time. The runtime flags of the R monolith become plain Julia composition, and the pieces line up one-to-one with EpiSewer.jl's components. This page shows a slice of that library — one renewal core, a swappable Rt process, and a swappable observation family — and transpiles three assembled configurations to three focused Stan programs. The arrays are build fixtures, not real data.

Companion port in BayesianRegressionModels.jl

BRM ports the same EpiSewer renewal core through its formula interface: Wastewater-based Rt inference (EpiSewer / ww-inference-model). That page frames the renewal model as a @brm kernel(...) term; this page shows the pure-StanBlocks @slic / @deffun library the same backend emits.

One renewal core, swappable Rt process and observation family ​

Two axes vary here, both by composition:

  • Rt process — a weekly random walk, or a sparse spline basis evaluated with the csr_matrix_times_vector primitive (the R model's R_model 0/1/3 spline processes, in compressed-row form).

  • Observation family — an observation submodel per family: obs_lod measures log-concentration with an LOD-censored normal, and obs_dpcr measures digital-PCR positive-partition counts, using to_int for the real→integer partition total. Each family pairs the same renewal core with a different measurement node.

Each model is one Base.merge of the core with an Rt fragment; the monolith's if-on-a-flag is gone.

julia
using StanBlocks

@deffun begin
    es_conv_at(x::vector[nx], pmf::vector[np], t::int)::real = begin
        acc = 0.0
        kmax = min(np, t - 1)
        for k in 1:kmax
            acc = acc + x[t - k] * pmf[k]
        end
        acc
    end
    es_conv(x::vector[nx], pmf::vector[np], out_n::int, off::int)::vector[out_n] = begin
        out::vector[out_n]
        for i in 1:out_n
            out[i] = es_conv_at(x, pmf, i + off)
        end
        out
    end
    es_rw(x0::real, sd::real, z::vector[m])::vector[m+1] = begin
        out::vector[m+1]
        out[1] = x0
        for i in 1:m
            out[i+1] = out[i] + sd * z[i]
        end
        out
    end
    es_renewal(log_i0::real, growth::real, log_rt::vector[nt],
               gen_int::vector[gmax], uot::int)::vector[nt] = begin
        infections::vector[nt]
        for t in 1:uot
            infections[t] = exp(log_i0 + growth * (t - 1))
        end
        for t in (uot + 1):nt
            infections[t] = exp(log_rt[t]) * es_conv_at(infections, gen_int, t)
        end
        infections
    end
end

# ── Swappable observation-family submodels (§3 Form A) ──
# Concentration measured with LOD-censored normal.
obs_lod = @slic begin
    log10_g ~ normal(9.0, 1.0)
    sigma ~ normal(0.0, 1.0; lower = 0.0)
    model_conc = log(10.0) * log10_g + log(signal + 1.0e-8) - log(mwpd)
    obs ~ censored(normal, model_conc, sigma; lower = lod)
    return obs
end
# Digital-PCR positive-partition counts: real->int partition total via `to_int`.
obs_dpcr = @slic begin
    phi ~ gamma(2.0, 0.1)
    n_part = to_int(round(n_partitions))              # real -> int (data-side)
    lambda = (1.0e5 * signal) * n_part / mwpd
    obs ~ neg_binomial_2(lambda, phi)
    return obs
end

es_data = (;
    gen_int     = [0.15, 0.30, 0.30, 0.15, 0.10],
    shed        = [0.05, 0.15, 0.25, 0.2, 0.15, 0.1, 0.07, 0.03],
    week_of_day = repeat(1:5, inner = 7),
    conc        = [3.4, 2.0, 4.1, 3.8, 2.0, 3.2, 4.4, 3.9, 2.0, 3.1, 4.0, 3.7,
                   3.6, 2.0, 4.2, 3.5, 3.9, 2.0, 4.0, 3.7, 3.3, 2.0, 4.1, 3.8,
                   2.0, 3.2, 4.4, 3.9],
    lod         = fill(2.0, 28),
    partitions_pos = fill(1200, 28),
    n_partitions   = 20000.0,
    spl_w = vcat([[0.5, 0.5] for _ in 1:28]...),
    spl_v = vcat([[min(i, 5), min(i, 5) + 1] for i in 1:28]...),
    spl_u = collect(1:2:(2*28 + 1)),
    n_basis = 6, mwpd = 1.0e6,
)

# ── Renewal cores: one per observation family, each leaving the Rt process
#    (`log_rt`) open. The concentration and dPCR families pair the SAME renewal
#    core with a different observation submodel. ──
es_base_conc = @slic es_data begin
    nt  = dims(week_of_day)[1]
    ot  = dims(conc)[1]
    uot = nt - ot
    log_i0 ~ normal(-13.0, 1.0)
    growth ~ normal(0.0, 0.05)
    infections = es_renewal(log_i0, growth, log_rt, gen_int, uot)   # log_rt UNBOUND
    signal = es_conv(infections, shed, ot, uot)
    conc ~ obs_lod(; signal, mwpd, lod)
end
es_base_dpcr = @slic es_data begin
    nt  = dims(week_of_day)[1]
    ot  = dims(partitions_pos)[1]
    uot = nt - ot
    log_i0 ~ normal(-13.0, 1.0)
    growth ~ normal(0.0, 0.05)
    infections = es_renewal(log_i0, growth, log_rt, gen_int, uot)   # log_rt UNBOUND
    signal = es_conv(infections, shed, ot, uot)
    partitions_pos ~ obs_dpcr(; signal, n_partitions, mwpd)
end

# ── Swappable Rt processes ──
rt_rw = quote
    log_r0 ~ normal(0.0, 0.2)
    eta_sd ~ normal(0.0, 0.1; lower = 0.0)
    n_weeks = maximum(week_of_day)
    w :: vector[n_weeks - 1] ~ std_normal()
    log_rt = es_rw(log_r0, eta_sd, w)[week_of_day]
end
rt_spline = quote
    beta :: vector[n_basis] ~ std_normal()
    log_rt = csr_matrix_times_vector(nt, n_basis, spl_w, spl_v, spl_u, beta)
end

# ── The composable library — swap the Rt process (and the observation family) ──
es_models = (;
    rw_lod     = Base.merge(es_base_conc, rt_rw),      # RW Rt,     LOD concentration
    spline_lod = Base.merge(es_base_conc, rt_spline),  # CSR spline Rt, LOD concentration
    rw_dpcr    = Base.merge(es_base_dpcr, rt_rw),      # RW Rt,     dPCR counts (to_int)
)

rw_lod ​

stan
functions {
vector es_rw(
    real x0,
    real sd,
    vector z
) {
    int m = dims(z)[1];
    vector[(m + 1)] out;
    out[1] = x0;
    for(i in 1:m) {
        out[(i + 1)] = (out[i] + (sd * z[i]));
    }
    return out;
}
vector es_renewal(
    real log_i0,
    real growth,
    vector log_rt,
    vector gen_int,
    int uot
) {
    int nt = dims(log_rt)[1];
    vector[nt] infections;
    for(t in 1:uot) {
        infections[t] = exp((log_i0 + (growth * (t - 1))));
    }
    for(t in (uot + 1):nt) {
        infections[t] = (exp(log_rt[t]) * es_conv_at(infections, gen_int, t));
    }
    return infections;
}
real es_conv_at(
    vector x,
    vector pmf,
    int t
) {
    int np = dims(pmf)[1];
    real acc = 0.0;
    int kmax = min(np, (t - 1));
    for(k in 1:kmax) {
        acc = (acc + (x[(t - k)] * pmf[k]));
    }
    return acc;
}
vector es_conv(
    vector x,
    vector pmf,
    int out_n,
    int off
) {
    vector[out_n] out;
    for(i in 1:out_n) {
        out[i] = es_conv_at(x, pmf, (i + off));
    }
    return out;
}
real lower_clamping_normal_lpdf(
    vector y,
    vector lo,
    vector args1,
    real args2
) {
    return sum(lower_clamping_normal_lpdfs(y, lo, args1, args2));
}
vector lower_clamping_normal_lpdfs(
    vector y,
    vector lo,
    vector args1,
    real args2
) {
    int n = dims(y)[1];
    return jbroadcasted_lower_clamping_lpdf_normal(y, lo, args1, args2);
}
vector jbroadcasted_lower_clamping_lpdf_normal(
    vector x1,
    vector x3,
    vector x4,
    real x5
) {
    int n = dims(x1)[1];
    vector[n] rv;
    for(i in 1:n) {
        rv[i] = lower_clamping_normal_lpdf(broadcasted_getindex(x1, i) | 
            broadcasted_getindex(x3, i),
            broadcasted_getindex(x4, i),
            x5
        );
    }
    return rv;
}
real lower_clamping_normal_lpdf(
    real y,
    real lo,
    real args1,
    real args2
) {
    array[1] real rv;
    rv[1] = negative_infinity();
    if((y == lo)) {
        rv[1] = normal_lcdf_stable(lo, args1, args2);
    } else {
        if((y > lo)) {
            rv[1] = normal_lpdf(y | args1, args2);
        }
    }
    return rv[1];
}
real normal_lcdf_stable(
    real x,
    real loc,
    real scale
) {
    return (log(erfc(((-(x - loc)) / (scale * sqrt(2.0))))) - log(2.0));
}
real broadcasted_getindex(vector x, int i) {
    return x[i];
}
vector lower_clamping_vector_normal_rng(
    int anontok__1,
    vector lo,
    vector args1,
    real args2
) {
    int n = anontok__1;
    return jbroadcasted_lower_clamping_cell_rng_normal_rng(rep_vector(0.0, n), lo, args1, args2);
}
vector jbroadcasted_lower_clamping_cell_rng_normal_rng(
    vector x1,
    vector x3,
    vector x4,
    real x5
) {
    int n = dims(x1)[1];
    vector[n] rv;
    for(i in 1:n) {
        rv[i] = lower_clamping_cell_normal_rng(
            broadcasted_getindex(x1, i),
            broadcasted_getindex(x3, i),
            broadcasted_getindex(x4, i),
            x5
        );
    }
    return rv;
}
real lower_clamping_cell_normal_rng(
    real dummy,
    real lo,
    real args1,
    real args2
) {
    return lower_clamping_normal_rng(lo, args1, args2);
}
real lower_clamping_normal_rng(
    real lo,
    real args1,
    real args2
) {
    vector[1] draw;
    draw[1] = normal_rng(args1, args2);
    if((draw[1] < lo)) {
        draw[1] = lo;
    }
    return draw[1];
}
}
data {
    int week_of_day_n;
    int conc_n;
    array[week_of_day_n] int week_of_day;
    int gen_int_n;
    vector[gen_int_n] gen_int;
    int shed_n;
    vector[shed_n] shed;
    real mwpd;
    vector[conc_n] conc;
    int lod_n;
    vector[lod_n] lod;
}
transformed data {
    int nt = week_of_day_n;
    int ot = conc_n;
    int uot = (nt - ot);
    int n_weeks = max(week_of_day);
}
parameters {
    real log_i0;
    real growth;
    real log_r0;
    real<lower=0.0> eta_sd;
    vector[(n_weeks - 1)] w;
    real conc_log10_g;
    real<lower=0.0> conc_sigma;
}
transformed parameters {
    vector[week_of_day_n] log_rt = es_rw(log_r0, eta_sd, w)[week_of_day];
    vector[week_of_day_n] infections = es_renewal(log_i0, growth, log_rt, gen_int, uot);
    vector[ot] signal = es_conv(infections, shed, ot, uot);
    vector[ot] conc_model_conc = (((log(10.0) * conc_log10_g) + log((signal + 1.0e-8))) - log(mwpd));
}
model {
    log_i0 ~ normal(-13.0, 1.0);
    growth ~ normal(0.0, 0.05);
    log_r0 ~ normal(0.0, 0.2);
    eta_sd ~ normal(0.0, 0.1);
    w ~ std_normal();
    conc_log10_g ~ normal(9.0, 1.0);
    conc_sigma ~ normal(0.0, 1.0);
    conc ~ lower_clamping_normal(lod, conc_model_conc, conc_sigma);
}
generated quantities {
    vector[conc_n] conc_likelihood = lower_clamping_normal_lpdfs(conc, lod, conc_model_conc, conc_sigma);
    vector[conc_n] conc_gen = lower_clamping_vector_normal_rng(conc_n, lod, conc_model_conc, conc_sigma);
}

spline_lod ​

stan
functions {
vector es_renewal(
    real log_i0,
    real growth,
    vector log_rt,
    vector gen_int,
    int uot
) {
    int nt = dims(log_rt)[1];
    vector[nt] infections;
    for(t in 1:uot) {
        infections[t] = exp((log_i0 + (growth * (t - 1))));
    }
    for(t in (uot + 1):nt) {
        infections[t] = (exp(log_rt[t]) * es_conv_at(infections, gen_int, t));
    }
    return infections;
}
real es_conv_at(
    vector x,
    vector pmf,
    int t
) {
    int np = dims(pmf)[1];
    real acc = 0.0;
    int kmax = min(np, (t - 1));
    for(k in 1:kmax) {
        acc = (acc + (x[(t - k)] * pmf[k]));
    }
    return acc;
}
vector es_conv(
    vector x,
    vector pmf,
    int out_n,
    int off
) {
    vector[out_n] out;
    for(i in 1:out_n) {
        out[i] = es_conv_at(x, pmf, (i + off));
    }
    return out;
}
real lower_clamping_normal_lpdf(
    vector y,
    vector lo,
    vector args1,
    real args2
) {
    return sum(lower_clamping_normal_lpdfs(y, lo, args1, args2));
}
vector lower_clamping_normal_lpdfs(
    vector y,
    vector lo,
    vector args1,
    real args2
) {
    int n = dims(y)[1];
    return jbroadcasted_lower_clamping_lpdf_normal(y, lo, args1, args2);
}
vector jbroadcasted_lower_clamping_lpdf_normal(
    vector x1,
    vector x3,
    vector x4,
    real x5
) {
    int n = dims(x1)[1];
    vector[n] rv;
    for(i in 1:n) {
        rv[i] = lower_clamping_normal_lpdf(broadcasted_getindex(x1, i) | 
            broadcasted_getindex(x3, i),
            broadcasted_getindex(x4, i),
            x5
        );
    }
    return rv;
}
real lower_clamping_normal_lpdf(
    real y,
    real lo,
    real args1,
    real args2
) {
    array[1] real rv;
    rv[1] = negative_infinity();
    if((y == lo)) {
        rv[1] = normal_lcdf_stable(lo, args1, args2);
    } else {
        if((y > lo)) {
            rv[1] = normal_lpdf(y | args1, args2);
        }
    }
    return rv[1];
}
real normal_lcdf_stable(
    real x,
    real loc,
    real scale
) {
    return (log(erfc(((-(x - loc)) / (scale * sqrt(2.0))))) - log(2.0));
}
real broadcasted_getindex(vector x, int i) {
    return x[i];
}
vector lower_clamping_vector_normal_rng(
    int anontok__1,
    vector lo,
    vector args1,
    real args2
) {
    int n = anontok__1;
    return jbroadcasted_lower_clamping_cell_rng_normal_rng(rep_vector(0.0, n), lo, args1, args2);
}
vector jbroadcasted_lower_clamping_cell_rng_normal_rng(
    vector x1,
    vector x3,
    vector x4,
    real x5
) {
    int n = dims(x1)[1];
    vector[n] rv;
    for(i in 1:n) {
        rv[i] = lower_clamping_cell_normal_rng(
            broadcasted_getindex(x1, i),
            broadcasted_getindex(x3, i),
            broadcasted_getindex(x4, i),
            x5
        );
    }
    return rv;
}
real lower_clamping_cell_normal_rng(
    real dummy,
    real lo,
    real args1,
    real args2
) {
    return lower_clamping_normal_rng(lo, args1, args2);
}
real lower_clamping_normal_rng(
    real lo,
    real args1,
    real args2
) {
    vector[1] draw;
    draw[1] = normal_rng(args1, args2);
    if((draw[1] < lo)) {
        draw[1] = lo;
    }
    return draw[1];
}
}
data {
    int week_of_day_n;
    int conc_n;
    int n_basis;
    int spl_w_n;
    vector[spl_w_n] spl_w;
    int spl_v_n;
    array[spl_v_n] int spl_v;
    int spl_u_n;
    array[spl_u_n] int spl_u;
    int gen_int_n;
    vector[gen_int_n] gen_int;
    int shed_n;
    vector[shed_n] shed;
    real mwpd;
    vector[conc_n] conc;
    int lod_n;
    vector[lod_n] lod;
}
transformed data {
    int nt = week_of_day_n;
    int ot = conc_n;
    int uot = (nt - ot);
}
parameters {
    real log_i0;
    real growth;
    vector[n_basis] beta;
    real conc_log10_g;
    real<lower=0.0> conc_sigma;
}
transformed parameters {
    vector[nt] log_rt = csr_matrix_times_vector(nt, n_basis, spl_w, spl_v, spl_u, beta);
    vector[nt] infections = es_renewal(log_i0, growth, log_rt, gen_int, uot);
    vector[ot] signal = es_conv(infections, shed, ot, uot);
    vector[ot] conc_model_conc = (((log(10.0) * conc_log10_g) + log((signal + 1.0e-8))) - log(mwpd));
}
model {
    log_i0 ~ normal(-13.0, 1.0);
    growth ~ normal(0.0, 0.05);
    beta ~ std_normal();
    conc_log10_g ~ normal(9.0, 1.0);
    conc_sigma ~ normal(0.0, 1.0);
    conc ~ lower_clamping_normal(lod, conc_model_conc, conc_sigma);
}
generated quantities {
    vector[conc_n] conc_likelihood = lower_clamping_normal_lpdfs(conc, lod, conc_model_conc, conc_sigma);
    vector[conc_n] conc_gen = lower_clamping_vector_normal_rng(conc_n, lod, conc_model_conc, conc_sigma);
}

rw_dpcr ​

stan
functions {
vector es_rw(
    real x0,
    real sd,
    vector z
) {
    int m = dims(z)[1];
    vector[(m + 1)] out;
    out[1] = x0;
    for(i in 1:m) {
        out[(i + 1)] = (out[i] + (sd * z[i]));
    }
    return out;
}
vector es_renewal(
    real log_i0,
    real growth,
    vector log_rt,
    vector gen_int,
    int uot
) {
    int nt = dims(log_rt)[1];
    vector[nt] infections;
    for(t in 1:uot) {
        infections[t] = exp((log_i0 + (growth * (t - 1))));
    }
    for(t in (uot + 1):nt) {
        infections[t] = (exp(log_rt[t]) * es_conv_at(infections, gen_int, t));
    }
    return infections;
}
real es_conv_at(
    vector x,
    vector pmf,
    int t
) {
    int np = dims(pmf)[1];
    real acc = 0.0;
    int kmax = min(np, (t - 1));
    for(k in 1:kmax) {
        acc = (acc + (x[(t - k)] * pmf[k]));
    }
    return acc;
}
vector es_conv(
    vector x,
    vector pmf,
    int out_n,
    int off
) {
    vector[out_n] out;
    for(i in 1:out_n) {
        out[i] = es_conv_at(x, pmf, (i + off));
    }
    return out;
}
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;
    if((n == 0)) {
        array[n] int rv;
        return rv;
    } else {
        return neg_binomial_2_rng(a, b);
    }
}
}
data {
    int week_of_day_n;
    int partitions_pos_n;
    array[week_of_day_n] int week_of_day;
    int gen_int_n;
    vector[gen_int_n] gen_int;
    int shed_n;
    vector[shed_n] shed;
    real n_partitions;
    real mwpd;
    array[partitions_pos_n] int partitions_pos;
}
transformed data {
    int nt = week_of_day_n;
    int ot = partitions_pos_n;
    int uot = (nt - ot);
    int n_weeks = max(week_of_day);
    int partitions_pos_n_part = to_int(round(n_partitions));
}
parameters {
    real log_i0;
    real growth;
    real log_r0;
    real<lower=0.0> eta_sd;
    vector[(n_weeks - 1)] w;
    real<lower=0.0> partitions_pos_phi;
}
transformed parameters {
    vector[week_of_day_n] log_rt = es_rw(log_r0, eta_sd, w)[week_of_day];
    vector[week_of_day_n] infections = es_renewal(log_i0, growth, log_rt, gen_int, uot);
    vector[ot] signal = es_conv(infections, shed, ot, uot);
    vector[ot] partitions_pos_lambda = (((100000.0 * signal) * partitions_pos_n_part) / mwpd);
}
model {
    log_i0 ~ normal(-13.0, 1.0);
    growth ~ normal(0.0, 0.05);
    log_r0 ~ normal(0.0, 0.2);
    eta_sd ~ normal(0.0, 0.1);
    w ~ std_normal();
    partitions_pos_phi ~ gamma(2.0, 0.1);
    partitions_pos ~ neg_binomial_2(partitions_pos_lambda, partitions_pos_phi);
}
generated quantities {
    vector[partitions_pos_n] partitions_pos_likelihood = neg_binomial_2_lpmfs(partitions_pos, partitions_pos_lambda, partitions_pos_phi);
    array[partitions_pos_n] int partitions_pos_gen = neg_binomial_2_int_rng(partitions_pos_n, partitions_pos_lambda, partitions_pos_phi);
}

What this exercises ​

  • to_int — the real→integer partition total in the digital-PCR observation family (to_int(round(n_partitions))), a data-side conversion that lands in transformed data where Stan's data-qualified-argument rule is satisfied.

  • csr_matrix_times_vector — the sparse compressed-row spline basis for the spline Rt process, exposed directly rather than densified.

  • Observation submodels — the concentration and dPCR measurement families as self-contained, swappable nodes.

  • Base.merge composition — the R monolith's runtime module flags become construction-time assembly, mirroring EpiSewer.jl's component design.

For the full joint renewal model — two coupled observation streams, a multi-subpopulation structure, tuple-returning renewal and per-site effects — see the wastewater renewal model case study.

You are viewing the dev branch. This branch may include code written with Claude Code with less human supervision. Only human-approved code is merged into main.