Parameter space

Fit a PyMC model, extract posterior draws with parameter_draws(), and plot them

parameter_draws() is the entry point for parameter-space plots: it pulls posterior draws out of an ArviZ DataTree into a tidy Polars DataFrame — one row per chain × draw × coordinate, each variable a column. This example starts with simulated observed data, fits a real PyMC model, and then uses the resulting posterior draws for densities, intervals, contrasts, and cross-parameter plots.

The string spec

A spec names a variable and, in brackets, the dimensions to spread over. The bracketed names must match coordinate names in the DataTree.

Spec	Meaning	Rows
`"sigma"`	scalar parameter	`chain × draw`
`"beta[groups]"`	one-dimensional array	`chain × draw × groups`
`"intercept[groups]"`	another group-level array	`chain × draw × groups`

Request several variables in one call and tidydraws joins them: variables sharing a dimension are inner-joined; a scalar is broadcast across an array’s dimensions.

Full PyMC workflow

The data are generated from four groups with deliberately different intercepts and slopes, including one negative slope. That separation makes the posterior plots diagnose real group differences rather than noise around one common line.

Code

from pathlib import Path
import sys

import numpy as np
import polars as pl
import pymc as pm
import tidydraws as td
from lets_plot import (
    LetsPlot,
    aes,
    facet_wrap,
    geom_density,
    geom_hline,
    geom_jitter,
    geom_line,
    geom_linerange,
    geom_point,
    geom_pointrange,
    geom_vline,
    ggplot,
    labs,
    scale_color_gradient,
)

for parent in [Path.cwd(), *Path.cwd().parents]:
    helper_dir = parent / "docs" / "examples"
    if (helper_dir / "_pymc_workflow.py").exists():
        sys.path.insert(0, str(helper_dir))
        break

from _pymc_workflow import interval_summary, simulate_grouped_regression

LetsPlot.setup_html()

workflow = simulate_grouped_regression(seed=2026)
observed = workflow.observed
truth = workflow.truth

The observed data and true generating lines show the group separation before fitting.

(
    ggplot(observed.sort(["groups", "x"]).to_pandas(), aes("x", "y"))
    + geom_point(aes(color="groups"), alpha=0.65, size=2.0)
    + geom_line(aes(y="mu_true", color="groups"), size=1.0)
    + labs(
        x="x", y="y", color="group", title="Observed data from known group differences"
    )
)

Simulated observed data with true group-specific regression lines.

Now we build our PyMC model and fit.

coords = {
    "groups": workflow.group_names,
    "obs_ind": observed.get_column("obs_ind").to_numpy(),
}

with pm.Model(coords=coords) as model:
    x = pm.Data("x", observed.get_column("x").to_numpy(), dims="obs_ind")
    group_idx = pm.Data(
        "group_idx",
        observed.get_column("group_idx").to_numpy().astype("int64"),
        dims="obs_ind",
    )
    intercept = pm.Normal("intercept", mu=0.0, sigma=2.0, dims="groups")
    beta = pm.Normal("beta", mu=0.0, sigma=1.5, dims="groups")
    sigma = pm.HalfNormal("sigma", sigma=1.0)
    mu = pm.Deterministic(
        "mu",
        intercept[group_idx] + beta[group_idx] * x,
        dims="obs_ind",
    )
    pm.Normal(
        "y",
        mu=mu,
        sigma=sigma,
        observed=observed.get_column("y").to_numpy(),
        dims="obs_ind",
    )
    dt = pm.sample(
        draws=400,
        tune=400,
        random_seed=2026,
    )

Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [intercept, beta, sigma]

Sampling 2 chains for 400 tune and 400 draw iterations (800 + 800 draws total) took 1 seconds.
We recommend running at least 4 chains for robust computation of convergence diagnostics
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details

dt

One parameter_draws() call extracts group slopes and intercepts from the fitted posterior:

beta_df = td.parameter_draws(dt, "beta[groups]", "intercept[groups]")
beta_df.head()

shape: (5, 5)

chain	draw	groups	beta	intercept
i64	i64	str	f64	f64
0	0	"North"	0.356443	-1.365803
0	0	"South"	0.647371	0.253135
0	0	"East"	1.384258	1.345093
0	0	"West"	-0.538669	2.264232
0	1	"North"	0.505888	-1.138709

Gallery

Posterior density by group

The dashed red lines are the true slopes used to simulate the data.

(
    ggplot(beta_df.to_pandas(), aes("beta", fill="groups"))
    + geom_density(alpha=0.45)
    + geom_vline(
        data=truth.to_pandas(),
        mapping=aes(xintercept="beta_true"),
        color="firebrick",
        linetype="dashed",
        size=0.8,
    )
    + facet_wrap(facets="groups", ncol=2)
    + labs(x="beta", y="density", fill="group", title="Posterior slopes by group")
)

Posterior density of group-specific slopes, with true values marked.

Forest plot

Summarise each group with one central 89% interval. The red points are the true slopes used to generate the observed data.

beta_forest = interval_summary(beta_df, "beta", "groups", [0.89])

