[plotting] Migrate plotting from matplotlib/xarray to tidydraws + plotnine (#988)

[drbenvincent]

Summary

Migrates CausalPy experiment plotting from imperative matplotlib/ArviZ composition to a declarative tidydraws + plotnine stack (#988).

The public contract remains stable: experiment plot() methods return (matplotlib.figure.Figure, matplotlib.axes.Axes | numpy.ndarray). Internally, Bayesian plot builders return PlotSpec, which contains a declarative ggplot and an optional narrow post-draw matplotlib overlay.

Architecture

model output
    ↓
prediction_draws / dataarray_draws       # canonical long-form posterior draws
    ↓
summarize_draws / spaghetti_draws         # reusable posterior transformations
    ↓
posterior_kind_layers                     # ribbon / histogram / spaghetti geoms
    ↓
PlotSpec                                  # declarative plot + compatibility overlay
    ↓
BaseExperiment._finalize_plot             # stable public (Figure, Axes) return

ITS, Synthetic Control, Synthetic DiD, and Piecewise ITS additionally share a semantic CausalPanelData → build_causal_panel_plot() path for their common three-panel outcome/effect/cumulative-effect view.

Plotnine is imported consistently through the p9 namespace. Matplotlib overlays remain only for compatibility behavior that is awkward to express declaratively, including date-axis formatting, custom legends required by legend_kwargs, and selected post-draw axis adjustments.

What changed

• Added plotnine>=0.15.7, tidydraws==0.4.3, and direct polars dependencies while retaining Python 3.11 support.
• Migrated Bayesian ribbon, histogram, and spaghetti uncertainty views to shared tidy-draw transformations and Plotnine geoms.
• Migrated RD, RK, DiD, PrePostNEGD, ITS, PiecewiseITS, Synthetic Control, Synthetic DiD, Staggered DiD, and PanelRegression plotting surfaces.
• Replaced ArviZ plotting calls with explicit summaries and declarative layers; production experiment plotting no longer calls az.plot_*.
• Preserved explicit keyword-only public plot signatures and stable matplotlib return values.
• Added centralized validation for kind, ci_kind, and num_samples.
• Ensured spaghetti paths are isolated by panel, series, chain, and draw identity.
• Added dtype-aware DiD x-axis handling for numeric, boolean, categorical, and datetime inputs.
• Preserved Patsy-based formula validation from current main across the affected experiments.
• Restored PrePostNEGD legends/axis labels, corrected PanelRegression OLS coefficient orientation, and fixed Staggered DiD shared-y behavior for three or more cohorts.

Scope

• In scope: migrated experiment plotting surfaces, including Bayesian kind ∈ {"ribbon", "histogram", "spaghetti"} views and PanelRegression/Staggered DiD auxiliary plots.
• Out of scope: IPW-specific plots and the IV plot() stub (NotImplementedError).
• Histogram heatmaps intentionally use Plotnine geom_bin_2d raw bin counts mapped to alpha for this pass.
• Cosmetic parity and broader theming refinements are intentionally deferred.

Size

Against the current PR base:

• Whole diff: +3,003 / −2,621, net +382 lines.
• Production plotting source: net +397 lines.
• Experiment modules including base.py: net −3 lines.
• plot_utils.py: net +400 lines, containing the shared extraction, transformation, composition, and rendering helpers.
• Tests: net −30 lines.

The remaining production increase is concentrated in the shared plotting layer rather than duplicated experiment implementations.

Test evidence

• Full local suite: 1,213 passed, 5 skipped.
• prek run --all-files.
• make test-patch-cov: 97% total coverage, 96% patch coverage.
• Public (Figure, Axes) contract coverage for every posterior kind.
• Regression coverage for invalid options, boolean DiD time, multi-series spaghetti grouping, PrePost legends/labels, OLS coefficient orientation, and 3+ cohort Staggered DiD facets.
• Semantic CausalPanelData coverage for ITS, Synthetic Control, and Synthetic DiD.
• Remote CI for the latest commit is currently running.

[read-the-docs-community]

Documentation build overview

📚 causalpy | 🛠️ Build #33595776 | 📁 Comparing 23eb421 against latest (642f177)

  🔍 Preview build  

145 files changed · + 34 added · ± 111 modified

+ Added

• notebooks/ancova-pymc.html
• notebooks/difference-in-differences-banks-pymc.html
• notebooks/difference-in-differences-pymc.html
• notebooks/difference-in-differences-sklearn.html
• notebooks/geolift-single-cell.html
• notebooks/instrumental-variables-pymc.html
• notebooks/instrumental-variables-variable-selection-priors.html
• notebooks/instrumental-variables-weak-instruments.html
• notebooks/interrupted-time-series-causalpy-vs-causalimpact.html
• notebooks/interrupted-time-series-comparative-pymc.html
• and 24 more...

± Modified

• 404.html
• knowledgebase/custom_pymc_models.html
• notebooks/ancova_pymc.html
• notebooks/did_pymc.html
• notebooks/did_pymc_banks.html
• notebooks/did_skl.html
• notebooks/geolift1.html
• notebooks/index.html
• notebooks/inv_prop_latent.html
• notebooks/inv_prop_pymc.html
• and 101 more...

[drbenvincent]

Cleanup after RD / PrePostNEGD migration

The earlier RD and PrePostNEGD commits added plotnine paths but left the full legacy matplotlib methods as _bayesian_plot_mpl, so those files only grew (Guardrail #2 anti-pattern).

This commit deletes both duplicate plotters (~218 net lines removed). Default ribbon (and RD spaghetti) stay on the single plotnine path. Unsupported kinds (histogram, and PrePost spaghetti) now raise ValueError pointing at #988 rather than carrying a second full implementation — we can revisit those later when there is a clean geom.

Next up: Guardrail #0 inventory for DifferenceInDifferences, then the component-wise migration from the issue’s worked example.

[drbenvincent]

Commit — RegressionKink (ribbon, default)

Full ggplot return (same shape as RD). Scatter + posterior ribbon/line + kink geom_vline. No matplotlib escape hatch; kind='histogram' raises until a plotnine path exists.

Before (matplotlib)

After (plotnine)

Review checklist

• HDI band centred on the correct data ✓
• Kink threshold at x=0.5 ✓
• Colours: black data, orange posterior mean+band, red threshold ✓
• Title (Bayesian R² + change-in-gradient CI) preserved ✓

Known cosmetic deltas (deferred to the final theming pass): legend moved to the right (plotnine default) and includes the threshold series; multi-line title left-aligned rather than centred.

Visual review approved.

[drbenvincent]

Commit — StaggeredDifferenceInDifferences event study (Bayesian plot())

Migrated the default event-study view only: plotnine geom_pointrange + pre-treatment geom_rect / hline / vline over cached att/att_lower/att_upper → .draw() → titles + matplotlib legend. Still returns (fig, [ax]). No tidydraws (HDI fixed at fit time). plot_group_time and OLS paths left on matplotlib for a separate assessment.

Before (matplotlib)

After (plotnine + matplotlib overlay)

Review checklist

• Placebo (grey squares) / ATT (blue circles) with HDI intervals ✓
• Pre-treatment shade + treatment onset vline ✓
• Zero reference hline ✓
• plot_group_time unchanged this commit ✓

Visual review approved. Next: Guardrail #0 for Bayesian plot_group_time.

[drbenvincent]

Commit — StaggeredDifferenceInDifferences plot_group_time (Bayesian)

Migrated Bayesian cohort trajectories: plotnine ribbons/lines/points (facet default + overlay) over cached ATT/HDI → .draw() → titles, shared-y, treatment vlines, legends. Fill scale matched to grey/blue (was default cyan/red). OLS paths and shared matplotlib helpers left alone. Deleted unused _plot_bayesian_group_time_segment.

Before (matplotlib)

After (plotnine + matplotlib overlay)

Review checklist

• Placebo (grey dashed) / ATT (blue solid) with matching ribbon fills ✓
• Treatment onset vline ✓
• Per-cohort titles + shared y ✓
• OLS plot_group_time unchanged ✓

Visual review approved. Next: Guardrail #0 for PanelRegression (assess methods independently; low priority).

[drbenvincent]

SC + SDiD — spaghetti & histogram (Step 1–2)

Commit: 81c89376

Migrated SyntheticControl and SyntheticDifferenceInDifferences off the matplotlib fallback for kind="spaghetti" and kind="histogram", following the ITS pattern:

• plotnine: facets, causal-impact ribbon, mean lines, observations, zero lines
• spaghetti: sampled posterior draws via tidydraws + geom_line (alpha 0.1) + mean overlay
• histogram: matplotlib pcolormesh overlay after .draw() (same fix as ITS — avoids geom_tile 1-unit height artifact)

Review images (migration_review/):

• sc_spaghetti_step1_after.png, sdid_spaghetti_step1_after.png
• sc_histogram_step2_after.png, sdid_histogram_step2_after.png

Still on matplotlib fallback: DiD, PiecewiseITS, PrePostNEGD (spaghetti/histogram); RD/RKink (histogram only).

Next: DiD / PiecewiseITS / PrePostNEGD spaghetti + histogram.

[drbenvincent]

DiD + PiecewiseITS + PrePostNEGD — spaghetti & histogram

Commit: dedf7468

All three experiments now support kind="spaghetti" and kind="histogram" on the plotnine path (same ITS/SC pattern):

Experiment	Layout	Spaghetti	Histogram
DiD	single panel	sampled draws per group + mean lines	pcolormesh overlay per trajectory (shared y-bins)
PiecewiseITS	3-panel facets	draws on fitted/counterfactual/effect/cumulative	pcolormesh per panel (top: shared bins for fitted + CF)
PrePostNEGD	2-panel facets	top facet only; bottom effect density unchanged	top facet pcolormesh; ribbon still returns ggplot

Review images (migration_review/):

• did_spaghetti_step1_after.png, did_histogram_step2_after.png
• piecewise_its_spaghetti_step1_after.png, piecewise_its_histogram_step2_after.png
• prepostnegd_spaghetti_step1_after.png, prepostnegd_histogram_step2_after.png

Remaining matplotlib fallback: RD/RKink histogram only.

Next: RD/RK histogram plotnine path, then Step 3 cleanup (_bayesian_plot_matplotlib removal).

[drbenvincent]

PrePostNEGD spaghetti fix — visual gate passed

Commit: 28fedf1c

Bug: geom_line(group="_draw_id") connected control and treatment posterior draws that share the same chain/draw id at each pretest x → vertical comb artifacts.

Fix: per-series _line_id + group="_line_id"; sort by [_draw_id, pre].

Review image: migration_review/prepostnegd_spaghetti_step1_after.png (approved)

Remaining: RD/RKink histogram plotnine path, then Step 3 cleanup.

[drbenvincent]

RD + RKink — histogram (final kind slice)

Commit: (see latest on branch)

Last matplotlib fallback removed: RegressionDiscontinuity and RegressionKink kind="histogram" now use plotnine points/lines/vlines + matplotlib pcolormesh overlay (same pattern as ITS).

Review images (migration_review/):

• rd_histogram_step2_after.png
• rkink_histogram_step2_after.png

Compare to existing baselines: rd_histogram_step0_baseline.png (in .scratch/plot_validation/ from Step 0).

All experiments now have plotnine paths for ribbon/spaghetti/histogram (where applicable). Next: Step 3 cleanup — remove _bayesian_plot_matplotlib, dedup helpers.

[drbenvincent]

RD + RKink histogram — visual gate passed

Commit: 65c4741c

Plotnine scatter/line/vline base + matplotlib pcolormesh overlay approved for both experiments.

Review images: migration_review/rd_histogram_step2_after.png, migration_review/rkink_histogram_step2_after.png

Migration slice complete: all Bayesian experiments with kind now on plotnine paths (ribbon/spaghetti/histogram). Ready for Step 3 cleanup (_bayesian_plot_matplotlib removal, helper dedup).

[drbenvincent]

Step 3 complete — dead matplotlib Bayesian paths removed

Removed _bayesian_plot_matplotlib from all 8 experiments that had plotnine parity (~1,390 lines). Unused plot_xY / _PlotXYStyle imports cleaned up.

Still intentional:

• plot_xY + _plot_histogram remain in plot_utils (tests + histogram pcolormesh overlay)
• kind="histogram" still returns (fig, ax) after plotnine layout + matplotlib density overlay
• plot() dual return handling in base.py updated to document that

Checks: prek on changed files; test_plot_utils, test_public_plot_signatures, test_piecewise_its (139 passed).

Migration slice for Bayesian kind plots is done. Optional follow-ups: dedupe _sample_draw_lines, notebook CI / tidydraws items if still open on this PR.

[codecov]

Codecov Report

❌ Patch coverage is 93.96157% with 66 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.42%. Comparing base (13af0ef) to head (23eb421).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
causalpy/plot_utils.py	90.61%	16 Missing and 4 partials ⚠️
causalpy/experiments/staggered_did.py	80.55%	8 Missing and 6 partials ⚠️
causalpy/experiments/base.py	79.31%	6 Missing and 6 partials ⚠️
causalpy/experiments/diff_in_diff.py	89.65%	4 Missing and 2 partials ⚠️
causalpy/experiments/panel_regression.py	96.38%	1 Missing and 2 partials ⚠️
causalpy/tests/test_input_validation.py	85.00%	1 Missing and 2 partials ⚠️
causalpy/experiments/regression_kink.py	95.12%	1 Missing and 1 partial ⚠️
causalpy/tests/test_plot_show_parameter.py	96.36%	1 Missing and 1 partial ⚠️
causalpy/tests/test_posterior_plot_contracts.py	98.40%	1 Missing and 1 partial ⚠️
causalpy/experiments/interrupted_time_series.py	95.45%	0 Missing and 1 partial ⚠️
... and 1 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #990      +/-   ##
==========================================
- Coverage   95.51%   95.42%   -0.10%     
==========================================
  Files          98       99       +1     
  Lines       15868    16080     +212     
  Branches      931      979      +48     
==========================================
+ Hits        15157    15345     +188     
- Misses        504      524      +20     
- Partials      207      211       +4     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[drbenvincent]

Notebook rendering checkpoint

Commit 23eb4218 contains the 28 documentation notebooks whose final Papermill executions completed successfully with no exception. The notebook output changes remain separate from the plotting code commits.

Not rerendered or not complete yet:

• interrupted-time-series-causalpy-vs-causalimpact.ipynb — skipped because it requires the optional causalimpact package.
• interrupted-time-series-placebo-in-time-analysis.ipynb — very slow; the Papermill run did not complete.
• inverse-propensity-latent.ipynb — slow model fit; the run was stopped before completion.
• synthetic-difference-in-differences-pymc.ipynb — slow model fit; the full run was stopped before completion.
• instrumental-variables-pymc.ipynb and instrumental-variables-weak-instruments.ipynb — excluded from this rendering batch because of their JAX/NumPyro environment requirement.

These notebooks should be rerendered on a faster machine when their dependencies and runtime allow it.