neal.jl

#=
Neal's funnel
https://mc-stan.org/docs/2_29/stan-users-guide/reparameterization.html#ref-papa-et-al:2007
=#
using Turing, StatsPlots, Distributions, DataFrames
gr() # set backend



@model function Neal()
    y ~ Normal(0,3)
    x ~ arraydist([Normal(0, exp(y/2)) for i in 1:9])
end

simple_chain = sample(Neal(), NUTS(), 10_000)

p = plot(layout = 2);
scatter!(simple_chain["x[1]"], simple_chain[:y], color = :red, opacity = 0.3, subplot = 1);
title!("Original Parametrization", subplot = 1)

################################################################################### REPARAM

@model function Neal2()
    # raw draws
    y_raw ~ Normal(0,1)
    x_raw ~ arraydist([Normal(0, 1) for i in 1:9])

    # transform:
    y = 3*y_raw
    x = exp.(y./2) .* x_raw

    # return:
    return [x; y]
end

raw_chain = Turing.MCMCChains.get_sections(sample(Neal2(), NUTS(), 10_000), :parameters)
reparam_chain = reduce(hcat, generated_quantities(Neal2(), raw_chain))

scatter!(reparam_chain[1, 1:end], reparam_chain[10, 1:end], color = :blue, opacity = 0.3, subplot= 2);
title!("Reparametrized", subplot = 2)