StanBlocks.jl

Appearance

Multilevel regression modeling (Radon data) ​

Indoor radon varies with both building characteristics and local geology. The original case study therefore relates log radon to a house-level predictor such as measurement floor while allowing the baseline level to vary by county. The main modelling choice is how much information the counties should share.

These compact examples use random values only as build fixtures; they are not the original radon observations. Open any comparison to inspect the exact generated Stan program beside the Julia source that produced it. That Stan pane is generated during the documentation build, so readers do not have to execute the example themselves.

The arguments supplied when a model is called (y, x, and, where needed, county) become Stan data. Names first introduced on the left of ~ become unknowns. StanBlocks infers scalar and vector shapes from the distribution and indexing expressions, while support such as lower = 0.0 becomes a Stan declaration constraint.

Complete pooling (radon_cp below) ​

Complete pooling assumes every county has the same intercept alpha. The single slope beta describes the relationship with x, and positive sigma is the residual standard deviation. County labels are deliberately absent: after conditioning on x, every observation is treated as coming from one population.

julia
using StanBlocks
y = x = randn(10)
county = rand(1:10, 10)
radon_cp = @slic begin 
    alpha ~ normal(0, 10)
    beta ~ normal(0, 10)
    sigma ~ normal(0, 10; lower=0.)
    y ~ normal(alpha + beta * x, sigma)
end
radon_cp_posterior = radon_cp(;y,x)
stan
functions {
vector normal_lpdfs(
    vector obs,
    vector loc,
    real scale
) {
    return jbroadcasted_normal_lpdfs(obs, loc, scale);
}
vector jbroadcasted_normal_lpdfs(
    vector x1,
    vector x2,
    real x3
) {
    int n = dims(x1)[1];
    vector[n] rv;
    for(i in 1:n) {
        rv[i] = normal_lpdfs(broadcasted_getindex(x1, i), broadcasted_getindex(x2, i), x3);
    }
    return rv;
}
real normal_lpdfs(
    real args1,
    real args2,
    real args3
) {
    return normal_lpdf(args1 | args2, args3);
}
real broadcasted_getindex(vector x, int i) {
    return x[i];
}
vector normal_vector_rng(
    int anontok__1,
    vector a,
    real b
) {
    int n = anontok__1;
    return to_vector(normal_rng(a, b));
}
}
data {
    int y_n;
    vector[y_n] y;
    int x_n;
    vector[x_n] x;
}
transformed data {
}
parameters {
    real alpha;
    real beta;
    real<lower=0.0> sigma;
}
transformed parameters {
}
model {
    alpha ~ normal(0, 10);
    beta ~ normal(0, 10);
    sigma ~ normal(0, 10);
    y ~ normal((alpha + (beta * x)), sigma);
}
generated quantities {
    vector[y_n] y_likelihood = normal_lpdfs(y, (alpha + (beta * x)), sigma);
    vector[y_n] y_gen = normal_vector_rng(y_n, (alpha + (beta * x)), sigma);
}

No pooling ​

At the other extreme, no pooling gives every county an unrelated intercept. n = n_counties makes alpha a vector, and alpha[county] selects the right entry for each observation. Its length is inferred from the largest county index supplied as data. The slope and residual scale remain shared, so “no pooling” refers specifically to the county baselines.

The model is written out in full to keep the statistical contrast visible: the substantive change from complete pooling is the vector-valued intercept and its indexed use in the likelihood.

julia
radon_np = @slic begin 
    n_counties = max(county)
    alpha ~ normal(0, 10; n=n_counties)
    beta ~ normal(0, 10)
    sigma ~ normal(0, 10; lower=0.)
    y ~ normal(alpha[county] + beta * x, sigma)
end
radon_np_posterior = radon_np(;y,x,county)
stan
functions {
vector normal_lpdfs(
    vector obs,
    vector loc,
    real scale
) {
    return jbroadcasted_normal_lpdfs(obs, loc, scale);
}
vector jbroadcasted_normal_lpdfs(
    vector x1,
    vector x2,
    real x3
) {
    int n = dims(x1)[1];
    vector[n] rv;
    for(i in 1:n) {
        rv[i] = normal_lpdfs(broadcasted_getindex(x1, i), broadcasted_getindex(x2, i), x3);
    }
    return rv;
}
real normal_lpdfs(
    real args1,
    real args2,
    real args3
) {
    return normal_lpdf(args1 | args2, args3);
}
real broadcasted_getindex(vector x, int i) {
    return x[i];
}
vector normal_vector_rng(
    int anontok__1,
    vector a,
    real b
) {
    int n = anontok__1;
    return to_vector(normal_rng(a, b));
}
}
data {
    int county_n;
    array[county_n] int county;
    int y_n;
    vector[y_n] y;
    int x_n;
    vector[x_n] x;
}
transformed data {
    int n_counties = max(county);
}
parameters {
    vector[n_counties] alpha;
    real beta;
    real<lower=0.0> sigma;
}
transformed parameters {
}
model {
    alpha ~ normal(0, 10);
    beta ~ normal(0, 10);
    sigma ~ normal(0, 10);
    y ~ normal((alpha[county] + (beta * x)), sigma);
}
generated quantities {
    vector[y_n] y_likelihood = normal_lpdfs(y, (alpha[county] + (beta * x)), sigma);
    vector[y_n] y_gen = normal_vector_rng(y_n, (alpha[county] + (beta * x)), sigma);
}

