Indexing
A FlexiChain stores data in a rich format: instead of just storing a raw matrix of data, it also includes information about the iteration numbers and chain numbers.
Additionally, FlexiSummary objects (which you get when performing any kind of summarisation on a chain, e.g. with summarystats) also sometimes store information about which summary functions were applied (especially when there are multiple of these).
This information is used when constructing the DimensionalData.DimArray outputs that you see when indexing into a FlexiChain or FlexiSummary object. But, on top of this, it also allows you to more surgically index into these objects using this information.
This page first begins with some illustrative examples, which might be the clearest way to demonstrate the indexing behaviour. If you prefer reading a fuller specification, the sections below that describe the exact behaviour in more detail.
Examples: chains
Let's first set up a chain:
using FlexiChains, Turing
@model function f()
x ~ MvNormal(zeros(2), I)
end
chn = sample(
f(),
MH(),
MCMCThreads(),
5,
2;
discard_initial=100,
chain_type=VNChain,
progress=false,
verbose=false,
)╭─FlexiChain (5 iterations, 2 chains) ─────────────────────────────────────────╮
│ ↓ iter = 101:105 │
│ → chain = 1:2 │
│ │
│ Parameters (1) ── VarName │
│ Vector{Float64} x (2,) │
│ │
│ Extras (4) │
│ Bool accepted │
│ Float64 logprior, loglikelihood, logjoint │
╰──────────────────────────────────────────────────────────────────────────────╯Notice how the iteration numbers here start from 101: that is because of the discard_initial argument.
# Picking out a single parameter; this returns a `DimMatrix`.
chn[@varname(x)]┌ 5×2 DimArray{Vector{Float64}, 2} Parameter(x) ┐
├───────────────────────────────────────────────┴───────── dims ┐
↓ iter Sampled{Int64} 101:105 ForwardOrdered Regular Points,
→ chain Sampled{Int64} 1:2 ForwardOrdered Regular Points
└───────────────────────────────────────────────────────────────┘
↓ → 1 2
101 [-0.244415, -0.669353] [-1.42033, -1.28706]
102 [-0.0427301, -2.22401] [0.63706, -1.24444]
103 [-0.0393705, -0.386514] [1.67414, -2.95014]
104 [0.415553, 1.01929] [0.371067, -1.14188]
105 [-0.405602, 1.31328] [-1.53292, 0.801801]# This picks out the first of the iterations (note: this has iteration number 101)
chn[@varname(x), iter=1]┌ 2-element DimArray{Vector{Float64}, 1} Parameter(x) ┐
├─────────────────────────────────────────────────────┴ dims ┐
↓ chain Sampled{Int64} 1:2 ForwardOrdered Regular Points
└────────────────────────────────────────────────────────────┘
1 [-0.244415, -0.669353]
2 [-1.42033, -1.28706]Keyword arguments to getindex
Note that keyword arguments when indexing with square brackets must be separated from positional arguments by a comma. Using a semicolon will lead to an error.
# You can also select a specific chain
chn[@varname(x), iter=1, chain=2]2-element Vector{Float64}:
-1.420330333068879
-1.2870636116009129# This picks out iteration number 101
chn[@varname(x), iter=At(101)]┌ 2-element DimArray{Vector{Float64}, 1} Parameter(x) ┐
├─────────────────────────────────────────────────────┴ dims ┐
↓ chain Sampled{Int64} 1:2 ForwardOrdered Regular Points
└────────────────────────────────────────────────────────────┘
1 [-0.244415, -0.669353]
2 [-1.42033, -1.28706]# This picks out iteration numbers 101 through 103
chn[@varname(x), iter=101..103]┌ 3×2 DimArray{Vector{Float64}, 2} Parameter(x) ┐
├───────────────────────────────────────────────┴───────── dims ┐
↓ iter Sampled{Int64} 101:103 ForwardOrdered Regular Points,
→ chain Sampled{Int64} 1:2 ForwardOrdered Regular Points
└───────────────────────────────────────────────────────────────┘
↓ → 1 2
101 [-0.244415, -0.669353] [-1.42033, -1.28706]
102 [-0.0427301, -2.22401] [0.63706, -1.24444]
103 [-0.0393705, -0.386514] [1.67414, -2.95014]# If you only want the first element of `x`:
chn[@varname(x[1])]┌ 5×2 DimArray{Float64, 2} Parameter(x[1]) ┐
├──────────────────────────────────────────┴────────────── dims ┐
↓ iter Sampled{Int64} 101:105 ForwardOrdered Regular Points,
→ chain Sampled{Int64} 1:2 ForwardOrdered Regular Points
└───────────────────────────────────────────────────────────────┘
↓ → 1 2
101 -0.244415 -1.42033
102 -0.0427301 0.63706
103 -0.0393705 1.67414
104 0.415553 0.371067
105 -0.405602 -1.53292# You can specify a vector of parameters
chn[[@varname(x[1]), :logjoint]]╭─FlexiChain (5 iterations, 2 chains) ─────────────────────────────────────────╮
│ ↓ iter = 101:105 │
│ → chain = 1:2 │
│ │
│ Parameters (1) ── VarName │
│ Float64 x[1] │
│ │
│ Extras (1) │
│ Float64 logjoint │
╰──────────────────────────────────────────────────────────────────────────────╯This last one returns a FlexiChain object, because multiple keys were specified. The data that we didn't care for, such as @varname(x[2]), are simply dropped.
Notice that this also gives us a way to 'flatten' a FlexiChain object such that all of its keys point to scalar values. We just need to find a full set of sub-VarNames, like the following. In practice you don't need to construct this set yourself: FlexiChains has an internal function, called _split_varnames, that will do this for you.
chn[[@varname(x[1]), @varname(x[2])]]╭─FlexiChain (5 iterations, 2 chains) ─────────────────────────────────────────╮
│ ↓ iter = 101:105 │
│ → chain = 1:2 │
│ │
│ Parameters (2) ── VarName │
│ Float64 x[1], x[2] │
│ │
│ Extras (0) │
│ (none) │
╰──────────────────────────────────────────────────────────────────────────────╯Examples: summaries
Calling any summary function such as mean, std, or summarystats on a FlexiChain object will return a FlexiSummary object.
sm = summarystats(chn)╭─FlexiSummary (9 statistics) ─────────────────────────────────────────────────╮
│ iter collapsed │
│ chain collapsed │
│ ↓ stat = [mean, std, mcse, ess_bulk, ess_tail, rhat, q5, q50, q95] │
│ │
│ Parameters (2) ── VarName │
│ Float64 x[1], x[2] │
│ │
│ Extras (4) │
│ Float64 accepted, logprior, loglikelihood, logjoint │
│ │
│ Summary │
│ param mean std mcse ess_bulk ess_tail rhat q5 … │
│ x[1] -0.0588 0.9475 0.2996 10.0000 10.0000 1.7894 -1.4823 … │
│ x[2] -0.6769 1.3976 0.5243 6.3147 10.0000 1.1071 -2.6234 … │
╰──────────────────────────────────────────────────────────────────────────────╯Notice two things: 2. This summary no longer has iter or chain dimensions, because the summary statistics have been calculated over all iterations and chains. However, it has a stat dimension, which we will need to use when accessing the data.
- The variable
xhas been broken up for you into its componentsx[1]andx[2].
Indexing is very similar as for chains, but there is an additional stat dimension which lets you specify which summary statistic you want to access.
sm[@varname(x[1]), stat=:mean]-0.058755369284811246Named statistic shorthand
Passing a Symbol to stat is shorthand for a named lookup: stat=:mean is equivalent to stat=At(:mean). Other selectors follow DimensionalData.jl's usual behaviour; for example, stat=1 selects the first statistic by position.
If only a single summary function was applied, e.g. via mean(chn), then the stat dimension will be automatically collapsed for you; you won't need to again specify stat when indexing.
sm_mean = mean(chn)
sm_mean[@varname(x[1])]-0.058755369284811246If you don't want to split the VarNames up, you can specify this as a keyword argument.
sm_mean_nosplit = mean(chn; split_varnames=false)
sm_mean_nosplit[@varname(x)]2-element Vector{Float64}:
-0.058755369284811246
-0.6769025989601005If you collapse only over iterations (for example), then you can specify the chain keyword argument (and likewise for iter if you collapse over chains).
sm_iter = mean(chn; dims=:iter)
sm_iter[@varname(x[1]), chain=2]-0.05419773091504143Positional arguments
When indexing into a FlexiChain (or FlexiSummary) object, you can use one optional positional argument. This positional argument can either be an object pointing to a single key, in which case a DimMatrix is returned; or it can be an object pointing to multiple keys, in which case a FlexiChain (or FlexiSummary) is returned.
To specify a single key, you can use:
a parameter name (e.g. for a
FlexiChain{T}, an object of typeT);a
VarNameor sub-VarName, for aFlexiChain{VarName}(i.e.VNChain);a
FlexiChains.Extrafor non-parameter keys;a
Symbol, as long as it refers to an unambiguous key;a
FlexiChains.Parameter(this is mentioned for completeness; as a user you probably don't need to do this)
On the other hand, you could specify multiple keys via:
a
Vectorcontaining any combination of the above;a colon
:, which refers to all keys in the chain or summary.
If a positional argument is not specified, it defaults to :.
Keyword arguments: chains
iter and chain
In addition to the positional argument, you can also specify the iter and chain keyword arguments when indexing into the FlexiChain object. (FlexiSummary objects are covered right below this.) Both of these are optional, and exist to allow you to select data from specific iterations and/or chains.
Keyword arguments to getindex
When indexing with square brackets, the keyword arguments must be separated from positional arguments by a comma, not a semicolon as is usual for other Julia functions. That is to say, you should use:
chn[param, iter=iter, chain=chain]
# this is also fine, albeit a bit wordy
getindex(chn, param; iter=iter, chain=chain)rather than
# this will error
# chn[param; iter=iter, chain=chain]The allowed values for these keyword arguments almost exactly mimic the behaviour of DimensionalData.jl. Suppose that you sampled a chain with 100 iterations, but with a thinning factor of 2. FlexiChains will record this information, and its iteration numbers will be 1:2:199 (i.e. 1, 3, 5, ..., 199).
For clarity, we will refer to the actual iteration numbers (1, 3, 5, ..., 199) as iteration numbers, and the entries in the chain (1st entry, 2nd entry, ..., 100th entry) as entries.
You can then specify, for example:
iter=... | Description |
|---|---|
5 | the fifth entry in the chain, i.e. iteration number 9 |
At(9) | iteration number 9 |
Not(5) | all entries except the fifth one, i.e. all iteration numbers except 9 |
Not(At(9)) | all entries except iteration number 9 |
6..30 | all iteration numbers between 6 and 30, i.e. all but the first entry |
2:10 | 2nd through 10th entries, i.e. iteration numbers 6 through 30 |
[At(9), At(30)] | this will get the entries corresponding to iteration numbers 9 and 30 |
: | all entries (i.e. all iteration numbers) |
For convenience, FlexiChains re-exports the DimensionalData.jl selectors such as Not, At, and ...
The same applies to the chain keyword argument, except that here you are selecting which chains to include. This is slightly less interesting because chains are always numbered consecutively starting from 1. Consequently, i and At(i) have the same meaning. Nonetheless, you can still use all the same selectors as described above, e.g. Not(2) to drop the second chain.
Finally, note that regular begin and end cannot be used with the iter and chain keyword arguments, because Julia's lowering pass does not handle those appropriately for keyword arguments. In place of this you can use Begin and End, which are semantically equivalent and re-exported from DimensionalData.
chn[@varname(x), iter=(Begin+1):End]┌ 4×2 DimArray{Vector{Float64}, 2} Parameter(x) ┐
├───────────────────────────────────────────────┴───────── dims ┐
↓ iter Sampled{Int64} 102:105 ForwardOrdered Regular Points,
→ chain Sampled{Int64} 1:2 ForwardOrdered Regular Points
└───────────────────────────────────────────────────────────────┘
↓ → 1 2
102 [-0.0427301, -2.22401] [0.63706, -1.24444]
103 [-0.0393705, -0.386514] [1.67414, -2.95014]
104 [0.415553, 1.01929] [0.371067, -1.14188]
105 [-0.405602, 1.31328] [-1.53292, 0.801801]For more information about DimensionalData's selectors, please see their docs.
stack
The stack keyword argument only affects parameters which are arrays; other parameters ignore it.
If stack=false (the default), then accessing an array-valued parameter will return a DimArray of Arrays. If stack=true, then the array-valued parameter will be stacked into a single DimArray. Note that the parameter's axes will be placed at the end of the returned DimArray. Thus, for example if the parameter :x is a Vector{T}, then chn[:x] will return a DimArray{T,3} with shape (iters, chains, length(x)).
DimArray
In the current version of FlexiChains, for DimArray-valued parameters the stacking happens by default. This will be removed in a future version: you will have to explicitly specify stack=true to get this behaviour.
Keyword arguments: summaries
Positional arguments
The positional argument when indexing into a FlexiSummary objects is exactly the same as for FlexiChain. Only keyword arguments behave differently.
There are two differences between a FlexiChain and a FlexiSummary in terms of their indexing behaviour:
FlexiSummaryobjects contain one additional dimension, calledstat.FlexiSummarydimensions may be collapsed, meaning that they cannot be indexed into.
Consequently, there are three possible keyword arguments: iter, chain, and stat; but depending on which dimensions have been collapsed, you may not be able to use them.
iter and chain
In general, if you apply a summary function like mean without specifying dimensions, then both iter and chain dimensions will be collapsed.
If you have performed the mean over a single dimension only, such as via summary = mean(chn; dims=:iter), then the iter dimension will be collapsed, but you can still index into the chain dimension using summary[key, chain=...].
stat
In general, the stat dimension is generally:
not collapsed if multiple summary functions were applied, e.g. via
summarystats(chn);collapsed if a single summary function was applied, e.g. via
mean(chn).
Unlike the iter and chain dimensions, the stat dimension's indices are Symbols instead of numbers. Thus, for example, if you have a summary that contains the mean and std of the chain, you could use:
stat=... | Description |
|---|---|
1 | the first statistic, i.e. :mean |
At(:mean) | the :mean statistic |
:mean | shorthand for At(:mean) |
Not(At(:mean)) | everything but the :mean statistic |
stack
The stack keyword argument behaves the same as for FlexiChain objects.