Matryoshka.jl ​

Matryoshka.jl builds Bayesian regression models from composable components and returns a MatryoshkaModel wrapping a DynamicPPL.Model.

Installation ​

This package is not registered with the General registry. You can install it via URL:

import Pkg
Pkg.add(url = "https://github.com/simonsteiger/Matryoshka.jl")

Get started ​

The example below shows how to use Matryoshka to fit a single categorical predictor to a continuous response variable:

using Matryoshka, Turing, FlexiChains
using PalmerPenguins, DataFrames, CategoricalArrays
using StatsBase: denserank

# We will z-standardise the outcome
standardise(x) = (x .- mean(x)) ./ std(x)

data = DataFrame(PalmerPenguins.load())

# There's one poor penguin with missings
dropmissing!(data)
transform!(data, :bill_length_mm => standardise => identity)
transform!(data, :species => categorical => identity)

# Bill length by species, Normal family
likelihood = @likelihood Normal bill_length_mm ~ species

# Standard Normal priors on intercept and species coefs
priors = @priors begin
    Intercept() ~ Normal(0, 1)
    FixedEffect() ~ Normal(0, 1)
end

# Now we create a `MatryoshkaModel`, then Turing takes over
bill_model = model(likelihood, priors, data)
chain = sample(bill_model, NUTS(), 1000)

# summarystats(chain) row names are selectors, e.g. `Intercept()` (Adelie),
# `FixedEffect(:species_Chinstrap)`, `FixedEffect(:species_Gentoo)` —
# Adelie have shortest bills (they're cute!)
summarystats(chain)

╭─FlexiSummary (9 statistics) ─────────────────────────────────────────────────╮
│   iter    collapsed                                                          │
│   chain   collapsed                                                          │
│ ↓ stat  = [mean, std, mcse, ess_bulk, ess_tail, rhat, q5, q50, q95]          │
│                                                                              │
│ Parameters (4) ── MatryoshkaParam                                            │
│  Float64  Intercept(), FixedEffect(:species_Chinstrap),                      │
│           FixedEffect(:species_Gentoo), Residuals()                          │
│                                                                              │
│ Extras (14)                                                                  │
│  Float64  n_steps, is_accept, acceptance_rate, log_density,                  │
│           hamiltonian_energy, hamiltonian_energy_error,                      │
│           max_hamiltonian_energy_error, tree_depth, numerical_error,         │
│           step_size, nom_step_size, logprior, loglikelihood, logjoint        │
│                                                                              │
│ Summary                                                                      │
│          param     mean     std    mcse  ess_bulk  ess_tail    rhat  …       │
│    Intercept()  -0.9348  0.0449  0.0019  570.4345  679.5142  1.0045  …       │
│   FixedEffect…   1.8108  0.0749  0.0029  673.9366  726.3871  1.0007  …       │
│   FixedEffect…   1.5853  0.0692  0.0028  590.9397  681.3907  0.9993  …       │
│    Residuals()   0.5447  0.0211  0.0008  619.6507  731.1019  1.0052  …       │
╰──────────────────────────────────────────────────────────────────────────────╯

Parameter selectors ​

sample(model, ...) returns a FlexiChain{MatryoshkaParam}: every parameter is keyed by a selector, a small value type that names what the parameter is rather than an opaque coefficient index. The same selectors are the @priors targets, so there is one vocabulary for "what am I putting a prior on" and "what am I pulling out of the chain":

A selector with unset fields (FixedEffect(), SD(), RandomEffect(), RandomEffect(:g)) is a class selector — it matches a family of parameters. A selector with every field set (FixedEffect(:x), RandomEffect(:g, :a)) is exact — it names one parameter and is a valid chain key. Intercept() and Residuals() are always exact.

Coefficient and level names are sanitized StatsModels names:

Rules: ": " becomes _; " & " becomes __; level strings are stripped to identifier-safe characters ("Very High" → VeryHigh). If two sanitized names collide, model() errors and asks you to rename the offending column or level. Use default_priors(lik, df) to list every target for your model and data.

@priors ​

@priors accepts only selector calls on the left of ~:

priors = @priors begin
    Intercept() ~ Normal(0, 1)
    FixedEffect() ~ Normal(0, 1)        # class selector: every coefficient
    FixedEffect(:x) ~ Normal(0, 5)      # exact selector: overrides the class above for `:x`
    SD(:g) ~ Exponential(1)
    Residuals() ~ Exponential(1)
end

Resolution order per parameter: exact target > class target > component default > family default. Old string/@varname-based targets are not accepted — @priors errors with a message naming the valid selector types.

Indexing a chain ​

An exact selector returns the raw draws matrix; a class selector returns a sub-chain; a class selector with stack = true returns a labeled DimArray stacking every matching parameter along a new axis:

chain[FixedEffect(:x)]                    # (iter, chain) draws matrix
chain[RandomEffect(:g)]                   # sub-chain: every level of group :g
chain[FixedEffect(), stack = true]        # DimArray, (iter, chain, coef)
chain[RandomEffect(:g), stack = true][g = At(:a)]   # DimensionalData indexing on the stacked axis

summarystats(chain) (from FlexiChains) displays one row per parameter using these same selector names — every displayed name is a valid index into the chain (chain[p] works for any p in summarystats(chain)'s parameters).

predict ​

DynamicPPL.predict(m_new, chain) returns a DimArray shaped (iter, chain, obs) — one predicted draw per posterior sample, per observation in m_new's data:

m_new = model(bill_model, newdata)   # no response column → predict-mode
preds = predict(m_new, chain)
preds[obs = At(1)]                   # per-observation slice

Escape hatches ​

• m.inner — the raw DynamicPPL.Model; use it with any DynamicPPL/Turing tooling Matryoshka does not wrap directly.

• rekey(m, vnchain) — convert a VarName-keyed chain (e.g. one sampled via m.inner directly, or with chain_type overridden) into the selector-keyed chain sample(m, ...) normally returns.

• to_vnchain(m, chain) — the inverse of rekey; reassemble the VarName-keyed chain DynamicPPL understands.

• chain_type keyword — pass chain_type = ... to sample(m, ...) to skip rekey entirely and get whatever chain type you asked for, unmodified.