Skip to content

feat: add plot_dist for posterior density plots - #467

Open
RussellSB wants to merge 5 commits into
pymc-labs:mainfrom
RussellSB:feat/323-add-plot-dist-wrapper
Open

feat: add plot_dist for posterior density plots#467
RussellSB wants to merge 5 commits into
pymc-labs:mainfrom
RussellSB:feat/323-add-plot-dist-wrapper

Conversation

@RussellSB

@RussellSB RussellSB commented Aug 21, 2026

Copy link
Copy Markdown

Summary

  • Add plot_dist() for posterior density plots on PathModel and the three draw-backed result objects: EffectResult, EstimandResult, and DoResult.
  • Resolve names in one call, whether labeled coefficients ("a"), defined parameters ("indirect"), or raw posterior variables ("sigma_Y"); coords= disambiguates multi-coordinate variables.
  • Add a public plot_density() renderer in pathmc.plotting (KDE curve + translucent fill + optional dashed reference line).

Issue links

Examples

Fitted mediation model (X -> M -> Y with a direct effect):

spec = "M ~ a*X\nY ~ b*M + c*X\nindirect := a*b\ntotal := c + a*b"
model_med = pathmc.model(spec, data=df)
model_med.fit()

1. Model: labeled coefficients

1.1 Plain

model_med.plot_dist(var=["a", "b", "c"])
image

1.2 Reference lines at the true values

model_med.plot_dist(var=["a", "b", "c"], ref=[0.5, 0.8, 0.3])
image

1.3 Plus custom colors

model_med.plot_dist(
    var=["a", "b", "c"],
    ref=[0.5, 0.8, 0.3],
    color=["lightblue", "lightcoral", "lightgreen"],
)
image

2. Model: defined parameters

model_med.plot_dist(var=["indirect", "total"], ref=[0.4, 0.7])
image

3. EffectResult: path effect

effect = model_med.effect("X -> M -> Y")
effect.plot_dist(ref=0.4)
image

4. EstimandResult: ATE

ate = model_med.ate("Y", "X", values=(0.0, 1.0))
ate.plot_dist(ref=0.65)
image

5. DoResult: interventional outcome

do_result = model_med.do(set={"X": 1})
do_result.plot_dist(var="Y")
image

Design decisions

  • var / ref naming: generalised from the issue's var_names / ref_val so one name works across all four entry points.
  • color optional: str | list[str] on the model; a single str on result objects.
  • plot_dist() returns None (calls plt.show() internally).
  • Shared axes + legend, no titles: multiple densities overlay on one plot.
  • plot_density() is public: exported from pathmc.

Questions

  • Is it OK that I made plot_density() public? It surfaces in the API reference and is directly usable for custom cases (comparing a posterior to a ground-truth distribution, for instance).
  • What would be your thoughts on a plot_params argument for figure-level knobs (figsize, title, axis labels)? A future idea to extend on this method's customisation. Deferred and not blocking; happy to follow up in a later PR if there's appetite.

Test plan

  • uv run pytest tests/test_plot_dist.py -x -v
  • make test-fast
  • make lint

Made with OpenCode Go (and a moderate amount of AI nagging)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant