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                                                    │
│  Vector{Float64}  z (2,)                                                     │
│  Int64            y                                                          │
│  Float64          x                                                          │
│                                                                              │
│ 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"#s22"}, FlexiChains.Extra} where var"#s22"<:VarName, Matrix} with 4 entries. Keys:
  Parameter(z)
  Parameter(y)
  Parameter(x)
  Extra(:something)

FlexiChains.parameters(chain)

3-element Vector{VarName}:
 z
 y
 x

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.150646   0.988284  0.319257
   2    0.707445   0.223128  0.545701
   3    0.451996   0.712241  0.685354
   4    0.356777   0.120077  0.36447
   ⋮                         
  97    0.253731   0.19106   0.96936
  98    0.882421   0.881372  0.362167
  99    0.0404113  0.732511  0.715824
 100    0.536547   0.450988  0.681791

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     [-1.09501, 0.0953908]       [-0.494889, -1.7073]
   2     [1.46972, 0.781811]         [-0.23304, 0.00330474]
   3     [0.525796, -1.85899]        [0.43996, 0.199913]
   4     [-0.795238, -0.646161]      [0.976783, 0.0378068]
   ⋮                             ⋱  
  97     [0.418158, 1.87304]         [-0.185688, 1.62558]
  98     [1.41462, -1.11359]         [-1.94139, -0.877165]
  99     [0.0684099, -0.687165]      [-1.72699, 1.37733]
 100     [1.04224, -1.16396]     …   [-2.56339, 0.307695]

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)).