What's in a FlexiChain?

A FlexiChain{T}, at its core, is a wrapper around a dictionary which maps keys to matrices of size (niters x nchains).

Let's start by setting up an example FlexiChain. Don't worry too much about the code to generate it; the important part is the object that is created.

using FlexiChains: FlexiChains, FlexiChain, VarName, @varname, Parameter, Extra

N_iters, N_chains = 100, 3
values = Dict(
    Parameter(@varname(x)) => rand(N_iters, N_chains),
    Parameter(@varname(y)) => rand(1:5, N_iters, N_chains),
    Parameter(@varname(z)) => [randn(2) for _ in 1:N_iters, _ in 1:N_chains],
    Extra(:something) => rand(N_iters, N_chains),
)
chain = FlexiChain{VarName}(N_iters, N_chains, values)

╭─FlexiChain (100 iterations, 3 chains) ───────────────────────────────────────╮
│ ↓ iter  = 1:100                                                              │
│ → chain = 1:3                                                                │
│                                                                              │
│ Parameters (3) ── VarName                                                    │
│  Float64          x                                                          │
│  Vector{Float64}  z (2,)                                                     │
│  Int64            y                                                          │
│                                                                              │
│ Extras (1)                                                                   │
│  Float64  something                                                          │
╰──────────────────────────────────────────────────────────────────────────────╯

Keys: parameters and extras

Each key in a chain can be either a Parameter{<:T} or an Extra. Each Parameter or Extra itself carries a name, which can be retrieved with FlexiChains.get_name. For a Parameter{T}, the name must be a T. For an Extra, the name can in theory be anything, but in practice is most commonly Symbol.

You can list all keys, all parameters, or all extras:

keys(chain)

KeySet for a OrderedCollections.OrderedDict{Union{FlexiChains.Parameter{var"#s35"}, FlexiChains.Extra} where var"#s35"<:VarName, Matrix} with 4 entries. Keys:
  Parameter(x)
  Extra(:something)
  Parameter(z)
  Parameter(y)

FlexiChains.parameters(chain)

3-element Vector{VarName}:
 x
 z
 y

FlexiChains.extras(chain)

1-element Vector{FlexiChains.Extra}:
 Extra(:something)

Values: matrices

Because each key maps to a different matrix, FlexiChains can store values of different types and sizes.

For example, x is stored as Float64, and z as a Vector{Float64}. Indexing into a chain returns an niters x nchains matrix of these samples:

chain[@varname(x)]

┌ 100×3 DimArray{Float64, 2} Parameter(x) ┐
├─────────────────────────────────────────┴───────────── dims ┐
  ↓ iter Sampled{Int64} 1:100 ForwardOrdered Regular Points,
  → chain Sampled{Int64} 1:3 ForwardOrdered Regular Points
└─────────────────────────────────────────────────────────────┘
   ↓ →  1          2         3
   1    0.125307   0.252942  0.407364
   2    0.162859   0.184877  0.738396
   3    0.166096   0.317212  0.715162
   4    0.976122   0.55471   0.789517
   ⋮                         
  97    0.0876886  0.408641  0.539017
  98    0.812473   0.420145  0.956818
  99    0.126165   0.586887  0.113687
 100    0.883634   0.991117  0.0657168

chain[@varname(z)]

┌ 100×3 DimArray{Vector{Float64}, 2} Parameter(z) ┐
├─────────────────────────────────────────────────┴───── dims ┐
  ↓ iter Sampled{Int64} 1:100 ForwardOrdered Regular Points,
  → chain Sampled{Int64} 1:3 ForwardOrdered Regular Points
└─────────────────────────────────────────────────────────────┘
   ↓ →  1                        …  3
   1     [0.466306, -1.23611]        [1.13695, -0.991442]
   2     [0.0901291, 0.851485]       [-0.0205137, -1.65892]
   3     [-0.149613, -1.58821]       [-0.493337, 0.394369]
   4     [-0.423979, -0.171413]      [0.760748, -0.330966]
   ⋮                             ⋱  
  97     [1.12618, -0.444655]        [0.713035, 0.566492]
  98     [-0.241728, 0.489731]       [0.832556, 1.94047]
  99     [0.97271, 0.67589]          [-0.248712, -0.711949]
 100     [0.211681, 0.379948]    …   [-0.687828, -0.527427]

To obtain the number of iterations or chains, you can use FlexiChains.niters and FlexiChains.nchains:

FlexiChains.niters(chain), FlexiChains.nchains(chain)

(100, 3)

Alternatively, size(chain) returns (niters, nchains):

size(chain)

(100, 3)

Notice that there is no measure of the number of parameters. Unlike other chains packages, such as MCMCChains.jl, FlexiChains does not have a third 'parameter axis'. If you want to get such a number, you can use length(keys(chain)) or length(FlexiChains.parameters(chain)).