(
    ggplot(beta_forest.to_pandas(), aes("groups", "median"))
    + geom_pointrange(
        aes(ymin="lower", ymax="upper"),
        color="steelblue",
        size=0.9,
    )
    + geom_point(
        data=truth.to_pandas(),
        mapping=aes("groups", "beta_true"),
        color="firebrick",
        size=2.4,
    )
    + geom_hline(yintercept=0, linetype="dashed", color="#888888")
    + labs(x="group", y="beta", title="Posterior slope forest plot")
)

Posterior 89% intervals of group-specific slopes, with true values in red.

Quantile dotplot

Each row below has 100 equally likely dots from the posterior slope distribution for a group. This is a frequency-format alternative to density plots.

quantiles = np.linspace(0.005, 0.995, 100)
group_levels = (
    beta_df.select("groups").unique().sort("groups").get_column("groups").to_list()
)
quantile_dots = pl.DataFrame([
    {
        "groups": group_name,
        "beta": float(
            beta_df
            .filter(pl.col("groups") == group_name)
            .get_column("beta")
            .quantile(q)
        ),
        "dot": dot,
    }
    for group_name in group_levels
    for dot, q in enumerate(quantiles, start=1)
])

(
    ggplot(quantile_dots.to_pandas(), aes("beta", "groups"))
    + geom_point(size=1.4, alpha=0.65, color="steelblue")
    + geom_point(
        data=truth.to_pandas(),
        mapping=aes("beta_true", "groups"),
        color="firebrick",
        size=2.4,
    )
    + labs(x="beta", y="group", title="100-dot posterior summaries")
)

Quantile dotplot of group-specific slopes: each dot represents 1% posterior mass.

Derived contrasts against a reference group

Because every row keeps its chain and draw, derived quantities are ordinary dataframe operations. Here each contrast is beta[group] - beta[North] within draw, summarised with one 89% interval.

reference_group = "North"
wide_beta = beta_df.pivot(index=["chain", "draw"], on="groups", values="beta")
contrast_draws = pl.concat([
    wide_beta.select(
        "chain",
        "draw",
        (pl.col(group_name) - pl.col(reference_group)).alias("contrast"),
        pl.lit(f"{group_name} - {reference_group}").alias("contrast_name"),
    )
    for group_name in group_levels
    if group_name != reference_group
])
contrast_forest = interval_summary(contrast_draws, "contrast", "contrast_name", [0.89])
reference_truth = truth.filter(pl.col("groups") == reference_group).get_column(
    "beta_true"
)[0]
truth_contrasts = truth.filter(pl.col("groups") != reference_group).select(
    (pl.col("groups") + " - " + pl.lit(reference_group)).alias("contrast_name"),
    (pl.col("beta_true") - reference_truth).alias("contrast_true"),
)

(
    ggplot(contrast_forest.to_pandas(), aes("contrast_name", "median"))
    + geom_pointrange(aes(ymin="lower", ymax="upper"), color="steelblue", size=0.9)
    + geom_point(
        data=truth_contrasts.to_pandas(),
        mapping=aes("contrast_name", "contrast_true"),
        color="firebrick",
        size=2.4,
    )
    + geom_hline(yintercept=0, linetype="dashed", color="#888888")
    + labs(
        x="contrast",
        y="beta difference",
        title="Posterior slope contrasts against North",
    )
)

Derived slope contrasts against North, with true contrasts in red.

Parameters against each other (cross-dim join)

Request beta[groups] and the scalar sigma together: sigma is broadcast onto every beta[groups] row, so you can colour one by the other directly.

mixed = td.parameter_draws(dt, "beta[groups]", "sigma")

(
    ggplot(mixed.to_pandas(), aes("groups", "beta"))
    + geom_jitter(aes(color="sigma"), width=0.15, alpha=0.15, size=0.8)
    + geom_point(
        data=truth.to_pandas(),
        mapping=aes("groups", "beta_true"),
        color="black",
        size=2.0,
    )
    + scale_color_gradient(low="steelblue", high="firebrick")
    + labs(
        x="group",
        y="beta",
        color="sigma",
        title="Slope draws coloured by residual scale",
    )
)

Cross-join detected between frame 0 and 1. Broadcasting scalar or differently-dimensioned variable on dims ['draw', 'chain'].

Posterior slopes by group, coloured by each draw’s residual sigma.

Same data, different backend

tidydraws returns a Polars DataFrame; .to_pandas() is the only bridge. The same forest-plot frame renders identically with plotnine as with lets-plot — swap the import, keep the geoms.

from plotnine import (
    ggplot,
    aes,
    geom_hline,
    geom_point,
    geom_pointrange,
    labs,
)

(
    ggplot(beta_forest.to_pandas(), aes("groups", "median"))
    + geom_pointrange(aes(ymin="lower", ymax="upper"), color="steelblue", size=0.8)
    + geom_point(
        data=truth.to_pandas(),
        mapping=aes("groups", "beta_true"),
        color="firebrick",
        size=2.0,
    )
    + geom_hline(yintercept=0, linetype="dashed", color="#888888")
    + labs(x="group", y="beta", title="Posterior slope forest plot (plotnine)")
)

The 89% slope forest plot above, rendered with plotnine.

Next, compare fitted posterior draws with model priors using compare_draws().