index.qmd

---
title: "tidydraws"
---


![](docs/assets/logo.jpg)
A tidybayes-inspired data layer for declarative Bayesian visualisation in Python.

`tidydraws` turns MCMC output (ArviZ) into tidy Polars frames that are ready to plot — one `.to_pandas()` away from any ggplot-like backend. It does no plotting itself. Four functions, four roles:

| Function | Role | Plot archetype |
| --- | --- | --- |
| [`parameter_draws()`](docs/examples/parameter_draws.qmd) | extract parameters | density, forest, scatter |
| [`prediction_draws()`](docs/examples/prediction_draws.qmd) | join predictions to covariates | ribbon + line, fit + data |
| [`compare_draws()`](docs/examples/compare_draws.qmd) | stack prior & posterior | prior vs posterior, intervals |
| [`point_interval()`](docs/examples/parameter_draws.qmd) | summarise draws | forest, ribbon, pointrange |

```{python}
# | code-fold: true
# | cache: true
# | output: false
from pathlib import Path
import sys

import lets_plot as lp
import plotnine as p9
import seaborn as sns
import seaborn.objects as so
from matplotlib.figure import Figure
import altair as alt
import pandas as pd
import polars as pl
import pymc as pm
import tidydraws as td

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 simulate_grouped_regression

lp.LetsPlot.setup_html()

workflow = simulate_grouped_regression(seed=2026)
observed = workflow.observed.with_columns(pl.col("groups").alias("group"))

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",
    )
    prior = pm.sample_prior_predictive(
        draws=800,
        random_seed=2029,
        var_names=["intercept", "beta", "sigma"],
    )
    dt = pm.sample(
        draws=400,
        tune=400,
        random_seed=2026,
        progressbar=False,
        compute_convergence_checks=False,
    )
    pm.sample_posterior_predictive(
        dt,
        var_names=["mu"],
        predictions=True,
        extend_inferencedata=True,
        random_seed=2028,
        progressbar=False,
    )

    # arviz 1.0+ uses .update(), 0.x uses .extend()
    try:
        dt.update(prior)
    except AttributeError:
        dt.extend(prior)
```

## Use tidydraws

```{python}
beta_df = td.parameter_draws(dt, "beta")
beta_summary = td.point_interval(beta_df, "beta", group_by="groups").sort("groups")
compare = td.compare_draws(dt, "beta")
pred = td.prediction_draws(dt, newdata=observed, var_name="mu")
pred_summary = td.point_interval(pred, "mu", group_by=["obs_ind", "x", "group"]).sort(
    "x"
)
```
## Plot
:::: {.panel-tabset}

#### lets-plot

```{python}
p_forest = (
    lp.ggplot(beta_summary.to_pandas(), lp.aes("groups", "beta"))
    + lp.geom_pointrange(lp.aes(ymin="beta_lower", ymax="beta_upper"), size=0.8)
    + lp.geom_hline(yintercept=0, linetype="dashed", color="#888888")
    + lp.labs(x="group", y="beta", title="Forest plot")
)
p_density = (
    lp.ggplot(beta_df.to_pandas(), lp.aes("beta", fill="groups"))
    + lp.geom_density(alpha=0.5)
    + lp.labs(x="beta", y="density", fill="group", title="Posterior density")
)
p_compare = (
    lp.ggplot(compare.to_pandas(), lp.aes("beta", fill="source"))
    + lp.geom_density(alpha=0.5)
    + lp.facet_wrap(facets="groups", ncol=2)
    + lp.labs(x="beta", y="density", fill="source", title="Prior vs posterior")
)
p_pred = (
    lp.ggplot(pred_summary.to_pandas(), lp.aes("x"))
    + lp.geom_ribbon(
        lp.aes(ymin="mu_lower", ymax="mu_upper", fill=lp.as_discrete("group")),
        alpha=0.25,
    )
    + lp.geom_line(lp.aes(y="mu", color=lp.as_discrete("group")), size=0.8)
    + lp.labs(x="x", y="mu", color="group", fill="group", title="Predictive fit")
)
g = lp.gggrid([p_forest, p_density, p_compare, p_pred], ncol=2)
lp.ggsave(g, "index-plot.png", path="../docs/assets", scale=1.5)
g
```

#### plotnine

```{python}
p9_forest = (
    p9.ggplot(beta_summary.to_pandas(), p9.aes("groups", "beta"))
    + p9.geom_pointrange(p9.aes(ymin="beta_lower", ymax="beta_upper"), size=0.8)
    + p9.geom_hline(yintercept=0, linetype="dashed", color="#888888")
    + p9.labs(x="group", y="beta", title="Forest plot")
)
p9_density = (
    p9.ggplot(beta_df.to_pandas(), p9.aes("beta", fill="groups"))
    + p9.geom_density(alpha=0.5)
    + p9.labs(x="beta", y="density", fill="group", title="Posterior density")
)
p9_compare = (
    p9.ggplot(compare.to_pandas(), p9.aes("beta", fill="source"))
    + p9.geom_density(alpha=0.5)
    + p9.facet_wrap("~groups", ncol=2)
    + p9.labs(x="beta", y="density", fill="source", title="Prior vs posterior")
)
p9_pred = (
    p9.ggplot(pred_summary.to_pandas(), p9.aes("x"))
    + p9.geom_ribbon(
        p9.aes(ymin="mu_lower", ymax="mu_upper", fill="group"),
        alpha=0.25,
    )
    + p9.geom_line(p9.aes(y="mu", color="group"), size=0.8)
    + p9.labs(x="x", y="mu", color="group", fill="group", title="Predictive fit")
)
(p9_forest | p9_density) / (p9_compare | p9_pred)
```


#### seaborn

```{python}
# | fig-cap: "Posterior forest plot, density, prior vs posterior, and predictive fit via seaborn objects."
so.Plot.config.theme.update(sns.axes_style("whitegrid"))

fig = Figure(figsize=(8, 6))
(sf1, sf2), (sf3, sf4) = fig.subfigures(2, 2)

# Forest plot
(
    so
    .Plot(beta_summary.to_pandas(), x="groups")
    .add(so.Range(alpha=0.7), ymin="beta_lower", ymax="beta_upper")
    .add(so.Dot(color="steelblue"), y="beta")
    .label(x="group", y="beta", title="Forest plot")
    .on(sf1)
    .plot()
)

# Density
(
    so
    .Plot(beta_df.to_pandas(), x="beta", color="groups")
    .add(so.Area(alpha=0.5), so.KDE(common_norm=False))
    .label(x="beta", y="density", title="Posterior density")
    .on(sf2)
    .plot()
)

# Prior vs posterior
(
    so
    .Plot(compare.to_pandas(), x="beta", color="source")
    .facet(col="groups", wrap=2)
    .add(so.Area(alpha=0.5), so.KDE(common_norm=False))
    .label(x="beta", y="density", title="Prior vs posterior")
    .on(sf3)
    .plot()
)

# Predictive fit
(
    so
    .Plot(pred_summary.to_pandas(), x="x", color="group")
    .add(so.Band(alpha=0.25), ymin="mu_lower", ymax="mu_upper")
    .add(so.Line(), y="mu")
    .label(x="x", y="mu", title="Predictive fit")
    .on(sf4)
    .plot()
)

fig
```

#### altair

```{python}
# | fig-cap: "Posterior forest plot, density, prior vs posterior, and predictive fit via Altair."
alt.data_transformers.enable("vegafusion")
# Forest plot
alt_forest = (
    alt
    .layer(
        alt
        .Chart(beta_summary.to_pandas())
        .mark_errorbar(ticks=True)
        .encode(x="groups:N", y="beta_lower:Q", y2="beta_upper:Q"),
        alt
        .Chart(beta_summary.to_pandas())
        .mark_point(color="steelblue")
        .encode(x="groups:N", y="beta:Q"),
        alt
        .Chart(pd.DataFrame({"y0": [0]}))
        .mark_rule(color="#888888", strokeDash=[4, 4])
        .encode(y="y0:Q"),
    )
    .resolve_scale(y="shared")
    .properties(title="Forest plot", width=200, height=200)
)

# Density
alt_density = (
    alt
    .Chart(beta_df.to_pandas())
    .transform_density("beta", groupby=["groups"], as_=["beta", "density"])
    .mark_area(opacity=0.5)
    .encode(x="beta:Q", y="density:Q", color="groups:N")
    .properties(title="Posterior density", width=200, height=200)
)
alt_compare = (
    alt
    .Chart(compare.to_pandas())
    .transform_density("beta", groupby=["groups", "source"], as_=["beta", "density"])
    .mark_area(opacity=0.5)
    .encode(x="beta:Q", y="density:Q", color="source:N")
    .properties(width=95, height=85)
    .facet("groups:N", columns=2)
    .properties(title="Prior vs posterior")
)
# Predictive fit
alt_pred = (
    alt
    .layer(
        alt
        .Chart(pred_summary.to_pandas())
        .mark_area(opacity=0.25)
        .encode(x="x:Q", y="mu_lower:Q", y2="mu_upper:Q", color="group:N"),
        alt
        .Chart(pred_summary.to_pandas())
        .mark_line()
        .encode(x="x:Q", y="mu:Q", color="group:N"),
    )
    .resolve_scale(y="shared")
    .properties(title="Predictive fit", width=200, height=200)
)

(
    alt.hconcat(alt_forest, alt_density) & alt.hconcat(alt_compare, alt_pred)
).resolve_scale(color="independent")
```
::::


## Install

::::: {.panel-tabset}

#### uv

```bash
uv add tidydraws
```

If you want the latest functionality merged into main but not yet released, install directly from GitHub:

```bash
uv add git+https://github.com/drbenvincent/tidydraws.git
```

#### pip

```bash
pip install tidydraws
```

If you want the latest functionality merged into main but not yet released, install directly from GitHub:

```bash
pip install git+https://github.com/drbenvincent/tidydraws.git
```

:::::

## Why tidydraws?

Plotting MCMC output in Python means manually slicing xarray dimensions, iterating groups, and aligning coordinates — imperative, verbose, error-prone. R's [`tidybayes`](https://github.com/mjskay/tidybayes) solved this with a data layer that respects parameter space vs prediction space. `tidydraws` brings that to Python on Polars.

`tidydraws` is plotting-backend-agnostic: it returns Polars `DataFrame`s, and `.to_pandas()` bridges to lets-plot, plotnine, or any library that takes pandas. The examples use [lets-plot](https://lets-plot.org/); see the [`parameter_draws()`](docs/examples/parameter_draws.qmd) page for the same plot in plotnine.

---

*Inspired by [tidybayes](https://github.com/mjskay/tidybayes) for R.*