using Turing
using DataFrames
using CSV
using Random
using StatisticalRethinking
using StatisticalRethinking: link
using StatisticalRethinkingPlots
using ParetoSmooth
using StatsPlots
using StatsBase
using FreqTables
using Logging
default(label=false)
Logging.disable_logging(Logging.Warn);
Code 11.1
d = DataFrame(CSV.File("data/chimpanzees.csv", delim=';'));
Code 11.2
d[!,:treatment] = 1 .+ d.prosoc_left .+ 2*d.condition;
Code 11.3
println(freqtable(d, :treatment, :prosoc_left, subset=d.condition .== 0))
println(freqtable(d, :treatment, :prosoc_left, subset=d.condition .== 1))
Code 11.4
@model function m11_1(pulled_left)
a ~ Normal(0, 10)
p = logistic(a) # inverse of the `logit` function
pulled_left ~ Binomial(1, p)
end
Code 11.5
Random.seed!(1999)
prior_chain = sample(m11_1(d.pulled_left), Prior(), 10000)
prior = DataFrame(prior_chain);
Code 11.6
p = logistic.(prior.a)
density(p, bandwidth=0.01)
Code 11.7
@model function m11_2(pulled_left, treatment)
a ~ Normal(0, 1.5)
treat_levels = length(levels(d.treatment))
b ~ MvNormal(zeros(treat_levels), 10)
p = @. logistic(a + b[treatment])
for i ∈ eachindex(pulled_left)
pulled_left[i] ~ Binomial(1, p[i])
end
end
Random.seed!(1999)
prior_chain = sample(m11_2(d.pulled_left, d.treatment), Prior(), 10000)
prior = DataFrame(prior_chain);
f = i -> @. logistic(prior.a + prior[!,"b[$i]"])
p = map(f, 1:4);
Code 11.8
density(abs.(p[1] .- p[2]), bandwidth=0.01)
Code 11.9
@model function m11_3(pulled_left, treatment)
a ~ Normal(0, 1.5)
treat_levels = length(levels(treatment))
b ~ MvNormal(zeros(treat_levels), 0.5)
p = @. logistic(a + b[treatment])
for i ∈ eachindex(pulled_left)
pulled_left[i] ~ Binomial(1, p[i])
end
end
Random.seed!(1999)
prior_chain = sample(m11_3(d.pulled_left, d.treatment), Prior(), 10000)
prior = DataFrame(prior_chain);
f = i -> @. logistic(prior.a + prior[!,"b[$i]"])
p = map(f, 1:4);
mean(abs.(p[1] .- p[2]))
Code 11.10
dat_list = d[!,[:pulled_left, :actor, :treatment]];
Code 11.11
@model function m11_4(actor, treatment, pulled_left)
act_levels = length(levels(actor))
a ~ MvNormal(zeros(act_levels), 1.5)
treat_levels = length(levels(treatment))
b ~ MvNormal(zeros(treat_levels), 0.5)
p = @. logistic(a[actor] + b[treatment])
for i ∈ eachindex(pulled_left)
pulled_left[i] ~ Binomial(1, p[i])
end
end
m11_4_chain = sample(m11_4(dat_list.actor, dat_list.treatment, dat_list.pulled_left), NUTS(), 1000)
m11_4_df = DataFrame(m11_4_chain)
precis(m11_4_df)
Code 11.12
p_left = DataFrame(map(i -> "V$i" => logistic.(m11_4_df[:,"a[$i]"]), 1:7)...);
coeftab_plot(p_left)
Code 11.13
names = ["R/N", "L/N", "R/P", "L/P"]
labs = DataFrame(map(i -> names[i] => m11_4_df[:,"b[$i]"], 1:4)...)
coeftab_plot(labs)
Code 11.14
diffs = DataFrame(
db13=m11_4_df."b[1]" .- m11_4_df."b[3]",
db24=m11_4_df."b[2]" .- m11_4_df."b[4]",
)
coeftab_plot(diffs)
Code 11.15
gd = groupby(d, [:actor, :treatment])
c = combine(gd, :pulled_left => mean => :val)
pl = unstack(c, :actor, :treatment, :val)
pl[1,:]
Code 11.16
names = ["R/N", "L/N", "R/P", "L/P"]
p = plot(ylims=(0, 1.1), ylab="proportion left lever", showaxis=:y, xticks=false)
hline!([0.5], c=:gray, s=:dash)
for actor in 1:7
ofs = (actor-1)*4
actor > 1 && vline!([ofs+0.5], c=:gray)
plot!([ofs+1,ofs+3], collect(pl[actor,["1","3"]]), lw=2, m=:o, c=:blue)
plot!([ofs+2,ofs+4], collect(pl[actor,["2","4"]]), lw=2, m=:o, c=:blue)
anns = [
(ofs+idx, pl[actor,string(idx)]+.04, (name, 8))
for (idx,name) ∈ enumerate(names)
]
actor != 2 && annotate!(anns)
end
annotate!([
(2.5 + (idx-1)*4, 1.1, ("actor $idx", 8))
for idx ∈ 1:7
])
p
Code 11.17
l_fun = (r, (ai, bi)) -> logistic(get(r, "a[$ai]", missing) + get(r, "b[$bi]", missing))
p_post = link(m11_4_df, l_fun, Iterators.product(1:7, 1:4))
p_μ = map(mean, p_post)
p_ci = map(PI, p_post);
# visualize mean and cred intervals of posterior distribution
# compute relative intervals
rel_ci = map(idx -> (p_μ[idx]-p_ci[idx][1], p_ci[idx][2]-p_μ[idx]), CartesianIndices(p_ci))
names = ["R/N", "L/N", "R/P", "L/P"]
p = plot(ylims=(0, 1.1), ylab="proportion left lever", title="Posterior", showaxis=:y, xticks=false)
hline!([0.5], c=:gray, s=:dash)
for actor in 1:7
ofs = (actor-1)*4
actor > 1 && vline!([ofs+0.5], c=:gray)
err = [rel_ci[actor,1], rel_ci[actor,3]]
plot!([ofs+1,ofs+3], collect(p_μ[actor,[1,3]]), err=err, lw=2, m=:o, c=:blue)
err = [rel_ci[actor,2], rel_ci[actor,4]]
plot!([ofs+2,ofs+4], collect(p_μ[actor,[2,4]]), err=err, lw=2, m=:o, c=:blue)
anns = [
(ofs+idx, p_μ[actor,idx]+.04, (name, 8))
for (idx,name) ∈ enumerate(names)
]
actor != 2 && annotate!(anns)
end
annotate!([
(2.5 + (idx-1)*4, 1.1, ("actor $idx", 8))
for idx ∈ 1:7
])
p
Code 11.18
d.side = d.prosoc_left .+ 1
d.cond = d.condition .+ 1;
Code 11.19
@model function m11_5(actor, side, cond, pulled_left)
act_levels = length(levels(actor))
a ~ MvNormal(zeros(act_levels), 1.5)
side_levels = length(levels(side))
bs ~ MvNormal(zeros(side_levels), 0.5)
cond_levels = length(levels(cond))
bc ~ MvNormal(zeros(cond_levels), 0.5)
p = @. logistic(a[actor] + bs[side] + bc[cond])
for i ∈ eachindex(pulled_left)
pulled_left[i] ~ Binomial(1, p[i])
end
end
m11_5_chain = sample(m11_5(d.actor, d.side, d.cond, d.pulled_left), NUTS(), 1000)
m11_5_df = DataFrame(m11_5_chain);
Code 11.20
l_fun = (r, (ai, bi, pull_left)) -> begin
p = logistic(get(r, "a[$ai]", 0) + get(r, "b[$bi]", 0))
binomlogpdf(1, p, pull_left)
end
m11_4_ll = link(m11_4_df, l_fun, zip(d.actor, d.treatment, d.pulled_left))
m11_4_ll = hcat(m11_4_ll...);
l_fun = (r, (ai, si, ci, pull_left)) -> begin
p = logistic(get(r, "a[$ai]", 0) + get(r, "bs[$si]", 0) + get(r, "bc[$ci]", 0))
binomlogpdf(1, p, pull_left)
end
m11_5_ll = link(m11_5_df, l_fun, zip(d.actor, d.side, d.cond, d.pulled_left))
m11_5_ll = hcat(m11_5_ll...);
compare([m11_5_ll, m11_4_ll], :psis, mnames=["m5", "m4"])
Code 11.21 and 11.22 were omitted, as they are stan-specific
Code 11.23
mean(@. exp(m11_4_df."b[4]" - m11_4_df."b[2]"))
Code 11.24
gb = groupby(d, [:treatment, :actor, :side, :cond])
d_aggregated = combine(gb, :pulled_left => sum => :left_pulls)
first(d_aggregated, 8)
Code 11.25
@model function m11_6(actor, treatment, left_pulls)
act_levels = length(levels(actor))
a ~ MvNormal(zeros(act_levels), 1.5)
treat_levels = length(levels(treatment))
b ~ MvNormal(zeros(treat_levels), 0.5)
p = @. logistic(a[actor] + b[treatment])
for i ∈ eachindex(left_pulls)
left_pulls[i] ~ Binomial(18, p[i])
end
end
m11_6_chain = sample(m11_6(d_aggregated.actor, d_aggregated.treatment, d_aggregated.left_pulls), NUTS(), 1000)
m11_6_df = DataFrame(m11_6_chain);
Code 11.26
l_fun = (r, (ai, bi, left_pulls)) -> begin
p = logistic(get(r, "a[$ai]", 0) + get(r, "b[$bi]", 0))
binomlogpdf(18, p, left_pulls)
end
m11_6_ll = link(m11_6_df, l_fun, zip(d_aggregated.actor, d_aggregated.treatment, d_aggregated.left_pulls))
m11_6_ll = hcat(m11_6_ll...)
try
compare([m11_6_ll, m11_4_ll], :psis, mnames=["m6", "m4"])
catch e
println(e)
end
Code 11.27
(
-2*binomlogpdf(9, 0.2, 6),
-2*sum(binomlogpdf.(1, 0.2, [1,1,1,1,1,1,0,0,0]))
)
Code 11.28
d = DataFrame(CSV.File("data/UCBadmit.csv", skipto=2,
header=[:id, :dept, :gender, :admit, :reject, :applications]))
Code 11.29
dat = d[!, [:admit, :applications]]
dat.gid = @. ifelse(d.gender == "male", 1, 2)
@model function m11_7(admit, applications, gid)
a ~ MvNormal([0, 0], 1.5)
p = @. logistic(a[gid])
for i ∈ eachindex(applications)
admit[i] ~ Binomial(applications[i], p[i])
end
end
m11_7_chain = sample(m11_7(dat.admit, dat.applications, dat.gid), NUTS(), 1000)
m11_7_df = DataFrame(m11_7_chain)
precis(m11_7_df)
Code 11.30
diff_a = m11_7_df."a[1]" .- m11_7_df."a[2]"
diff_p = @. logistic(m11_7_df."a[1]") - logistic(m11_7_df."a[2]")
precis(DataFrame(diff_a=diff_a, diff_p=diff_p))
Code 11.31
Random.seed!(1)
fun = (r, (apps, gid)) -> begin
p = logistic(get(r, "a[$gid]", 0))
rand(Binomial(apps, p))
end
admit_pred = link(m11_7_df, fun, zip(dat.applications, dat.gid))
admit_rate = @. admit_pred ./ dat.applications
# plot
xrange = 1:12
p = plot(yrange=(0, 1), title="Posterior validation check", xlab="case", ylab="admit")
scatter!(xrange, mean.(admit_rate), yerr=std.(admit_rate))
pi_rate = PI.(admit_rate)
scatter!(xrange, first.(pi_rate), shape=:cross, c=:black)
scatter!(xrange, last.(pi_rate), shape=:cross, c=:black)
dat_rates = dat.admit ./ dat.applications
for idx in 1:2:12
r = [dat_rates[idx], dat_rates[idx+1]]
plot!([idx, idx+1], r, mark=:o, lw=2, c=:blue)
annotate!([(idx+0.5, mean(r)+0.03, (d.dept[idx], 8))])
end
p
Code 11.32
dat.dept_id = reshape(hcat(1:6, 1:6)', 12)
@model function m11_8(admit, applications, gid, dept_id)
a ~ MvNormal([0, 0], 1.5)
delta ~ MvNormal(zeros(6), 1.5)
p = @. logistic(a[gid] + delta[dept_id])
for i ∈ eachindex(applications)
admit[i] ~ Binomial(applications[i], p[i])
end
end
m11_8_chain = sample(m11_8(dat.admit, dat.applications, dat.gid, dat.dept_id), NUTS(), 1000)
m11_8_df = DataFrame(m11_8_chain)
precis(m11_8_df)
Code 11.33
diff_a = m11_8_df."a[1]" .- m11_8_df."a[2]"
diff_p = logistic.(m11_8_df."a[1]") .- logistic.(m11_8_df."a[2]")
precis(DataFrame(diff_a=diff_a, diff_p=diff_p))
Code 11.34
gb = groupby(d, :dept)
fun = (as,gs) -> [(gender=g, a_ratio=a/sum(as)) for (a,g) in zip(as,gs)]
c = combine(gb, [:applications, :gender] => fun => AsTable)
unstack(c, :gender, :dept, :a_ratio)
Code 11.35
Random.seed!(1)
y = rand(Binomial(1000, 1/1000), 10^5)
mean(y), var(y)
Code 11.36
d = DataFrame(CSV.File("data/Kline.csv"))
Code 11.37
d.P = standardize(ZScoreTransform, log.(d.population))
d.contact_id = ifelse.(d.contact .== "high", 2, 1);
Code 11.38
p = plot(xlims=(-1, 100), ylims=(0, 0.08),
xlab="mean number of tools", ylab="Density")
r = 0:0.5:100
plot!(r, LogNormal(0, 10), label=L"a \sim \ln \mathcal{N}(0, 10)")
plot!(r, LogNormal(3, 0.5), label=L"a \sim \ln \mathcal{N}(3, 0.5)")
Code 11.39
a = rand(Normal(0, 10), 10^4)
λ = exp.(a)
mean(λ)
Code 11.40 - joined with 11.38
Code 11.41
Random.seed!(1)
N = 100
a = rand(Normal(3, 0.5), N)
b = rand(Normal(0, 10), N)
p = plot(xlims=(-2, 2), ylims=(0, 100))
r = -2:0.05:2
for (aᵢ, bᵢ) ∈ zip(a, b)
plot!(r, x -> exp(aᵢ + bᵢ * x), c=:gray)
end
p
Code 11.42
Random.seed!(1)
N = 100
a = rand(Normal(3, 0.5), N)
b = rand(Normal(0, .2), N)
p = plot(xlims=(-2, 2), ylims=(0, 100))
r = -2:0.05:2
for (aᵢ, bᵢ) ∈ zip(a, b)
plot!(r, x -> exp(aᵢ + bᵢ * x), c=:gray)
end
p
Code 11.43
x_seq = range(log(100), log(200_000), length=100)
λ = map(x -> (@. exp(a + b * x)), x_seq)
λ = hcat(λ...)
p = plot(xlims=extrema(x_seq), ylims=(0, 500), xlab="log population", ylab="total tools")
for λᵢ ∈ eachrow(λ)
plot!(x_seq, λᵢ, c=:gray)
end
p
Code 11.44
p = plot(xlims=extrema(exp.(x_seq)), ylims=(0, 500),
xlab="population", ylab="total tools")
for λᵢ ∈ eachrow(λ)
plot!(exp.(x_seq), λᵢ, c=:gray)
end
p
Code 11.45
# intercept only
@model function m11_9(T)
a ~ Normal(3, 0.5)
λ = exp(a)
T ~ Poisson(λ)
end
m11_9_ch = sample(m11_9(d.total_tools), NUTS(), 1000)
m11_9_df = DataFrame(m11_9_ch)
# interaction model
@model function m11_10(T, cid, P)
a ~ MvNormal([3, 3], 0.5)
b ~ MvNormal([0, 0], 0.2)
λ = @. exp(a[cid] + b[cid]*P)
for i ∈ eachindex(T)
T[i] ~ Poisson(λ[i])
end
end
m11_10_ch = sample(m11_10(d.total_tools, d.contact_id, d.P), NUTS(), 1000)
m11_10_df = DataFrame(m11_10_ch);
Code 11.46
f = (r, T) -> poislogpdf(exp(r.a), T)
ll_m9 = link(m11_9_df, f, d.total_tools)
ll_m9 = hcat(ll_m9...);
f = (r, (T, cid, P)) -> poislogpdf(exp(get(r, "a[$cid]", 0) + get(r, "b[$cid]", 0)*P), T)
ll_m10 = link(m11_10_df, f, zip(d.total_tools, d.contact_id, d.P))
ll_m10 = hcat(ll_m10...);
compare([ll_m9, ll_m10], :psis, mnames=["m9", "m10"])
Code 11.47
# get psis K values
t = ll_m10'
m10_t = collect(reshape(t, size(t)..., 1))
PSIS_m10 = psis_loo(m10_t)
k = PSIS_m10.pointwise(:pareto_k);
# plot
mark_size = standardize(ZScoreTransform, k).+5
mark_color = ifelse.(d.contact_id .== 1, :white, :blue)
scatter(d.P, d.total_tools, xlab="log population (std)", ylab="total tools", ylim=(0, 75),
ms=mark_size, mc=mark_color, msw=2)
ns = 100
P_seq = range(-1.4, 3, length=ns)
# cid=1 predictions
f = (r,p) -> exp(r."a[1]" + r."b[1]"*p)
λ = link(m11_10_df, f, P_seq)
λ = hcat(λ...);
λ1_mean = mean.(eachcol(λ))
λ1_ci = PI.(eachcol(λ))
λ1_ci = vcat(λ1_ci'...)
plot!(P_seq, [λ1_mean λ1_mean], fillrange=λ1_ci, fillalpha=0.2, ls=:dash, c=:black, lw=1.5)
# cid=2 predictions
f = (r,p) -> exp(r."a[2]" + r."b[2]"*p)
λ = link(m11_10_df, f, P_seq)
λ = hcat(λ...);
λ2_mean = mean.(eachcol(λ))
λ2_ci = PI.(eachcol(λ))
λ2_ci = vcat(λ2_ci'...)
plot!(P_seq, [λ2_mean λ2_mean], fillrange=λ2_ci, fillalpha=0.2, c=:black, lw=1.5)
Code 11.48
scatter(d.population, d.total_tools, xlab="population", ylab="total tools", ylim=(0, 75),
ms=mark_size, mc=mark_color, msw=2)
P_seq = range(-5, 3, length=ns)
# 1.53 is sd of log(population)
# 9 is mean of log(population)
pop_seq = @. exp(P_seq * 1.53 + 9)
plot!(pop_seq, [λ1_mean λ1_mean], fillrange=λ1_ci, fillalpha=0.2, ls=:dash, c=:black, lw=1.5)
plot!(pop_seq, [λ2_mean λ2_mean], fillrange=λ2_ci, fillalpha=0.2, c=:black, lw=1.5)
Code 11.49
@model function m11_11(T, P, cid)
a ~ MvNormal([1,1], 1)
b₁ ~ Exponential(1)
b₂ ~ Exponential(1)
b = [b₁, b₂]
g ~ Exponential(1)
λ = @. exp(a[cid]) * P ^ b[cid] / g
for i ∈ eachindex(T)
T[i] ~ Poisson(λ[i])
end
end
m11_11_ch = sample(m11_11(d.total_tools, d.population, d.contact_id), NUTS(), 1000)
m11_11_df = DataFrame(m11_11_ch)
precis(m11_11_df)
Code 11.50
Random.seed!(1)
num_days = 30
y = rand(Poisson(1.5), num_days);
Code 11.51
Random.seed!(2)
num_weeks = 4
y_new = rand(Poisson(0.5*7), num_weeks);
Code 11.52
y_all = [y; y_new]
exposure = [repeat([1], 30); repeat([7], 4)]
monastery = [repeat([0], 30); repeat([1], 4)]
d = DataFrame(y=y_all, days=exposure, monastery=monastery);
Code 11.53
Random.seed!(1)
d.log_days = log.(d.days)
@model function m11_12(y, log_days, monastery)
a ~ Normal()
b ~ Normal()
λ = @. exp(log_days + a + b*monastery)
@. y ~ Poisson(λ)
end
m11_12_chain = sample(m11_12(d.y, d.log_days, d.monastery), NUTS(), 1000)
m11_12_df = DataFrame(m11_12_chain);
Code 11.54
Note that values are slightly different from the book. This is due to non-Normal distributions (you can check this yourself with optimize)
λ_old = exp.(m11_12_df.a)
λ_new = @. exp(m11_12_df.a + m11_12_df.b)
precis(DataFrame(λ_old=λ_old, λ_new=λ_new))
Code 11.55
# simulate career choice among 500 individuals
N = 500
income = [1,2,5]
c_score = 0.5*income
p = softmax(c_score)
# simulate choice
Random.seed!(34302)
w = Weights(p)
career = [sample(w) for _ in 1:N];
Code 11.56
@model function m11_13(career, income)
a ~ MvNormal([0, 0], 1)
b ~ TruncatedNormal(0, 0.5, 0, Inf)
p = softmax([a[1] + b*income[1], a[2] + b*income[2], 0])
career ~ Categorical(p)
end
Code 11.57
Random.seed!(121)
m11_13_chain = sample(m11_13(career, income), NUTS(), 5000, n_chains=4)
m11_13_df = DataFrame(m11_13_chain)
precis(m11_13_df)
Code 11.58
# logit scores
s1 = m11_13_df."a[1]" + m11_13_df.b * income[1]
s2_orig = m11_13_df."a[2]" + m11_13_df.b * income[2]
s2_new = m11_13_df."a[2]" + m11_13_df.b * income[2] * 2
# compute probabilities for original and counterfactual
p_orig = softmax.(eachrow(hcat(s1, s2_orig, zeros(length(s1)))))
p_orig = hcat(p_orig...)'
p_new = softmax.(eachrow(hcat(s1, s2_new, zeros(length(s1)))))
p_new = hcat(p_new...)'
p_diff = p_new[:, 2] - p_orig[:, 2]
precis(DataFrame(p_diff=p_diff))
Code 11.59
Random.seed!(1)
N = 500
family_income = rand(Uniform(), N)
b = [-2, 0, 2]
career = [
begin
score = (0.5 * 1:3) + b * inc
p = softmax(score)
sample(Weights(p))
end
for inc in family_income
]
@model function m11_14(career, family_income)
a ~ MvNormal([0, 0], 1.5)
b ~ MvNormal([0, 0], 1)
for i ∈ eachindex(career)
income = family_income[i]
s = [a[1] + b[1] * income, a[2] + b[2] * income, 0]
p = softmax(s)
career[i] ~ Categorical(p)
end
end
m11_14_chain = sample(m11_14(career, family_income), NUTS(), 1000)
m11_14_df = DataFrame(m11_14_chain)
precis(m11_14_df)
Code 11.60
d = DataFrame(CSV.File("data/UCBadmit.csv", skipto=2,
header=[:id, :dept, :gender, :admit, :reject, :applications]))
Code 11.61
@model function m_binom(admit, applications)
a ~ Normal(0, 1.5)
p = logistic(a)
@. admit ~ Binomial(applications, p)
end
m_binom_chain = sample(m_binom(d.admit, d.applications), NUTS(), 1000)
m_binom_df = DataFrame(m_binom_chain);
@model function m_pois(admit, reject)
a ~ MvNormal([0, 0], 1.5)
λ = exp.(a)
admit ~ Poisson(λ[1])
reject ~ Poisson(λ[2])
end
m_pois_chain = sample(m_pois(d.admit, d.reject), NUTS(), 1000)
m_pois_df = DataFrame(m_pois_chain);
Code 11.62
m_binom_df.a .|> logistic |> mean
Code 11.63
a1 = m_pois_df."a[1]" |> mean
a2 = m_pois_df."a[2]" |> mean
exp(a1)/(exp(a1)+exp(a2))