Skip to content

Instantly share code, notes, and snippets.

@cpfiffer
Created August 3, 2019 17:27
Show Gist options
  • Save cpfiffer/70734e617c35bd2e2b68da0b5da4aae2 to your computer and use it in GitHub Desktop.
Save cpfiffer/70734e617c35bd2e2b68da0b5da4aae2 to your computer and use it in GitHub Desktop.
Test ESS draws
using CmdStan, DynamicHMC
using StatsPlots, Random, MCMCDiagnostics
using Revise
using Turing, AdvancedHMC; const AHMC = AdvancedHMC
Random.seed!(1239911)
ProjDir = @__DIR__
cd(ProjDir)
Nsamples = 2000
Nadapt = 1000
N = 5
normstanmodel = begin
"""
data {
int<lower=0> N;
vector[N] y;
}
parameters {
real mu;
real<lower=0> sigma;
}
model {
mu ~ normal(0,1);
sigma ~ cauchy(0,5);
y ~ normal(mu,sigma);
}
"""
end
@model model(y) = begin
μ ~ Normal(0,1)
σ ~ Truncated(Cauchy(0,5),0,Inf)
for n = 1:length(y)
y[n] ~ Normal(μ,σ)
end
end
stanmodel = Stanmodel(
name = "normstanmodel", model = normstanmodel, nchains = 4,
Sample(num_samples = Nsamples-Nadapt,
num_warmup = Nadapt,
# adapt = CmdStan.Adapt(engaged=false),
adapt = CmdStan.Adapt(delta=0.8),
save_warmup = false,
algorithm=CmdStan.Hmc(;
engine=CmdStan.Nuts(),
# engine=CmdStan.Static(; int_time=0.1),
# metric=CmdStan.unit_e,
# metric=CmdStan.diag_e,
metric=CmdStan.dense_e,
stepsize=1.0,
# stepsize=0.01,
stepsize_jitter=0.0)
)
);
initOutput() = DataFrame(μ=Float64[],σ=Float64[])
# Collect ess dfs
ess_array_cmdstan = initOutput()
ess_array_MCMCChains_cmdstan = initOutput()
ess_array_MCMCDiagnostics_cmdstan = initOutput()
ess_array_MCMCChains_turing = initOutput()
ess_array_MCMCDiagnostics_turing = initOutput()
# Collect rhat dfs
rhat_array_cmdstan = initOutput()
rhat_array_MCMCChains_turing = initOutput()
rhat_array_MCMCChains_cmdstan = initOutput()
rhat_array_MCMCDiagnostics_turing = initOutput()
rhat_array_MCMCDiagnostics_cmdstan = initOutput()
# Collect stepsize dfs
ϵ_array_cmdstan = initOutput()
ϵ_array_turing = initOutput()
N_runs = 50
for i in 1:N_runs
println("\n Loop $i\n")
data = Dict("y" => rand(Normal(0,1),N), "N" => N)
chn = mapreduce(x->sample(model(data["y"]),
Turing.NUTS(
Nsamples, Nadapt, 0.8;
max_depth = 10,
init_ϵ = 1.0,
# metricT=AHMC.UnitEuclideanMetric,
# metricT=AHMC.DiagEuclideanMetric,
metricT=AHMC.DenseEuclideanMetric,
),
# Turing.HMCDA(
# Nsamples, Nadapt, 0.8, 0.1;
# init_ϵ=1.0,
# # metricT=AHMC.UnitEuclideanMetric,
# metricT=AHMC.DenseEuclideanMetric,
# ),
# Turing.HMC(
# Nsamples, 0.01, 10
# ),
# Turing.DynamicNUTS(
# Nsamples - Nadapt
# ),
),
chainscat, 1:4)
dft = describe(chn)[1]
push!(ess_array_MCMCChains_turing, dft[:ess])
push!(rhat_array_MCMCChains_turing, dft[:r_hat])
global rc, chns, cnames = stan(stanmodel,data, summary=true, ProjDir)
dfc = describe(chns)[1]
push!(ess_array_MCMCChains_cmdstan, dfc[:ess])
push!(rhat_array_MCMCChains_cmdstan, dfc[:r_hat])
summary_df = read_summary(stanmodel)
push!(ess_array_cmdstan, summary_df[[:mu, :sigma], :ess])
push!(rhat_array_cmdstan, summary_df[[:mu, :sigma], :r_hat])
ac = DataFrame(chns);
acs = DataFrame(chns, append_chains=false)
at = DataFrame(chn);
push!(ess_array_MCMCDiagnostics_cmdstan,
[effective_sample_size(ac[:, :mu]),
effective_sample_size(ac[:, :sigma])])
push!(ess_array_MCMCDiagnostics_turing,
[effective_sample_size(at[:, :μ]),
effective_sample_size(at[:, :σ])])
acs_mu = hcat([acs[i][:mu] for i in 1:4])
acs_sigma = hcat([acs[i][:sigma] for i in 1:4])
push!(rhat_array_MCMCDiagnostics_cmdstan,
[potential_scale_reduction(acs_mu...),
potential_scale_reduction(acs_sigma...)])
push!(ϵ_array_cmdstan,
[summary_df[:stepsize__, :mean][1],
summary_df[:stepsize__, :std][1]]
)
# push!(ϵ_array_turing, [0.01, 0.0])
end
# ess plots
p = Array{Plots.Plot{Plots.GRBackend}}(undef, 3);
p[1] = plot(convert(Vector{Float64}, ess_array_cmdstan[:, :μ]),
lab="CmdStan", xlim=(0, N_runs), title=":mu ess")
p[1] = plot!(convert(Vector{Float64}, ess_array_MCMCChains_cmdstan[:, :μ]),
lab="MCMCChains/CmdStan ess")
p[1] = plot!(convert(Vector{Float64}, ess_array_MCMCDiagnostics_cmdstan[:, :μ]),
lab="MCMCDiagnostics/CmdStan ess")
p[1] = plot!(convert(Vector{Float64}, ess_array_MCMCChains_turing[:, :μ]),
lab="MCMCChains/Turing ess")
p[1] = plot!(convert(Vector{Float64}, ess_array_MCMCDiagnostics_turing[:, :μ]),
lab="MCMCDiagnostics/Turing")
p[2] = plot(convert(Vector{Float64}, ess_array_cmdstan[:, :μ]),
lab="CmdStan", xlim=(0, N_runs))
p[2] = plot!(convert(Vector{Float64}, ess_array_MCMCChains_cmdstan[:, :μ]),
lab="MCMCChains/CmdStan ess")
p[2] = plot!(convert(Vector{Float64}, ess_array_MCMCDiagnostics_cmdstan[:, :μ]),
lab="MCMCDiagnostics/CmdStan ess")
p[3] = plot(convert(Vector{Float64}, ess_array_MCMCChains_turing[:, :μ]),
lab="MCMCChains/Turing ess",
xlim=(0, N_runs))
p[3] = plot!(convert(Vector{Float64}, ess_array_MCMCDiagnostics_turing[:, :μ]),
lab="MCMCDiagnostics/Turing ess")
plot(p..., layout=(3,1), legend=false)
savefig("ess_mu__estimates_plot.pdf")
p = Array{Plots.Plot{Plots.GRBackend}}(undef, 3);
p[1] = plot(convert(Vector{Float64}, ess_array_cmdstan[:, :σ]),
lab="CmdStan", xlim=(0, N_runs), title=":sigma ess")
p[1] = plot!(convert(Vector{Float64}, ess_array_MCMCChains_cmdstan[:, :σ]),
lab="MCMCChains/CmdStan ess")
p[1] = plot!(convert(Vector{Float64}, ess_array_MCMCDiagnostics_cmdstan[:, :σ]),
lab="MCMCDiagnostics/CmdStan ess")
p[1] = plot!(convert(Vector{Float64}, ess_array_MCMCChains_turing[:, :σ]),
lab="MCMCChains/Turing ess")
p[1] = plot!(convert(Vector{Float64}, ess_array_MCMCDiagnostics_turing[:, :σ]),
lab="MCMCDiagnostics/Turing")
p[2] = plot(convert(Vector{Float64}, ess_array_cmdstan[:, :σ]),
lab="CmdStan", xlim=(0, N_runs))
p[2] = plot!(convert(Vector{Float64}, ess_array_MCMCChains_cmdstan[:, :σ]),
lab="MCMCChains/CmdStan ess")
p[2] = plot!(convert(Vector{Float64}, ess_array_MCMCDiagnostics_cmdstan[:, :σ]),
lab="MCMCDiagnostics/CmdStan ess")
p[3] = plot(convert(Vector{Float64}, ess_array_MCMCChains_turing[:, :σ]),
lab="MCMCChains/Turing ess",
xlim=(0, N_runs))
p[3] = plot!(convert(Vector{Float64}, ess_array_MCMCDiagnostics_turing[:, :σ]),
lab="MCMCDiagnostics/Turing ess")
plot(p..., layout=(3,1))
savefig("ess_sigma__estimates_plot.pdf")
# rhat plots
q = Array{Plots.Plot{Plots.GRBackend}}(undef, 2);
q[1] = plot(convert(Vector{Float64}, rhat_array_cmdstan[:, :μ]),
lab="CmdStan r_hat", xlim=(0, N_runs), title=":mu r_hat")
q[1] = plot!(convert(Vector{Float64}, rhat_array_MCMCChains_cmdstan[:, :μ]),
line=(:dash), lab="MCMCChains/CmdStan r_hat")
q[1] = plot!(convert(Vector{Float64}, rhat_array_MCMCChains_turing[:, :μ]),
lab="MCMCChains/Turing r_hat")
q[2] = plot(convert(Vector{Float64}, rhat_array_cmdstan[:, :μ]),
lab="CmdStan r_hat", xlim=(0, N_runs))
q[2] = plot!(convert(Vector{Float64}, rhat_array_MCMCChains_cmdstan[:, :μ]),
lab="MCMCChains/CmdStan r_hat")
q[2] = plot!(convert(Vector{Float64}, rhat_array_MCMCDiagnostics_cmdstan[:, :μ]),
line=(:dot), lab="MCMCDiagnostics/CmdStan r_hat")
plot(q..., layout=(2,1))
savefig("rhat_mu__estimates_plot.pdf")
q = Array{Plots.Plot{Plots.GRBackend}}(undef, 2);
q[1] = plot(convert(Vector{Float64}, rhat_array_cmdstan[:, :σ]),
lab="CmdStan r_hat", xlim=(0, N_runs), title=":sigma r_hat")
q[1] = plot!(convert(Vector{Float64}, rhat_array_MCMCChains_cmdstan[:, :σ]),
line=(:dash), lab="MCMCChains/CmdStan r_hat")
q[1] = plot!(convert(Vector{Float64}, rhat_array_MCMCChains_turing[:, :σ]),
lab="MCMCChains/Turing r_hat")
q[2] = plot(convert(Vector{Float64}, rhat_array_cmdstan[:, :σ]),
lab="CmdStan r_hat", xlim=(0, N_runs))
q[2] = plot!(convert(Vector{Float64}, rhat_array_MCMCChains_cmdstan[:, :σ]),
line=(:dash), lab="MCMCChains/CmdStan r_hat")
q[2] = plot!(convert(Vector{Float64}, rhat_array_MCMCDiagnostics_cmdstan[:, :σ]),
line=(:dot), lab="MCMCDiagnostics/CmdStan r_hat")
plot(q..., layout=(2,1))
savefig("rhat_sigma_estimates_plot.pdf")
# nuts plots
q = Array{Plots.Plot{Plots.GRBackend}}(undef, 2);
q[1] = plot(convert(Vector{Float64}, ϵ_array_cmdstan[:, :μ]),
lab="CmdStan stepsize__", xlim=(0, N_runs), title="mean stepsize__")
q[1] = plot!(convert(Vector{Float64}, ϵ_array_turing[:, :μ]),
lab="Turing epsilon")
q[2] = plot(convert(Vector{Float64}, ϵ_array_cmdstan[:, :σ]),
lab="CmdStan stepsize__", xlim=(0, N_runs), title="std stepsize__")
q[2] = plot!(convert(Vector{Float64}, ϵ_array_turing[:, :σ]),
lab="Turing epsilon")
plot(q..., layout=(2,1))
savefig("stepsize_sigma_estimates_plot.pdf")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment