[plotting] Allow user to specify type of posterior data visualisation

[AllenDowney]

Summary

This draft PR extends the plotting capabilities of CausalPy to support multiple visualization types for posterior data. The purpose of this draft PR is to discuss the API design and approach, not to finalize implementation or testing.

Currently, CausalPy only supports CI ribbon visualizations using Highest Density Intervals (HDI). This PR adds:

• Support for Equal-Tailed Intervals (ETI) in addition to HDI
• Support for histogram visualizations
• Support for spaghetti plot visualizations (individual posterior sample trajectories)

The API design aligns with ArviZ's naming conventions (ci_prob, ci_kind) while maintaining backward compatibility with existing code.

Fixes #671

Changes

• Extended plot_xY() function (causalpy/plot_utils.py):

  • Added kind parameter: "ribbon", "histogram", or "spaghetti" (default: "ribbon")
  • Added ci_kind parameter: "hdi" or "eti" (default: "hdi" to match current behavior)
  • Added ci_prob parameter (default: 0.94 to match current behavior)
  • Added num_samples parameter for spaghetti plots (default: 50)
  • Maintained backward compatibility with hdi_prob parameter
  • Implemented three visualization types:
    • Ribbon: HDI/ETI intervals using ArviZ functions
    • Histogram: Marginal histograms at key time points (basic implementation -- will be extended after API review)
    • Spaghetti: Random posterior sample trajectories with mean overlay

• Updated BaseExperiment.plot() method (causalpy/experiments/base.py):

  • Added new parameters to method signature
  • Parameters are passed through to _bayesian_plot() and _ols_plot() methods
  • Maintains backward compatibility (all parameters optional with defaults)

Testing

• Manual testing completed in marimo notebook (causalpy_test.py) demonstrating all visualization types
• All visualization types verified to render correctly
• Pre-commit checks pass (linting, formatting, type checking)
• Note: Unit and integration tests are deferred until after API review and feedback

API Design Rationale

1. kind parameter: Uses familiar naming convention from seaborn/pandas
2. ci_prob and ci_kind naming: Aligns with ArviZ's naming conventions for ecosystem consistency
3. Default values: Match current behavior (0.94, "hdi") to maintain backward compatibility
4. Backward compatibility: Existing code using hdi_prob continues to work unchanged

Open Questions for Reviewers

1. Any comments or suggestions on the proposed API?
2. Should we proceed with updating all experiment classes (_bayesian_plot() methods) to explicitly accept and pass through these parameters, or is passing via **kwargs sufficient?
3. Are the default values (0.94, "hdi") appropriate, or should we consider ArviZ's new defaults (0.89, "eti")?

Checklist

• Pre-commit checks pass
• Code follows project conventions
• All tests pass (deferred until after API review)
• Documentation updated (deferred until after API review)
• Example notebooks created (deferred until after API review)

---

📚 Documentation preview 📚: https://causalpy--680.org.readthedocs.build/en/680/

[drbenvincent]

Thanks @AllenDowney ! Hoping to look at this very soon. Just flagging up that there might be some conflicts to resolve if #643 gets merged first. That one allows the user to decide if they are looking at things related to the posterior expectation or the posterior predictive. So these PR's are nicely complementary, but it's possible there could be some code overlap/conflicts.

[drbenvincent]

Thanks for this — the API direction looks promising. I reviewed this mainly from a maintainability/reviewability perspective and have a few blocking + process suggestions.

Blocking issues

• CI currently fails early due to import error in causalpy/plot_utils.py:
  • ImportError: cannot import name 'eti' from 'arviz_stats'
• Because this import fails at module load, reviewers cannot run plotting tests or examples yet.

Requested updates before final review

1. Fix ETI import/implementation path so test/docs/sdist are green.
2. Add visual review artifacts to the PR description/comments, since GitHub diff does not show rendered plots:
   • kind="ribbon" with ci_kind="hdi"
   • kind="ribbon" with ci_kind="eti"
   • kind="histogram"
   • kind="spaghetti"
3. Include one minimal reproducible script/snippet (not notebook-only) to generate mock posterior data and produce all plot kinds. This allows reviewers to validate quickly without rerunning full notebooks.

API/compatibility checks to confirm

• plot_xY() now returns either (Line2D, PolyCollection) or (list[Line2D], None) depending on kind.
• Please confirm experiment-level .plot() legend handling remains consistent for non-ribbon kinds, since several call sites still assume tuple-style handles.
• Please clarify whether hdi_prob deprecation timeline is planned, or if it remains indefinitely as compatibility alias.

Efficient review request

Could you add 3–5 static PNGs (before/after where useful) generated from the same mock posterior dataset and style settings? That will make visual review possible without manual notebook execution.

Suggested mock posterior script for reviewers

import numpy as np
import pandas as pd
import xarray as xr
import matplotlib.pyplot as plt

from causalpy.plot_utils import plot_xY

# Reproducible synthetic posterior: (chain, draw, obs_ind)
rng = np.random.default_rng(42)
n_chains, n_draws, n_t = 2, 200, 40
x = pd.date_range("2022-01-01", periods=n_t, freq="D")

trend = 10 + 0.05 * np.arange(n_t) + 0.01 * np.arange(n_t) ** 2
samples = np.empty((n_chains, n_draws, n_t))
for c in range(n_chains):
    for d in range(n_draws):
        draw_mean = trend + rng.normal(0, 0.4, n_t)
        samples[c, d, :] = draw_mean + rng.normal(0, 0.8, n_t)

Y = xr.DataArray(
    samples,
    dims=["chain", "draw", "obs_ind"],
    coords={"chain": np.arange(n_chains), "draw": np.arange(n_draws), "obs_ind": x},
)

