BRM feature atlas
Every executable example in the user guide has the same four semantic views: the BRM declaration, the StanBlocks model BRM emitted, the complete Stan source emitted from that model, and the Turing model selected directly from the BRM. Choose a tab for one view, or Compare and select any two or more views for a side-by-side reading. If a backend does not support an example, its pane remains present and shows the exact construction error produced by the current build.
Nothing below is copied output. The docs build evaluates the displayed source and derives all three emission panes from it.
Gaussian population model and effect priors
brm-comparison
Gaussian population modeljulia
gaussian = (@brm begin
sigma ~ Exponential(2)
mu ~ 1 + x
effect(:, :) ~ Normal(0, 2)
effect(mu, x) ~ Normal(0, 0.25)
y ~ Normal(mu, sigma)
end)((; x=[-1.0, 0.5, 2.0], y=[0.2, 1.1, -0.4]))julia
BRMI:
sigma ~ Exponential(2)
x: data (eltype=Float64, n=3)
mu ~ 1 + x
effect(:, :) ~ Normal(0, 2)
effect(mu, x) ~ Normal(0, 0.25)
y ~ Normal(mu, sigma)julia
SBBRMI with data keys = [:x, :y]
emitted @slic body:
begin
sigma ~ exponential(1.0 ./ 2)
X_mu = hcat(rep_vector(1.0, num_elements(x)), x)
pop_mu ~ _popefs_normal_coefs(; X = X_mu, beta_loc = [0, 0], beta_scale = [2, 0.25])
y ~ normal_id_glm(X_mu, 0.0, pop_mu, sigma)
mu = X_mu * pop_mu
endstan
functions {
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
vector normal_id_glm_lpdfs(
vector y,
matrix X,
real alpha,
vector beta,
real sigma
) {
int n = dims(y)[1];
if (dims(X)[1] != n) reject("normal_id_glm_lpdfs: dim mismatch — `X` dim 1 (= ", dims(X)[1], ") does not match `n` (= ", n, "), inferred from `y` dim 1. `n` sizes: `y` dim 1 (= ", dims(y)[1], "), `X` dim 1 (= ", dims(X)[1], ").");
vector[n] rv;
for(i in 1:n) {
rv[i] = normal_lpdf(y[i] | (alpha + (X[i, :] * beta)), sigma);
}
return rv;
}
vector normal_id_glm_vector_rng(
int anontok__1,
matrix X,
real alpha,
vector beta,
real sigma
) {
int m = anontok__1;
if (dims(X)[1] != m) reject("normal_id_glm_rng: dim mismatch — `X` dim 1 (= ", dims(X)[1], ") does not match `m` (= ", m, "), inferred from `anontok__1` dim 1. `m` sizes: `anontok__1` dim 1 (= ", anontok__1, "), `X` dim 1 (= ", dims(X)[1], ").");
if((m == 0)) {
vector[m] rv;
return rv;
} else {
return normal_id_glm_rng(X, alpha, beta, sigma);
}
}
vector normal_id_glm_rng(
matrix X,
real alpha,
vector beta,
real sigma
) {
int m = dims(X)[1];
if((m == 0)) {
vector[m] rv;
return rv;
} else {
return to_vector(normal_rng((rep_vector(alpha, m) + (X * beta)), sigma));
}
}
}
data {
int x_n;
vector[x_n] x;
int y_n;
vector[y_n] y;
}
transformed data {
matrix[x_n, 2] X_mu = hcat(rep_vector(1.0, num_elements(x)), x);
int pop_mu_n_covariates = 2;
}
parameters {
real<lower=0.0> sigma;
vector[pop_mu_n_covariates] pop_mu_beta_pop;
}
transformed parameters {
vector[pop_mu_n_covariates] pop_mu = pop_mu_beta_pop;
}
model {
sigma ~ exponential((1.0 ./ 2));
pop_mu_beta_pop ~ normal([0, 0]', [2, 0.25]');
y ~ normal_id_glm(X_mu, 0.0, pop_mu, sigma);
}
generated quantities {
vector[x_n] y_likelihood = normal_id_glm_lpdfs(y, X_mu, 0.0, pop_mu, sigma);
vector[x_n] y_gen = normal_id_glm_vector_rng(y_n, X_mu, 0.0, pop_mu, sigma);
vector[x_n] mu = (X_mu * pop_mu);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu)
sigma ~ Distributions.Exponential(2)
beta_pop ~ Distributions.product_distribution([Distributions.Normal(0, 2), Distributions.Normal(0, 0.25)])
eta_mu = X_mu * beta_pop
mu = eta_mu
begin
for i = Base.eachindex(y)
y[i] ~ Distributions.Normal(mu[i], sigma)
end
end
(; mu = mu, sigma = sigma, response = y)
end)Partly missing Gaussian response
brm-comparison
Partly missing Gaussian responsejulia
missing_gaussian = (@brm begin
sigma ~ Exponential(2)
mu ~ 1 + x
mi(y) ~ Normal(mu, sigma)
end)((;
x=[-1.0, 0.5, 2.0, 0.25],
y=Union{Missing,Float64}[0.2, missing, -0.4, missing],
))julia
BRMI:
sigma ~ Exponential(2)
x: data (eltype=Float64, n=4)
mu ~ 1 + x
mi(y) ~ Normal(mu, sigma)julia
SBBRMI with data keys = [:Jmis_y, :Jobs_y, :x, :y_obs]
configured submodels:
_sb_mi_response_configured_1 = Base.merge(BayesianRegressionModels._sb_mi_response, quote
y_mis::vector[n_mis] ~ normal(maybe_index(mi_arg_1, Jmis), maybe_index(mi_arg_2, Jmis))
y_obs ~ normal(maybe_index(mi_arg_1, Jobs), maybe_index(mi_arg_2, Jobs))
end)
emitted @slic body:
begin
sigma ~ exponential(1.0 ./ 2)
X_mu = hcat(rep_vector(1.0, num_elements(x)), x)
pop_mu ~ popefs(; X = X_mu)
mu = pop_mu
y ~ _sb_mi_response_configured_1(; y_obs = y_obs, Jobs = Jobs_y, Jmis = Jmis_y, mi_arg_1 = mu, mi_arg_2 = sigma)
endstan
functions {
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
vector normal_vector_rng(
int anontok__1,
vector a,
real b
) {
int n = anontok__1;
if((n == 0)) {
vector[n] rv;
return rv;
} else {
return to_vector(normal_rng(a, 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 mi_merge(
vector y_obs,
vector y_mis,
array[] int Jobs,
array[] int Jmis,
int n
) {
int n_obs = dims(y_obs)[1];
int n_mis = dims(y_mis)[1];
if (dims(Jobs)[1] != n_obs) reject("mi_merge: dim mismatch — `Jobs` dim 1 (= ", dims(Jobs)[1], ") does not match `n_obs` (= ", n_obs, "), inferred from `y_obs` dim 1. `n_obs` sizes: `y_obs` dim 1 (= ", dims(y_obs)[1], "), `Jobs` dim 1 (= ", dims(Jobs)[1], ").");
if (dims(Jmis)[1] != n_mis) reject("mi_merge: dim mismatch — `Jmis` dim 1 (= ", dims(Jmis)[1], ") does not match `n_mis` (= ", n_mis, "), inferred from `y_mis` dim 1. `n_mis` sizes: `y_mis` dim 1 (= ", dims(y_mis)[1], "), `Jmis` dim 1 (= ", dims(Jmis)[1], ").");
vector[n] rv = rep_vector(0.0, n);
if((n_obs > 0)) {
for(i in 1:n_obs) {
rv[Jobs[i]] = y_obs[i];
}
}
if((n_mis > 0)) {
for(i in 1:n_mis) {
rv[Jmis[i]] = y_mis[i];
}
}
return rv;
}
}
data {
int x_n;
vector[x_n] x;
int Jmis_y_n;
array[Jmis_y_n] int Jmis_y;
int y_obs_n;
vector[y_obs_n] y_obs;
int Jobs_y_n;
array[Jobs_y_n] int Jobs_y;
}
transformed data {
matrix[x_n, 2] X_mu = hcat(rep_vector(1.0, num_elements(x)), x);
int pop_mu_n_covariates = 2;
int y_n_mis = num_elements(Jmis_y);
}
parameters {
real<lower=0.0> sigma;
vector[pop_mu_n_covariates] pop_mu_beta_pop;
}
transformed parameters {
vector[x_n] pop_mu = (X_mu * pop_mu_beta_pop);
vector[x_n] mu = pop_mu;
}
model {
sigma ~ exponential((1.0 ./ 2));
pop_mu_beta_pop ~ std_normal();
y_obs ~ normal(mu[Jobs_y], sigma);
}
generated quantities {
vector[y_n_mis] y_y_mis = normal_vector_rng(y_n_mis, mu[Jmis_y], sigma);
vector[y_obs_n] y_obs_likelihood = normal_lpdfs(y_obs, mu[Jobs_y], sigma);
vector[y_obs_n] y_obs_gen = normal_vector_rng(y_obs_n, mu[Jobs_y], sigma);
vector[(num_elements(Jobs_y) + y_n_mis)] y = mi_merge(y_obs, y_y_mis, Jobs_y, Jmis_y, (num_elements(Jobs_y) + y_n_mis));
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu)
sigma ~ Distributions.Exponential(2)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_mu = X_mu * beta_pop
mu = eta_mu
begin
for i = Base.eachindex(y)
y[i] ~ Distributions.Normal(mu[i], sigma)
end
end
(; mu = mu, sigma = sigma, response = y)
end)Shared and independent multi-response model
brm-comparison
Shared and independent responsesjulia
multi_response = (@brm begin
sigma ~ Exponential(2)
mu ~ 1 + x
y ~ Normal(mu, sigma)
log(rate) ~ 1 + z
count ~ Poisson(rate)
y_replicate ~ Normal(mu, sigma)
end)((;
x=[-1.0, 0.5, 2.0],
z=[0.25, -0.5, 1.0],
y=[0.2, 1.1, -0.4],
count=[0, 2, 4],
y_replicate=[-0.1, 0.3, 0.7],
))julia
BRMI:
sigma ~ Exponential(2)
x: data (eltype=Float64, n=3)
mu ~ 1 + x
y ~ Normal(mu, sigma)
z: data (eltype=Float64, n=3)
log(rate) ~ 1 + z
count ~ Poisson(rate)
y_replicate ~ Normal(mu, sigma)julia
SBBRMI with data keys = [:count, :x, :y, :y_replicate, :z]
emitted @slic body:
begin
sigma ~ exponential(1.0 ./ 2)
X_mu = hcat(rep_vector(1.0, num_elements(x)), x)
pop_mu ~ popefs(; X = X_mu)
mu = pop_mu
y ~ normal(mu, sigma)
X_log_rate = hcat(rep_vector(1.0, num_elements(z)), z)
pop_log_rate ~ popefs(; X = X_log_rate)
log_rate = pop_log_rate
rate = exp(log_rate)
count ~ poisson(rate)
y_replicate ~ normal(mu, sigma)
endstan
functions {
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
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;
if((n == 0)) {
vector[n] rv;
return rv;
} else {
return to_vector(normal_rng(a, b));
}
}
vector poisson_lpmfs(
array[] int obs,
vector lambda
) {
return jbroadcasted_poisson_lpmfs(obs, lambda);
}
vector jbroadcasted_poisson_lpmfs(
array[] int x1,
vector x2
) {
int n = dims(x1)[1];
vector[n] rv;
for(i in 1:n) {
rv[i] = poisson_lpmfs(broadcasted_getindex(x1, i), broadcasted_getindex(x2, i));
}
return rv;
}
real poisson_lpmfs(int args1, real args2) {
return poisson_lpmf(args1 | args2);
}
int broadcasted_getindex(array[] int x, int i) {
return x[i];
}
array[] int poisson_int_rng(
int anontok__1,
vector p
) {
int n = anontok__1;
return poisson_rng(p);
}
}
data {
int x_n;
vector[x_n] x;
int y_n;
vector[y_n] y;
int z_n;
vector[z_n] z;
int count_n;
array[count_n] int count;
int y_replicate_n;
vector[y_replicate_n] y_replicate;
}
transformed data {
matrix[x_n, 2] X_mu = hcat(rep_vector(1.0, num_elements(x)), x);
int pop_mu_n_covariates = 2;
matrix[z_n, 2] X_log_rate = hcat(rep_vector(1.0, num_elements(z)), z);
int pop_log_rate_n_covariates = 2;
}
parameters {
real<lower=0.0> sigma;
vector[pop_mu_n_covariates] pop_mu_beta_pop;
vector[pop_log_rate_n_covariates] pop_log_rate_beta_pop;
}
transformed parameters {
vector[x_n] pop_mu = (X_mu * pop_mu_beta_pop);
vector[x_n] mu = pop_mu;
vector[z_n] pop_log_rate = (X_log_rate * pop_log_rate_beta_pop);
vector[z_n] log_rate = pop_log_rate;
vector[z_n] rate = exp(log_rate);
}
model {
sigma ~ exponential((1.0 ./ 2));
pop_mu_beta_pop ~ std_normal();
y ~ normal(mu, sigma);
pop_log_rate_beta_pop ~ std_normal();
count ~ poisson(rate);
y_replicate ~ normal(mu, sigma);
}
generated quantities {
vector[y_n] y_likelihood = normal_lpdfs(y, mu, sigma);
vector[y_n] y_gen = normal_vector_rng(y_n, mu, sigma);
vector[count_n] count_likelihood = poisson_lpmfs(count, rate);
array[count_n] int count_gen = poisson_int_rng(count_n, rate);
vector[y_replicate_n] y_replicate_likelihood = normal_lpdfs(y_replicate, mu, sigma);
vector[y_replicate_n] y_replicate_gen = normal_vector_rng(y_replicate_n, mu, sigma);
}julia
#= line 0 =# Turing.@model(function brm_multi_model(y_1, y_2, y_3, X_mu, X_rate)
sigma ~ Distributions.Exponential(2)
beta_pop_mu ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_mu = X_mu * beta_pop_mu
mu = eta_mu
beta_pop_rate ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_rate = X_rate * beta_pop_rate
rate = Base.exp.(eta_rate)
begin
for i = Base.eachindex(y_1)
y_1[i] ~ Distributions.Normal(mu[i], sigma)
end
end
begin
for i = Base.eachindex(y_2)
y_2[i] ~ Distributions.Poisson(rate[i])
end
end
begin
for i = Base.eachindex(y_3)
y_3[i] ~ Distributions.Normal(mu[i], sigma)
end
end
(; responses = ((; mu = mu, rate = rate, sigma = sigma, response = y_1), (; mu = mu, rate = rate, sigma = sigma, response = y_2), (; mu = mu, rate = rate, sigma = sigma, response = y_3)))
end)Correlated Gaussian multi-response model
The vector response is one row-wise multivariate likelihood. Outcome order is the order written on both sides, and L_res is the estimated residual covariance Cholesky factor (marginal scales plus an LKJ correlation factor).
brm-comparison
Correlated Gaussian responsesjulia
correlated_response = (@brm begin
L_res ~ LKJCovarianceFactor(
2; scale_prior=Exponential(1), shape=2,
)
shared_log_rate ~ Normal(0, 1)
concentration_mu = exp(-exp(shared_log_rate) * time)
effect_mu = 1.0 - concentration_mu
[concentration, effect] ~ MvNormalCholesky(
[concentration_mu, effect_mu], L_res)
end)((;
time=[0.0, 1.0, 2.0, 4.0],
concentration=[1.0, 0.72, 0.51, 0.27],
effect=[0.03, 0.18, 0.43, 0.79],
))julia
BRMI:
L_res ~ LKJCovarianceFactor(2; scale_prior=Exponential(1), shape=2)
shared_log_rate ~ Normal(0, 1)
time: data (eltype=Float64, n=4)
:concentration_mu = exp((-(exp(shared_log_rate)) * time))
:effect_mu = -(1.0, concentration_mu)
[concentration, effect] ~ MvNormalCholesky(NamedColumn{Symbol}[concentration_mu, effect_mu], L_res)julia
SBBRMI with data keys = [:L_res_n, :brm_joint_concentration__effect_n, :brm_joint_concentration__effect_observed, :time]
emitted @slic body:
begin
L_res_scales ~ exponential(1.0; n = L_res_n)
L_res_L_corr::cholesky_factor_corr[L_res_n] ~ lkj_corr_cholesky(2.0)
L_res = diag_pre_multiply(L_res_scales, L_res_L_corr)
shared_log_rate ~ normal(0, 1)
concentration_mu = (exp)((-)((exp)(shared_log_rate)) .* time)
effect_mu = (-)(1.0, concentration_mu)
brm_joint_concentration__effect_means ~ plate(brm_joint_concentration__effect_observed, concentration_mu, effect_mu; outer = (brm_joint_concentration__effect_n,)) do brm_joint_concentration__effect_observed_cell, brm_joint_concentration__effect_mean_cell_1, brm_joint_concentration__effect_mean_cell_2
brm_joint_concentration__effect_mean_vector = 0.0 .* brm_joint_concentration__effect_observed_cell + [brm_joint_concentration__effect_mean_cell_1, brm_joint_concentration__effect_mean_cell_2]
brm_joint_concentration__effect_mean_vector
end
brm_joint_concentration__effect_observed ~ multi_normal_cholesky(brm_joint_concentration__effect_means, L_res)
endstan
functions {
int ragged_end(array[] int ends, int i) {
return ends[i];
}
int ragged_start(
array[] int ends,
int i
) {
if((i == 1)) {
return 1;
} else {
return (1 + ends[(i - 1)]);
}
}
int num_elements_RaggedVector(tuple(vector, array[] int) rv) {
return size(rv.2);
}
real multi_normal_cholesky_lpdfs(
vector args1,
vector args2,
matrix args3
) {
return multi_normal_cholesky_lpdf(args1 | args2, args3);
}
vector multi_normal_cholesky_vector_rng(
int anontok__1,
vector loc,
matrix scale
) {
int n = anontok__1;
if (dims(loc)[1] != n) reject("multi_normal_cholesky_rng: dim mismatch — `loc` dim 1 (= ", dims(loc)[1], ") does not match `n` (= ", n, "), inferred from `anontok__1` dim 1. `n` sizes: `anontok__1` dim 1 (= ", anontok__1, "), `loc` dim 1 (= ", dims(loc)[1], ").");
return multi_normal_cholesky_rng(loc, scale);
}
int ragged_end_RaggedVector(tuple(vector, array[] int) x, int i) {
return x.2[i];
}
int ragged_start_RaggedVector(
tuple(vector, array[] int) x,
int i
) {
if((i == 1)) {
return 1;
} else {
return (1 + x.2[(i - 1)]);
}
}
vector getindex_RaggedVector(
tuple(vector, array[] int) rv,
int i
) {
return rv.1[ragged_start_RaggedVector(rv, i):ragged_end_RaggedVector(rv, i)];
}
}
data {
int L_res_n;
int time_n;
vector[time_n] time;
int brm_joint_concentration__effect_n;
int brm_joint_concentration__effect_observed_ends_n;
int brm_joint_concentration__effect_observed_mem_n;
tuple(
vector[brm_joint_concentration__effect_observed_mem_n],
array[brm_joint_concentration__effect_observed_ends_n] int
) brm_joint_concentration__effect_observed;
}
transformed data {
array[brm_joint_concentration__effect_n] int brm_joint_concentration__effect_means_brm_joint_concentration__effect_mean_vector__pl_len_1;
array[brm_joint_concentration__effect_n] int brm_joint_concentration__effect_means__pl_len_1;
for(plate_i__pl_1 in 1:brm_joint_concentration__effect_n) {
brm_joint_concentration__effect_means_brm_joint_concentration__effect_mean_vector__pl_len_1[
plate_i__pl_1
] = (
1 +
(
ragged_end(brm_joint_concentration__effect_observed.2, plate_i__pl_1) -
ragged_start(brm_joint_concentration__effect_observed.2, plate_i__pl_1)
)
);
brm_joint_concentration__effect_means__pl_len_1[plate_i__pl_1] = (
1 +
(
ragged_end(brm_joint_concentration__effect_observed.2, plate_i__pl_1) -
ragged_start(brm_joint_concentration__effect_observed.2, plate_i__pl_1)
)
);
}
array[brm_joint_concentration__effect_n] int brm_joint_concentration__effect_means_brm_joint_concentration__effect_mean_vector__pl_end_1 = cumulative_sum(
brm_joint_concentration__effect_means_brm_joint_concentration__effect_mean_vector__pl_len_1
);
array[brm_joint_concentration__effect_n] int brm_joint_concentration__effect_means__pl_end_1 = cumulative_sum(brm_joint_concentration__effect_means__pl_len_1);
}
parameters {
vector<lower=0.0>[L_res_n] L_res_scales;
cholesky_factor_corr[L_res_n] L_res_L_corr;
real shared_log_rate;
}
transformed parameters {
matrix[L_res_n, L_res_n] L_res = diag_pre_multiply(L_res_scales, L_res_L_corr);
vector[time_n] concentration_mu = exp(((-exp(shared_log_rate)) .* time));
vector[time_n] effect_mu = (1.0 - concentration_mu);
vector[sum(brm_joint_concentration__effect_means_brm_joint_concentration__effect_mean_vector__pl_len_1)] brm_joint_concentration__effect_means_brm_joint_concentration__effect_mean_vector__pl_mem_1;
vector[sum(brm_joint_concentration__effect_means__pl_len_1)] brm_joint_concentration__effect_means__pl_mem_1;
for(plate_i__pl_1 in 1:brm_joint_concentration__effect_n) {
brm_joint_concentration__effect_means_brm_joint_concentration__effect_mean_vector__pl_mem_1[
ragged_start(
brm_joint_concentration__effect_means_brm_joint_concentration__effect_mean_vector__pl_end_1,
plate_i__pl_1
):ragged_end(
brm_joint_concentration__effect_means_brm_joint_concentration__effect_mean_vector__pl_end_1,
plate_i__pl_1
)
] = (
(
0.0 .*
brm_joint_concentration__effect_observed.1[
ragged_start(brm_joint_concentration__effect_observed.2, plate_i__pl_1):ragged_end(brm_joint_concentration__effect_observed.2, plate_i__pl_1)
]
) +
[concentration_mu[plate_i__pl_1], effect_mu[plate_i__pl_1]]'
);
brm_joint_concentration__effect_means__pl_mem_1[
ragged_start(brm_joint_concentration__effect_means__pl_end_1, plate_i__pl_1):ragged_end(brm_joint_concentration__effect_means__pl_end_1, plate_i__pl_1)
] = brm_joint_concentration__effect_means_brm_joint_concentration__effect_mean_vector__pl_mem_1[
ragged_start(
brm_joint_concentration__effect_means_brm_joint_concentration__effect_mean_vector__pl_end_1,
plate_i__pl_1
):ragged_end(
brm_joint_concentration__effect_means_brm_joint_concentration__effect_mean_vector__pl_end_1,
plate_i__pl_1
)
];
}
}
model {
L_res_scales ~ exponential(1.0);
L_res_L_corr ~ lkj_corr_cholesky(2.0);
shared_log_rate ~ normal(0, 1);
for(g__ro_2 in 1:num_elements_RaggedVector(brm_joint_concentration__effect_observed)) {
getindex_RaggedVector(brm_joint_concentration__effect_observed, g__ro_2) ~ multi_normal_cholesky(
brm_joint_concentration__effect_means__pl_mem_1[
ragged_start(brm_joint_concentration__effect_means__pl_end_1, g__ro_2):ragged_end(brm_joint_concentration__effect_means__pl_end_1, g__ro_2)
],
L_res
);
}
}
generated quantities {
vector[num_elements(brm_joint_concentration__effect_observed.1)] brm_joint_concentration__effect_observed_gen;
vector[num_elements_RaggedVector(brm_joint_concentration__effect_observed)] brm_joint_concentration__effect_observed_likelihood;
for(g__rq_2 in 1:num_elements_RaggedVector(brm_joint_concentration__effect_observed)) {
brm_joint_concentration__effect_observed_gen[
ragged_start(brm_joint_concentration__effect_observed.2, g__rq_2):ragged_end(brm_joint_concentration__effect_observed.2, g__rq_2)
] = multi_normal_cholesky_vector_rng(
(1 + (ragged_end_RaggedVector(brm_joint_concentration__effect_observed, g__rq_2) - ragged_start_RaggedVector(brm_joint_concentration__effect_observed, g__rq_2))),
brm_joint_concentration__effect_means__pl_mem_1[
ragged_start(brm_joint_concentration__effect_means__pl_end_1, g__rq_2):ragged_end(brm_joint_concentration__effect_means__pl_end_1, g__rq_2)
],
L_res
);
brm_joint_concentration__effect_observed_likelihood[g__rq_2] = multi_normal_cholesky_lpdf(getindex_RaggedVector(brm_joint_concentration__effect_observed, g__rq_2) |
brm_joint_concentration__effect_means__pl_mem_1[
ragged_start(brm_joint_concentration__effect_means__pl_end_1, g__rq_2):ragged_end(brm_joint_concentration__effect_means__pl_end_1, g__rq_2)
],
L_res
);
}
}julia
#= line 0 =# Turing.@model(function brm_model(y, time)
L_res ~ DynamicPPL.to_submodel(BayesianRegressionModels._brm_turing_covariance_prior(2; scale_prior = Distributions.Exponential(1), shape = 2))
shared_log_rate ~ Distributions.Normal(0, 1)
concentration_mu = [Base.exp(-(Base.exp(shared_log_rate)) * time[i]) for i = Base.eachindex(y)]
effect_mu = [1.0 - concentration_mu[i] for i = Base.eachindex(y)]
begin
for i = Base.eachindex(y)
y[i] ~ BayesianRegressionModelsTuringExt._brm_mvn_cholesky(Base.vect(concentration_mu[i], effect_mu[i]), L_res)
end
end
(; shared_log_rate = shared_log_rate, L_res = L_res, concentration_mu = concentration_mu, effect_mu = effect_mu, response = y)
end)Canonical Binomial-logit model
brm-comparison
Binomial logitjulia
using LogExpFunctions: logit
binomial = (@brm begin
logit(p) ~ 1 + x
successes ~ Binomial(trials, p)
end)((;
x=[-1.0, 0.5, 2.0, 0.25],
trials=[2, 4, 6, 3],
successes=[0, 2, 5, 1],
))julia
BRMI:
x: data (eltype=Float64, n=4)
logit(p) ~ 1 + x
trials: data (eltype=Int64, n=4)
successes ~ Binomial(trials, p)julia
SBBRMI with data keys = [:successes, :trials, :x]
emitted @slic body:
begin
X_logit_p = hcat(rep_vector(1.0, num_elements(x)), x)
pop_logit_p ~ popefs(; X = X_logit_p)
logit_p = pop_logit_p
p = inv_logit(logit_p)
successes ~ binomial(trials, p)
endstan
functions {
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
vector binomial_lpmfs(
array[] int y,
array[] int args1,
vector args2
) {
return jbroadcasted_binomial_lpmfs(y, args1, args2);
}
vector jbroadcasted_binomial_lpmfs(
array[] int x1,
array[] int x2,
vector x3
) {
int n = dims(x1)[1];
vector[n] rv;
for(i in 1:n) {
rv[i] = binomial_lpmfs(
broadcasted_getindex(x1, i),
broadcasted_getindex(x2, i),
broadcasted_getindex(x3, i)
);
}
return rv;
}
real binomial_lpmfs(
int args1,
int args2,
real args3
) {
return binomial_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 binomial_int_rng(
int anontok__1,
array[] int N,
vector p
) {
int n = anontok__1;
if (dims(N)[1] != n) reject("binomial_rng: dim mismatch — `N` dim 1 (= ", dims(N)[1], ") does not match `n` (= ", n, "), inferred from `anontok__1` dim 1. `n` sizes: `anontok__1` dim 1 (= ", anontok__1, "), `N` dim 1 (= ", dims(N)[1], ").");
if((n == 0)) {
array[n] int rv;
return rv;
} else {
return binomial_rng(N, p);
}
}
}
data {
int x_n;
vector[x_n] x;
int successes_n;
array[successes_n] int successes;
int trials_n;
array[trials_n] int trials;
}
transformed data {
matrix[x_n, 2] X_logit_p = hcat(rep_vector(1.0, num_elements(x)), x);
int pop_logit_p_n_covariates = 2;
}
parameters {
vector[pop_logit_p_n_covariates] pop_logit_p_beta_pop;
}
transformed parameters {
vector[x_n] pop_logit_p = (X_logit_p * pop_logit_p_beta_pop);
vector[x_n] logit_p = pop_logit_p;
vector[x_n] p = inv_logit(logit_p);
}
model {
pop_logit_p_beta_pop ~ std_normal();
successes ~ binomial(trials, p);
}
generated quantities {
vector[successes_n] successes_likelihood = binomial_lpmfs(successes, trials, p);
array[trials_n] int successes_gen = binomial_int_rng(successes_n, trials, p);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_p, trials)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_p = X_p * beta_pop
p = LogExpFunctions.logistic.(eta_p)
begin
for i = Base.eachindex(y)
y[i] ~ Distributions.Binomial(trials[i], p[i])
end
end
(; p = p, response = y)
end)Poisson model with transformed data and an offset
brm-comparison
Poisson exposure modeljulia
poisson_offset = (@brm begin
log(rate) ~ 1 + center(x) + offset(log(exposure))
counts ~ Poisson(rate)
end)((;
x=[1.0, 2.0, 4.0, 8.0],
exposure=[2.0, 4.0, 8.0, 16.0],
counts=[0, 2, 5, 7],
))julia
BRMI:
x: data (eltype=Float64, n=4)
exposure: data (eltype=Float64, n=4)
log(rate) ~ 1 + center(x) + offset(log(exposure))
counts ~ Poisson(rate)julia
SBBRMI with data keys = [:center_x, :counts, :exposure, :x]
emitted @slic body:
begin
X_log_rate = hcat(rep_vector(1.0, num_elements(x)), center_x)
pop_log_rate ~ popefs(; X = X_log_rate)
log_rate = pop_log_rate + (log)(exposure)
rate = exp(log_rate)
counts ~ poisson(rate)
endstan
functions {
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
vector poisson_lpmfs(
array[] int obs,
vector lambda
) {
return jbroadcasted_poisson_lpmfs(obs, lambda);
}
vector jbroadcasted_poisson_lpmfs(
array[] int x1,
vector x2
) {
int n = dims(x1)[1];
vector[n] rv;
for(i in 1:n) {
rv[i] = poisson_lpmfs(broadcasted_getindex(x1, i), broadcasted_getindex(x2, i));
}
return rv;
}
real poisson_lpmfs(int args1, real args2) {
return poisson_lpmf(args1 | args2);
}
int broadcasted_getindex(array[] int x, int i) {
return x[i];
}
real broadcasted_getindex(vector x, int i) {
return x[i];
}
array[] int poisson_int_rng(
int anontok__1,
vector p
) {
int n = anontok__1;
return poisson_rng(p);
}
}
data {
int center_x_n;
int x_n;
vector[x_n] x;
vector[center_x_n] center_x;
int exposure_n;
vector[exposure_n] exposure;
int counts_n;
array[counts_n] int counts;
}
transformed data {
matrix[center_x_n, 2] X_log_rate = hcat(rep_vector(1.0, num_elements(x)), center_x);
int pop_log_rate_n_covariates = 2;
}
parameters {
vector[pop_log_rate_n_covariates] pop_log_rate_beta_pop;
}
transformed parameters {
vector[center_x_n] pop_log_rate = (X_log_rate * pop_log_rate_beta_pop);
vector[center_x_n] log_rate = (pop_log_rate + log(exposure));
vector[center_x_n] rate = exp(log_rate);
}
model {
pop_log_rate_beta_pop ~ std_normal();
counts ~ poisson(rate);
}
generated quantities {
vector[counts_n] counts_likelihood = poisson_lpmfs(counts, rate);
array[counts_n] int counts_gen = poisson_int_rng(counts_n, rate);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_rate, offset_rate)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_rate = X_rate * beta_pop + offset_rate
rate = Base.exp.(eta_rate)
begin
for i = Base.eachindex(y)
y[i] ~ Distributions.Poisson(rate[i])
end
end
(; rate = rate, response = y)
end)Distributional count model
brm-comparison
Negative-binomial mean and precisionjulia
negative_binomial = (@brm begin
log(mu) ~ 1 + x
log(phi) ~ 1
y ~ NegativeBinomial2(mu, phi)
end)((; x=[-1.0, 0.5, 2.0, 0.25], y=[0, 2, 5, 1]))julia
BRMI:
x: data (eltype=Float64, n=4)
log(mu) ~ 1 + x
log(phi) ~ 1
y ~ NegativeBinomial2(mu, phi)julia
SBBRMI with data keys = [:x, :y]
emitted @slic body:
begin
X_log_mu = hcat(rep_vector(1.0, num_elements(x)), x)
pop_log_mu ~ popefs(; X = X_log_mu)
log_mu = pop_log_mu
mu = exp(log_mu)
X_log_phi = hcat(rep_vector(1.0, num_elements(y)))
pop_log_phi ~ popefs(; X = X_log_phi)
log_phi = pop_log_phi
phi = exp(log_phi)
y ~ neg_binomial_2(mu, phi)
endstan
functions {
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
matrix hcat(vector x) {
int n = dims(x)[1];
return to_matrix(x, n, 1);
}
vector neg_binomial_2_lpmfs(
array[] int obs,
vector mu,
vector phi
) {
return jbroadcasted_neg_binomial_2_lpmfs(obs, mu, phi);
}
vector jbroadcasted_neg_binomial_2_lpmfs(
array[] int x1,
vector x2,
vector 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),
broadcasted_getindex(x3, i)
);
}
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,
vector b
) {
int n = anontok__1;
if((n == 0)) {
array[n] int rv;
return rv;
} else {
return neg_binomial_2_rng(a, b);
}
}
}
data {
int x_n;
vector[x_n] x;
int y_n;
array[y_n] int y;
}
transformed data {
matrix[x_n, 2] X_log_mu = hcat(rep_vector(1.0, num_elements(x)), x);
int pop_log_mu_n_covariates = 2;
matrix[num_elements(y), 1] X_log_phi = hcat(rep_vector(1.0, num_elements(y)));
int pop_log_phi_n_covariates = 1;
}
parameters {
vector[pop_log_mu_n_covariates] pop_log_mu_beta_pop;
vector[pop_log_phi_n_covariates] pop_log_phi_beta_pop;
}
transformed parameters {
vector[x_n] pop_log_mu = (X_log_mu * pop_log_mu_beta_pop);
vector[x_n] log_mu = pop_log_mu;
vector[x_n] mu = exp(log_mu);
vector[num_elements(y)] pop_log_phi = (X_log_phi * pop_log_phi_beta_pop);
vector[num_elements(y)] log_phi = pop_log_phi;
vector[num_elements(y)] phi = exp(log_phi);
}
model {
pop_log_mu_beta_pop ~ std_normal();
pop_log_phi_beta_pop ~ std_normal();
y ~ neg_binomial_2(mu, phi);
}
generated quantities {
vector[y_n] y_likelihood = neg_binomial_2_lpmfs(y, mu, phi);
array[y_n] int y_gen = neg_binomial_2_int_rng(y_n, mu, phi);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu, X_phi)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_mu = X_mu * beta_pop
mu = Base.exp.(eta_mu)
beta_pop_phi ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 1))
eta_phi = X_phi * beta_pop_phi
phi = Base.exp.(eta_phi)
begin
for i = Base.eachindex(y)
y[i] ~ BayesianRegressionModels.NegativeBinomial2(mu[i], phi[i])
end
end
(; mu = mu, phi = phi, response = y)
end)Distributional count model with independent group slopes
brm-comparison
Independent mean and precision group slopesjulia
grouped_negative_binomial = (@brm begin
log(mu) ~ 1 + x + (1 + x || subject)
log(phi) ~ 1 + z + (0 + z || batch)
y ~ NegativeBinomial2(mu, phi)
end)((;
x=[-1.0, 0.5, 2.0, 0.25],
z=[0.0, 1.0, -0.5, 0.75],
subject=[2, 1, 2, 3],
batch=[2, 1, 1, 2],
y=[0, 2, 5, 1],
))julia
BRMI:
x: data (eltype=Float64, n=4)
subject: data (eltype=Int64, n=4)
log(mu) ~ 1 + x + ((1 + x) || subject)
z: data (eltype=Float64, n=4)
batch: data (eltype=Int64, n=4)
log(phi) ~ 1 + z + ((0 + z) || batch)
y ~ NegativeBinomial2(mu, phi)julia
SBBRMI with data keys = [:batch, :subject, :total_A_mu, :total_A_phi, :total_group_mu, :total_group_phi, :total_location_mu, :total_location_phi, :total_ng_mu, :total_ng_phi, :total_nk_mu, :total_nk_phi, :total_np_mu, :total_np_phi, :total_precision_mu, :total_precision_phi, :x, :y, :z]
configured submodels:
_brm_total_scales_configured_1 = Base.merge(BayesianRegressionModels._brm_total_scales, quote
tau ~ (ValueFamily(brm_vector_prior_40534eb26bf4642a))(0.0, 1.0, 0.0, 1.0; n = 2)
end)
_brm_total_scales_configured_2 = Base.merge(BayesianRegressionModels._brm_total_scales, quote
tau ~ (ValueFamily(brm_vector_prior_ca4b8a1c1bc116d6))(0.0, 1.0; n = 1)
end)
emitted @slic body:
begin
total_scale_mu ~ _brm_total_scales_configured_1(; n = total_nk_mu)
total_mu::matrix[total_ng_mu, total_nk_mu] ~ brm_total(total_scale_mu, total_A_mu, total_location_mu, total_precision_mu)
population_mu = brm_total_recover_rng(total_mu, total_scale_mu, total_A_mu, total_location_mu, total_precision_mu)
deviation_mu = brm_total_deviations(total_mu, total_A_mu * population_mu)
total_Z_mu = hcat(rep_vector(1.0, num_elements(total_group_mu)), x)
log_mu = rows_dot_product(total_mu[total_group_mu, :], total_Z_mu)
mu = exp(log_mu)
total_scale_phi ~ _brm_total_scales_configured_2(; n = total_nk_phi)
total_phi::matrix[total_ng_phi, total_nk_phi] ~ brm_total(total_scale_phi, total_A_phi, total_location_phi, total_precision_phi)
population_phi = brm_total_recover_rng(total_phi, total_scale_phi, total_A_phi, total_location_phi, total_precision_phi)
deviation_phi = brm_total_deviations(total_phi, total_A_phi * population_phi)
total_Z_phi = hcat(z)
X_log_phi = hcat(rep_vector(1.0, num_elements(total_group_phi)))
pop_log_phi ~ popefs(; X = X_log_phi)
log_phi = rows_dot_product(total_phi[total_group_phi, :], total_Z_phi) + pop_log_phi
phi = exp(log_phi)
y ~ neg_binomial_2(mu, phi)
endstan
functions {
// value UDF brm_vector_prior_40534eb26bf4642a_lpdf
real brm_vector_prior_40534eb26bf4642a_lpdf(
vector x,
real arg_1,
real arg_2,
real arg_3,
real arg_4
) {
if((x[1] < 0.0)) {
return negative_infinity();
}
if((x[2] < 0.0)) {
return negative_infinity();
}
return (lognormal_lpdf(x[1] | arg_1, arg_2) + normal_lpdf(x[2] | arg_3, arg_4));
}
real brm_total_lpdf(
matrix total,
vector tau,
matrix A,
vector location,
vector precision
) {
int j = dims(total)[1];
int k = dims(total)[2];
int p = dims(A)[2];
if (dims(tau)[1] != k) reject("brm_total_lpdf: dim mismatch — `tau` dim 1 (= ", dims(tau)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
if (dims(A)[1] != k) reject("brm_total_lpdf: dim mismatch — `A` dim 1 (= ", dims(A)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
if (dims(location)[1] != p) reject("brm_total_lpdf: dim mismatch — `location` dim 1 (= ", dims(location)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], ").");
if (dims(precision)[1] != p) reject("brm_total_lpdf: dim mismatch — `precision` dim 1 (= ", dims(precision)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], ").");
matrix[dims(precision)[1], dims(precision)[1]] Q = brm_total_precision(tau, A, precision, j);
vector[dims(total)[2]] average = brm_total_mean(total);
vector[dims(precision)[1]] beta = brm_total_conditional_mean(total, tau, A, location, precision, Q);
vector[dims(total)[2]] residual = (average - (A * beta));
real quadratic = 0.0;
real lp = ((-0.5 * ((j * k) - p) * 1.8378770664093453) - (j * sum(log(tau))));
for(c in 1:k) {
quadratic += ((j * square(residual[c])) / square(tau[c]));
for(g in 1:j) {
quadratic += (square((total[g, c] - average[c])) / square(tau[c]));
}
}
for(a in 1:p) {
if((precision[a] > 0.0)) {
lp += (0.5 * (log(precision[a]) - 1.8378770664093453));
quadratic += (precision[a] * square((beta[a] - location[a])));
}
}
return (lp - (0.5 * (log_determinant(Q) + quadratic)));
}
matrix brm_total_precision(
vector tau,
matrix A,
vector precision,
int n_groups
) {
int k = dims(tau)[1];
int p = dims(A)[2];
if (dims(A)[1] != k) reject("brm_total_precision: dim mismatch — `A` dim 1 (= ", dims(A)[1], ") does not match `k` (= ", k, "), inferred from `tau` dim 1. `k` sizes: `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
if (dims(precision)[1] != p) reject("brm_total_precision: dim mismatch — `precision` dim 1 (= ", dims(precision)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `precision` dim 1 (= ", dims(precision)[1], ").");
matrix[dims(precision)[1], dims(precision)[1]] out = diag_matrix(precision);
for(a in 1:p) {
for(b in 1:p) {
for(c in 1:k) {
out[a, b] += ((n_groups * A[c, a] * A[c, b]) / square(tau[c]));
}
}
}
return out;
}
vector brm_total_mean(
matrix total
) {
int j = dims(total)[1];
int k = dims(total)[2];
vector[k] out = rep_vector(0.0, k);
for(c in 1:k) {
out[c] = (sum(total[:, c]) / j);
}
return out;
}
vector brm_total_conditional_mean(
matrix total,
vector tau,
matrix A,
vector location,
vector precision,
matrix Q
) {
int j = dims(total)[1];
int k = dims(total)[2];
int p = dims(A)[2];
if (dims(tau)[1] != k) reject("brm_total_conditional_mean: dim mismatch — `tau` dim 1 (= ", dims(tau)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
if (dims(A)[1] != k) reject("brm_total_conditional_mean: dim mismatch — `A` dim 1 (= ", dims(A)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
if (dims(location)[1] != p) reject("brm_total_conditional_mean: dim mismatch — `location` dim 1 (= ", dims(location)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], "), `Q` dim 1 (= ", dims(Q)[1], "), `Q` dim 2 (= ", dims(Q)[2], ").");
if (dims(precision)[1] != p) reject("brm_total_conditional_mean: dim mismatch — `precision` dim 1 (= ", dims(precision)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], "), `Q` dim 1 (= ", dims(Q)[1], "), `Q` dim 2 (= ", dims(Q)[2], ").");
if (dims(Q)[1] != p) reject("brm_total_conditional_mean: dim mismatch — `Q` dim 1 (= ", dims(Q)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], "), `Q` dim 1 (= ", dims(Q)[1], "), `Q` dim 2 (= ", dims(Q)[2], ").");
if (dims(Q)[2] != p) reject("brm_total_conditional_mean: dim mismatch — `Q` dim 2 (= ", dims(Q)[2], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], "), `Q` dim 1 (= ", dims(Q)[1], "), `Q` dim 2 (= ", dims(Q)[2], ").");
vector[dims(total)[2]] average = brm_total_mean(total);
vector[dims(location)[1]] natural = (precision .* location);
for(a in 1:p) {
for(c in 1:k) {
natural[a] += ((j * A[c, a] * average[c]) / square(tau[c]));
}
}
return mdivide_left_spd(Q, natural);
}
vector brm_total_recover_rng(
matrix total,
vector tau,
matrix A,
vector location,
vector precision
) {
int j = dims(total)[1];
int k = dims(total)[2];
int p = dims(A)[2];
if (dims(tau)[1] != k) reject("brm_total_recover_rng: dim mismatch — `tau` dim 1 (= ", dims(tau)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
if (dims(A)[1] != k) reject("brm_total_recover_rng: dim mismatch — `A` dim 1 (= ", dims(A)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
if (dims(location)[1] != p) reject("brm_total_recover_rng: dim mismatch — `location` dim 1 (= ", dims(location)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], ").");
if (dims(precision)[1] != p) reject("brm_total_recover_rng: dim mismatch — `precision` dim 1 (= ", dims(precision)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], ").");
matrix[dims(precision)[1], dims(precision)[1]] Q = brm_total_precision(tau, A, precision, j);
vector[dims(precision)[1]] beta = brm_total_conditional_mean(total, tau, A, location, precision, Q);
return multi_normal_rng(beta, inverse_spd(Q));
}
matrix brm_total_deviations(
matrix total,
vector mu
) {
int j = dims(total)[1];
int k = dims(total)[2];
if (dims(mu)[1] != k) reject("brm_total_deviations: dim mismatch — `mu` dim 1 (= ", dims(mu)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `mu` dim 1 (= ", dims(mu)[1], ").");
matrix[dims(total)[1], dims(total)[2]] out = total;
for(c in 1:k) {
out[:, c] = (total[:, c] - rep_vector(mu[c], j));
}
return out;
}
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
// value UDF brm_vector_prior_ca4b8a1c1bc116d6_lpdf
real brm_vector_prior_ca4b8a1c1bc116d6_lpdf(
vector x,
real arg_1,
real arg_2
) {
if((x[1] < 0.0)) {
return negative_infinity();
}
return normal_lpdf(x[1] | arg_1, arg_2);
}
matrix hcat(vector x) {
int n = dims(x)[1];
return to_matrix(x, n, 1);
}
vector neg_binomial_2_lpmfs(
array[] int obs,
vector mu,
vector phi
) {
return jbroadcasted_neg_binomial_2_lpmfs(obs, mu, phi);
}
vector jbroadcasted_neg_binomial_2_lpmfs(
array[] int x1,
vector x2,
vector 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),
broadcasted_getindex(x3, i)
);
}
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,
vector b
) {
int n = anontok__1;
if((n == 0)) {
array[n] int rv;
return rv;
} else {
return neg_binomial_2_rng(a, b);
}
}
}
data {
int total_ng_mu;
int total_nk_mu;
int total_A_mu_m;
int total_A_mu_n;
matrix[total_A_mu_m, total_A_mu_n] total_A_mu;
int total_location_mu_n;
vector[total_location_mu_n] total_location_mu;
int total_precision_mu_n;
vector[total_precision_mu_n] total_precision_mu;
int x_n;
int total_group_mu_n;
array[total_group_mu_n] int total_group_mu;
vector[x_n] x;
int total_ng_phi;
int total_nk_phi;
int total_A_phi_m;
int total_A_phi_n;
matrix[total_A_phi_m, total_A_phi_n] total_A_phi;
int total_location_phi_n;
vector[total_location_phi_n] total_location_phi;
int total_precision_phi_n;
vector[total_precision_phi_n] total_precision_phi;
int z_n;
vector[z_n] z;
int total_group_phi_n;
array[total_group_phi_n] int total_group_phi;
int y_n;
array[y_n] int y;
}
transformed data {
matrix[x_n, 2] total_Z_mu = hcat(rep_vector(1.0, num_elements(total_group_mu)), x);
matrix[z_n, 1] total_Z_phi = hcat(z);
matrix[num_elements(total_group_phi), 1] X_log_phi = hcat(rep_vector(1.0, num_elements(total_group_phi)));
int pop_log_phi_n_covariates = 1;
}
parameters {
vector<lower=0.0>[2] total_scale_mu_tau;
matrix[total_ng_mu, total_nk_mu] total_mu;
vector<lower=0.0>[1] total_scale_phi_tau;
matrix[total_ng_phi, total_nk_phi] total_phi;
vector[pop_log_phi_n_covariates] pop_log_phi_beta_pop;
}
transformed parameters {
vector<lower=0.0>[2] total_scale_mu = total_scale_mu_tau;
vector[x_n] log_mu = rows_dot_product(total_mu[total_group_mu, :], total_Z_mu);
vector[x_n] mu = exp(log_mu);
vector<lower=0.0>[1] total_scale_phi = total_scale_phi_tau;
vector[num_elements(total_group_phi)] pop_log_phi = (X_log_phi * pop_log_phi_beta_pop);
vector[z_n] log_phi = (rows_dot_product(total_phi[total_group_phi, :], total_Z_phi) + pop_log_phi);
vector[z_n] phi = exp(log_phi);
}
model {
total_scale_mu_tau ~ brm_vector_prior_40534eb26bf4642a(0.0, 1.0, 0.0, 1.0);
total_mu ~ brm_total(total_scale_mu, total_A_mu, total_location_mu, total_precision_mu);
total_scale_phi_tau ~ brm_vector_prior_ca4b8a1c1bc116d6(0.0, 1.0);
total_phi ~ brm_total(total_scale_phi, total_A_phi, total_location_phi, total_precision_phi);
pop_log_phi_beta_pop ~ std_normal();
y ~ neg_binomial_2(mu, phi);
}
generated quantities {
vector[total_precision_mu_n] population_mu = brm_total_recover_rng(total_mu, total_scale_mu, total_A_mu, total_location_mu, total_precision_mu);
matrix[total_ng_mu, total_A_mu_m] deviation_mu = brm_total_deviations(total_mu, (total_A_mu * population_mu));
vector[total_precision_phi_n] population_phi = brm_total_recover_rng(
total_phi,
total_scale_phi,
total_A_phi,
total_location_phi,
total_precision_phi
);
matrix[total_ng_phi, total_A_phi_m] deviation_phi = brm_total_deviations(total_phi, (total_A_phi * population_phi));
vector[y_n] y_likelihood = neg_binomial_2_lpmfs(y, mu, phi);
array[y_n] int y_gen = neg_binomial_2_int_rng(y_n, mu, phi);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu, group_effects_mu_1, X_phi, group_effects_phi_1)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_mu = X_mu * beta_pop
group_effect_1 = Base.zeros(Base.length(y))
group_1_1 ~ DynamicPPL.to_submodel(BRM.turing_group_effect(group_effects_mu_1, (nothing, nothing), nothing))
group_effect_1 = group_effect_1 + group_1_1.effect
eta_mu = eta_mu + group_effect_1
mu = Base.exp.(eta_mu)
beta_pop_phi ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_phi = X_phi * beta_pop_phi
group_effect_2 = Base.zeros(Base.length(y))
group_2_1 ~ DynamicPPL.to_submodel(BRM.turing_group_effect(group_effects_phi_1, (nothing,), nothing))
group_effect_2 = group_effect_2 + group_2_1.effect
eta_phi = eta_phi + group_effect_2
phi = Base.exp.(eta_phi)
begin
for i = Base.eachindex(y)
y[i] ~ BayesianRegressionModels.NegativeBinomial2(mu[i], phi[i])
end
end
(; mu = mu, phi = phi, response = y)
end)Shared distributional group covariance
The repeated joint ID makes the mean and precision slopes slices of one four-dimensional covariance block for each subject. Removing the ID would make the two predictor blocks independent.
brm-comparison
Shared mean and precision group covariancejulia
shared_group_negative_binomial = (@brm begin
log(mu) ~ 1 + x + (1 + x | joint | subject)
log(phi) ~ 1 + z + (1 + z | joint | subject)
y ~ NegativeBinomial2(mu, phi)
end)((;
x=[-1.0, 0.5, 2.0, 0.25, -0.75, 1.25],
z=[0.0, 1.0, -0.5, 0.75, 0.25, -1.0],
subject=[2, 1, 2, 3, 1, 3],
y=[0, 2, 5, 1, 3, 4],
))julia
BRMI:
x: data (eltype=Float64, n=6)
subject: data (eltype=Int64, n=6)
log(mu) ~ 1 + x + ((1 + x) | joint | subject)
z: data (eltype=Float64, n=6)
log(phi) ~ 1 + z + ((1 + z) | joint | subject)
y ~ NegativeBinomial2(mu, phi)julia
SBBRMI with data keys = [:col_idx_log_mu_joint_subject, :col_idx_log_phi_joint_subject, :n_subject, :n_terms_joint_subject, :subject_idx, :x, :y, :z]
emitted @slic body:
begin
b_joint_subject ~ ranef_correlated_draws(; group_idx = subject_idx, n_groups = n_subject, n_terms = n_terms_joint_subject)
X_log_mu = hcat(rep_vector(1.0, num_elements(x)), x)
pop_log_mu ~ popefs(; X = X_log_mu)
Z_log_mu_joint_subject = hcat(rep_vector(1.0, num_elements(x)), x)
r_log_mu_joint_subject = rows_dot_product(Z_log_mu_joint_subject, b_joint_subject[subject_idx, col_idx_log_mu_joint_subject])
log_mu = pop_log_mu + r_log_mu_joint_subject
mu = exp(log_mu)
X_log_phi = hcat(rep_vector(1.0, num_elements(z)), z)
pop_log_phi ~ popefs(; X = X_log_phi)
Z_log_phi_joint_subject = hcat(rep_vector(1.0, num_elements(z)), z)
r_log_phi_joint_subject = rows_dot_product(Z_log_phi_joint_subject, b_joint_subject[subject_idx, col_idx_log_phi_joint_subject])
log_phi = pop_log_phi + r_log_phi_joint_subject
phi = exp(log_phi)
y ~ neg_binomial_2(mu, phi)
endstan
functions {
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
vector neg_binomial_2_lpmfs(
array[] int obs,
vector mu,
vector phi
) {
return jbroadcasted_neg_binomial_2_lpmfs(obs, mu, phi);
}
vector jbroadcasted_neg_binomial_2_lpmfs(
array[] int x1,
vector x2,
vector 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),
broadcasted_getindex(x3, i)
);
}
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,
vector b
) {
int n = anontok__1;
if((n == 0)) {
array[n] int rv;
return rv;
} else {
return neg_binomial_2_rng(a, b);
}
}
}
data {
int n_terms_joint_subject;
int n_subject;
int x_n;
vector[x_n] x;
int subject_idx_n;
int col_idx_log_mu_joint_subject_n;
array[subject_idx_n] int subject_idx;
array[col_idx_log_mu_joint_subject_n] int col_idx_log_mu_joint_subject;
int z_n;
vector[z_n] z;
int col_idx_log_phi_joint_subject_n;
array[col_idx_log_phi_joint_subject_n] int col_idx_log_phi_joint_subject;
int y_n;
array[y_n] int y;
}
transformed data {
matrix[x_n, 2] X_log_mu = hcat(rep_vector(1.0, num_elements(x)), x);
int pop_log_mu_n_covariates = 2;
matrix[x_n, 2] Z_log_mu_joint_subject = hcat(rep_vector(1.0, num_elements(x)), x);
matrix[z_n, 2] X_log_phi = hcat(rep_vector(1.0, num_elements(z)), z);
int pop_log_phi_n_covariates = 2;
matrix[z_n, 2] Z_log_phi_joint_subject = hcat(rep_vector(1.0, num_elements(z)), z);
}
parameters {
cholesky_factor_corr[n_terms_joint_subject] b_joint_subject_L;
vector<lower=0.0>[n_terms_joint_subject] b_joint_subject_tau;
vector[(n_terms_joint_subject * n_subject)] b_joint_subject_z_flat;
vector[pop_log_mu_n_covariates] pop_log_mu_beta_pop;
vector[pop_log_phi_n_covariates] pop_log_phi_beta_pop;
}
transformed parameters {
matrix[n_terms_joint_subject, n_subject] b_joint_subject_z = to_matrix(b_joint_subject_z_flat, n_terms_joint_subject, n_subject);
matrix[n_subject, n_terms_joint_subject] b_joint_subject = ((diag_pre_multiply(b_joint_subject_tau, b_joint_subject_L) * b_joint_subject_z)');
vector[x_n] pop_log_mu = (X_log_mu * pop_log_mu_beta_pop);
vector[subject_idx_n] r_log_mu_joint_subject = rows_dot_product(Z_log_mu_joint_subject, b_joint_subject[subject_idx, col_idx_log_mu_joint_subject]);
vector[x_n] log_mu = (pop_log_mu + r_log_mu_joint_subject);
vector[x_n] mu = exp(log_mu);
vector[z_n] pop_log_phi = (X_log_phi * pop_log_phi_beta_pop);
vector[subject_idx_n] r_log_phi_joint_subject = rows_dot_product(
Z_log_phi_joint_subject,
b_joint_subject[subject_idx, col_idx_log_phi_joint_subject]
);
vector[z_n] log_phi = (pop_log_phi + r_log_phi_joint_subject);
vector[z_n] phi = exp(log_phi);
}
model {
b_joint_subject_L ~ lkj_corr_cholesky(1.0);
b_joint_subject_tau ~ std_normal();
b_joint_subject_z_flat ~ std_normal();
pop_log_mu_beta_pop ~ std_normal();
pop_log_phi_beta_pop ~ std_normal();
y ~ neg_binomial_2(mu, phi);
}
generated quantities {
vector[y_n] y_likelihood = neg_binomial_2_lpmfs(y, mu, phi);
array[y_n] int y_gen = neg_binomial_2_int_rng(y_n, mu, phi);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu, X_phi, group_effects_mu_1, group_effects_phi_1)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_mu = X_mu * beta_pop
group_effect_1 = Base.zeros(Base.length(y))
eta_mu = eta_mu + group_effect_1
beta_pop_phi ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_phi = X_phi * beta_pop_phi
group_effect_2 = Base.zeros(Base.length(y))
eta_phi = eta_phi + group_effect_2
shared_group_1 ~ DynamicPPL.to_submodel(BayesianRegressionModelsTuringExt._brm_shared_group_effect_model((group_effects_mu_1, group_effects_phi_1), (nothing, nothing, nothing, nothing), (nothing, nothing)))
eta_mu = eta_mu + shared_group_1.effects[1]
eta_phi = eta_phi + shared_group_1.effects[2]
mu = Base.exp.(eta_mu)
phi = Base.exp.(eta_phi)
begin
for i = Base.eachindex(y)
y[i] ~ BayesianRegressionModels.NegativeBinomial2(mu[i], phi[i])
end
end
(; mu = mu, phi = phi, response = y)
end)Multi-axis population PK kernel
The subject frame has one row per person, while the observation frame has one row per concentration measurement. Those axes deliberately have different lengths and the observation rows are interleaved. ragged(x, group) joins the flat observation columns to the subject axis; kernel(...) then evaluates one structural-model cell per subject. In the generated StanBlocks pane, that public BRM kernel lowers to a plate.
This is a deliberately small one-compartment IV-bolus model, C(t) = dose / V * exp(-(CL / V)t). The shared pk ID gives CL and V one correlated between-subject variability block.
brm-comparison
Multi-axis population PK kerneljulia
population_pk = (@brm begin
sigma ~ Exponential(1)
log(CL) ~ 1 + (1 | pk | subject)
log(V) ~ 1 + (1 | pk | subject)
predicted_concentration ~ kernel(
ragged(time, obs_subject), dose, CL, V,
) do ts, d, cl, volume
d / volume * exp((-cl / volume) * ts)
end
ragged(concentration, obs_subject) ~
Normal(predicted_concentration, sigma)
end)((;
# Subject axis: one row per subject.
subject=["alice", "bob"],
dose=[100.0, 80.0],
# Observation axis: one row per sample, interleaved by subject.
obs_subject=["alice", "bob", "alice", "bob", "alice"],
time=[0.5, 0.25, 1.5, 1.0, 3.0],
concentration=[8.1, 7.6, 5.2, 4.9, 2.1],
))julia
BRMI:
sigma ~ Exponential(1)
subject: data (eltype=String, n=2)
log(CL) ~ 1 + (1 | pk | subject)
log(V) ~ 1 + (1 | pk | subject)
time: data (eltype=Float64, n=5)
obs_subject: data (eltype=String, n=5)
dose: data (eltype=Float64, n=2)
predicted_concentration ~ kernel((ts, d, cl, volume)->begin
#= brm-docs-example.jl:9 =#
(d / volume) * exp((-cl / volume) * ts)
end, ragged(time, obs_subject), dose, CL, V)
ragged(concentration, obs_subject) ~ Normal(predicted_concentration, sigma)julia
SBBRMI with data keys = [:concentration, :dose, :kernel_nsub_predicted_concentration, :kernel_predicted_concentration_time_ragged, :n_subject, :n_terms_pk_subject, :subject_idx, :time]
emitted @slic body:
begin
b_pk_subject ~ ranef_correlated_draws(; group_idx = subject_idx, n_groups = n_subject, n_terms = n_terms_pk_subject)
sigma ~ exponential(1.0 ./ 1)
X_log_CL = hcat(rep_vector(1.0, num_elements(subject_idx)))
pop_log_CL ~ popefs(; X = X_log_CL)
r_log_CL_pk_subject = b_pk_subject[subject_idx, 1]
log_CL = pop_log_CL + r_log_CL_pk_subject
CL = exp(log_CL)
X_log_V = hcat(rep_vector(1.0, num_elements(subject_idx)))
pop_log_V ~ popefs(; X = X_log_V)
r_log_V_pk_subject = b_pk_subject[subject_idx, 2]
log_V = pop_log_V + r_log_V_pk_subject
V = exp(log_V)
predicted_concentration ~ plate(kernel_predicted_concentration_time_ragged, dose, CL, V; outer = (kernel_nsub_predicted_concentration,)) do ts, d, cl, volume
#= brm-docs-example.jl:9 =#
(d / volume) * exp((-cl / volume) * ts)
end
concentration ~ normal(predicted_concentration, sigma)
endstan
functions {
matrix hcat(vector x) {
int n = dims(x)[1];
return to_matrix(x, n, 1);
}
int ragged_end(array[] int ends, int i) {
return ends[i];
}
int ragged_start(
array[] int ends,
int i
) {
if((i == 1)) {
return 1;
} else {
return (1 + ends[(i - 1)]);
}
}
int num_elements_RaggedVector(tuple(vector, array[] int) rv) {
return size(rv.2);
}
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;
if((n == 0)) {
vector[n] rv;
return rv;
} else {
return to_vector(normal_rng(a, b));
}
}
int ragged_end_RaggedVector(tuple(vector, array[] int) x, int i) {
return x.2[i];
}
int ragged_start_RaggedVector(
tuple(vector, array[] int) x,
int i
) {
if((i == 1)) {
return 1;
} else {
return (1 + x.2[(i - 1)]);
}
}
vector getindex_RaggedVector(
tuple(vector, array[] int) rv,
int i
) {
return rv.1[ragged_start_RaggedVector(rv, i):ragged_end_RaggedVector(rv, i)];
}
}
data {
int n_terms_pk_subject;
int n_subject;
int subject_idx_n;
array[subject_idx_n] int subject_idx;
int kernel_nsub_predicted_concentration;
int kernel_predicted_concentration_time_ragged_ends_n;
int kernel_predicted_concentration_time_ragged_mem_n;
tuple(
vector[kernel_predicted_concentration_time_ragged_mem_n],
array[kernel_predicted_concentration_time_ragged_ends_n] int
) kernel_predicted_concentration_time_ragged;
int dose_n;
vector[dose_n] dose;
int concentration_mem_n;
int concentration_ends_n;
tuple(vector[concentration_mem_n], array[concentration_ends_n] int) concentration;
}
transformed data {
matrix[num_elements(subject_idx), 1] X_log_CL = hcat(rep_vector(1.0, num_elements(subject_idx)));
int pop_log_CL_n_covariates = 1;
matrix[num_elements(subject_idx), 1] X_log_V = hcat(rep_vector(1.0, num_elements(subject_idx)));
int pop_log_V_n_covariates = 1;
array[kernel_nsub_predicted_concentration] int predicted_concentration__pl_len_1;
for(plate_i__pl_1 in 1:kernel_nsub_predicted_concentration) {
predicted_concentration__pl_len_1[plate_i__pl_1] = (
1 +
(
ragged_end(kernel_predicted_concentration_time_ragged.2, plate_i__pl_1) -
ragged_start(kernel_predicted_concentration_time_ragged.2, plate_i__pl_1)
)
);
}
array[kernel_nsub_predicted_concentration] int predicted_concentration__pl_end_1 = cumulative_sum(predicted_concentration__pl_len_1);
}
parameters {
cholesky_factor_corr[n_terms_pk_subject] b_pk_subject_L;
vector<lower=0.0>[n_terms_pk_subject] b_pk_subject_tau;
vector[(n_terms_pk_subject * n_subject)] b_pk_subject_z_flat;
real<lower=0.0> sigma;
vector[pop_log_CL_n_covariates] pop_log_CL_beta_pop;
vector[pop_log_V_n_covariates] pop_log_V_beta_pop;
}
transformed parameters {
matrix[n_terms_pk_subject, n_subject] b_pk_subject_z = to_matrix(b_pk_subject_z_flat, n_terms_pk_subject, n_subject);
matrix[n_subject, n_terms_pk_subject] b_pk_subject = ((diag_pre_multiply(b_pk_subject_tau, b_pk_subject_L) * b_pk_subject_z)');
vector[num_elements(subject_idx)] pop_log_CL = (X_log_CL * pop_log_CL_beta_pop);
vector[subject_idx_n] r_log_CL_pk_subject = b_pk_subject[subject_idx, 1];
vector[num_elements(subject_idx)] log_CL = (pop_log_CL + r_log_CL_pk_subject);
vector[num_elements(subject_idx)] CL = exp(log_CL);
vector[num_elements(subject_idx)] pop_log_V = (X_log_V * pop_log_V_beta_pop);
vector[subject_idx_n] r_log_V_pk_subject = b_pk_subject[subject_idx, 2];
vector[num_elements(subject_idx)] log_V = (pop_log_V + r_log_V_pk_subject);
vector[num_elements(subject_idx)] V = exp(log_V);
vector[sum(predicted_concentration__pl_len_1)] predicted_concentration__pl_mem_1;
for(plate_i__pl_1 in 1:kernel_nsub_predicted_concentration) {
predicted_concentration__pl_mem_1[
ragged_start(predicted_concentration__pl_end_1, plate_i__pl_1):ragged_end(predicted_concentration__pl_end_1, plate_i__pl_1)
] = (
(dose[plate_i__pl_1] / V[plate_i__pl_1]) *
exp(
(
((-CL[plate_i__pl_1]) / V[plate_i__pl_1]) *
kernel_predicted_concentration_time_ragged.1[
ragged_start(kernel_predicted_concentration_time_ragged.2, plate_i__pl_1):ragged_end(kernel_predicted_concentration_time_ragged.2, plate_i__pl_1)
]
)
)
);
}
}
model {
b_pk_subject_L ~ lkj_corr_cholesky(1.0);
b_pk_subject_tau ~ std_normal();
b_pk_subject_z_flat ~ std_normal();
sigma ~ exponential((1.0 ./ 1));
pop_log_CL_beta_pop ~ std_normal();
pop_log_V_beta_pop ~ std_normal();
for(g__ro_2 in 1:num_elements_RaggedVector(concentration)) {
getindex_RaggedVector(concentration, g__ro_2) ~ normal(
predicted_concentration__pl_mem_1[
ragged_start(predicted_concentration__pl_end_1, g__ro_2):ragged_end(predicted_concentration__pl_end_1, g__ro_2)
],
sigma
);
}
}
generated quantities {
vector[num_elements(concentration.1)] concentration_gen;
vector[num_elements_RaggedVector(concentration)] concentration_likelihood;
for(g__rq_2 in 1:num_elements_RaggedVector(concentration)) {
concentration_gen[ragged_start(concentration.2, g__rq_2):ragged_end(concentration.2, g__rq_2)] = normal_vector_rng(
(1 + (ragged_end_RaggedVector(concentration, g__rq_2) - ragged_start_RaggedVector(concentration, g__rq_2))),
predicted_concentration__pl_mem_1[
ragged_start(predicted_concentration__pl_end_1, g__rq_2):ragged_end(predicted_concentration__pl_end_1, g__rq_2)
],
sigma
);
concentration_likelihood[g__rq_2] = normal_lpdf(getindex_RaggedVector(concentration, g__rq_2) |
predicted_concentration__pl_mem_1[
ragged_start(predicted_concentration__pl_end_1, g__rq_2):ragged_end(predicted_concentration__pl_end_1, g__rq_2)
],
sigma
);
}
}julia
Turing unsupported for this BRM example
Turing backend: response decorators other than `mi(response)` and response links are not yet supportedVerified public Warfarin PK/PD reproduction
The repository includes a complete executable translation of Sebastian Weber's public StanCon 2018 Warfarin programs: a first-stage one-compartment oral PK model with lag, allometry, four independent subject effects, and Gamma overdispersion; followed by a turnover PD ODE conditioned on the public PK posterior medians, with three independent PD subject effects and the same observation family.
The reproduction script contains the typed ODE and likelihood definitions, a public two-subject data slice, both @brm declarations, stanc checks, and finite BridgeStan density and gradient checks. The accompanying audit notes map every retained prior, structural equation, solver tolerance, and known generated-quantity difference back to the public source.
The Warfarin PK/PD page also renders the separately specified joint model: one posterior in which shared subject-level PK effects feed both the PK and PD likelihoods, so PK uncertainty propagates into PD and PD data can update PK quantities.
julia
include("research/warfarin/reproduce.jl")
models = warfarin_sbbrmis()This is the strongest identifiable public match to the Warfarin model mentioned in brms issue #1509, but the issue itself provides no equations or citation, so identity with the commenter's private working model cannot be proved. The public StanCon model is two separate scalar-Gamma stages; it does not need the correlated Gaussian outcome surface above.
Categorical population terms
brm-comparison
Categorical population designjulia
categorical = (@brm begin
sigma ~ Exponential(2)
mu ~ 1 + factor(group; ref=3) + x + x & group
effect(mu, group) ~ Normal(0, 0.5)
y ~ Normal(mu, sigma)
end)((;
group=[1, 2, 3, 1, 2, 3],
x=[-1.0, -0.5, 0.0, 0.5, 1.0, 1.5],
y=[-2.4, -2.2, -2.0, -1.8, -1.7, -1.5],
))julia
BRMI:
sigma ~ Exponential(2)
group: data (eltype=Int64, n=6)
x: data (eltype=Float64, n=6)
mu ~ 1 + factor(group; ref=3) + x + &(x, group)
effect(mu, group) ~ Normal(0, 0.5)
y ~ Normal(mu, sigma)julia
SBBRMI with data keys = [:group, :group__ref_3_idx, :group__ref_3_n_levels, :group_lvl_2, :group_lvl_3, :int_x_x_group_lvl_2, :int_x_x_group_lvl_3, :x, :y]
emitted @slic body:
begin
sigma ~ exponential(1.0 ./ 2)
X_mu = hcat(rep_vector(1.0, num_elements(x)), x, int_x_x_group_lvl_2, int_x_x_group_lvl_3)
pop_mu ~ popefs(; X = X_mu)
cat_mu_group__ref_3 ~ _sb_cat_normal(; x = group__ref_3_idx, n_levels = group__ref_3_n_levels, beta_loc = 0, beta_scale = 0.5)
mu = pop_mu + cat_mu_group__ref_3
y ~ normal(mu, sigma)
endstan
functions {
matrix hcat(vector x, vector y, vector z, vector args1) {
return hcat(hcat(x, y), z, args1);
}
matrix hcat(matrix x, vector y, vector z) {
return hcat(hcat(x, y), z);
}
matrix hcat(
matrix x,
vector y
) {
int m = dims(x)[1];
int n = dims(x)[2];
if (dims(y)[1] != m) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `m` (= ", m, "), inferred from `x` dim 1. `m` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
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;
if((n == 0)) {
vector[n] rv;
return rv;
} else {
return to_vector(normal_rng(a, b));
}
}
}
data {
int int_x_x_group_lvl_3_n;
int x_n;
vector[x_n] x;
int int_x_x_group_lvl_2_n;
vector[int_x_x_group_lvl_2_n] int_x_x_group_lvl_2;
vector[int_x_x_group_lvl_3_n] int_x_x_group_lvl_3;
int group__ref_3_n_levels;
int group__ref_3_idx_n;
array[group__ref_3_idx_n] int group__ref_3_idx;
int y_n;
vector[y_n] y;
}
transformed data {
matrix[int_x_x_group_lvl_3_n, ((2 + 1) + 1)] X_mu = hcat(rep_vector(1.0, num_elements(x)), x, int_x_x_group_lvl_2, int_x_x_group_lvl_3);
int pop_mu_n_covariates = ((2 + 1) + 1);
}
parameters {
real<lower=0.0> sigma;
vector[pop_mu_n_covariates] pop_mu_beta_pop;
vector[(group__ref_3_n_levels - 1)] cat_mu_group__ref_3_beta;
}
transformed parameters {
vector[int_x_x_group_lvl_3_n] pop_mu = (X_mu * pop_mu_beta_pop);
vector[group__ref_3_idx_n] cat_mu_group__ref_3 = append_row(0.0, cat_mu_group__ref_3_beta)[group__ref_3_idx];
vector[int_x_x_group_lvl_3_n] mu = (pop_mu + cat_mu_group__ref_3);
}
model {
sigma ~ exponential((1.0 ./ 2));
pop_mu_beta_pop ~ std_normal();
cat_mu_group__ref_3_beta ~ normal(0, 0.5);
y ~ normal(mu, sigma);
}
generated quantities {
vector[y_n] y_likelihood = normal_lpdfs(y, mu, sigma);
vector[y_n] y_gen = normal_vector_rng(y_n, mu, sigma);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu)
sigma ~ Distributions.Exponential(2)
beta_pop ~ Distributions.product_distribution([Distributions.Normal(0, 1), Distributions.Normal(0, 0.5), Distributions.Normal(0, 0.5), Distributions.Normal(0, 1), Distributions.Normal(0, 1), Distributions.Normal(0, 1)])
eta_mu = X_mu * beta_pop
mu = eta_mu
begin
for i = Base.eachindex(y)
y[i] ~ Distributions.Normal(mu[i], sigma)
end
end
(; mu = mu, sigma = sigma, response = y)
end)Correlated group effects
brm-comparison
Correlated random slopesjulia
grouped = (@brm begin
sigma ~ Exponential(2)
mu ~ 1 + x + (1 + x | subject)
outcome ~ Normal(mu, sigma)
end)((;
x=[-1.0, 0.5, 2.0, 0.25],
subject=[2, 1, 2, 3],
outcome=[0.2, 1.1, -0.4, 0.7],
))julia
BRMI:
sigma ~ Exponential(2)
x: data (eltype=Float64, n=4)
subject: data (eltype=Int64, n=4)
mu ~ 1 + x + ((1 + x) | subject)
outcome ~ Normal(mu, sigma)julia
SBBRMI with data keys = [:n_subject, :n_terms_mu_subject, :outcome, :subject, :subject_idx, :x]
emitted @slic body:
begin
sigma ~ exponential(1.0 ./ 2)
X_mu = hcat(rep_vector(1.0, num_elements(x)), x)
pop_mu ~ popefs(; X = X_mu)
Z_mu_subject = hcat(rep_vector(1.0, num_elements(x)), x)
r_mu_subject ~ ranef_correlated(; Z = Z_mu_subject, group_idx = subject_idx, n_groups = n_subject, n_terms = n_terms_mu_subject)
mu = pop_mu + r_mu_subject
outcome ~ normal(mu, sigma)
endstan
functions {
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
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;
if((n == 0)) {
vector[n] rv;
return rv;
} else {
return to_vector(normal_rng(a, b));
}
}
}
data {
int x_n;
vector[x_n] x;
int n_terms_mu_subject;
int n_subject;
int subject_idx_n;
array[subject_idx_n] int subject_idx;
int outcome_n;
vector[outcome_n] outcome;
}
transformed data {
matrix[x_n, 2] X_mu = hcat(rep_vector(1.0, num_elements(x)), x);
int pop_mu_n_covariates = 2;
matrix[x_n, 2] Z_mu_subject = hcat(rep_vector(1.0, num_elements(x)), x);
}
parameters {
real<lower=0.0> sigma;
vector[pop_mu_n_covariates] pop_mu_beta_pop;
cholesky_factor_corr[n_terms_mu_subject] r_mu_subject_L;
vector<lower=0.0>[n_terms_mu_subject] r_mu_subject_tau;
vector[(n_terms_mu_subject * n_subject)] r_mu_subject_z_flat;
}
transformed parameters {
vector[x_n] pop_mu = (X_mu * pop_mu_beta_pop);
matrix[n_terms_mu_subject, n_subject] r_mu_subject_z = to_matrix(r_mu_subject_z_flat, n_terms_mu_subject, n_subject);
matrix[n_subject, n_terms_mu_subject] r_mu_subject_b = ((diag_pre_multiply(r_mu_subject_tau, r_mu_subject_L) * r_mu_subject_z)');
vector[subject_idx_n] r_mu_subject = rows_dot_product(Z_mu_subject, r_mu_subject_b[subject_idx, :]);
vector[x_n] mu = (pop_mu + r_mu_subject);
}
model {
sigma ~ exponential((1.0 ./ 2));
pop_mu_beta_pop ~ std_normal();
r_mu_subject_L ~ lkj_corr_cholesky(1.0);
r_mu_subject_tau ~ std_normal();
r_mu_subject_z_flat ~ std_normal();
outcome ~ normal(mu, sigma);
}
generated quantities {
vector[outcome_n] outcome_likelihood = normal_lpdfs(outcome, mu, sigma);
vector[outcome_n] outcome_gen = normal_vector_rng(outcome_n, mu, sigma);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu, group_matrix_mu_1, group_indices_mu_1, group_levels_mu_1)
sigma ~ Distributions.Exponential(2)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_mu = X_mu * beta_pop
group_effect_1 = Base.zeros(Base.length(y))
group_1_1 ~ DynamicPPL.to_submodel(BRM.turing_default_correlated_group(group_matrix_mu_1, group_indices_mu_1, group_levels_mu_1, 1.0))
group_effect_1 = group_effect_1 + group_1_1.effect
eta_mu = eta_mu + group_effect_1
mu = eta_mu
begin
for i = Base.eachindex(y)
y[i] ~ Distributions.Normal(mu[i], sigma)
end
end
(; mu = mu, sigma = sigma, response = y)
end)Random-effect scale and correlation priors
The p identifier gives the group block a stable prior address. The shared half-Normal sets both scales, then the more-specific Exponential override replaces it only on the x margin.
brm-comparison
Addressed group scale and correlation priorsjulia
grouped_priors = (@brm begin
sigma ~ Exponential(2)
mu ~ 1 + x + (1 + x | p | subject)
sd(:, p) ~ Normal(0, 0.5)
sd(mu, p, x) ~ Exponential(0.25)
cor(:, p) ~ LKJCholesky(2, 2.5)
outcome ~ Normal(mu, sigma)
end)((;
x=[-1.0, 0.5, 2.0, 0.25],
subject=[2, 1, 2, 3],
outcome=[0.2, 1.1, -0.4, 0.7],
))julia
BRMI:
sigma ~ Exponential(2)
x: data (eltype=Float64, n=4)
subject: data (eltype=Int64, n=4)
mu ~ 1 + x + ((1 + x) | p | subject)
effect(sd, p) ~ Normal(0, 0.5)
effect(sd, p, mu, x) ~ Exponential(0.25)
effect(cor, p) ~ LKJCholesky(2, 2.5)
outcome ~ Normal(mu, sigma)julia
SBBRMI with data keys = [:col_idx_mu_p_subject, :n_subject, :n_terms_p_subject, :outcome, :subject_idx, :x]
configured submodels:
ranef_correlated_draws_generic_configured_1 = Base.merge(BayesianRegressionModels.ranef_correlated_draws_generic, quote
tau ~ (ValueFamily(brm_vector_prior_5cc114893928eff4))(0.0, 0.5, 4.0; n = n_terms)
end)
emitted @slic body:
begin
b_p_subject ~ ranef_correlated_draws_generic_configured_1(; group_idx = subject_idx, n_groups = n_subject, n_terms = n_terms_p_subject, lkj_eta = 2.5)
sigma ~ exponential(1.0 ./ 2)
X_mu = hcat(rep_vector(1.0, num_elements(x)), x)
pop_mu ~ popefs(; X = X_mu)
Z_mu_p_subject = hcat(rep_vector(1.0, num_elements(x)), x)
r_mu_p_subject = rows_dot_product(Z_mu_p_subject, b_p_subject[subject_idx, col_idx_mu_p_subject])
mu = pop_mu + r_mu_p_subject
outcome ~ normal(mu, sigma)
endstan
functions {
// value UDF brm_vector_prior_5cc114893928eff4_lpdf
real brm_vector_prior_5cc114893928eff4_lpdf(
vector x,
real arg_1,
real arg_2,
real arg_3
) {
if((x[1] < 0.0)) {
return negative_infinity();
}
if((x[2] < 0.0)) {
return negative_infinity();
}
return (normal_lpdf(x[1] | arg_1, arg_2) + exponential_lpdf(x[2] | arg_3));
}
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
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;
if((n == 0)) {
vector[n] rv;
return rv;
} else {
return to_vector(normal_rng(a, b));
}
}
}
data {
int n_terms_p_subject;
int n_subject;
int x_n;
vector[x_n] x;
int subject_idx_n;
int col_idx_mu_p_subject_n;
array[subject_idx_n] int subject_idx;
array[col_idx_mu_p_subject_n] int col_idx_mu_p_subject;
int outcome_n;
vector[outcome_n] outcome;
}
transformed data {
matrix[x_n, 2] X_mu = hcat(rep_vector(1.0, num_elements(x)), x);
int pop_mu_n_covariates = 2;
matrix[x_n, 2] Z_mu_p_subject = hcat(rep_vector(1.0, num_elements(x)), x);
}
parameters {
cholesky_factor_corr[n_terms_p_subject] b_p_subject_L;
vector<lower=0.0>[n_terms_p_subject] b_p_subject_tau;
vector[(n_terms_p_subject * n_subject)] b_p_subject_z_flat;
real<lower=0.0> sigma;
vector[pop_mu_n_covariates] pop_mu_beta_pop;
}
transformed parameters {
matrix[n_terms_p_subject, n_subject] b_p_subject_z = to_matrix(b_p_subject_z_flat, n_terms_p_subject, n_subject);
matrix[n_subject, n_terms_p_subject] b_p_subject = ((diag_pre_multiply(b_p_subject_tau, b_p_subject_L) * b_p_subject_z)');
vector[x_n] pop_mu = (X_mu * pop_mu_beta_pop);
vector[subject_idx_n] r_mu_p_subject = rows_dot_product(Z_mu_p_subject, b_p_subject[subject_idx, col_idx_mu_p_subject]);
vector[x_n] mu = (pop_mu + r_mu_p_subject);
}
model {
b_p_subject_L ~ lkj_corr_cholesky(2.5);
b_p_subject_tau ~ brm_vector_prior_5cc114893928eff4(0.0, 0.5, 4.0);
b_p_subject_z_flat ~ std_normal();
sigma ~ exponential((1.0 ./ 2));
pop_mu_beta_pop ~ std_normal();
outcome ~ normal(mu, sigma);
}
generated quantities {
vector[outcome_n] outcome_likelihood = normal_lpdfs(outcome, mu, sigma);
vector[outcome_n] outcome_gen = normal_vector_rng(outcome_n, mu, sigma);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu, group_effects_mu_1)
sigma ~ Distributions.Exponential(2)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_mu = X_mu * beta_pop
group_effect_1 = Base.zeros(Base.length(y))
group_1_1 ~ DynamicPPL.to_submodel(BRM.turing_group_effect(group_effects_mu_1, (Distributions.Normal(0, 0.5), Distributions.Exponential(0.25)), nothing))
group_effect_1 = group_effect_1 + group_1_1.effect
eta_mu = eta_mu + group_effect_1
mu = eta_mu
begin
for i = Base.eachindex(y)
y[i] ~ Distributions.Normal(mu[i], sigma)
end
end
(; mu = mu, sigma = sigma, response = y)
end)Stratified group covariance
gr(subject, by=arm) fits a separate scale/correlation frame in each arm while retaining one pooled subject coordinate.
brm-comparison
Group covariance stratified by armjulia
stratified_groups = (@brm begin
sigma ~ Exponential(2)
mu ~ 1 + x + (1 + x | gr(subject, by=arm))
outcome ~ Normal(mu, sigma)
end)((;
x=[-1.0, 0.5, 2.0, 0.25, -0.5, 1.25],
subject=[1, 1, 2, 3, 3, 4],
arm=[1, 1, 1, 2, 2, 2],
outcome=[0.2, 1.1, -0.4, 0.7, -0.2, 0.5],
))julia
BRMI:
sigma ~ Exponential(2)
x: data (eltype=Float64, n=6)
subject: data (eltype=Int64, n=6)
arm: data (eltype=Int64, n=6)
mu ~ 1 + x + ((1 + x) | gr(subject; by=arm))
outcome ~ Normal(mu, sigma)julia
SBBRMI with data keys = [:arm, :n_strata_subject__by__arm, :n_subject__by__arm, :n_terms_mu_subject__by__arm, :outcome, :subject, :subject__by__arm_idx, :subject__by__arm_stratum_idx, :x]
emitted @slic body:
begin
sigma ~ exponential(1.0 ./ 2)
X_mu = hcat(rep_vector(1.0, num_elements(x)), x)
pop_mu ~ popefs(; X = X_mu)
Z_mu_subject__by__arm = hcat(rep_vector(1.0, num_elements(x)), x)
r_mu_subject__by__arm ~ ranef_correlated_by(; Z = Z_mu_subject__by__arm, group_idx = subject__by__arm_idx, n_groups = n_subject__by__arm, n_terms = n_terms_mu_subject__by__arm, stratum_idx = subject__by__arm_stratum_idx, n_strata = n_strata_subject__by__arm)
mu = pop_mu + r_mu_subject__by__arm
outcome ~ normal(mu, sigma)
endstan
functions {
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
real lkj_corr_cholesky_lpdfs(matrix L, real x) {
return lkj_corr_cholesky_lpdf(L | x);
}
matrix lkj_corr_cholesky_cholesky_factor_corr_rng(
int anontok__1,
real eta
) {
int n = anontok__1;
return lkj_corr_cholesky_rng(n, eta);
}
vector lower_conditioning_vector_normal_rng(
int anontok__1,
real lo,
real args1,
real args2
) {
int n = anontok__1;
return jbroadcasted_lower_conditioning_cell_rng_normal_rng(rep_vector(0.0, n), lo, args1, args2);
}
vector jbroadcasted_lower_conditioning_cell_rng_normal_rng(
vector x1,
real x3,
real x4,
real x5
) {
int n = dims(x1)[1];
vector[n] rv;
for(i in 1:n) {
rv[i] = lower_conditioning_cell_normal_rng(broadcasted_getindex(x1, i), x3, x4, x5);
}
return rv;
}
real lower_conditioning_cell_normal_rng(
real dummy,
real lo,
real args1,
real args2
) {
return lower_conditioning_normal_rng(lo, args1, args2);
}
real lower_conditioning_normal_rng(
real lo,
real args1,
real args2
) {
vector[1] draw;
array[1] int attempts;
draw[1] = normal_rng(args1, args2);
attempts[1] = 1;
while((draw[1] < lo)) {
if((attempts[1] >= 100000)) {
reject("truncated: lower-tail rejection sampler exceeded 100000 draws");
}
draw[1] = normal_rng(args1, args2);
attempts[1] = (attempts[1] + 1);
}
return draw[1];
}
real broadcasted_getindex(vector x, int i) {
return x[i];
}
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));
}
}
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);
}
vector normal_vector_rng(
int anontok__1,
vector a,
real b
) {
int n = anontok__1;
if((n == 0)) {
vector[n] rv;
return rv;
} else {
return to_vector(normal_rng(a, b));
}
}
}
data {
int x_n;
vector[x_n] x;
int n_strata_subject__by__arm;
int n_terms_mu_subject__by__arm;
int n_subject__by__arm;
int subject__by__arm_stratum_idx_n;
array[subject__by__arm_stratum_idx_n] int subject__by__arm_stratum_idx;
int subject__by__arm_idx_n;
array[subject__by__arm_idx_n] int subject__by__arm_idx;
int outcome_n;
vector[outcome_n] outcome;
}
transformed data {
matrix[x_n, 2] X_mu = hcat(rep_vector(1.0, num_elements(x)), x);
int pop_mu_n_covariates = 2;
matrix[x_n, 2] Z_mu_subject__by__arm = hcat(rep_vector(1.0, num_elements(x)), x);
}
parameters {
real<lower=0.0> sigma;
vector[pop_mu_n_covariates] pop_mu_beta_pop;
array[n_strata_subject__by__arm] cholesky_factor_corr[n_terms_mu_subject__by__arm] r_mu_subject__by__arm_L_s_L;
matrix<lower=0.0>[n_terms_mu_subject__by__arm, n_strata_subject__by__arm] r_mu_subject__by__arm_tau_s_tau;
matrix[n_terms_mu_subject__by__arm, n_subject__by__arm] r_mu_subject__by__arm_b_T_z_g;
}
transformed parameters {
vector[x_n] pop_mu = (X_mu * pop_mu_beta_pop);
array[n_strata_subject__by__arm] cholesky_factor_corr[n_terms_mu_subject__by__arm] r_mu_subject__by__arm_L_s;
for(r_mu_subject__by__arm_plate_i__pl_1 in 1:n_strata_subject__by__arm) {
r_mu_subject__by__arm_L_s[r_mu_subject__by__arm_plate_i__pl_1] = r_mu_subject__by__arm_L_s_L[r_mu_subject__by__arm_plate_i__pl_1];
}
matrix<lower=0.0>[n_terms_mu_subject__by__arm, n_strata_subject__by__arm] r_mu_subject__by__arm_tau_s;
for(r_mu_subject__by__arm_plate_i__pl_2 in 1:n_strata_subject__by__arm) {
r_mu_subject__by__arm_tau_s[:, r_mu_subject__by__arm_plate_i__pl_2] = r_mu_subject__by__arm_tau_s_tau[:, r_mu_subject__by__arm_plate_i__pl_2];
}
array[n_subject__by__arm] cholesky_factor_corr[n_terms_mu_subject__by__arm] r_mu_subject__by__arm_b_T_L_g;
matrix[n_terms_mu_subject__by__arm, n_subject__by__arm] r_mu_subject__by__arm_b_T_tau_g;
matrix[n_terms_mu_subject__by__arm, n_subject__by__arm] r_mu_subject__by__arm_b_T;
for(r_mu_subject__by__arm_plate_i__pl_3 in 1:n_subject__by__arm) {
r_mu_subject__by__arm_b_T_L_g[r_mu_subject__by__arm_plate_i__pl_3] = r_mu_subject__by__arm_L_s[subject__by__arm_stratum_idx[r_mu_subject__by__arm_plate_i__pl_3]];
r_mu_subject__by__arm_b_T_tau_g[:, r_mu_subject__by__arm_plate_i__pl_3] = r_mu_subject__by__arm_tau_s[:, subject__by__arm_stratum_idx[r_mu_subject__by__arm_plate_i__pl_3]];
r_mu_subject__by__arm_b_T[:, r_mu_subject__by__arm_plate_i__pl_3] = (
diag_pre_multiply(
r_mu_subject__by__arm_b_T_tau_g[:, r_mu_subject__by__arm_plate_i__pl_3],
r_mu_subject__by__arm_b_T_L_g[r_mu_subject__by__arm_plate_i__pl_3]
) *
r_mu_subject__by__arm_b_T_z_g[:, r_mu_subject__by__arm_plate_i__pl_3]
);
}
matrix[n_subject__by__arm, n_terms_mu_subject__by__arm] r_mu_subject__by__arm_b = (r_mu_subject__by__arm_b_T');
vector[subject__by__arm_idx_n] r_mu_subject__by__arm = rows_dot_product(Z_mu_subject__by__arm, r_mu_subject__by__arm_b[subject__by__arm_idx, :]);
vector[x_n] mu = (pop_mu + r_mu_subject__by__arm);
}
model {
sigma ~ exponential((1.0 ./ 2));
pop_mu_beta_pop ~ std_normal();
for(r_mu_subject__by__arm_plate_i__pl_1 in 1:n_strata_subject__by__arm) {
r_mu_subject__by__arm_L_s_L[r_mu_subject__by__arm_plate_i__pl_1] ~ lkj_corr_cholesky(1.0);
}
for(r_mu_subject__by__arm_plate_i__pl_2 in 1:n_strata_subject__by__arm) {
r_mu_subject__by__arm_tau_s_tau[:, r_mu_subject__by__arm_plate_i__pl_2] ~ std_normal();
}
for(r_mu_subject__by__arm_plate_i__pl_3 in 1:n_subject__by__arm) {
r_mu_subject__by__arm_b_T_z_g[:, r_mu_subject__by__arm_plate_i__pl_3] ~ std_normal();
}
outcome ~ normal(mu, sigma);
}
generated quantities {
vector[outcome_n] outcome_likelihood = normal_lpdfs(outcome, mu, sigma);
vector[outcome_n] outcome_gen = normal_vector_rng(outcome_n, mu, sigma);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu, group_effects_mu_1)
sigma ~ Distributions.Exponential(2)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_mu = X_mu * beta_pop
group_effect_1 = Base.zeros(Base.length(y))
group_1_1 ~ DynamicPPL.to_submodel(BRM.turing_group_effect(group_effects_mu_1, (nothing, nothing), nothing))
group_effect_1 = group_effect_1 + group_1_1.effect
eta_mu = eta_mu + group_effect_1
mu = eta_mu
begin
for i = Base.eachindex(y)
y[i] ~ Distributions.Normal(mu[i], sigma)
end
end
(; mu = mu, sigma = sigma, response = y)
end)Multi-membership group effects
Each row belongs to two groups. The row weights are normalized before the two group contributions are pooled; normalize=false is the explicit raw-weight variant.
brm-comparison
Weighted multi-membership random slopesjulia
multi_membership = (@brm begin
sigma ~ Exponential(2)
mu ~ 1 + x + (1 + x | mm(g1, g2; weights=(w1, w2)))
outcome ~ Normal(mu, sigma)
end)((;
x=[-1.0, 0.5, 2.0, 0.25],
g1=[1, 1, 2, 3],
g2=[2, 3, 3, 1],
w1=[2.0, 1.0, 0.5, 3.0],
w2=[1.0, 1.0, 1.5, 1.0],
outcome=[0.2, 1.1, -0.4, 0.7],
))julia
BRMI:
sigma ~ Exponential(2)
x: data (eltype=Float64, n=4)
w1: data (eltype=Float64, n=4)
w2: data (eltype=Float64, n=4)
g1: data (eltype=Int64, n=4)
g2: data (eltype=Int64, n=4)
mu ~ 1 + x + ((1 + x) | mm(g1, g2; weights=(w1, w2)))
outcome ~ Normal(mu, sigma)julia
SBBRMI with data keys = [:mm__g1__g2__w__w1__w2_idx, :mm__g1__g2__w__w1__w2_weights, :n_memberships_mm__g1__g2__w__w1__w2, :n_mm__g1__g2__w__w1__w2, :n_obs_mm__g1__g2__w__w1__w2, :n_terms_mu_mm__g1__g2__w__w1__w2, :outcome, :x]
emitted @slic body:
begin
sigma ~ exponential(1.0 ./ 2)
X_mu = hcat(rep_vector(1.0, num_elements(x)), x)
pop_mu ~ popefs(; X = X_mu)
Z_mu_mm__g1__g2__w__w1__w2 = hcat(rep_vector(1.0, num_elements(x)), x)
b_mu_mm__g1__g2__w__w1__w2 ~ ranef_correlated_draws(; Z = Z_mu_mm__g1__g2__w__w1__w2, group_idx = mm__g1__g2__w__w1__w2_idx, n_groups = n_mm__g1__g2__w__w1__w2, n_terms = n_terms_mu_mm__g1__g2__w__w1__w2)
r_mu_mm__g1__g2__w__w1__w2 = multi_membership_correlated(Z_mu_mm__g1__g2__w__w1__w2, b_mu_mm__g1__g2__w__w1__w2, mm__g1__g2__w__w1__w2_idx, mm__g1__g2__w__w1__w2_weights, n_obs_mm__g1__g2__w__w1__w2, n_memberships_mm__g1__g2__w__w1__w2)
mu = pop_mu + r_mu_mm__g1__g2__w__w1__w2
outcome ~ normal(mu, sigma)
endstan
functions {
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
vector multi_membership_correlated(
matrix Z,
matrix b,
array[] int group_idx,
vector weights,
int n_obs,
int n_memberships
) {
int n_terms = dims(Z)[2];
int n_mm = dims(group_idx)[1];
if (dims(b)[2] != n_terms) reject("multi_membership_correlated: dim mismatch — `b` dim 2 (= ", dims(b)[2], ") does not match `n_terms` (= ", n_terms, "), inferred from `Z` dim 2. `n_terms` sizes: `Z` dim 2 (= ", dims(Z)[2], "), `b` dim 2 (= ", dims(b)[2], ").");
if (dims(weights)[1] != n_mm) reject("multi_membership_correlated: dim mismatch — `weights` dim 1 (= ", dims(weights)[1], ") does not match `n_mm` (= ", n_mm, "), inferred from `group_idx` dim 1. `n_mm` sizes: `group_idx` dim 1 (= ", dims(group_idx)[1], "), `weights` dim 1 (= ", dims(weights)[1], ").");
vector[n_obs] rv = rep_vector(0.0, n_obs);
for(i in 1:n_obs) {
for(m in 1:n_memberships) {
int j = (((i - 1) * n_memberships) + m);
for(k in 1:n_terms) {
rv[i] += (weights[j] * Z[i, k] * b[group_idx[j], k]);
}
}
}
return rv;
}
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;
if((n == 0)) {
vector[n] rv;
return rv;
} else {
return to_vector(normal_rng(a, b));
}
}
}
data {
int x_n;
vector[x_n] x;
int n_terms_mu_mm__g1__g2__w__w1__w2;
int n_mm__g1__g2__w__w1__w2;
int n_obs_mm__g1__g2__w__w1__w2;
int mm__g1__g2__w__w1__w2_idx_n;
array[mm__g1__g2__w__w1__w2_idx_n] int mm__g1__g2__w__w1__w2_idx;
int mm__g1__g2__w__w1__w2_weights_n;
vector[mm__g1__g2__w__w1__w2_weights_n] mm__g1__g2__w__w1__w2_weights;
int n_memberships_mm__g1__g2__w__w1__w2;
int outcome_n;
vector[outcome_n] outcome;
}
transformed data {
matrix[x_n, 2] X_mu = hcat(rep_vector(1.0, num_elements(x)), x);
int pop_mu_n_covariates = 2;
matrix[x_n, 2] Z_mu_mm__g1__g2__w__w1__w2 = hcat(rep_vector(1.0, num_elements(x)), x);
}
parameters {
real<lower=0.0> sigma;
vector[pop_mu_n_covariates] pop_mu_beta_pop;
cholesky_factor_corr[n_terms_mu_mm__g1__g2__w__w1__w2] b_mu_mm__g1__g2__w__w1__w2_L;
vector<lower=0.0>[n_terms_mu_mm__g1__g2__w__w1__w2] b_mu_mm__g1__g2__w__w1__w2_tau;
vector[(n_terms_mu_mm__g1__g2__w__w1__w2 * n_mm__g1__g2__w__w1__w2)] b_mu_mm__g1__g2__w__w1__w2_z_flat;
}
transformed parameters {
vector[x_n] pop_mu = (X_mu * pop_mu_beta_pop);
matrix[n_terms_mu_mm__g1__g2__w__w1__w2, n_mm__g1__g2__w__w1__w2] b_mu_mm__g1__g2__w__w1__w2_z = to_matrix(
b_mu_mm__g1__g2__w__w1__w2_z_flat,
n_terms_mu_mm__g1__g2__w__w1__w2,
n_mm__g1__g2__w__w1__w2
);
matrix[n_mm__g1__g2__w__w1__w2, n_terms_mu_mm__g1__g2__w__w1__w2] b_mu_mm__g1__g2__w__w1__w2 = ((
diag_pre_multiply(b_mu_mm__g1__g2__w__w1__w2_tau, b_mu_mm__g1__g2__w__w1__w2_L) *
b_mu_mm__g1__g2__w__w1__w2_z
)');
vector[n_obs_mm__g1__g2__w__w1__w2] r_mu_mm__g1__g2__w__w1__w2 = multi_membership_correlated(
Z_mu_mm__g1__g2__w__w1__w2,
b_mu_mm__g1__g2__w__w1__w2,
mm__g1__g2__w__w1__w2_idx,
mm__g1__g2__w__w1__w2_weights,
n_obs_mm__g1__g2__w__w1__w2,
n_memberships_mm__g1__g2__w__w1__w2
);
vector[x_n] mu = (pop_mu + r_mu_mm__g1__g2__w__w1__w2);
}
model {
sigma ~ exponential((1.0 ./ 2));
pop_mu_beta_pop ~ std_normal();
b_mu_mm__g1__g2__w__w1__w2_L ~ lkj_corr_cholesky(1.0);
b_mu_mm__g1__g2__w__w1__w2_tau ~ std_normal();
b_mu_mm__g1__g2__w__w1__w2_z_flat ~ std_normal();
outcome ~ normal(mu, sigma);
}
generated quantities {
vector[outcome_n] outcome_likelihood = normal_lpdfs(outcome, mu, sigma);
vector[outcome_n] outcome_gen = normal_vector_rng(outcome_n, mu, sigma);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu, group_matrix_mu_1, group_indices_mu_1, group_levels_mu_1)
sigma ~ Distributions.Exponential(2)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_mu = X_mu * beta_pop
group_effect_1 = Base.zeros(Base.length(y))
group_1_1 ~ DynamicPPL.to_submodel(BRM.turing_default_correlated_group(group_matrix_mu_1, group_indices_mu_1, group_levels_mu_1, 1.0))
group_effect_1 = group_effect_1 + group_1_1.effect
eta_mu = eta_mu + group_effect_1
mu = eta_mu
begin
for i = Base.eachindex(y)
y[i] ~ Distributions.Normal(mu[i], sigma)
end
end
(; mu = mu, sigma = sigma, response = y)
end)Multiple crossed group effects
brm-comparison
Crossed correlated and independent group blocksjulia
crossed_groups = (@brm begin
sigma ~ Exponential(2)
mu ~ 1 + x + (1 + x | subject) + (1 | item) + (0 + x || site)
outcome ~ Normal(mu, sigma)
end)((;
x=[-1.0, 0.5, 2.0, 0.25],
subject=[2, 1, 2, 3],
item=[2, 1, 1, 2],
site=[10, 10, 20, 20],
outcome=[0.2, 1.1, -0.4, 0.7],
))julia
BRMI:
sigma ~ Exponential(2)
x: data (eltype=Float64, n=4)
subject: data (eltype=Int64, n=4)
item: data (eltype=Int64, n=4)
site: data (eltype=Int64, n=4)
mu ~ 1 + x + ((1 + x) | subject) + (1 | item) + ((0 + x) || site)
outcome ~ Normal(mu, sigma)julia
SBBRMI with data keys = [:item, :item_idx, :n_item, :n_site__nocor__1, :n_subject, :n_terms_mu_subject, :outcome, :site, :site__nocor__1_idx, :subject, :subject_idx, :x]
emitted @slic body:
begin
sigma ~ exponential(1.0 ./ 2)
X_mu = hcat(rep_vector(1.0, num_elements(x)), x)
pop_mu ~ popefs(; X = X_mu)
Z_mu_subject = hcat(rep_vector(1.0, num_elements(x)), x)
r_mu_subject ~ ranef_correlated(; Z = Z_mu_subject, group_idx = subject_idx, n_groups = n_subject, n_terms = n_terms_mu_subject)
r_mu_item ~ ranef_intercept(; group_idx = item_idx, n_groups = n_item)
Z_mu_site__nocor__1 = hcat(x)
r_mu_site__nocor__1 ~ ranef_slope(; Z = Z_mu_site__nocor__1, group_idx = site__nocor__1_idx, n_groups = n_site__nocor__1)
mu = pop_mu + r_mu_subject + r_mu_item + r_mu_site__nocor__1
outcome ~ normal(mu, sigma)
endstan
functions {
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
matrix hcat(vector x) {
int n = dims(x)[1];
return to_matrix(x, n, 1);
}
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;
if((n == 0)) {
vector[n] rv;
return rv;
} else {
return to_vector(normal_rng(a, b));
}
}
}
data {
int x_n;
vector[x_n] x;
int n_terms_mu_subject;
int n_subject;
int subject_idx_n;
array[subject_idx_n] int subject_idx;
int n_item;
int item_idx_n;
array[item_idx_n] int item_idx;
int n_site__nocor__1;
int site__nocor__1_idx_n;
array[site__nocor__1_idx_n] int site__nocor__1_idx;
int outcome_n;
vector[outcome_n] outcome;
}
transformed data {
matrix[x_n, 2] X_mu = hcat(rep_vector(1.0, num_elements(x)), x);
int pop_mu_n_covariates = 2;
matrix[x_n, 2] Z_mu_subject = hcat(rep_vector(1.0, num_elements(x)), x);
matrix[x_n, 1] Z_mu_site__nocor__1 = hcat(x);
}
parameters {
real<lower=0.0> sigma;
vector[pop_mu_n_covariates] pop_mu_beta_pop;
cholesky_factor_corr[n_terms_mu_subject] r_mu_subject_L;
vector<lower=0.0>[n_terms_mu_subject] r_mu_subject_tau;
vector[(n_terms_mu_subject * n_subject)] r_mu_subject_z_flat;
real r_mu_item_log_scale;
vector[n_item] r_mu_item_xi;
vector<lower=0.0>[1] r_mu_site__nocor__1_tau;
vector[n_site__nocor__1] r_mu_site__nocor__1_xi;
}
transformed parameters {
vector[x_n] pop_mu = (X_mu * pop_mu_beta_pop);
matrix[n_terms_mu_subject, n_subject] r_mu_subject_z = to_matrix(r_mu_subject_z_flat, n_terms_mu_subject, n_subject);
matrix[n_subject, n_terms_mu_subject] r_mu_subject_b = ((diag_pre_multiply(r_mu_subject_tau, r_mu_subject_L) * r_mu_subject_z)');
vector[subject_idx_n] r_mu_subject = rows_dot_product(Z_mu_subject, r_mu_subject_b[subject_idx, :]);
vector[item_idx_n] r_mu_item = (exp(r_mu_item_log_scale) * r_mu_item_xi[item_idx]);
vector[x_n] r_mu_site__nocor__1 = (
r_mu_site__nocor__1_tau[1] *
(r_mu_site__nocor__1_xi[site__nocor__1_idx] .* Z_mu_site__nocor__1[:, 1])
);
vector[x_n] mu = (pop_mu + r_mu_subject + r_mu_item + r_mu_site__nocor__1);
}
model {
sigma ~ exponential((1.0 ./ 2));
pop_mu_beta_pop ~ std_normal();
r_mu_subject_L ~ lkj_corr_cholesky(1.0);
r_mu_subject_tau ~ std_normal();
r_mu_subject_z_flat ~ std_normal();
r_mu_item_log_scale ~ std_normal();
r_mu_item_xi ~ std_normal();
r_mu_site__nocor__1_tau ~ std_normal();
r_mu_site__nocor__1_xi ~ std_normal();
outcome ~ normal(mu, sigma);
}
generated quantities {
vector[outcome_n] outcome_likelihood = normal_lpdfs(outcome, mu, sigma);
vector[outcome_n] outcome_gen = normal_vector_rng(outcome_n, mu, sigma);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu, group_matrix_mu_1, group_indices_mu_1, group_levels_mu_1, group_effects_mu_2, group_effects_mu_3)
sigma ~ Distributions.Exponential(2)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_mu = X_mu * beta_pop
group_effect_1 = Base.zeros(Base.length(y))
group_1_1 ~ DynamicPPL.to_submodel(BRM.turing_default_correlated_group(group_matrix_mu_1, group_indices_mu_1, group_levels_mu_1, 1.0))
group_effect_1 = group_effect_1 + group_1_1.effect
group_1_2 ~ DynamicPPL.to_submodel(BRM.turing_group_effect(group_effects_mu_2, (nothing,), nothing))
group_effect_1 = group_effect_1 + group_1_2.effect
group_1_3 ~ DynamicPPL.to_submodel(BRM.turing_group_effect(group_effects_mu_3, (nothing,), nothing))
group_effect_1 = group_effect_1 + group_1_3.effect
eta_mu = eta_mu + group_effect_1
mu = eta_mu
begin
for i = Base.eachindex(y)
y[i] ~ Distributions.Normal(mu[i], sigma)
end
end
(; mu = mu, sigma = sigma, response = y)
end)Interval-censored evidence
brm-comparison
Interval-censored Normal evidencejulia
interval_normal = (@brm begin
sigma ~ Exponential(2)
mu ~ 1 + x
y_lower ~ interval_censored(Normal(mu, sigma); upper=y_upper)
end)((;
x=[-1.0, 0.0, 1.0],
y_lower=[-0.4, 0.1, 0.8],
y_upper=[-0.1, 0.4, 1.2],
))julia
BRMI:
sigma ~ Exponential(2)
x: data (eltype=Float64, n=3)
mu ~ 1 + x
y_upper: data (eltype=Float64, n=3)
y_lower ~ interval_censored(Normal(mu, sigma); upper=y_upper)julia
SBBRMI with data keys = [:x, :y_lower, :y_upper]
emitted @slic body:
begin
sigma ~ exponential(1.0 ./ 2)
X_mu = hcat(rep_vector(1.0, num_elements(x)), x)
pop_mu ~ popefs(; X = X_mu)
mu = pop_mu
y_lower ~ interval_censored(normal, y_lower, y_upper, mu, sigma)
endstan
functions {
matrix hcat(
vector x,
vector y
) {
int n = dims(x)[1];
if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
return append_col(x, y);
}
real interval_evidence_impl_normal_lpdf(
vector y,
vector lo,
vector hi,
vector args1,
real args2
) {
return sum(interval_evidence_impl_normal_lpdfs(y, lo, hi, args1, args2));
}
vector interval_evidence_impl_normal_lpdfs(
vector y,
vector lo,
vector hi,
vector args1,
real args2
) {
int n = dims(y)[1];
return jbroadcasted_interval_evidence_impl_lpdf_normal(y, lo, hi, args1, args2);
}
vector jbroadcasted_interval_evidence_impl_lpdf_normal(
vector x1,
vector x3,
vector x4,
vector x5,
real x6
) {
int n = dims(x1)[1];
vector[n] rv;
for(i in 1:n) {
rv[i] = interval_evidence_impl_normal_lpdf(broadcasted_getindex(x1, i) |
broadcasted_getindex(x3, i),
broadcasted_getindex(x4, i),
broadcasted_getindex(x5, i),
x6
);
}
return rv;
}
real interval_evidence_impl_normal_lpdf(
real y,
real lo,
real hi,
real args1,
real args2
) {
array[1] real rv;
if((lo >= hi)) {
reject("interval_censored: lower bound must be less than upper bound");
}
rv[1] = log_diff_exp(normal_lcdf_stable(hi, args1, args2), normal_lcdf_stable(lo, 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 interval_evidence_impl_vector_normal_rng(
int anontok__1,
vector lo,
vector hi,
vector args1,
real args2
) {
int n = anontok__1;
return normal_vector_rng(n, args1, args2);
}
vector normal_vector_rng(
int anontok__1,
vector a,
real b
) {
int n = anontok__1;
if((n == 0)) {
vector[n] rv;
return rv;
} else {
return to_vector(normal_rng(a, b));
}
}
}
data {
int x_n;
vector[x_n] x;
int y_lower_n;
vector[y_lower_n] y_lower;
int y_upper_n;
vector[y_upper_n] y_upper;
}
transformed data {
matrix[x_n, 2] X_mu = hcat(rep_vector(1.0, num_elements(x)), x);
int pop_mu_n_covariates = 2;
}
parameters {
real<lower=0.0> sigma;
vector[pop_mu_n_covariates] pop_mu_beta_pop;
}
transformed parameters {
vector[x_n] pop_mu = (X_mu * pop_mu_beta_pop);
vector[x_n] mu = pop_mu;
}
model {
sigma ~ exponential((1.0 ./ 2));
pop_mu_beta_pop ~ std_normal();
y_lower ~ interval_evidence_impl_normal(y_lower, y_upper, mu, sigma);
}
generated quantities {
vector[y_lower_n] y_lower_likelihood = interval_evidence_impl_normal_lpdfs(y_lower, y_lower, y_upper, mu, sigma);
vector[y_lower_n] y_lower_gen = interval_evidence_impl_vector_normal_rng(y_lower_n, y_lower, y_upper, mu, sigma);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu, upper_y_lower)
sigma ~ Distributions.Exponential(2)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 2))
eta_mu = X_mu * beta_pop
mu = eta_mu
begin
for i = Base.eachindex(y)
y[i] ~ BayesianRegressionModelsTuringExt._brm_interval_evidence(Distributions.Normal(mu[i], sigma), upper_y_lower[i])
end
end
(; mu = mu, sigma = sigma, response = y)
end)Spline term
This smooth example is supported by both backends. The Turing pane shows the native spline submodel selected directly from the shared BRMI preparation.
brm-comparison
Spline termjulia
smooth = (@brm begin
sigma ~ Exponential(1)
mu ~ 1 + s(x)
y ~ Normal(mu, sigma)
end)((;
x=collect(range(-2, 2; length=20)),
y=sin.(collect(range(-2, 2; length=20))),
))julia
BRMI:
sigma ~ Exponential(1)
x: data (eltype=Float64, n=20)
mu ~ 1 + s(x)
y ~ Normal(mu, sigma)julia
SBBRMI with data keys = [:Xnull_x, :Zpen_x, :x, :y]
emitted @slic body:
begin
sigma ~ exponential(1.0 ./ 1)
X_mu = hcat(rep_vector(1.0, num_elements(x)))
pop_mu ~ popefs(; X = X_mu)
s_x ~ BayesianRegressionModels._sb_s_generic(; Xnull = Xnull_x, Zpen = Zpen_x)
mu = pop_mu + s_x
y ~ normal(mu, sigma)
endstan
functions {
matrix hcat(vector x) {
int n = dims(x)[1];
return to_matrix(x, n, 1);
}
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;
if((n == 0)) {
vector[n] rv;
return rv;
} else {
return to_vector(normal_rng(a, b));
}
}
}
data {
int x_n;
vector[x_n] x;
int Zpen_x_n;
int Xnull_x_m;
int Xnull_x_n;
matrix[Xnull_x_m, Xnull_x_n] Xnull_x;
int Zpen_x_m;
matrix[Zpen_x_m, Zpen_x_n] Zpen_x;
int y_n;
vector[y_n] y;
}
transformed data {
matrix[num_elements(x), 1] X_mu = hcat(rep_vector(1.0, num_elements(x)));
int pop_mu_n_covariates = 1;
int s_x_n_pen = Zpen_x_n;
}
parameters {
real<lower=0.0> sigma;
vector[pop_mu_n_covariates] pop_mu_beta_pop;
vector[2] s_x_b_fixed;
vector<lower=0.0>[1] s_x_sd_pen;
vector[s_x_n_pen] s_x_b_pen_raw;
}
transformed parameters {
vector[num_elements(x)] pop_mu = (X_mu * pop_mu_beta_pop);
vector[s_x_n_pen] s_x_b_pen = (s_x_sd_pen[1] * s_x_b_pen_raw);
vector[Xnull_x_m] s_x = ((Xnull_x * s_x_b_fixed) + (Zpen_x * s_x_b_pen));
vector[num_elements(x)] mu = (pop_mu + s_x);
}
model {
sigma ~ exponential((1.0 ./ 1));
pop_mu_beta_pop ~ std_normal();
s_x_sd_pen ~ std_normal();
s_x_b_pen_raw ~ std_normal();
y ~ normal(mu, sigma);
}
generated quantities {
vector[y_n] y_likelihood = normal_lpdfs(y, mu, sigma);
vector[y_n] y_gen = normal_vector_rng(y_n, mu, sigma);
}julia
#= line 0 =# Turing.@model(function brm_model(y, X_mu, terms_mu_1, x)
sigma ~ Distributions.Exponential(1)
beta_pop ~ Distributions.product_distribution(Base.fill(Distributions.Normal(), 1))
eta_mu = X_mu * beta_pop
term_mu_1 ~ DynamicPPL.to_submodel(BRM.turing_term_model(terms_mu_1, Base.length(y), NamedTuple{$(QuoteNode((:sd,)))}((Distributions.Normal(0.0, 1.0),)), NamedTuple{$(QuoteNode((:x,)))}((x,))))
eta_mu = eta_mu + term_mu_1.effect
mu = eta_mu
begin
for i = Base.eachindex(y)
y[i] ~ Distributions.Normal(mu[i], sigma)
end
end
(; mu = mu, sigma = sigma, response = y)
end)