Working with InferenceData

using ArviZ, DimensionalData, Statistics

Here we present a collection of common manipulations you can use while working with InferenceData.

Let's load one of ArviZ's example datasets. posterior, posterior_predictive, etc are the groups stored in idata, and they are stored as Datasets. In this HTML view, you can click a group name to expand a summary of the group.

idata = load_example_data("centered_eight")
InferenceData
posterior
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 3 layers:
  :mu    Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :theta Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)
  :tau   Float64 dims: Dim{:chain}, Dim{:draw} (4×500)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
posterior_predictive
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.489022"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
sample_stats
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 12 layers:
  :tune             Bool dims: Dim{:chain}, Dim{:draw} (4×500)
  :depth            Int64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :tree_size        Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :lp               Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :energy_error     Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :step_size_bar    Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :max_energy_error Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :energy           Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :mean_tree_accept Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :step_size        Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :diverging        Bool dims: Dim{:chain}, Dim{:draw} (4×500)
  :log_likelihood   Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.485802"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
prior
Dataset with dimensions: 
  Dim{:chain} Sampled StepRangeLen(0.0, 0.0, 1) ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 5 layers:
  :tau       Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :tau_log__ Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :mu        Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :theta     Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)
  :obs       Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.490387"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
observed_data
Dataset with dimensions: 
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:school} (8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.491909"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
Info

Datasets are DimensionalData.AbstractDimStacks and can be used identically. The variables a Dataset contains are called "layers", and dimensions of the same name that appear in more than one layer within a Dataset must have the same indices.

InferenceData behaves like a NamedTuple and can be used similarly. Note that unlike a NamedTuple, the groups always appear in a specific order.

length(idata) # number of groups
5
keys(idata) # group names
(:posterior, :posterior_predictive, :sample_stats, :prior, :observed_data)

Get the dataset corresponding to a single group

Group datasets can be accessed both as properties or as indexed items.

post = idata.posterior
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 3 layers:
  :mu    Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :theta Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)
  :tau   Float64 dims: Dim{:chain}, Dim{:draw} (4×500)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at                => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library         => "pymc3"

post is the dataset itself, so this is a non-allocating operation.

idata[:posterior] === post
true

InferenceData supports a more advanced indexing syntax, which we'll see later.

Getting a new InferenceData with a subset of groups

We can index by a collection of group names to get a new InferenceData with just those groups. This is also non-allocating.

idata_sub = idata[(:posterior, :posterior_predictive)]
InferenceData
posterior
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 3 layers:
  :mu    Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :theta Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)
  :tau   Float64 dims: Dim{:chain}, Dim{:draw} (4×500)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
posterior_predictive
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.489022"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"

Adding groups to an InferenceData

InferenceData is immutable, so to add or replace groups we use merge to create a new object.

merge(idata_sub, idata[(:observed_data, :prior)])
InferenceData
posterior
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 3 layers:
  :mu    Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :theta Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)
  :tau   Float64 dims: Dim{:chain}, Dim{:draw} (4×500)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
posterior_predictive
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.489022"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
prior
Dataset with dimensions: 
  Dim{:chain} Sampled StepRangeLen(0.0, 0.0, 1) ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 5 layers:
  :tau       Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :tau_log__ Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :mu        Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :theta     Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)
  :obs       Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.490387"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
observed_data
Dataset with dimensions: 
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:school} (8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.491909"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"

We can also use Base.setindex to out-of-place add or replace a single group.

Base.setindex(idata_sub, idata.prior, :prior)
InferenceData
posterior
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 3 layers:
  :mu    Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :theta Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)
  :tau   Float64 dims: Dim{:chain}, Dim{:draw} (4×500)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
posterior_predictive
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.489022"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
prior
Dataset with dimensions: 
  Dim{:chain} Sampled StepRangeLen(0.0, 0.0, 1) ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 5 layers:
  :tau       Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :tau_log__ Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :mu        Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :theta     Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)
  :obs       Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.490387"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"

Add a new variable

Dataset is also immutable. So while the values within the underlying data arrays can be mutated, layers cannot be added or removed from Datasets, and groups cannot be added/removed from InferenceData.

Instead, we do this out-of-place also using merge.

merge(post, (log_tau=log.(post[:tau]),))
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 4 layers:
  :mu      Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :theta   Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)
  :tau     Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :log_tau Float64 dims: Dim{:chain}, Dim{:draw} (4×500)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at                => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library         => "pymc3"

Obtain an array for a given parameter

Let’s say we want to get the values for mu as an array. Parameters can be accessed with either property or index syntax.

post.tau
4×500 DimArray{Float64,2} tau with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points
    0        1        2         3        …  497        498        499
 0  3.7301   2.07538  3.70299   4.14612      10.1079     8.07999    7.72886
 1  1.19333  1.19333  1.19333   3.0369       13.922      8.86992    4.76318
 2  5.13725  4.26438  2.14143   1.44099       2.81184   12.1797     4.45297
 3  0.50007  0.50007  0.902267  1.17612       8.34563    7.71079    5.4068
post[:tau] === post.tau
true

To remove the dimensions, just use parent to retrieve the underlying array.

parent(post.tau)
4×500 Matrix{Float64}:
 3.7301   2.07538  3.70299   4.14612  …  10.1079    8.07999  7.72886
 1.19333  1.19333  1.19333   3.0369      13.922     8.86992  4.76318
 5.13725  4.26438  2.14143   1.44099      2.81184  12.1797   4.45297
 0.50007  0.50007  0.902267  1.17612      8.34563   7.71079  5.4068

Get the dimension lengths

Let’s check how many groups are in our hierarchical model.

size(idata.observed_data, :school)
8

Get coordinate/index values

What are the names of the groups in our hierarchical model? You can access them from the coordinate name school in this case.

DimensionalData.index(idata.observed_data, :school)
8-element Vector{String}:
 "Choate"
 "Deerfield"
 "Phillips Andover"
 "Phillips Exeter"
 "Hotchkiss"
 "Lawrenceville"
 "St. Paul's"
 "Mt. Hermon"

Get a subset of chains

Let’s keep only chain 0 here. For the subset to take effect on all relevant InferenceData groups – posterior, sample_stats, log_likelihood, and posterior_predictive – we will index InferenceData instead of Dataset.

Here we use DimensionalData's At selector. Its other selectors are also supported.

idata[chain=At(0)]
InferenceData
posterior
Dataset with dimensions: 
  Dim{:chain} Sampled Int64[0] ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 3 layers:
  :mu    Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :theta Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)
  :tau   Float64 dims: Dim{:chain}, Dim{:draw} (1×500)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
posterior_predictive
Dataset with dimensions: 
  Dim{:chain} Sampled Int64[0] ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.489022"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
sample_stats
Dataset with dimensions: 
  Dim{:chain} Sampled Int64[0] ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 12 layers:
  :tune             Bool dims: Dim{:chain}, Dim{:draw} (1×500)
  :depth            Int64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :tree_size        Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :lp               Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :energy_error     Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :step_size_bar    Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :max_energy_error Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :energy           Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :mean_tree_accept Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :step_size        Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :diverging        Bool dims: Dim{:chain}, Dim{:draw} (1×500)
  :log_likelihood   Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.485802"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
prior
Dataset with dimensions: 
  Dim{:chain} Sampled Float64[0.0] ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 5 layers:
  :tau       Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :tau_log__ Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :mu        Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :theta     Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)
  :obs       Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.490387"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
observed_data
Dataset with dimensions: 
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:school} (8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.491909"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"

Note that in this case, prior only has a chain of 0. If it also had the other chains, we could have passed chain=At([0, 2]) to subset by chains 0 and 2.

Warning

If we used idata[chain=[0, 2]] without the At selector, this is equivalent to idata[chain=DimensionalData.index(idata.posterior, :chain)[0, 2]], that is, [0, 2] indexes an array of dimension indices, which here would error. But if we had requested idata[chain=[1, 2]] we would not hit an error, but we would index the wrong chains. So it's important to always use a selector to index by values of dimension indices.

Remove the first $n$ draws (burn-in)

Let’s say we want to remove the first 100 draws from all the chains and all InferenceData groups with draws. To do this we use the .. syntax from IntervalSets.jl, which is exported by DimensionalData.

idata[draw=100 .. Inf]
InferenceData
posterior
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 100:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 3 layers:
  :mu    Float64 dims: Dim{:chain}, Dim{:draw} (4×400)
  :theta Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×400×8)
  :tau   Float64 dims: Dim{:chain}, Dim{:draw} (4×400)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
posterior_predictive
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 100:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×400×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.489022"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
sample_stats
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 100:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 12 layers:
  :tune             Bool dims: Dim{:chain}, Dim{:draw} (4×400)
  :depth            Int64 dims: Dim{:chain}, Dim{:draw} (4×400)
  :tree_size        Float64 dims: Dim{:chain}, Dim{:draw} (4×400)
  :lp               Float64 dims: Dim{:chain}, Dim{:draw} (4×400)
  :energy_error     Float64 dims: Dim{:chain}, Dim{:draw} (4×400)
  :step_size_bar    Float64 dims: Dim{:chain}, Dim{:draw} (4×400)
  :max_energy_error Float64 dims: Dim{:chain}, Dim{:draw} (4×400)
  :energy           Float64 dims: Dim{:chain}, Dim{:draw} (4×400)
  :mean_tree_accept Float64 dims: Dim{:chain}, Dim{:draw} (4×400)
  :step_size        Float64 dims: Dim{:chain}, Dim{:draw} (4×400)
  :diverging        Bool dims: Dim{:chain}, Dim{:draw} (4×400)
  :log_likelihood   Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×400×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.485802"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
prior
Dataset with dimensions: 
  Dim{:chain} Sampled StepRangeLen(0.0, 0.0, 1) ForwardOrdered Regular Points,
  Dim{:draw} Sampled 100:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 5 layers:
  :tau       Float64 dims: Dim{:chain}, Dim{:draw} (1×400)
  :tau_log__ Float64 dims: Dim{:chain}, Dim{:draw} (1×400)
  :mu        Float64 dims: Dim{:chain}, Dim{:draw} (1×400)
  :theta     Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×400×8)
  :obs       Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×400×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.490387"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
observed_data
Dataset with dimensions: 
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:school} (8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.491909"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"

If you check the object you will see that the groups posterior, posterior_predictive, prior, and sample_stats have 400 draws compared to idata, which has 500. The group observed_data has not been affected because it does not have the draw dimension.

Alternatively, you can change a subset of groups by combining indexing styles with merge. Here we use this to build a new InferenceData where we have discarded the first 100 draws only from posterior.

merge(idata, idata[(:posterior,), draw=100 .. Inf])
InferenceData
posterior
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 100:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 3 layers:
  :mu    Float64 dims: Dim{:chain}, Dim{:draw} (4×400)
  :theta Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×400×8)
  :tau   Float64 dims: Dim{:chain}, Dim{:draw} (4×400)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
posterior_predictive
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.489022"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
sample_stats
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 12 layers:
  :tune             Bool dims: Dim{:chain}, Dim{:draw} (4×500)
  :depth            Int64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :tree_size        Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :lp               Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :energy_error     Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :step_size_bar    Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :max_energy_error Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :energy           Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :mean_tree_accept Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :step_size        Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :diverging        Bool dims: Dim{:chain}, Dim{:draw} (4×500)
  :log_likelihood   Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.485802"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
prior
Dataset with dimensions: 
  Dim{:chain} Sampled StepRangeLen(0.0, 0.0, 1) ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 5 layers:
  :tau       Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :tau_log__ Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :mu        Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :theta     Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)
  :obs       Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.490387"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
observed_data
Dataset with dimensions: 
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:school} (8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.491909"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"

Compute posterior mean values along draw and chain dimensions

To compute the mean value of the posterior samples, do the following:

mean(post)
(mu = 4.092610850912027, theta = 4.56047268323059, tau = 4.088982928754772)

This computes the mean along all dimensions, discarding all dimensions and returning the result as a NamedTuple. This may be what you wanted for mu and tau, which have only two dimensions (chain and draw), but maybe not what you expected for theta, which has one more dimension school.

You can specify along which dimension you want to compute the mean (or other functions), which instead returns a Dataset.

mean(post; dims=(:chain, :draw))
Dataset with dimensions: 
  Dim{:chain} Sampled 1.5:4.0:1.5 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 249.5:500.0:249.5 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 3 layers:
  :mu    Float64 dims: Dim{:chain}, Dim{:draw} (1×1)
  :theta Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×1×8)
  :tau   Float64 dims: Dim{:chain}, Dim{:draw} (1×1)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at                => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library         => "pymc3"

The singleton dimensions of chain and draw now contain meaningless indices, so you may want to discard them, which you can do with dropdims.

dropdims(mean(post; dims=(:chain, :draw)); dims=(:chain, :draw))
Dataset with dimensions: 
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 3 layers:
  :mu    Float64 dims: 
  :theta Float64 dims: Dim{:school} (8)
  :tau   Float64 dims: 

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at                => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library         => "pymc3"

Renaming a dimension

We can rename a dimension in a Dataset using DimensionalData's set method:

theta_bis = set(post.theta; school=:school_bis)
4×500×8 DimArray{Float64,3} theta with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school_bis} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
[:, :, 1]
     0         1        2        …  497       498        499
 0   1.66865  -6.23936  2.1951       21.6306    9.29298   11.7154
 1   8.09621   8.09621  8.09621      15.2759   14.7355    -4.83704
 2  14.5709   12.6867   9.66618       2.6685    5.36165   13.4391
 3   4.32639   4.32639  2.99078      14.1863   -1.42095   -0.0501594
[and 7 more slices...]

We can use this, for example, to broadcast functions across multiple arrays, automatically matching up shared dimensions, using DimensionalData.broadcast_dims.

theta_school_diff = broadcast_dims(-, post.theta, theta_bis)
4×500×8×8 DimArray{Float64,4} theta with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered,
  Dim{:school_bis} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
[:, :, 1, 1]
    0    1    2    3    4    5    6    …  495    496    497    498    499
 0  0.0  0.0  0.0  0.0  0.0  0.0  0.0       0.0    0.0    0.0    0.0    0.0
 1  0.0  0.0  0.0  0.0  0.0  0.0  0.0       0.0    0.0    0.0    0.0    0.0
 2  0.0  0.0  0.0  0.0  0.0  0.0  0.0       0.0    0.0    0.0    0.0    0.0
 3  0.0  0.0  0.0  0.0  0.0  0.0  0.0       0.0    0.0    0.0    0.0    0.0
[and 63 more slices...]

Compute and store posterior pushforward quantities

We use “posterior pushfoward quantities” to refer to quantities that are not variables in the posterior but deterministic computations using posterior variables.

You can compute these pushforward operations and store them as a new variable in a copy of the posterior group.

Here we'll create a new InferenceData with theta_school_diff in the posterior:

idata_new = Base.setindex(idata, merge(post, (; theta_school_diff)), :posterior)
InferenceData
posterior
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered,
  Dim{:school_bis} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 4 layers:
  :mu                Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :theta             Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)
  :tau               Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :theta_school_diff Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school}, Dim{:school_bis} (4×500×8×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
posterior_predictive
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.489022"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
sample_stats
Dataset with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 12 layers:
  :tune             Bool dims: Dim{:chain}, Dim{:draw} (4×500)
  :depth            Int64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :tree_size        Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :lp               Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :energy_error     Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :step_size_bar    Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :max_energy_error Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :energy           Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :mean_tree_accept Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :step_size        Float64 dims: Dim{:chain}, Dim{:draw} (4×500)
  :diverging        Bool dims: Dim{:chain}, Dim{:draw} (4×500)
  :log_likelihood   Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (4×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.485802"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
prior
Dataset with dimensions: 
  Dim{:chain} Sampled StepRangeLen(0.0, 0.0, 1) ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 5 layers:
  :tau       Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :tau_log__ Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :mu        Float64 dims: Dim{:chain}, Dim{:draw} (1×500)
  :theta     Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)
  :obs       Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (1×500×8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.490387"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"
observed_data
Dataset with dimensions: 
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 1 layer:
  :obs Float64 dims: Dim{:school} (8)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.491909"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"

Once you have these pushforward quantities in an InferenceData, you’ll then be able to plot them with ArviZ functions, calculate stats and diagnostics on them, or save and share the InferenceData object with the pushforward quantities included.

Here we compute the mcse of theta_school_diff:

mcse(idata_new.posterior).theta_school_diff
8×8 DimArray{Float64,2} theta_school_diff with dimensions: 
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered,
  Dim{:school_bis} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
                       "Choate"  …   "St. Paul's"   "Mt. Hermon"
  "Choate"            0.0           0.111869       0.175078
  "Deerfield"         0.149711      0.15666        0.11984
  "Phillips Andover"  0.276861      0.23738        0.163744
  "Phillips Exeter"   0.163212      0.167966       0.135588
  "Hotchkiss"         0.299984   …  0.260274       0.165326
  "Lawrenceville"     0.223715      0.209422       0.146585
  "St. Paul's"        0.111869      0.0            0.155986
  "Mt. Hermon"        0.175078      0.155986       0.0

Advanced subsetting

To select the value corresponding to the difference between the Choate and Deerfield schools do:

school_idx = ["Choate", "Hotchkiss", "Mt. Hermon"]
school_bis_idx = ["Deerfield", "Choate", "Lawrenceville"]
theta_school_diff[school=At(school_idx), school_bis=At(school_bis_idx)]
4×500×3×3 DimArray{Float64,4} theta with dimensions: 
  Dim{:chain} Sampled 0:3 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Hotchkiss, Mt. Hermon] Unordered,
  Dim{:school_bis} Categorical String[Deerfield, Choate, Lawrenceville] Unordered
[:, :, 1, 1]
     0          1         2         …  497         498        499
 0  10.2061    -7.31077   5.11594       11.8025     -4.39806    7.22325
 1   0.339696   0.339696  0.339696      10.027       7.18936  -13.3384
 2  -0.458749   5.00749   0.540794      -0.935327    2.57992    3.82478
 3  -0.872076  -0.872076  0.168416       2.42358     2.61346   -0.113698
[and 8 more slices...]

Add new chains using concat

Suppose after checking the mcse and realizing you need more samples, you rerun the model with two chains and obtain an idata_rerun object.

idata_rerun = InferenceData(; posterior=set(post[chain=At([0, 1])]; chain=[4, 5]))
InferenceData
posterior
Dataset with dimensions: 
  Dim{:chain} Sampled Int64[4, 5] ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 3 layers:
  :mu    Float64 dims: Dim{:chain}, Dim{:draw} (2×500)
  :theta Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (2×500×8)
  :tau   Float64 dims: Dim{:chain}, Dim{:draw} (2×500)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"

You can combine the two using concat.

concat(idata[[:posterior]], idata_rerun; dim=:chain)
InferenceData
posterior
Dataset with dimensions: 
  Dim{:chain} Sampled 0:5 ForwardOrdered Regular Points,
  Dim{:draw} Sampled 0:499 ForwardOrdered Regular Points,
  Dim{:school} Categorical String[Choate, Deerfield, …, St. Paul's, Mt. Hermon] Unordered
and 3 layers:
  :mu    Float64 dims: Dim{:chain}, Dim{:draw} (6×500)
  :theta Float64 dims: Dim{:chain}, Dim{:draw}, Dim{:school} (6×500×8)
  :tau   Float64 dims: Dim{:chain}, Dim{:draw} (6×500)

with metadata OrderedCollections.OrderedDict{Symbol, Any} with 3 entries:
  :created_at => "2019-06-21T17:36:34.398087"
  :inference_library_version => "3.7"
  :inference_library => "pymc3"