Skip to content

Plotting 1.1 — Phase 0: shared plotnine draw-and-finalize seam #1156

Description

@drbenvincent

Child of #1155 (1.1 declarative plotting migration). Nothing else in this program starts until this lands.

Workflow

Branch from pymc6_and_pymcmarketing1_migration and open the PR against pymc6_and_pymcmarketing1_migration. Not main. See the parent issue for why.

Problem

There are two independent, unrelated plotnine seams on the migration branch:

  1. RegressionDiscontinuity._plot (from Migrate RegressionDiscontinuity base rendering to Plotnine #1135) builds a ggplot, calls .draw(), then applies fig.set_size_inches(figsize), grabs fig.axes[0], adds axvline rules, calls the matplotlib posterior helper, suppresses plotnine's own colour guide with guides(color=False), hand-builds Line2D proxy handles, and merges them with ax.get_legend_handles_labels().
  2. causalpy/checks/placebo_in_time.py::_draw draws a plotnine plot or composition, stamps a suptitle, and restores the caller's figure size because plotnine's layout engine ignores subplots_adjust.

Every subsequent experiment migration needs the same post-.draw() work. If we migrate one more experiment before factoring this out, we have three copies of it.

Scope

Add one shared draw-and-finalize seam to causalpy/plot_utils.py covering:

  • .draw() plus figsize application, in a way that survives plotnine's layout pass (the placebo_in_time fix already knows how).
  • Data-axes discovery for faceted plots. ax.get_subplotspec() is not None filters out plotnine's internal strip-label axes; a caller should never have to know that.
  • Optional suptitle stamping (absorbing placebo_in_time._draw).
  • Legend reconciliation: suppress plotnine's guide, build proxy handles from a {label: colour} mapping, and merge with handles already on the Axes, so BaseExperiment._render_plot's in-place legend_kwargs mutation keeps working and tuple handles like (Line2D, PolyCollection) are preserved.
  • A helper for reserved renderer column names, so the __causalpy_plot_x / __causalpy_plot_y / __causalpy_plot_series convention from Migrate RegressionDiscontinuity base rendering to Plotnine #1135 is defined once rather than restated per experiment (guardrail G7 in the parent).
  • A documented return shape that lines up with what _render_plot expects: (Figure, Axes) or (Figure, ndarray[Axes]).

Also decide and document, in this PR, how plotnine themes coexist with _render_plot's plt.style.context("arviz-darkgrid"). Today that style context wraps _plot but does not theme plotnine output, so RD's plotnine layers and its matplotlib overlays are styled by different systems. Pick one answer (a shared theme_causalpy(), or keep the style context and match it, or drop the style context) and write it down.

Files likely touched

  • causalpy/plot_utils.py (new shared helpers)
  • causalpy/checks/placebo_in_time.py (adopt the shared seam, delete the local _draw)
  • causalpy/experiments/regression_discontinuity.py — only if adopting the seam is a pure substitution; otherwise leave it to child C
  • causalpy/experiments/base.py — only if the theming decision requires it
  • Tests for the new helpers

Acceptance criteria

  • One shared seam exists; placebo_in_time uses it and its local _draw is gone.
  • Unit tests cover: figsize survives the layout pass, data-axes discovery ignores strip axes in a faceted plot, proxy-legend merge preserves pre-existing Axes handles, reserved column names do not collide with user columns of the same name.
  • No public signature or return type changes anywhere.
  • The theming decision is documented in the module docstring or ARCHITECTURE.md.
  • prek run --all-files passes; make test-patch-cov passes.

Out of scope

  • Migrating any experiment's plot (children C–F).
  • Posterior uncertainty layers and tidy draw extraction (child B).
  • Removing plot_posterior_over_x or any other plot_utils function (child G).

Metadata

Metadata

Assignees

No one assigned

    Labels

    agentsAgent related issues specifically for use by developersplottingImprove or fix plottingrefactorRefactor, clean up, or improvement with no visible changes to the user

    Type

    No type

    Projects

    No projects

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions