ArviZ.jl Quickstart

Note

This tutorial is adapted from ArviZ's quickstart.

using ArviZ
using PyPlot

# ArviZ ships with style sheets!
ArviZ.use_style("arviz-darkgrid")

Get started with plotting

ArviZ.jl is designed to be used with libraries like CmdStan, Turing.jl, and Soss.jl but works fine with raw arrays.

using Random

rng = Random.MersenneTwister(37772)
plot_posterior(randn(rng, 100_000));
gcf()

Plotting a dictionary of arrays, ArviZ.jl will interpret each key as the name of a different random variable. Each row of an array is treated as an independent series of draws from the variable, called a chain. Below, we have 10 chains of 50 draws each for four different distributions.

using Distributions

s = (10, 50)
plot_forest(
    Dict(
        "normal" => randn(rng, s),
        "gumbel" => rand(rng, Gumbel(), s),
        "student t" => rand(rng, TDist(6), s),
        "exponential" => rand(rng, Exponential(), s),
    ),
);
gcf()

Plotting with MCMCChains.jl's `Chains` objects produced by Turing.jl

ArviZ is designed to work well with high dimensional, labelled data. Consider the eight schools model, which roughly tries to measure the effectiveness of SAT classes at eight different schools. To show off ArviZ's labelling, I give the schools the names of a different eight schools.

This model is small enough to write down, is hierarchical, and uses labelling. Additionally, a centered parameterization causes divergences (which are interesting for illustration).

First we create our data and set some sampling parameters.

J = 8
y = [28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0]
σ = [15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0]
schools = [
    "Choate",
    "Deerfield",
    "Phillips Andover",
    "Phillips Exeter",
    "Hotchkiss",
    "Lawrenceville",
    "St. Paul's",
    "Mt. Hermon",
];

nwarmup, nsamples, nchains = 1000, 1000, 4;

Now we write and run the model using Turing:

using Turing

Turing.@model function turing_model(y, σ, J=length(y))
    μ ~ Normal(0, 5)
    τ ~ truncated(Cauchy(0, 5), 0, Inf)
    θ ~ filldist(Normal(μ, τ), J)
    for i in 1:J
        y[i] ~ Normal(θ[i], σ[i])
    end
end

param_mod = turing_model(y, σ)
sampler = NUTS(nwarmup, 0.8)

rng = Random.MersenneTwister(16653)
turing_chns = sample(
    rng, param_mod, sampler, MCMCThreads(), nsamples, nchains; progress=false
);

┌ Info: Found initial step size
└   ϵ = 0.8
┌ Info: Found initial step size
└   ϵ = 0.8
┌ Info: Found initial step size
└   ϵ = 0.30000000000000004
┌ Info: Found initial step size
└   ϵ = 0.4

Most ArviZ functions work fine with Chains objects from Turing:

plot_autocorr(turing_chns; var_names=["μ", "τ"]);
gcf()

Convert to `InferenceData`

For much more powerful querying, analysis and plotting, we can use built-in ArviZ utilities to convert Chains objects to xarray datasets. Note we are also giving some information about labelling.

ArviZ is built to work with InferenceData (a netcdf datastore that loads data into xarray datasets), and the more groups it has access to, the more powerful analyses it can perform.

idata = from_mcmcchains(
    turing_chns;
    coords=Dict("school" => schools),
    dims=Dict("y" => ["school"], "σ" => ["school"], "θ" => ["school"]),
    library="Turing",
)

InferenceData

Each group is an ArviZ.Dataset (a thinly wrapped xarray.Dataset). We can view a summary of the dataset.

idata.posterior

Dataset (xarray.Dataset)
Dimensions:  (chain: 4, draw: 1000, school: 8)
Coordinates:
  * chain    (chain) int64 0 1 2 3
  * draw     (draw) int64 0 1 2 3 4 5 6 7 8 ... 992 993 994 995 996 997 998 999
  * school   (school) <U16 'Choate' 'Deerfield' ... "St. Paul's" 'Mt. Hermon'
