Modifying data

While FlexiChain and FlexiSummary are intended to be immutable, there are scenarios where you may want to modify the data that they contain.

FlexiChains provides generic, high-level functions that allow you to modify either the keys or the values stored inside a FlexiChain or FlexiSummary.

All of these functions return a new chain (or summary) and avoid mutating the original chain (or summary).

Modifying keys

Let's load our example chain first:

using FlexiChainschain = FlexiChains.EXAMPLE_CHAIN

╭─FlexiChain (10 iterations, 2 chains) ────────────────────────────────────────╮
│ ↓ iter  = 101:110                                                            │
│ → chain = 1:2                                                                │
│                                                                              │
│ Parameters (3) ── VarName                                                    │
│  Float64          mu, tau                                                    │
│  Vector{Float64}  theta (8,)                                                 │
│                                                                              │
│ Extras (2)                                                                   │
│  Float64  logjoint                                                           │
│  Bool     is_divergent                                                       │
╰──────────────────────────────────────────────────────────────────────────────╯

If you want to change the keys stored inside a FlexiChain or FlexiSummary, you can use the FlexiChains.map_keys function.

using FlexiChains: get_name, map_keys, Parameter, Extra, VarNameusing VarNames: getsym, getoptic# Define a function that takes the old key and returns a new key.function f(p::Parameter{<:VarName})    # Get the name of the parameter (e.g. @vn(x))    name = get_name(p)    # Get the components of that VarName.    sym, optic = getsym(name), getoptic(name)    # Add a prefix to the symbol.    new_sym = Symbol("new_", sym)    # Return a new Parameter with the new name    return Parameter{VarName}(VarName{new_sym}(optic))end# Don't change extras!f(p::Extra) = pmap_keys(f, chain)

╭─FlexiChain (10 iterations, 2 chains) ────────────────────────────────────────╮
│ ↓ iter  = 101:110                                                            │
│ → chain = 1:2                                                                │
│                                                                              │
│ Parameters (3) ── VarName                                                    │
│  Float64          new_mu, new_tau                                            │
│  Vector{Float64}  new_theta (8,)                                             │
│                                                                              │
│ Extras (2)                                                                   │
│  Float64  logjoint                                                           │
│  Bool     is_divergent                                                       │
╰──────────────────────────────────────────────────────────────────────────────╯

This example is admittedly quite contrived. However, it is intended to be general enough to cover any use case you might have. For example, if you want to rename a single key, f can simply be a function that looks it up in a dictionary:

mapping = Dict(old_key => new_key)f(p) = get(mapping, p, p)  # Return the new key if it exists, otherwise return the old key.

Often for the most part you will only want to modify Parameters, in which case you can use FlexiChains.map_parameters. This has the same behaviour as map_keys, except that you can skip the wrapping/unwrapping in Parameter.

Note that both map_keys and map_parameters work with FlexiSummary as well.

Modifying values

It is also possible to modify the values stored inside a FlexiChain (but not a FlexiSummary). This is done with the transform_values function:

using FlexiChains: transform_valueschain2 = transform_values(    chain,    @vn(mu) => (i -> i + 1),    @vn(tau) => (i -> i * 2) => @vn(new_tau))

╭─FlexiChain (10 iterations, 2 chains) ────────────────────────────────────────╮
│ ↓ iter  = 101:110                                                            │
│ → chain = 1:2                                                                │
│                                                                              │
│ Parameters (4) ── VarName                                                    │
│  Float64          mu, tau, new_tau                                           │
│  Vector{Float64}  theta (8,)                                                 │
│                                                                              │
│ Extras (2)                                                                   │
│  Float64  logjoint                                                           │
│  Bool     is_divergent                                                       │
╰──────────────────────────────────────────────────────────────────────────────╯

The above example:

• adds 1 to each element of chain[@vn(mu)]; and
• multiplies each element of chain[@vn(tau)] by 2, and stores the result in a new key @vn(new_tau).

chain2[@vn(mu)] .- chain[@vn(mu)]   # Should be all 1.

┌ 10×2 DimArray{Float64, 2}  ┐
├────────────────────────────┴──────────────────────────── dims ┐
  ↓ iter Sampled{Int64} 101:110 ForwardOrdered Regular Points,
  → chain Sampled{Int64} 1:2 ForwardOrdered Regular Points
└───────────────────────────────────────────────────────────────┘
   ↓ →  1    2
 101    1.0  1.0
 102    1.0  1.0
 103    1.0  1.0
 104    1.0  1.0
   ⋮         
 107    1.0  1.0
 108    1.0  1.0
 109    1.0  1.0
 110    1.0  1.0

You can pass as many transformations as you like. The syntax is designed to be similar to that of DataFrames.transform, but has some slight differences: notably, the function being applied acts on individual draws from chain[@vn(x)] rather than the matrix of all draws. (In DataFrames, the functions act on entire columns of data.)

You can also pass binary (or n-ary) functions to transform_values to combine multiple keys. Again, this is similar to DataFrames.transform, but the function combines individual draws from chain[@vn(mu)] and chain[@vn(tau)] rather than the matrices themselves.

chain3 = transform_values(chain, [@vn(mu), @vn(tau)] => (+) => @vn(mu_plus_tau))