Partial pooling ​

Partial pooling replaces the independent, zero-centred intercept priors with a population distribution. mu_alpha learns the overall county level and positive sigma_alpha learns the between-county variation. Each county can still move toward its own data, but a county with little information is shrunk toward mu_alpha rather than estimated in isolation. This is the multilevel compromise between the first two models.

julia
radon_pp = @slic begin 
    n_counties = max(county)
    mu_alpha ~ normal(0, 10)
    sigma_alpha ~ normal(0, 10; lower=0)
    alpha ~ normal(mu_alpha, sigma_alpha; n=n_counties)
    beta ~ normal(0, 10)
    sigma ~ normal(0, 10; lower=0.)
    y ~ normal(alpha[county] + beta * x, sigma)
end
radon_pp_posterior = radon_pp(;y,x,county)
stan
functions {
vector normal_lpdfs(
    vector obs,
    vector loc,
    real scale
) {
    return jbroadcasted_normal_lpdfs(obs, loc, scale);
}
vector jbroadcasted_normal_lpdfs(
    vector x1,
    vector x2,
    real x3
) {
    int n = dims(x1)[1];
    vector[n] rv;
    for(i in 1:n) {
        rv[i] = normal_lpdfs(broadcasted_getindex(x1, i), broadcasted_getindex(x2, i), x3);
    }
    return rv;
}
real normal_lpdfs(
    real args1,
    real args2,
    real args3
) {
    return normal_lpdf(args1 | args2, args3);
}
real broadcasted_getindex(vector x, int i) {
    return x[i];
}
vector normal_vector_rng(
    int anontok__1,
    vector a,
    real b
) {
    int n = anontok__1;
    return to_vector(normal_rng(a, b));
}
}
data {
    int county_n;
    array[county_n] int county;
    int y_n;
    vector[y_n] y;
    int x_n;
    vector[x_n] x;
}
transformed data {
    int n_counties = max(county);
}
parameters {
    real mu_alpha;
    real<lower=0> sigma_alpha;
    vector[n_counties] alpha;
    real beta;
    real<lower=0.0> sigma;
}
transformed parameters {
}
model {
    mu_alpha ~ normal(0, 10);
    sigma_alpha ~ normal(0, 10);
    alpha ~ normal(mu_alpha, sigma_alpha);
    beta ~ normal(0, 10);
    sigma ~ normal(0, 10);
    y ~ normal((alpha[county] + (beta * x)), sigma);
}
generated quantities {
    vector[y_n] y_likelihood = normal_lpdfs(y, (alpha[county] + (beta * x)), sigma);
    vector[y_n] y_gen = normal_vector_rng(y_n, (alpha[county] + (beta * x)), sigma);
}

Cross-validation-aware county labels ​

The final model has exactly the same partial-pooling likelihood, but wraps the county index with StanBlocks.stan.maybecv. In an ordinary fit this behaves like the original data. In StanBlocks' leave-one-group-out workflow, the wrapper marks the county dimension as cross-validation-aware: the held-out contribution is omitted from the fitted likelihood, and quantities depending on that county's intercept can be regenerated from the population model in generated quantities. Population-level parameters are still learned from the retained counties.

The annotation belongs at the data boundary, so the statistical model need not be duplicated. The generated-Stan pane makes the resulting parameter/generated-quantities split explicit.

julia
radon_pp_cv_posterior = radon_pp(;y,x,county=StanBlocks.stan.maybecv(:county, county))
stan
functions {
vector normal_vector_rng(
    int anontok__1,
    real a,
    real b
) {
    int n = anontok__1;
    return to_vector(normal_rng(rep_vector(a, n), b));
}
vector normal_lpdfs(
    vector obs,
    vector loc,
    real scale
) {
    return jbroadcasted_normal_lpdfs(obs, loc, scale);
}
vector jbroadcasted_normal_lpdfs(
    vector x1,
    vector x2,
    real x3
) {
    int n = dims(x1)[1];
    vector[n] rv;
    for(i in 1:n) {
        rv[i] = normal_lpdfs(broadcasted_getindex(x1, i), broadcasted_getindex(x2, i), x3);
    }
    return rv;
}
real normal_lpdfs(
    real args1,
    real args2,
    real args3
) {
    return normal_lpdf(args1 | args2, args3);
}
real broadcasted_getindex(vector x, int i) {
    return x[i];
}
vector normal_vector_rng(
    int anontok__1,
    vector a,
    real b
) {
    int n = anontok__1;
    return to_vector(normal_rng(a, b));
}
}
data {
    int county_n;
    array[county_n] int county;
    int y_n;
    vector[y_n] y;
    int x_n;
    vector[x_n] x;
}
transformed data {
    int n_counties = max(county);
}
parameters {
    real mu_alpha;
    real<lower=0> sigma_alpha;
    real beta;
    real<lower=0.0> sigma;
}
transformed parameters {
}
model {
    mu_alpha ~ normal(0, 10);
    sigma_alpha ~ normal(0, 10);
    beta ~ normal(0, 10);
    sigma ~ normal(0, 10);
}
generated quantities {
    vector[n_counties] alpha = normal_vector_rng(n_counties, mu_alpha, sigma_alpha);
    vector[y_n] y_likelihood = normal_lpdfs(y, (alpha[county] + (beta * x)), sigma);
    vector[y_n] y_gen = normal_vector_rng(y_n, (alpha[county] + (beta * x)), sigma);
}
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.