Data variables:
    μ        (chain, draw) float64 3.598 3.844 3.727 1.557 ... 7.414 3.701 3.061
    τ        (chain, draw) float64 1.341 1.253 0.9012 ... 2.304 4.701 3.185
    θ        (chain, draw, school) float64 3.826 4.778 2.754 ... 5.444 7.008
Attributes:
    created_at:         2021-05-11T17:20:41.996609
    arviz_version:      0.11.2
    start_time:         1620753599.753838
    stop_time:          1620753617.297758
    inference_library:  Turing

Here is a plot of the trace. Note the intelligent labels.

plot_trace(idata);
gcf()

We can also generate summary stats

summarystats(idata)

10 rows × 10 columns

	variable	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
	String	Float64	Float64	Float64	Float64	Float64	Float64	Float64	Float64	Float64
1	μ	4.377	3.433	-2.29	10.685	0.126	0.092	743.0	881.0	1.01
2	τ	3.649	3.0	0.407	8.922	0.184	0.13	161.0	106.0	1.01
3	θ[1]	6.233	5.668	-3.119	17.324	0.206	0.145	787.0	1324.0	1.01
4	θ[2]	4.968	4.668	-3.795	13.939	0.138	0.105	1087.0	1834.0	1.01
5	θ[3]	3.945	5.228	-5.885	13.905	0.15	0.113	1126.0	1586.0	1.01
6	θ[4]	4.667	4.889	-4.779	13.732	0.147	0.115	1042.0	1677.0	1.0
7	θ[5]	3.542	4.797	-5.586	12.766	0.135	0.117	1192.0	1511.0	1.01
8	θ[6]	4.044	4.768	-4.807	13.581	0.145	0.108	1024.0	1348.0	1.01
9	θ[7]	6.296	5.07	-2.371	16.512	0.189	0.133	727.0	1258.0	1.01
10	θ[8]	4.766	5.289	-6.104	14.186	0.154	0.127	1047.0	1611.0	1.01

and examine the energy distribution of the Hamiltonian sampler

plot_energy(idata);
gcf()

Additional information in Turing.jl

With a few more steps, we can use Turing to compute additional useful groups to add to the InferenceData.

To sample from the prior, one simply calls sample but with the Prior sampler:

prior = sample(param_mod, Prior(), nsamples; progress=false)

Chains MCMC chain (1000×11×1 Array{Float64, 3}):

Start time        = 2021-05-11T17:20:52.494
Stop time         = 2021-05-11T17:20:52.760
Wall duration     = 0.27 seconds
Iterations        = 1:1000
Thinning interval = 1
Chains            = 1
Samples per chain = 1000
parameters        = θ[1], θ[2], θ[3], θ[4], θ[5], θ[6], θ[7], θ[8], μ, τ
internals         = lp

Summary Statistics
  parameters      mean       std   naive_se      mcse         ess      rhat    ⋯
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64    ⋯

        θ[1]    0.9495   67.9183     2.1478    1.8926    556.5266    1.0001    ⋯
        θ[2]   -0.6053   48.9281     1.5472    1.0120   1024.0944    1.0000    ⋯
        θ[3]   -1.1761   60.8375     1.9238    1.5571    999.5217    0.9991    ⋯
        θ[4]    0.1684   52.6024     1.6634    1.8087    811.6064    1.0025    ⋯
        θ[5]    1.7505   71.2369     2.2527    2.2212    918.8990    0.9993    ⋯
        θ[6]    2.0359   74.2804     2.3490    2.5844   1004.3982    1.0023    ⋯
        θ[7]    1.1503   53.3496     1.6871    1.6193   1044.4259    0.9991    ⋯
        θ[8]    0.7439   58.6209     1.8538    1.9887    965.1255    1.0011    ⋯
           μ    0.0522    4.8552     0.1535    0.1401   1030.5527    1.0003    ⋯
           τ   18.2545   68.3487     2.1614    2.4279    638.0227    0.9996    ⋯
                                                                1 column omitted