╭─FlexiChain (10 iterations, 2 chains) ────────────────────────────────────────╮
│ ↓ iter  = 101:110                                                            │
│ → chain = 1:2                                                                │
│                                                                              │
│ Parameters (4) ── VarName                                                    │
│  Float64          mu, tau, mu_plus_tau                                       │
│  Vector{Float64}  theta (8,)                                                 │
│                                                                              │
│ Extras (2)                                                                   │
│  Float64  logjoint                                                           │
│  Bool     is_divergent                                                       │
╰──────────────────────────────────────────────────────────────────────────────╯

chain3[@vn(mu_plus_tau)] == chain[@vn(mu)] .+ chain[@vn(tau)]

true

Attaching labels to data

A common use case for transform_values is to attach labels to data, for example, converting a Vector of parameters into a DimVector.

In our example chain above, theta has one entry per school, but is stored as a plain Vector. While this is fine for many applications, attaching labels can help make it easier to analyse the data.

using DimensionalData: DimArray, Dimschool_names = [    "Choate",    "Deerfield",    "Phillips Andover",    "Phillips Exeter",    "Hotchkiss",    "Lawrenceville",    "St. Paul's",    "Mt. Hermon",]add_labels(v::Vector{<:Real}) = DimArray(v, Dim{:school}(school_names))labelled_chain = transform_values(chain, @vn(theta) => add_labels)

╭─FlexiChain (10 iterations, 2 chains) ────────────────────────────────────────╮
│ ↓ iter  = 101:110                                                            │
│ → chain = 1:2                                                                │
│                                                                              │
│ Parameters (3) ── VarName                                                    │
│  Float64                                   mu, tau                           │
│  DimensionalData.DimVector{Float64, Tupl…  theta (8,)                        │
│                                                                              │
│ Extras (2)                                                                   │
│  Float64  logjoint                                                           │
│  Bool     is_divergent                                                       │
╰──────────────────────────────────────────────────────────────────────────────╯

For example, DimVector parameters get special labels when plotting:

using CairoMakieFlexiChains.Makie.traceplot(labelled_chain, @vn(theta); layout=(4, 2))

Having the labels also benefits downstream analysis of any data extracted from the chain. For example, you can obtain the same functionality as tidybayes' spread_draws:

using DataFrames# Subset to the first three iterations, or else the output will be very long!labelled_chain = labelled_chain[iter=1:3]# Create a 3D DimArray (iter x chain x school)dimarr = labelled_chain[@vn(theta), stack=true]# Convert it to a DataFrame (long-format by default)df = DataFrame(dimarr)

If you want a wide-format table, you can tap into the functionality in DimensionalData.jl (please see their docs for full info):

using DimensionalData: DimTabledf = DataFrame(DimTable(dimarr; layersfrom=:school))

Standardisation

Another use case for transforming values is to (un)standardise parameter values.

For example, consider a case where we have some model data y which is standardised before being passed to a model:

y = randn(100) .* 4.0 .+ 2.0using StatsBase: fit, ZScoreTransform, transformzs = fit(ZScoreTransform, y)y_scaled = transform(zs, y)mean(y_scaled), std(y_scaled)  # ≈ 0 and 1

(6.88338275267597e-17, 0.9999999999999998)

Then, suppose that we pass y_scaled to a model and sample a parameter m_scaled which is the mean of the standardised y. We'll construct one such chain for demonstration purposes:

using FlexiChains: FlexiChain, VarName, Parameterniters, nchains = 1000, 2fake_m_data = [randn(100) .+ 1 for _ in 1:niters, _ in 1:nchains]fake_chain =    FlexiChain{VarName}(niters, nchains, Dict(Parameter(@vn(m_scaled)) => fake_m_data))

╭─FlexiChain (1000 iterations, 2 chains) ──────────────────────────────────────╮
│ ↓ iter  = 1:1000                                                             │
│ → chain = 1:2                                                                │
│                                                                              │
│ Parameters (1) ── VarName                                                    │
│  Vector{Float64}  m_scaled (100,)                                            │
│                                                                              │
│ Extras (0)                                                                   │
│  (none)                                                                      │
╰──────────────────────────────────────────────────────────────────────────────╯

The mean of this m_scaled should be approximately 1:

mean(fake_chain[@vn(m_scaled), stack=true])

0.9963935660290663

We can unstandardise the m_scaled values by applying the inverse transformation:

using StatsBase: reconstructfake_chain =    transform_values(fake_chain, @vn(m_scaled) => (i -> reconstruct(zs, i)) => @vn(m))

╭─FlexiChain (1000 iterations, 2 chains) ──────────────────────────────────────╮
│ ↓ iter  = 1:1000                                                             │
│ → chain = 1:2                                                                │
│                                                                              │
│ Parameters (2) ── VarName                                                    │
│  Vector{Float64}  m_scaled (100,), m (100,)                                  │
│                                                                              │
│ Extras (0)                                                                   │
│  (none)                                                                      │
╰──────────────────────────────────────────────────────────────────────────────╯

The mean of the reconstructed m should be approximately (4 * 1) + 2 = 6:

mean(fake_chain[@vn(m), stack=true])

6.116315588145866

Docstrings

FlexiChains.map_keys(f, cs::ChainOrSummary{T})::ChainOrSummary{S} where {T,S}

FlexiChains.map_parameters(f, cs::ChainOrSummary{T})::ChainOrSummary{S} where {T,S}

FlexiChains.transform_values(chn::FlexiChain{T}, args...)