fig, axes = plt.subplots(2, 2, figsize=(12, 8), sharex=True)

plot_xY(x, Y, ax=axes[0, 0], kind="ribbon", ci_kind="hdi", ci_prob=0.94, label="HDI")
axes[0, 0].set_title("Ribbon (HDI)")

plot_xY(x, Y, ax=axes[0, 1], kind="ribbon", ci_kind="eti", ci_prob=0.94, label="ETI")
axes[0, 1].set_title("Ribbon (ETI)")

plot_xY(x, Y, ax=axes[1, 0], kind="histogram", label="Histogram")
axes[1, 0].set_title("Histogram")

plot_xY(x, Y, ax=axes[1, 1], kind="spaghetti", num_samples=60, label="Spaghetti")
axes[1, 1].set_title("Spaghetti")

for ax in axes.ravel():
    ax.legend(loc="best")
    ax.grid(alpha=0.3)

plt.tight_layout()
plt.show()

Conflict check with #643

I checked this directly with a branch-to-branch merge simulation.

There are real merge conflicts between #643 and #680 in these files:

• causalpy/experiments/base.py
• causalpy/plot_utils.py
• causalpy/tests/test_plot_utils.py

So yes, conflict risk is concrete, not just theoretical.

Best sequencing: merge #643 first, then rebase/update #680 on top and re-run tests.

[drbenvincent]

Including images is optional - especially if there's a temp script to generate plots for dev/review purposes.

If the conflicts become gnarly when #643 is merged, then I'm happy to give that a stab - I feel guilty about that kind of thing :)

[drbenvincent]

Thanks for the nudge — adding a minimal reproducible script here to speed visual/API review without needing notebook execution.

This script creates synthetic posterior draws and renders all current plot kinds:

• kind="ribbon" with ci_kind="hdi"
• kind="ribbon" with ci_kind="eti"
• kind="histogram"
• kind="spaghetti"

import numpy as np
import pandas as pd
import xarray as xr
import matplotlib.pyplot as plt

from causalpy.plot_utils import plot_xY

# Reproducible synthetic posterior: (chain, draw, obs_ind)
rng = np.random.default_rng(42)
n_chains, n_draws, n_t = 2, 200, 40
x = pd.date_range("2022-01-01", periods=n_t, freq="D")

trend = 10 + 0.05 * np.arange(n_t) + 0.01 * np.arange(n_t) ** 2
samples = np.empty((n_chains, n_draws, n_t))
for c in range(n_chains):
    for d in range(n_draws):
        draw_mean = trend + rng.normal(0, 0.4, n_t)
        samples[c, d, :] = draw_mean + rng.normal(0, 0.8, n_t)

Y = xr.DataArray(
    samples,
    dims=["chain", "draw", "obs_ind"],
    coords={"chain": np.arange(n_chains), "draw": np.arange(n_draws), "obs_ind": x},
)

fig, axes = plt.subplots(2, 2, figsize=(12, 8), sharex=True)

plot_xY(x, Y, ax=axes[0, 0], kind="ribbon", ci_kind="hdi", ci_prob=0.94, label="HDI")
axes[0, 0].set_title("Ribbon (HDI)")

plot_xY(x, Y, ax=axes[0, 1], kind="ribbon", ci_kind="eti", ci_prob=0.94, label="ETI")
axes[0, 1].set_title("Ribbon (ETI)")

plot_xY(x, Y, ax=axes[1, 0], kind="histogram", label="Histogram")
axes[1, 0].set_title("Histogram")

plot_xY(x, Y, ax=axes[1, 1], kind="spaghetti", num_samples=60, label="Spaghetti")
axes[1, 1].set_title("Spaghetti")

for ax in axes.ravel():
    ax.legend(loc="best")
    ax.grid(alpha=0.3)

plt.tight_layout()
plt.show()

If useful, I can follow up with static PNGs from this same dataset/style setup.

[drbenvincent]

Looks cool!

[codecov]

Codecov Report

❌ Patch coverage is 96.00000% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.65%. Comparing base (deb8774) to head (9397732).
⚠️ Report is 50 commits behind head on main.

Files with missing lines	Patch %	Lines
causalpy/plot_utils.py	89.65%	7 Missing and 5 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #680      +/-   ##
==========================================
+ Coverage   94.60%   94.65%   +0.04%     
==========================================
  Files          80       80              
  Lines       12764    13058     +294     
  Branches      770      796      +26     
==========================================
+ Hits        12076    12360     +284     
- Misses        485      491       +6     
- Partials      203      207       +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.

[read-the-docs-community]

Documentation build overview

📚 causalpy | 🛠️ Build #32439580 | 📁 Comparing 9397732 against latest (78ec91d)

  🔍 Preview build  

28 files changed · ± 28 modified

± Modified

• 404.html
• index.html
• api/generated/causalpy.experiments.base.BaseExperiment.html
• api/generated/causalpy.experiments.base.BaseExperiment.plot.html
• api/generated/causalpy.experiments.diff_in_diff.DifferenceInDifferences.html
• api/generated/causalpy.experiments.diff_in_diff.DifferenceInDifferences.plot.html
• api/generated/causalpy.experiments.interrupted_time_series.InterruptedTimeSeries.get_plot_data_bayesian.html
• api/generated/causalpy.experiments.interrupted_time_series.InterruptedTimeSeries.html
• api/generated/causalpy.experiments.interrupted_time_series.InterruptedTimeSeries.plot.html
• api/generated/causalpy.experiments.inverse_propensity_weighting.InversePropensityWeighting.html
• and 18 more...

[AllenDowney]

@drbenvincent Ready for review. Please take a look and let me know if there's anything else we should do.