Quantiles
  parameters       2.5%     25.0%     50.0%     75.0%      97.5%
      Symbol    Float64   Float64   Float64   Float64    Float64

        θ[1]   -53.1021   -5.5494    0.0282    5.4924    42.0920
        θ[2]   -47.9095   -4.9937    0.2079    6.0771    38.2048
        θ[3]   -48.4412   -5.0144    0.3014    5.8416    39.0567
        θ[4]   -36.2510   -4.9584    0.4786    6.0296    46.7597
        θ[5]   -41.1661   -5.4494   -0.3680    5.0671    53.3500
        θ[6]   -40.9711   -5.4135    0.2481    5.6112    47.2051
        θ[7]   -39.5905   -5.2035    0.4159    6.4169    50.1237
        θ[8]   -49.2875   -4.9863    0.2812    5.8763    55.9611
           μ    -9.7057   -3.3349   -0.0951    3.3130     9.5634
           τ     0.1994    2.0832    4.9683   11.2756   112.9419

To draw from the prior and posterior predictive distributions we can instantiate a "predictive model", i.e. a Turing model but with the observations set to missing, and then calling predict on the predictive model and the previously drawn samples:

# Instantiate the predictive model
param_mod_predict = turing_model(similar(y, Missing), σ)
# and then sample!
prior_predictive = predict(param_mod_predict, prior)
posterior_predictive = predict(param_mod_predict, turing_chns)

Chains MCMC chain (1000×8×4 Array{Float64, 3}):

Iterations        = 1:1000
Thinning interval = 1
Chains            = 1, 2, 3, 4
Samples per chain = 1000
parameters        = y[1], y[2], y[3], y[4], y[5], y[6], y[7], y[8]
internals         = 

Summary Statistics
  parameters      mean       std   naive_se      mcse         ess      rhat
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64

        y[1]    6.1940   15.8365     0.2504    0.3187   2598.9319    0.9996
        y[2]    5.1315   11.0163     0.1742    0.2020   2972.2101    1.0004
        y[3]    4.1674   16.8618     0.2666    0.2583   3529.7049    1.0003
        y[4]    4.7150   12.1285     0.1918    0.2096   2859.1765    0.9996
        y[5]    3.5240   10.2518     0.1621    0.2029   2745.2054    1.0010
        y[6]    4.1669   11.9747     0.1893    0.2232   3018.5444    0.9996
        y[7]    6.5947   11.0310     0.1744    0.2484   2290.3377    1.0000
        y[8]    5.0596   18.4197     0.2912    0.2796   3201.2609    0.9999

Quantiles
  parameters       2.5%     25.0%     50.0%     75.0%     97.5%
      Symbol    Float64   Float64   Float64   Float64   Float64

        y[1]   -25.4038   -4.0195    6.3500   16.6845   37.8474
        y[2]   -16.5638   -2.5435    5.3023   12.6178   26.2114
        y[3]   -29.6595   -7.0289    4.1721   15.3849   37.5304
        y[4]   -19.2035   -3.6940    4.5058   12.9910   28.7925
        y[5]   -16.5995   -3.2405    3.5815   10.3291   23.3610
        y[6]   -18.4540   -3.8768    4.2179   12.0992   27.4589
        y[7]   -14.3394   -0.7633    6.5626   13.9443   28.7613
        y[8]   -30.7903   -7.2315    5.0240   17.7800   41.0833

And to extract the pointwise log-likelihoods, which is useful if you want to compute metrics such as loo,

loglikelihoods = Turing.pointwise_loglikelihoods(
    param_mod, MCMCChains.get_sections(turing_chns, :parameters)
)

