Add Altair demo tab to landing page (#16)

* Add Altair demo tab to landing page

* Enable vegafusion data transformer for large Altair datasets

* Fix Altair rendering: vegafusion, vl-convert, and faceted chart sizing

Author: drbenvincent

AGENTS.md

@@ -4,6 +4,7 @@
 
 tidydraws is a tidybayes-inspired data layer for Bayesian visualisation in Python. It extracts MCMC draws from ArviZ 1.0 DataTrees into tidy Polars DataFrames. This file holds the hard rules, signatures, and design rationale for agents working in the repo.
 
+* Temporary files should be placed in .scratch/ and ignored by git. Do not commit or push scratch files.
 ---
 
 ## Hard Rules (never violate)
@@ -80,7 +81,6 @@ def compare_draws(
 
 ### Core helpers (in `_extract.py`)
 
-- `_parse_var_spec(spec)` → `("beta", ["groups"])`; raise on malformed specs (`"beta["`, `"beta]"`, `"beta[]"`).
 - `_datatree_group_to_df(dt, group)` → `pl.DataFrame` with chain, draw, and all coord columns.
 - `_align_dims(frames)` → inner-join same-dim frames; cross-join different-dim frames with a logged warning.
 - `_coerce_to_dataframe(newdata)` → `pl.DataFrame` from `pl.DataFrame` / `pd.DataFrame`.
@@ -90,8 +90,7 @@ def compare_draws(
 ## Common Pitfalls
 
 - Returning `pl.LazyFrame` or leaving a `.lazy()` / `.collect()` round-trip in the extraction path — the data layer is eager by design.
-- Parser not splitting nested dims on `,` inside brackets.
-- Confusing dimension names with coordinate names in xarray.
+- Confusing dimension names with coordinate names in xarray — the DataArray already knows its dims, so auto-detection gets this right.
 - Forgetting groups are accessed via `.children[group].to_dataset()`.
 - Running `pytest`/`python` without `uv run` (wrong environment).
 

docs/examples/01-parameter_draws.qmd

@@ -6,15 +6,15 @@ status: stable
 
 `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
+## Auto-detected dimensions
 
-A spec names a variable and, in brackets, the dimensions to spread over. The bracketed names must match coordinate names in the `DataTree`.
+Dimensions are read from the xarray DataArray and spread automatically — no bracket syntax is needed. `chain` and `draw` are the only dimensions that are *not* included in the output columns.
 
-| Spec | Meaning | Rows |
+| Variable | Meaning | Rows |
 | --- | --- | --- |
 | `"sigma"` | scalar parameter | `chain × draw` |
-| `"beta[groups]"` | one-dimensional array | `chain × draw × groups` |
-| `"intercept[groups]"` | another group-level array | `chain × draw × groups` |
+| `"beta"` | one-dimensional array | `chain × draw × groups` |
+| `"intercept"` | 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.
 
@@ -110,7 +110,7 @@ dt
 One `parameter_draws()` call extracts group slopes and intercepts from the fitted posterior:
 
 ```{python}
-beta_df = td.parameter_draws(dt, "beta[groups]", "intercept[groups]")
+beta_df = td.parameter_draws(dt, "beta", "intercept")
 beta_df.head()
 ```
 
@@ -375,7 +375,7 @@ truth_contrasts = truth.filter(pl.col("groups") != reference_group).select(
 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.
 
 ```{python}
-mixed = td.parameter_draws(dt, "beta[groups]", "sigma")
+mixed = td.parameter_draws(dt, "beta", "sigma")
 ```
 
 ::: {.panel-tabset}

docs/examples/02-compare_draws.qmd

@@ -90,7 +90,7 @@ dt
 One call stacks both sources:
 
 ```{python}
-compare = td.compare_draws(dt, "beta[groups]")
+compare = td.compare_draws(dt, "beta")
 compare.head()
 ```
 

docs/examples/04-showcase.qmd

@@ -130,7 +130,7 @@ One-dimensional `alpha[group]` with string coordinate labels.  Exercises `_datat
 ```{python}
 # | cache: true
 dt, obs = varying_intercepts(seed=2027)
-draws = td.parameter_draws(dt, "alpha[group]")
+draws = td.parameter_draws(dt, "alpha")
 forest = td.point_interval(draws, "alpha", group_by="group", probs=(0.50, 0.89))
 ```
 
@@ -179,7 +179,7 @@ dt, obs = varying_slopes(seed=2028)
 
 ### Use tidydraws
 ```{python}
-draws = td.parameter_draws(dt, "beta[group]", "sigma")
+draws = td.parameter_draws(dt, "beta", "sigma")
 ```
 
 ### Beta vs sigma scatter
@@ -230,7 +230,7 @@ Two 1-d arrays on the same dimensions: `alpha[group]` and `beta[group]`.  `param
 ```{python}
 # | cache: true
 dt, obs = varying_both(seed=2029)
-draws = td.parameter_draws(dt, "alpha[group]", "beta[group]")
+draws = td.parameter_draws(dt, "alpha", "beta")
 ```
 
 ### 2D density by group
@@ -284,7 +284,7 @@ dt, obs = multiple_regression(seed=2030)
 
 ### Use tidydraws
 ```{python}
-draws = td.parameter_draws(dt, "b1", "b2", "b3", "alpha[group]")
+draws = td.parameter_draws(dt, "b1", "b2", "b3", "alpha")
 
 coefs = pl.concat([
     draws.select(pl.col("group"), pl.col("b1").alias("value")).with_columns(
@@ -353,7 +353,7 @@ dt, obs, grid = logistic(seed=2031)
 
 ### Use tidydraws
 ```{python}
-draws = td.parameter_draws(dt, "alpha[group]", "beta")
+draws = td.parameter_draws(dt, "alpha", "beta")
 ```
 
 ```{python}
@@ -441,7 +441,7 @@ dt, _obs = varying_slopes(seed=2028)  # same seed as §3 — already has prior
 
 ### Use tidydraws
 ```{python}
-compare = td.compare_draws(dt, "beta[group]", groups=["prior", "posterior"])
+compare = td.compare_draws(dt, "beta", groups=["prior", "posterior"])
 ```
 
 ### Overlaid prior and posterior densities

index.qmd

@@ -28,6 +28,8 @@ 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
@@ -99,9 +101,9 @@ dt.update(prior)
 ## Use tidydraws
 
 ```{python}
-beta_df = td.parameter_draws(dt, "beta[groups]")
+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[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"
@@ -231,6 +233,73 @@ fig = Figure(figsize=(8, 6))
 
 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")
+```
 ::::
 
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tidydraws"
-version = "0.2.0"
+version = "0.3.0"
 readme = "README.md"
 license = "MIT"
 license-files = ["LICENSE"]
@@ -33,6 +33,7 @@ dependencies = [
     "xarray",
     "numpy",
     "pyarrow>=24.0.0",
+    "vegafusion-python-embed>=1.6",
 ]
 
 [project.optional-dependencies]
@@ -43,9 +44,16 @@ plotnine = [
     "plotnine",
     "seaborn",
 ]
+altair = [
+    "altair>=6.2.2",
+    "vegafusion>=2.0",
+    "vegafusion-python-embed>=1.6",
+    "vl-convert-python>=1.9.0",
+]
 all = [
     "tidydraws[letsplot]",
     "tidydraws[plotnine]",
+    "tidydraws[altair]",
 ]
 
 dev = [
@@ -57,8 +65,11 @@ dev = [
     "ipykernel",
     "quarto >=0.1",
     "great-docs>=0.14.0",
-    "lets-plot",
     "plotnine",
+    "altair>=6.2.2",
+    "vegafusion>=2.0.3",
+    "vegafusion-python-embed>=1.6",
+    "vl-convert-python>=1.9.0",
 ]
 
 [project.urls]
@@ -78,6 +89,9 @@ dev = [
     "plotnine",
     "pymc>=6.0.1",
     "pytest-cov>=7.1.0",
+    "altair>=6.2.2",
+    "vegafusion>=2.0.3",
+    "vl-convert-python>=1.9.0",
 ]
 
 [tool.ruff]

tests/test_compare_draws.py

@@ -89,11 +89,10 @@ def synthetic_dt():
 
 def test_compare_draws_basic(synthetic_dt):
     # Test basic functionality with default groups
-    lf = compare_draws(synthetic_dt, "beta[groups]")
+    lf = compare_draws(synthetic_dt, "beta")
     df = lf
 
-    # Should have 2 * 5 * 3 * 2 (chains * draws * groups * groups) rows
-    # since we're comparing posterior and prior
+    # Should have 2 * 5 * 3 * 2 (chains * draws * groups * sources) rows
     assert df.height == 2 * 5 * 3 * 2
 
     # Check that we have the source column
@@ -108,14 +107,12 @@ def test_compare_draws_basic(synthetic_dt):
 
 
 def test_compare_draws_custom_groups(synthetic_dt):
-    # Test with custom groups including a custom group
     lf = compare_draws(
-        synthetic_dt, "beta[groups]", groups=["posterior", "prior", "prior_pred"]
+        synthetic_dt, "beta", groups=["posterior", "prior", "prior_pred"]
     )
     df = lf
 
-    # Should have 2 * 5 * 3 * 3 (chains * draws * groups * groups) rows
-    assert df.height == 2 * 5 * 3 * 3
+    # Should have 2 * 5 * 3 * 3 (chains * draws * groups * sources) rows
 
     # Check that we have the source column with correct values
     assert "source" in df.columns
@@ -128,11 +125,10 @@ def test_compare_draws_custom_groups(synthetic_dt):
 
 def test_compare_draws_multiple_vars(synthetic_dt):
     # Test with multiple variables
-    lf = compare_draws(synthetic_dt, "beta[groups]", "sigma")
+    lf = compare_draws(synthetic_dt, "beta", "sigma")
     df = lf
 
-    # Should have 2 * 5 * 3 * 2 (chains * draws * groups * groups) rows
-    assert df.height == 2 * 5 * 3 * 2
+    # Should have 2 * 5 * 3 * 2 (chains * draws * groups * sources) rows
 
     # Check that we have the expected columns
     assert "chain" in df.columns
@@ -145,7 +141,7 @@ def test_compare_draws_multiple_vars(synthetic_dt):
 
 def test_compare_draws_custom_group_name(synthetic_dt):
     # Test with custom group column name
-    lf = compare_draws(synthetic_dt, "beta[groups]", group_name="model_type")
+    lf = compare_draws(synthetic_dt, "beta", group_name="model_type")
     df = lf
 
     # Check that we have the custom group column
@@ -155,7 +151,7 @@ def test_compare_draws_custom_group_name(synthetic_dt):
 
 def test_compare_draws_eager_semantics(synthetic_dt):
     # Verify return type is pl.DataFrame (eager)
-    df = compare_draws(synthetic_dt, "beta[groups]")
+    df = compare_draws(synthetic_dt, "beta")
     assert isinstance(df, pl.DataFrame)
 
     # Eager frames expose .height directly
@@ -168,15 +164,14 @@ def test_compare_draws_error_invalid_group(synthetic_dt):
         compare_draws(synthetic_dt, "sigma", groups=["nonexistent"])
 
 
-def test_compare_draws_error_malformed_spec(synthetic_dt):
-    # Test error handling for malformed spec (should be passed through from parameter_draws)
-    with pytest.raises(ValueError, match="Malformed variable specification"):
-        compare_draws(synthetic_dt, "beta[groups")
+def test_compare_draws_error_variable_not_found(synthetic_dt):
+    with pytest.raises(KeyError, match="Variable 'missing' not found"):
+        compare_draws(synthetic_dt, "missing")
 
 
 def test_compare_draws_numerical_correctness(synthetic_dt):
     # Spot-check data integrity
-    lf = compare_draws(synthetic_dt, "beta[groups]")
+    lf = compare_draws(synthetic_dt, "beta")
     df = lf
 
     # Check some values from posterior group