Wastewater renewal model
This example ports the CDC ww-inference-model (inst/stan/wwinference.stan) — a joint model that infers latent infections from two noisy signals at once: hospital admissions and pathogen concentration in wastewater. It is a renewal model: today's infections are generated from recent infections and a time-varying reproduction number, rather than from a mechanistic compartmental system.
The published program is a single monolith with a runtime flag for each modeling choice. The port takes the shape StanBlocks makes natural instead: a small library of composable pieces, assembled at construction time. This page climbs a modeling ladder — a base model, then rungs that swap one component at a time — and ends at the full multi-subpopulation model. Every rung is evaluated at documentation-build time, so the displayed Julia is the exact source that produced the Stan beside it. The small arrays are build fixtures, not real surveillance data.
Companion port in BayesianRegressionModels.jl
The same CDC ww-inference-model is also ported in BRM, onto its StanBlocks backend (@slic + @deffun): The full CDC ww-inference-model. Both target the identical upstream model — this page as a standalone StanBlocks modeling ladder, the BRM page from the @brm / StanBlocks layer boundary.
Two observation streams, two removable submodels
The joint model observes two data streams. Each is an observation submodel (data ~ submodel(...)): a self-contained node that carries its own parameters and its own likelihood, added or dropped as a single line.
admissions_streamcarries the infection-hospitalization ratio, the negative-binomial dispersion, and asimplex[7]day-of-week reporting effect, and observes admissions withneg_binomial_2.wastewater_streamcarries the genome-scaling and measurement noise, and observes log-concentration withcensored(normal, …; lower = lod)— so below-limit-of-detection samples contribute the log-CDF automatically.
The renewal core produces the latent infection curve and leaves two things open: the Rt process (an unbound log_rt) and the shedding kernel (a shed_kernel that defaults to a data pmf). Each observation stream is one line; each ladder rung below binds log_rt (or swaps shed_kernel) through Base.merge, without touching the core or the streams.
The Rt ladder — one core, three reproduction-number processes
The same core becomes three models by merging in a different Rt fragment. The CDC source selects between these with an integer flag; here each is a value you compose in. spline_csr uses the sparse csr_matrix_times_vector primitive on a compressed-row spline basis; rw_parametric additionally swaps the data shedding pmf for a parametric viral-shedding trajectory driven by peak-timing parameters.
julia
using StanBlocks
@deffun begin
cdc_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
cdc_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] = cdc_conv_at(x, pmf, i + off)
end
out
end
cdc_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
"diff-AR(1): the weekly *differences* follow an AR(1) around 0; the level is
x0 + cumulative sum of the differences (the CDC reference-Rt process)."
cdc_diff_ar1(x0::real, ar::real, sd::real, z::vector[m])::vector[m+1] = begin
diffs::vector[m]
diffs[1] = sd * z[1]
for t in 2:m
diffs[t] = ar * diffs[t-1] + sd * z[t]
end
lvl::vector[m+1]
lvl[1] = x0
for t in 1:m
lvl[t+1] = lvl[t] + diffs[t]
end
lvl
end
"Parametric viral-shedding trajectory: a triangular-on-log profile peaking at
day t_peak, declining to zero at dur_shed, normalised to sum 1."
cdc_vl_trajectory(t_peak::real, viral_peak::real, dur_shed::real, n::int)::vector[n] = begin
v::vector[n]
for k in 1:n
tk = k * 1.0
frac = tk <= t_peak ? tk / t_peak : (dur_shed - tk) / (dur_shed - t_peak)
height = frac > 0.0 ? frac : 0.0
v[k] = exp(viral_peak * height)
end
v / sum(v)
end
cdc_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]) * cdc_conv_at(infections, gen_int, t)
end
infections
end
end
# ── Observation submodels (§3 Form A): one removable node per stream ──
admissions_stream = @slic begin
ihr ~ beta(2.0, 20.0)
phi ~ gamma(2.0, 0.1)
wday_effect :: simplex[7] ~ dirichlet(rep_vector(1.0, 7))
wday = 7.0 * wday_effect[day_of_week]
mu = (state_pop * ihr) * exp_hosp_pc .* wday
obs ~ neg_binomial_2(mu, phi)
return obs
end
wastewater_stream = @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(conc_pc + 1.0e-8) - log(mwpd)
obs ~ censored(normal, model_conc, sigma; lower = lod)
return obs
end
cdc_data = (;
week_of_day = repeat(1:5, inner = 7),
day_of_week = repeat(1:7, outer = 4),
hosp = fill(20, 28),
gen_int = [0.15, 0.30, 0.30, 0.15, 0.10],
inf_to_hosp = [0.02, 0.1, 0.2, 0.25, 0.2, 0.12, 0.07, 0.04],
shed = [0.05, 0.15, 0.25, 0.2, 0.15, 0.1, 0.07, 0.03],
ww_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],
ww_lod = fill(2.0, 28),
# sparse spline basis (28 obs days x 6 basis fns), banded CSR triple:
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, n_shed = 8,
state_pop = 1.0e5, mwpd = 1.0e6,
)
# ── Renewal core with an UNBOUND `log_rt` (Rt axis) and a swappable shedding
# kernel (`shed_kernel`, data by default). Each observation stream is one line. ──
cdc_core = @slic cdc_data begin
nt = dims(week_of_day)[1]
ot = dims(hosp)[1]
uot = nt - ot
n_weeks = maximum(week_of_day)
log_i0 ~ normal(-13.0, 1.0)
growth ~ normal(0.0, 0.05)
infections = cdc_renewal(log_i0, growth, log_rt, gen_int, uot) # log_rt UNBOUND
exp_hosp_pc = cdc_conv(infections, inf_to_hosp, ot, uot)
shed_kernel = shed # shedding axis (data default)
conc_pc = cdc_conv(infections, shed_kernel, ot, uot)
hosp ~ admissions_stream(; day_of_week, state_pop, exp_hosp_pc)
ww_conc ~ wastewater_stream(; conc_pc, mwpd, lod = ww_lod)
end
# ── Swappable Rt-process fragments (each binds `log_rt`) ──
rt_rw = quote
log_r0 ~ normal(0.0, 0.2)
eta_sd ~ normal(0.0, 0.1; lower = 0.0)
w :: vector[n_weeks - 1] ~ std_normal()
log_rt = cdc_rw(log_r0, eta_sd, w)[week_of_day]
end
rt_diff_ar1 = quote
log_r0 ~ normal(0.0, 0.2)
ar ~ beta(2.0, 2.0)
eta_sd ~ normal(0.0, 0.1; lower = 0.0)
z :: vector[n_weeks - 1] ~ std_normal()
log_rt = cdc_diff_ar1(log_r0, ar, eta_sd, z)[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
# ── Swappable shedding fragment (parametric VL trajectory) ──
shed_parametric = quote
t_peak ~ normal(3.0, 1.0; lower = 0.0)
viral_peak ~ normal(2.0, 1.0)
dur_shed ~ normal(8.0, 2.0; lower = 0.0)
shed_kernel = cdc_vl_trajectory(t_peak, viral_peak, dur_shed, n_shed)
end
# ── The modeling ladder — each rung a composition on the previous ──
ladder = (;
rw = Base.merge(cdc_core, rt_rw),
diff_ar1 = Base.merge(cdc_core, rt_diff_ar1),
spline_csr = Base.merge(cdc_core, rt_spline),
rw_parametric = Base.merge(cdc_core, rt_rw, shed_parametric),
)rw
stan
functions {
vector cdc_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 cdc_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]) * cdc_conv_at(infections, gen_int, t));
}
return infections;
}
real cdc_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 cdc_conv(
vector x,
vector pmf,
int out_n,
int off
) {
vector[out_n] out;
for(i in 1:out_n) {
out[i] = cdc_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);
}
}
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));
}
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 hosp_n;
array[week_of_day_n] int week_of_day;
int gen_int_n;
vector[gen_int_n] gen_int;
int inf_to_hosp_n;
vector[inf_to_hosp_n] inf_to_hosp;
int shed_n;
vector[shed_n] shed;
int day_of_week_n;
array[day_of_week_n] int day_of_week;
real state_pop;
array[hosp_n] int hosp;
real mwpd;
int ww_conc_n;
vector[ww_conc_n] ww_conc;
int ww_lod_n;
vector[ww_lod_n] ww_lod;
}
transformed data {
int nt = week_of_day_n;
int ot = hosp_n;
int uot = (nt - ot);
int n_weeks = max(week_of_day);
vector[shed_n] shed_kernel = shed;
}
parameters {
real log_i0;
real growth;
real log_r0;
real<lower=0.0> eta_sd;
vector[(n_weeks - 1)] w;
real<lower=0, upper=1> hosp_ihr;
real<lower=0.0> hosp_phi;
simplex[7] hosp_wday_effect;
real ww_conc_log10_g;
real<lower=0.0> ww_conc_sigma;
}
transformed parameters {
vector[week_of_day_n] log_rt = cdc_rw(log_r0, eta_sd, w)[week_of_day];
vector[week_of_day_n] infections = cdc_renewal(log_i0, growth, log_rt, gen_int, uot);
vector[ot] exp_hosp_pc = cdc_conv(infections, inf_to_hosp, ot, uot);
vector[ot] conc_pc = cdc_conv(infections, shed_kernel, ot, uot);
vector[day_of_week_n] hosp_wday = (7.0 * hosp_wday_effect[day_of_week]);
vector[day_of_week_n] hosp_mu = (((state_pop * hosp_ihr) * exp_hosp_pc) .* hosp_wday);
vector[ot] ww_conc_model_conc = (((log(10.0) * ww_conc_log10_g) + log((conc_pc + 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();
hosp_ihr ~ beta(2.0, 20.0);
hosp_phi ~ gamma(2.0, 0.1);
hosp_wday_effect ~ dirichlet(rep_vector(1.0, 7));
hosp ~ neg_binomial_2(hosp_mu, hosp_phi);
ww_conc_log10_g ~ normal(9.0, 1.0);
ww_conc_sigma ~ normal(0.0, 1.0);
ww_conc ~ lower_clamping_normal(ww_lod, ww_conc_model_conc, ww_conc_sigma);
}
generated quantities {
vector[hosp_n] hosp_likelihood = neg_binomial_2_lpmfs(hosp, hosp_mu, hosp_phi);
array[hosp_n] int hosp_gen = neg_binomial_2_int_rng(hosp_n, hosp_mu, hosp_phi);
vector[ww_conc_n] ww_conc_likelihood = lower_clamping_normal_lpdfs(ww_conc, ww_lod, ww_conc_model_conc, ww_conc_sigma);
vector[ww_conc_n] ww_conc_gen = lower_clamping_vector_normal_rng(ww_conc_n, ww_lod, ww_conc_model_conc, ww_conc_sigma);
}diff_ar1
stan
functions {
// diff-AR(1): the weekly *differences* follow an AR(1) around 0; the level is
// x0 + cumulative sum of the differences (the CDC reference-Rt process).
vector cdc_diff_ar1(
real x0,
real ar,
real sd,
vector z
) {
int m = dims(z)[1];
vector[m] diffs;
diffs[1] = (sd * z[1]);
for(t in 2:m) {
diffs[t] = ((ar * diffs[(t - 1)]) + (sd * z[t]));
}
vector[(m + 1)] lvl;
lvl[1] = x0;
for(t in 1:m) {
lvl[(t + 1)] = (lvl[t] + diffs[t]);
}
return lvl;
}
vector cdc_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]) * cdc_conv_at(infections, gen_int, t));
}
return infections;
}
real cdc_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 cdc_conv(
vector x,
vector pmf,
int out_n,
int off
) {
vector[out_n] out;
for(i in 1:out_n) {
out[i] = cdc_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);
}
}
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));
}
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 hosp_n;
array[week_of_day_n] int week_of_day;
int gen_int_n;
vector[gen_int_n] gen_int;
int inf_to_hosp_n;
vector[inf_to_hosp_n] inf_to_hosp;
int shed_n;
vector[shed_n] shed;
int day_of_week_n;
array[day_of_week_n] int day_of_week;
real state_pop;
array[hosp_n] int hosp;
real mwpd;
int ww_conc_n;
vector[ww_conc_n] ww_conc;
int ww_lod_n;
vector[ww_lod_n] ww_lod;
}
transformed data {
int nt = week_of_day_n;
int ot = hosp_n;
int uot = (nt - ot);
int n_weeks = max(week_of_day);
vector[shed_n] shed_kernel = shed;
}
parameters {
real log_i0;
real growth;
real log_r0;
real<lower=0, upper=1> ar;
real<lower=0.0> eta_sd;
vector[(n_weeks - 1)] z;
real<lower=0, upper=1> hosp_ihr;
real<lower=0.0> hosp_phi;
simplex[7] hosp_wday_effect;
real ww_conc_log10_g;
real<lower=0.0> ww_conc_sigma;
}
transformed parameters {
vector[week_of_day_n] log_rt = cdc_diff_ar1(log_r0, ar, eta_sd, z)[week_of_day];
vector[week_of_day_n] infections = cdc_renewal(log_i0, growth, log_rt, gen_int, uot);
vector[ot] exp_hosp_pc = cdc_conv(infections, inf_to_hosp, ot, uot);
vector[ot] conc_pc = cdc_conv(infections, shed_kernel, ot, uot);
vector[day_of_week_n] hosp_wday = (7.0 * hosp_wday_effect[day_of_week]);
vector[day_of_week_n] hosp_mu = (((state_pop * hosp_ihr) * exp_hosp_pc) .* hosp_wday);
vector[ot] ww_conc_model_conc = (((log(10.0) * ww_conc_log10_g) + log((conc_pc + 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);
ar ~ beta(2.0, 2.0);
eta_sd ~ normal(0.0, 0.1);
z ~ std_normal();
hosp_ihr ~ beta(2.0, 20.0);
hosp_phi ~ gamma(2.0, 0.1);
hosp_wday_effect ~ dirichlet(rep_vector(1.0, 7));
hosp ~ neg_binomial_2(hosp_mu, hosp_phi);
ww_conc_log10_g ~ normal(9.0, 1.0);
ww_conc_sigma ~ normal(0.0, 1.0);
ww_conc ~ lower_clamping_normal(ww_lod, ww_conc_model_conc, ww_conc_sigma);
}
generated quantities {
vector[hosp_n] hosp_likelihood = neg_binomial_2_lpmfs(hosp, hosp_mu, hosp_phi);
array[hosp_n] int hosp_gen = neg_binomial_2_int_rng(hosp_n, hosp_mu, hosp_phi);
vector[ww_conc_n] ww_conc_likelihood = lower_clamping_normal_lpdfs(ww_conc, ww_lod, ww_conc_model_conc, ww_conc_sigma);
vector[ww_conc_n] ww_conc_gen = lower_clamping_vector_normal_rng(ww_conc_n, ww_lod, ww_conc_model_conc, ww_conc_sigma);
}spline_csr
stan
functions {
vector cdc_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]) * cdc_conv_at(infections, gen_int, t));
}
return infections;
}
real cdc_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 cdc_conv(
vector x,
vector pmf,
int out_n,
int off
) {
vector[out_n] out;
for(i in 1:out_n) {
out[i] = cdc_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);
}
}
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));
}
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 hosp_n;
array[week_of_day_n] int week_of_day;
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 inf_to_hosp_n;
vector[inf_to_hosp_n] inf_to_hosp;
int shed_n;
vector[shed_n] shed;
int day_of_week_n;
array[day_of_week_n] int day_of_week;
real state_pop;
array[hosp_n] int hosp;
real mwpd;
int ww_conc_n;
vector[ww_conc_n] ww_conc;
int ww_lod_n;
vector[ww_lod_n] ww_lod;
}
transformed data {
int nt = week_of_day_n;
int ot = hosp_n;
int uot = (nt - ot);
int n_weeks = max(week_of_day);
vector[shed_n] shed_kernel = shed;
}
parameters {
real log_i0;
real growth;
vector[n_basis] beta;
real<lower=0, upper=1> hosp_ihr;
real<lower=0.0> hosp_phi;
simplex[7] hosp_wday_effect;
real ww_conc_log10_g;
real<lower=0.0> ww_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 = cdc_renewal(log_i0, growth, log_rt, gen_int, uot);
vector[ot] exp_hosp_pc = cdc_conv(infections, inf_to_hosp, ot, uot);
vector[ot] conc_pc = cdc_conv(infections, shed_kernel, ot, uot);
vector[day_of_week_n] hosp_wday = (7.0 * hosp_wday_effect[day_of_week]);
vector[day_of_week_n] hosp_mu = (((state_pop * hosp_ihr) * exp_hosp_pc) .* hosp_wday);
vector[ot] ww_conc_model_conc = (((log(10.0) * ww_conc_log10_g) + log((conc_pc + 1.0e-8))) - log(mwpd));
}
model {
log_i0 ~ normal(-13.0, 1.0);
growth ~ normal(0.0, 0.05);
beta ~ std_normal();
hosp_ihr ~ beta(2.0, 20.0);
hosp_phi ~ gamma(2.0, 0.1);
hosp_wday_effect ~ dirichlet(rep_vector(1.0, 7));
hosp ~ neg_binomial_2(hosp_mu, hosp_phi);
ww_conc_log10_g ~ normal(9.0, 1.0);
ww_conc_sigma ~ normal(0.0, 1.0);
ww_conc ~ lower_clamping_normal(ww_lod, ww_conc_model_conc, ww_conc_sigma);
}
generated quantities {
vector[hosp_n] hosp_likelihood = neg_binomial_2_lpmfs(hosp, hosp_mu, hosp_phi);
array[hosp_n] int hosp_gen = neg_binomial_2_int_rng(hosp_n, hosp_mu, hosp_phi);
vector[ww_conc_n] ww_conc_likelihood = lower_clamping_normal_lpdfs(ww_conc, ww_lod, ww_conc_model_conc, ww_conc_sigma);
vector[ww_conc_n] ww_conc_gen = lower_clamping_vector_normal_rng(ww_conc_n, ww_lod, ww_conc_model_conc, ww_conc_sigma);
}rw_parametric
stan
functions {
vector cdc_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 cdc_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]) * cdc_conv_at(infections, gen_int, t));
}
return infections;
}
real cdc_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 cdc_conv(
vector x,
vector pmf,
int out_n,
int off
) {
vector[out_n] out;
for(i in 1:out_n) {
out[i] = cdc_conv_at(x, pmf, (i + off));
}
return out;
}
// Parametric viral-shedding trajectory: a triangular-on-log profile peaking at
// day t_peak, declining to zero at dur_shed, normalised to sum 1.
vector cdc_vl_trajectory(
real t_peak,
real viral_peak,
real dur_shed,
int n
) {
vector[n] v;
for(k in 1:n) {
real tk = (k * 1.0);
real frac = ((tk <= t_peak) ? (tk / t_peak) : ((dur_shed - tk) / (dur_shed - t_peak)));
real height = ((frac > 0.0) ? frac : 0.0);
v[k] = exp((viral_peak * height));
}
return (v / sum(v));
}
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);
}
}
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));
}
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 hosp_n;
array[week_of_day_n] int week_of_day;
int gen_int_n;
vector[gen_int_n] gen_int;
int inf_to_hosp_n;
vector[inf_to_hosp_n] inf_to_hosp;
int n_shed;
int day_of_week_n;
array[day_of_week_n] int day_of_week;
real state_pop;
array[hosp_n] int hosp;
real mwpd;
int ww_conc_n;
vector[ww_conc_n] ww_conc;
int ww_lod_n;
vector[ww_lod_n] ww_lod;
}
transformed data {
int nt = week_of_day_n;
int ot = hosp_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<lower=0.0> t_peak;
real viral_peak;
real<lower=0.0> dur_shed;
real<lower=0, upper=1> hosp_ihr;
real<lower=0.0> hosp_phi;
simplex[7] hosp_wday_effect;
real ww_conc_log10_g;
real<lower=0.0> ww_conc_sigma;
}
transformed parameters {
vector[week_of_day_n] log_rt = cdc_rw(log_r0, eta_sd, w)[week_of_day];
vector[week_of_day_n] infections = cdc_renewal(log_i0, growth, log_rt, gen_int, uot);
vector[ot] exp_hosp_pc = cdc_conv(infections, inf_to_hosp, ot, uot);
vector[n_shed] shed_kernel = cdc_vl_trajectory(t_peak, viral_peak, dur_shed, n_shed);
vector[ot] conc_pc = cdc_conv(infections, shed_kernel, ot, uot);
vector[day_of_week_n] hosp_wday = (7.0 * hosp_wday_effect[day_of_week]);
vector[day_of_week_n] hosp_mu = (((state_pop * hosp_ihr) * exp_hosp_pc) .* hosp_wday);
vector[ot] ww_conc_model_conc = (((log(10.0) * ww_conc_log10_g) + log((conc_pc + 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();
t_peak ~ normal(3.0, 1.0);
viral_peak ~ normal(2.0, 1.0);
dur_shed ~ normal(8.0, 2.0);
hosp_ihr ~ beta(2.0, 20.0);
hosp_phi ~ gamma(2.0, 0.1);
hosp_wday_effect ~ dirichlet(rep_vector(1.0, 7));
hosp ~ neg_binomial_2(hosp_mu, hosp_phi);
ww_conc_log10_g ~ normal(9.0, 1.0);
ww_conc_sigma ~ normal(0.0, 1.0);
ww_conc ~ lower_clamping_normal(ww_lod, ww_conc_model_conc, ww_conc_sigma);
}
generated quantities {
vector[hosp_n] hosp_likelihood = neg_binomial_2_lpmfs(hosp, hosp_mu, hosp_phi);
array[hosp_n] int hosp_gen = neg_binomial_2_int_rng(hosp_n, hosp_mu, hosp_phi);
vector[ww_conc_n] ww_conc_likelihood = lower_clamping_normal_lpdfs(ww_conc, ww_lod, ww_conc_model_conc, ww_conc_sigma);
vector[ww_conc_n] ww_conc_gen = lower_clamping_vector_normal_rng(ww_conc_n, ww_lod, ww_conc_model_conc, ww_conc_sigma);
}The full multi-subpopulation model
The published model is not single-population: it runs several coupled sub-epidemics and aggregates them to the state level. The capstone below carries that structure faithfully and showcases the remaining features:
a
plateover sub-epidemics — each cell samples that subpopulation's own initial size, growth, and weekly AR(1) innovations and returns its infection column; thevector[nt]cell outputs collect into amatrix[nt, S]. The unavoidable within-subpop time recurrence (a scan) lives in a@deffunthe cell calls; the subpopulation axis is the plate. This matches the companion BRM port, which expresses the same subpop axis as aplate;per-subpopulation
Rt— a shared reference weekly random walk plus a stationary AR(1) deviation per subpopulation (its own innovations sampled inside the plate cell), and infection feedback in the renewal recurrence;aggregation to admissions by a matrix-vector product
M * pop_frac, and per-subpopulation wastewater gathered at the sampled(subpopulation, day)pairs;a per-site wastewater stream — its own site offset and per-lab-site noise, indexed by each sample's site.
The two observation submodels are reused unchanged; only the wastewater stream gains the per-site index.
julia
using StanBlocks
@deffun begin
cdc_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
cdc_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] = cdc_conv_at(x, pmf, i + off)
end
out
end
cdc_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
cdc_ar1(ac::real, sd::real, z::vector[m])::vector[m] = begin
out::vector[m]
out[1] = sd / sqrt(1.0 - ac * ac) * z[1]
for t in 2:m
out[t] = ac * out[t-1] + sd * z[t]
end
out
end
"Single sub-epidemic renewal for ONE `plate` cell: the subpop's daily log-Rt
(shared reference weekly RW + that subpop's AR(1) deviation) drives the
feedback-renewal scan. Returns the subpop's daily infection vector[nt]. The
subpop LOOP is the `plate` below; this @deffun owns only the unavoidable
within-subpop TIME recurrence (a scan, which a plate cell cannot express)."
cdc_subpop_cell(log_i0::real, growth::real, log_rt_ref::vector[nw], z::vector[nz],
ac::real, sd::real, week_of_day::int[nt],
gen_int::vector[gmax], fb_pmf::vector[fl], feedback::real, uot::int)::vector[nt] = begin
dev = cdc_ar1(ac, sd, z)
M::vector[nt]
for t in 1:uot
M[t] = exp(log_i0 + growth * (t - 1))
end
for t in (uot + 1):nt
infness = 0.0
kmax = min(gmax, t - 1)
for k in 1:kmax
infness = infness + M[t - k] * gen_int[k]
end
fbk = 0.0
fmax = min(fl, t - 1)
for k in 1:fmax
fbk = fbk + M[t - k] * fb_pmf[k]
end
wk = week_of_day[t]
rt_eff = exp((log_rt_ref[wk] + dev[wk]) - feedback * fbk)
M[t] = rt_eff * infness
end
M
end
"Per-subpop net shedding convolution -> nt x S log-scale matrix (the log10_g
genome-scaling is added downstream in the wastewater stream)."
cdc_shed_conc(M::matrix[nt, S], shed::vector[ns], mwpd::real)::matrix[nt, S] = begin
C::matrix[nt, S]
for s in 1:S
for t in 1:nt
net = 0.0
kmax = min(ns, t)
for k in 1:kmax
net = net + M[t - k + 1, s] * shed[k]
end
C[t, s] = log(net + 1.0e-8) - log(mwpd)
end
end
C
end
"Gather per-subpop concentration at observation (subpop, day) pairs."
cdc_gather(conc::matrix[nt, S], subpop::int[nobs], day::int[nobs])::vector[nobs] = begin
out::vector[nobs]
for i in 1:nobs
out[i] = conc[day[i], subpop[i]]
end
out
end
end
# Observation stream submodels (reused, per §3 Form A).
admissions_stream = @slic begin
ihr ~ beta(2.0, 20.0)
phi ~ gamma(2.0, 0.1)
wday_effect :: simplex[7] ~ dirichlet(rep_vector(1.0, 7))
wday = 7.0 * wday_effect[day_of_week]
mu = (state_pop * ihr) * exp_hosp_pc .* wday
obs ~ neg_binomial_2(mu, phi)
return obs
end
# Per-site wastewater stream: its own log10_g + a per-lab-site sigma, indexed by
# the observation's site (ww_site_mod / per-site variance = faithful detail).
wastewater_site_stream = @slic begin
log10_g ~ normal(9.0, 1.0)
site_sigma :: vector[n_sites] ~ normal(0.0, 1.0; lower = 0.0)
site_mod :: vector[n_sites] ~ normal(0.0, 0.5)
model_conc = log(10.0) * log10_g + model_log_conc + site_mod[ww_site]
obs ~ censored(normal, model_conc, site_sigma[ww_site]; lower = lod)
return obs
end
cdc_capstone = @slic (;
week_of_day = repeat(1:5, inner = 7),
day_of_week = repeat(1:7, outer = 4),
hosp = fill(20, 28),
gen_int = [0.15, 0.30, 0.30, 0.15, 0.10],
fb_pmf = [0.5, 0.3, 0.2],
inf_to_hosp = [0.02, 0.1, 0.2, 0.25, 0.2, 0.12, 0.07, 0.04],
shed = [0.05, 0.15, 0.25, 0.2, 0.15, 0.1, 0.07, 0.03],
pop_frac = [0.6, 0.4],
ww_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],
ww_lod = fill(2.0, 12),
ww_day = [3, 7, 10, 14, 17, 21, 24, 5, 9, 13, 19, 23],
ww_subpop = [1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2],
ww_site = [1, 1, 2, 2, 3, 3, 3, 1, 2, 2, 3, 3],
n_sites = 3, n_shed = 8,
state_pop = 1.0e5, mwpd = 1.0e6,
) begin
nt = dims(week_of_day)[1]
ot = dims(hosp)[1]
uot = nt - ot
n_weeks = maximum(week_of_day)
S = dims(pop_frac)[1]
# shared reference weekly log-Rt (random walk), captured by every subpop cell
log_r0 ~ normal(0.0, 0.2)
eta_sd ~ normal(0.0, 0.1; lower = 0.0)
w_ref :: vector[n_weeks - 1] ~ std_normal()
autoreg ~ beta(2.0, 2.0)
sigma_subpop ~ normal(0.0, 0.1; lower = 0.0)
feedback ~ normal(0.0, 0.01; lower = 0.0)
log_rt_ref = cdc_rw(log_r0, eta_sd, w_ref)
# per-subpop renewal as a `plate` over the S sub-epidemics: each cell samples
# its OWN initial size / growth / weekly AR(1) innovations and returns that
# subpop's infection column. Shared pieces (log_rt_ref, feedback, autoreg, …)
# are captured; the vector[nt] cell outputs collect into I_mat :: matrix[nt, S].
I_mat :: matrix[nt, S] ~ plate(; outer = (S,)) do s
log_i0_s ~ normal(-13.0, 1.0)
growth_s ~ normal(0.0, 0.05)
z_s :: vector[n_weeks] ~ std_normal()
cdc_subpop_cell(log_i0_s, growth_s, log_rt_ref, z_s, autoreg, sigma_subpop,
week_of_day, gen_int, fb_pmf, feedback, uot)
end
state_inf = I_mat * pop_frac # aggregate per-capita infections
exp_hosp_pc = cdc_conv(state_inf, inf_to_hosp, ot, uot)
hosp ~ admissions_stream(; day_of_week, state_pop, exp_hosp_pc)
conc_matrix = cdc_shed_conc(I_mat, shed, mwpd) # net shed; log10_g added in the stream
model_log_conc = cdc_gather(conc_matrix, ww_subpop, ww_day)
ww_conc ~ wastewater_site_stream(; model_log_conc, ww_site, n_sites, lod = ww_lod)
endstan
functions {
vector cdc_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;
}
real normal_lpdfs(
real args1,
real args2,
real args3
) {
return normal_lpdf(args1 | args2, args3);
}
vector std_normal_vector_rng(
int anontok__1
) {
int n = anontok__1;
if((n == 0)) {
vector[n] rv;
return rv;
} else {
return to_vector(normal_rng(rep_vector(0, n), 1));
}
}
// Single sub-epidemic renewal for ONE `plate` cell: the subpop's daily log-Rt
// (shared reference weekly RW + that subpop's AR(1) deviation) drives the
// feedback-renewal scan. Returns the subpop's daily infection vector[nt]. The
// subpop LOOP is the `plate` below; this @deffun owns only the unavoidable
// within-subpop TIME recurrence (a scan, which a plate cell cannot express).
vector cdc_subpop_cell(
real log_i0,
real growth,
vector log_rt_ref,
vector z,
real ac,
real sd,
array[] int week_of_day,
vector gen_int,
vector fb_pmf,
real feedback,
int uot
) {
int nt = dims(week_of_day)[1];
int gmax = dims(gen_int)[1];
int fl = dims(fb_pmf)[1];
vector[dims(z)[1]] dev = cdc_ar1(ac, sd, z);
vector[nt] M;
for(t in 1:uot) {
M[t] = exp((log_i0 + (growth * (t - 1))));
}
for(t in (uot + 1):nt) {
real infness = 0.0;
int kmax = min(gmax, (t - 1));
for(k in 1:kmax) {
infness = (infness + (M[(t - k)] * gen_int[k]));
}
real fbk = 0.0;
int fmax = min(fl, (t - 1));
for(k in 1:fmax) {
fbk = (fbk + (M[(t - k)] * fb_pmf[k]));
}
int wk = week_of_day[t];
real rt_eff = exp(((log_rt_ref[wk] + dev[wk]) - (feedback * fbk)));
M[t] = (rt_eff * infness);
}
return M;
}
vector cdc_ar1(
real ac,
real sd,
vector z
) {
int m = dims(z)[1];
vector[m] out;
out[1] = ((sd / sqrt((1.0 - (ac * ac)))) * z[1]);
for(t in 2:m) {
out[t] = ((ac * out[(t - 1)]) + (sd * z[t]));
}
return out;
}
vector cdc_conv(
vector x,
vector pmf,
int out_n,
int off
) {
vector[out_n] out;
for(i in 1:out_n) {
out[i] = cdc_conv_at(x, pmf, (i + off));
}
return out;
}
real cdc_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 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);
}
}
// Per-subpop net shedding convolution -> nt x S log-scale matrix (the log10_g
// genome-scaling is added downstream in the wastewater stream).
matrix cdc_shed_conc(
matrix M,
vector shed,
real mwpd
) {
int nt = dims(M)[1];
int S = dims(M)[2];
int ns = dims(shed)[1];
matrix[nt, S] C;
for(s in 1:S) {
for(t in 1:nt) {
real net = 0.0;
int kmax = min(ns, t);
for(k in 1:kmax) {
net = (net + (M[((t - k) + 1), s] * shed[k]));
}
C[t, s] = (log((net + 1.0e-8)) - log(mwpd));
}
}
return C;
}
// Gather per-subpop concentration at observation (subpop, day) pairs.
vector cdc_gather(
matrix conc,
array[] int subpop,
array[] int day
) {
int nobs = dims(subpop)[1];
if (dims(day)[1] != nobs) reject("cdc_gather: dim mismatch — `day` dim 1 (= ", dims(day)[1], ") does not match `nobs` (= ", nobs, "), inferred from `subpop` dim 1. `nobs` sizes: `subpop` dim 1 (= ", dims(subpop)[1], "), `day` dim 1 (= ", dims(day)[1], ").");
vector[nobs] out;
for(i in 1:nobs) {
out[i] = conc[day[i], subpop[i]];
}
return out;
}
real lower_clamping_normal_lpdf(
vector y,
vector lo,
vector args1,
vector args2
) {
return sum(lower_clamping_normal_lpdfs(y, lo, args1, args2));
}
vector lower_clamping_normal_lpdfs(
vector y,
vector lo,
vector args1,
vector 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,
vector 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),
broadcasted_getindex(x5, i)
);
}
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));
}
vector lower_clamping_vector_normal_rng(
int anontok__1,
vector lo,
vector args1,
vector 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,
vector 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),
broadcasted_getindex(x5, i)
);
}
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 hosp_n;
array[week_of_day_n] int week_of_day;
int pop_frac_n;
int gen_int_n;
vector[gen_int_n] gen_int;
int fb_pmf_n;
vector[fb_pmf_n] fb_pmf;
vector[pop_frac_n] pop_frac;
int inf_to_hosp_n;
vector[inf_to_hosp_n] inf_to_hosp;
int day_of_week_n;
array[day_of_week_n] int day_of_week;
real state_pop;
array[hosp_n] int hosp;
int shed_n;
vector[shed_n] shed;
real mwpd;
int ww_day_n;
int ww_subpop_n;
array[ww_subpop_n] int ww_subpop;
array[ww_day_n] int ww_day;
int n_sites;
int ww_site_n;
array[ww_site_n] int ww_site;
int ww_conc_n;
vector[ww_conc_n] ww_conc;
int ww_lod_n;
vector[ww_lod_n] ww_lod;
}
transformed data {
int nt = week_of_day_n;
int ot = hosp_n;
int uot = (nt - ot);
int n_weeks = max(week_of_day);
int S = pop_frac_n;
}
parameters {
real log_r0;
real<lower=0.0> eta_sd;
vector[(n_weeks - 1)] w_ref;
real<lower=0, upper=1> autoreg;
real<lower=0.0> sigma_subpop;
real<lower=0.0> feedback;
vector[S] I_mat_growth_s;
matrix[n_weeks, S] I_mat_z_s;
vector[S] I_mat_log_i0_s;
real<lower=0, upper=1> hosp_ihr;
real<lower=0.0> hosp_phi;
simplex[7] hosp_wday_effect;
real ww_conc_log10_g;
vector<lower=0.0>[n_sites] ww_conc_site_sigma;
vector[n_sites] ww_conc_site_mod;
}
transformed parameters {
vector[((n_weeks - 1) + 1)] log_rt_ref = cdc_rw(log_r0, eta_sd, w_ref);
matrix[nt, S] I_mat;
for(plate_i__pl_1 in 1:S) {
I_mat[:, plate_i__pl_1] = cdc_subpop_cell(
I_mat_log_i0_s[plate_i__pl_1],
I_mat_growth_s[plate_i__pl_1],
log_rt_ref,
I_mat_z_s[:, plate_i__pl_1],
autoreg,
sigma_subpop,
week_of_day,
gen_int,
fb_pmf,
feedback,
uot
);
}
vector[nt] state_inf = (I_mat * pop_frac);
vector[ot] exp_hosp_pc = cdc_conv(state_inf, inf_to_hosp, ot, uot);
vector[day_of_week_n] hosp_wday = (7.0 * hosp_wday_effect[day_of_week]);
vector[day_of_week_n] hosp_mu = (((state_pop * hosp_ihr) * exp_hosp_pc) .* hosp_wday);
matrix[nt, S] conc_matrix = cdc_shed_conc(I_mat, shed, mwpd);
vector[ww_day_n] model_log_conc = cdc_gather(conc_matrix, ww_subpop, ww_day);
vector[ww_day_n] ww_conc_model_conc = ((log(10.0) * ww_conc_log10_g) + model_log_conc + ww_conc_site_mod[ww_site]);
}
model {
log_r0 ~ normal(0.0, 0.2);
eta_sd ~ normal(0.0, 0.1);
w_ref ~ std_normal();
autoreg ~ beta(2.0, 2.0);
sigma_subpop ~ normal(0.0, 0.1);
feedback ~ normal(0.0, 0.01);
for(plate_i__pl_1 in 1:S) {
I_mat_log_i0_s[plate_i__pl_1] ~ normal(-13.0, 1.0);
I_mat_growth_s[plate_i__pl_1] ~ normal(0.0, 0.05);
I_mat_z_s[:, plate_i__pl_1] ~ std_normal();
}
hosp_ihr ~ beta(2.0, 20.0);
hosp_phi ~ gamma(2.0, 0.1);
hosp_wday_effect ~ dirichlet(rep_vector(1.0, 7));
hosp ~ neg_binomial_2(hosp_mu, hosp_phi);
ww_conc_log10_g ~ normal(9.0, 1.0);
ww_conc_site_sigma ~ normal(0.0, 1.0);
ww_conc_site_mod ~ normal(0.0, 0.5);
ww_conc ~ lower_clamping_normal(ww_lod, ww_conc_model_conc, ww_conc_site_sigma[ww_site]);
}
generated quantities {
vector[hosp_n] hosp_likelihood = neg_binomial_2_lpmfs(hosp, hosp_mu, hosp_phi);
array[hosp_n] int hosp_gen = neg_binomial_2_int_rng(hosp_n, hosp_mu, hosp_phi);
vector[ww_conc_n] ww_conc_likelihood = lower_clamping_normal_lpdfs(ww_conc, ww_lod, ww_conc_model_conc, ww_conc_site_sigma[ww_site]);
vector[ww_conc_n] ww_conc_gen = lower_clamping_vector_normal_rng(ww_conc_n, ww_lod, ww_conc_model_conc, ww_conc_site_sigma[ww_site]);
}What this exercises
Observation submodels (
data ~ submodel(...)) — the two data streams as self-contained, removable nodes, each with its own parameters and likelihood.Composition —
Base.mergeswaps the Rt process and the shedding kernel into one shared core, dissolving the source's runtime-flag monolith into a construction-time modeling ladder.plateover sub-epidemics — per-subpopulation~parameters (initial size, growth, weekly AR(1) innovations) introduced inside the cell, with a shared reference-Rt random walk captured from the enclosing scope; thevector[nt]cell outputs collect into amatrix[nt, S], matching the BRM companion port's subpop-axis structure.The sparse CSR primitive
csr_matrix_times_vectorfor the spline Rt process, and a parametric shedding kernel built in a@deffun.@deffunscans for the renewal recurrence, the random-walk and AR(1) / diff-AR(1) Rt processes, and the shedding convolutions.censored(normal, …)for below-LOD wastewater, whose emitted log-CDF routes through StanBlocks'erfc-stable helper — a more accurate lower-tail gradient than the hand-writtentarget += normal_lcdfin the original source.Fancy indexing, matrix-vector aggregation, and a gather at the sampled
(subpopulation, day, site)observation points.