Dict{String, Matrix{Float64}} with 8 entries:
  "y[6]" => [-3.38973 -3.51189 -3.43534 -3.73188; -3.32076 -3.32165 -3.38096 -3…
  "y[2]" => [-3.27342 -3.22381 -3.30609 -3.35272; -3.36404 -3.42648 -3.23499 -3…
  "y[1]" => [-4.92564 -4.6233 -4.15682 -4.81507; -4.94726 -5.12105 -5.63477 -3.…
  "y[5]" => [-3.17195 -3.55913 -3.55442 -3.64785; -3.34779 -3.17499 -3.21129 -3…
  "y[8]" => [-3.92493 -3.82086 -4.23278 -3.88106; -3.9033 -3.92951 -4.07239 -3.…
  "y[7]" => [-4.20172 -3.96703 -3.49247 -4.77806; -4.1926 -4.5937 -4.66271 -3.4…
  "y[3]" => [-3.75618 -3.87783 -3.69645 -4.18097; -3.80406 -3.74826 -3.70795 -3…
  "y[4]" => [-3.38347 -3.32362 -3.32345 -3.31684; -3.33491 -3.50791 -3.35802 -3…

This can then be included in the from_mcmcchains call from above:

using LinearAlgebra
# Ensure the ordering of the loglikelihoods matches the ordering of `posterior_predictive`
ynames = string.(keys(posterior_predictive))
loglikelihoods_vals = getindex.(Ref(loglikelihoods), ynames)
# Reshape into `(nchains, nsamples, size(y)...)`
loglikelihoods_arr = permutedims(cat(loglikelihoods_vals...; dims=3), (2, 1, 3))

idata = from_mcmcchains(
    turing_chns;
    posterior_predictive=posterior_predictive,
    log_likelihood=Dict("y" => loglikelihoods_arr),
    prior=prior,
    prior_predictive=prior_predictive,
    observed_data=Dict("y" => y),
    coords=Dict("school" => schools),
    dims=Dict("y" => ["school"], "σ" => ["school"], "θ" => ["school"]),
    library="Turing",
)

InferenceData

Then we can for example compute the expected leave-one-out (LOO) predictive density, which is an estimate of the out-of-distribution predictive fit of the model:

loo(idata) # higher is better

1 rows × 7 columns

	loo	loo_se	p_loo	n_samples	n_data_points	warning	loo_scale
	Float64	Float64	Float64	Int64	Int64	Bool	String
1	-30.707	1.3569	0.861043	4000	8	0	log

If the model is well-calibrated, i.e. it replicates the true generative process well, the CDF of the pointwise LOO values should be similarly distributed to a uniform distribution. This can be inspected visually:

plot_loo_pit(idata; y="y", ecdf=true);
gcf()

Plotting with CmdStan.jl outputs

CmdStan.jl and StanSample.jl also default to producing Chains outputs, and we can easily plot these chains.

Here is the same centered eight schools model:

using CmdStan, MCMCChains

schools_code = """
data {
  int<lower=0> J;
  real y[J];
  real<lower=0> sigma[J];
}

parameters {
  real mu;
  real<lower=0> tau;
  real theta[J];
}

model {
  mu ~ normal(0, 5);
  tau ~ cauchy(0, 5);
  theta ~ normal(mu, tau);
  y ~ normal(theta, sigma);
}

generated quantities {
    vector[J] log_lik;
    vector[J] y_hat;
    for (j in 1:J) {
        log_lik[j] = normal_lpdf(y[j] | theta[j], sigma[j]);
        y_hat[j] = normal_rng(theta[j], sigma[j]);
    }
}
"""

schools_dat = Dict("J" => J, "y" => y, "sigma" => σ)
stan_model = Stanmodel(;
    model=schools_code,
    name="schools",
    nchains=nchains,
    num_warmup=nwarmup,
    num_samples=nsamples,
    output_format=:mcmcchains,
    random=CmdStan.Random(28983),
)
_, stan_chns, _ = stan(stan_model, schools_dat; summary=false);

File /home/runner/work/ArviZ.jl/ArviZ.jl/docs/build/tmp/schools.stan will be updated.

plot_density(stan_chns; var_names=["mu", "tau"]);
gcf()

Again, converting to InferenceData, we can get much richer labelling and mixing of data. Note that we're using the same from_cmdstan function used by ArviZ to process cmdstan output files, but through the power of dispatch in Julia, if we pass a Chains object, it instead uses ArviZ.jl's overloads, which forward to from_mcmcchains.

idata = from_cmdstan(
    stan_chns;
    posterior_predictive="y_hat",
    observed_data=Dict("y" => schools_dat["y"]),
    log_likelihood="log_lik",
    coords=Dict("school" => schools),
    dims=Dict(
        "y" => ["school"],
        "sigma" => ["school"],
        "theta" => ["school"],
        "log_lik" => ["school"],
        "y_hat" => ["school"],
    ),
)

InferenceData

Here is a plot showing where the Hamiltonian sampler had divergences:

plot_pair(
    idata;
    coords=Dict("school" => ["Choate", "Deerfield", "Phillips Andover"]),
    divergences=true,
);
gcf()

Plotting with Soss.jl outputs

With Soss, we can define our model for the posterior and easily use it to draw samples from the prior, prior predictive, posterior, and posterior predictive distributions.

First we define our model:

using Soss, NamedTupleTools

mod = Soss.@model (J, σ) begin
    μ ~ Normal(0, 5)
    τ ~ HalfCauchy(5)
    θ ~ iid(J)(Normal(μ, τ))
    y ~ For(1:J) do j
        Normal(θ[j], σ[j])
    end
end

constant_data = (J=J, σ=σ)
param_mod = mod(; constant_data...)

Joint Distribution
    Bound arguments: [J, σ]
    Variables: [τ, μ, θ, y]

@model (J, σ) begin
        τ ~ HalfCauchy(5)
        μ ~ Normal(0, 5)
        θ ~ (iid(J))(Normal(μ, τ))
        y ~ For(1:J) do j
                Normal(θ[j], σ[j])
            end
    end

Then we draw from the prior and prior predictive distributions.

Random.seed!(5298)
prior_priorpred = [
    map(1:(nchains * nsamples)) do _
        draw = rand(param_mod)
        return delete(draw, keys(constant_data))
    end,
];

Next, we draw from the posterior using DynamicHMC.jl.

post = map(1:nchains) do _
    dynamicHMC(param_mod, (y=y,), nsamples)
end;

Finally, we update the posterior samples with draws from the posterior predictive distribution.

pred = predictive(mod, :μ, :τ, :θ)
post_postpred = map(post) do post_draws
    map(post_draws) do post_draw
        pred_draw = rand(pred(post_draw))
        pred_draw = delete(pred_draw, keys(constant_data))
        return merge(pred_draw, post_draw)
    end
end;

Each Soss draw is a NamedTuple. We can plot the rank order statistics of the posterior to identify poor convergence:

plot_rank(post; var_names=["μ", "τ"]);
gcf()

Now we combine all of the samples to an InferenceData:

idata = from_namedtuple(
    post_postpred;
    posterior_predictive=[:y],
    prior=prior_priorpred,
    prior_predictive=[:y],
    observed_data=(y=y,),
    constant_data=constant_data,
    coords=Dict("school" => schools),
    dims=Dict("y" => ["school"], "σ" => ["school"], "θ" => ["school"]),
    library=Soss,
)

InferenceData

We can compare the prior and posterior predictive distributions:

plot_density(
    [idata.posterior_predictive, idata.prior_predictive];
    data_labels=["Post-pred", "Prior-pred"],
    var_names=["y"],
)
gcf()

Environment

using Pkg
Pkg.status()

using InteractiveUtils
versioninfo()

Julia Version 1.6.1
Commit 6aaedecc44 (2021-04-23 05:59 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Xeon(R) CPU E5-2673 v4 @ 2.30GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-11.0.1 (ORCJIT, broadwell)
Environment:
  JULIA_CMDSTAN_HOME = /home/runner/work/ArviZ.jl/ArviZ.jl/.cmdstan//cmdstan-2.25.0/
  JULIA_NUM_THREADS = 2