refactor: bundle experiment attributes into shared design/result types

[drbenvincent]

Summary

This refactor tackles the experiment-class “god object” attribute sprawl identified in the attribute audit. It does not aim to shrink line count — the PR is ~+330 net LOC (new _design.py / _results.py modules, guards, and test churn) — but it does collapse ~50 scattered instance attributes into a small set of shared bundles.

Core changes:

• DesignTransform / PatsyDesignTransform (_design.py): patsy DesignInfo is no longer stored on every formula experiment; rebuilds go through transform_x / transform_y. Future patsy → formulaic swap is a one-file change.
• CausalResult (_results.py): ITS, Synthetic Control, SDiD, and PiecewiseITS expose pre/post predictions, impacts, and score via self.result instead of six duplicate top-level attrs. reporting.py and placebo checks read this typed contract.
• Three-period ITS: eight conditional attributes (intervention_pred, post_intervention_impact, etc.) become optional intervention_result / post_intervention_result (CausalResult | None).
• Scenario: DiD/RD/RK prediction (input, output) pairs bundled for internal plotting.
• PriorCalibration: IV OLS first-stage/second-stage/naive fits moved off the public surface to _prior_calibration (public priors unchanged).
• Deleted transient *_raw attrs: constructor scratch matrices are locals, not self + del.

Attribute wins (approximate): ~50 duplicated/scattered attrs → ~15 bundled fields + 4 shared types; largest single-class gain is ITS (~14 fewer attrs in three-period mode, stable object shape).

Migration (pre-1.0 attribute renames)

Old attribute	New location	Experiments
pre_pred	result.predictions_pre	ITS, SC, SDiD, PiecewiseITS
post_pred	result.predictions_post	ITS, SC, SDiD, PiecewiseITS
pre_impact	result.impact_pre	ITS, SC, SDiD
post_impact	result.impact_post	ITS, SC, SDiD, PiecewiseITS
post_impact_cumulative	result.impact_post_cumulative	ITS, SC, SDiD
score (top-level)	result.score	SC, PiecewiseITS
intervention_pred / intervention_impact / …	intervention_result.*	Three-period ITS
post_intervention_pred / post_intervention_impact / …	post_intervention_result.*	Three-period ITS
ols_beta_params, first_stage_reg, second_stage_reg, naive_reg	_prior_calibration.beta_*, _prior_calibration.first_stage, …	IV
DiD/RD/RK prediction (inputs, output) pairs	_control, _treatment, _counterfactual (Scenario)	DiD, RD, RK

Public priors on IV is unchanged. experiment.result is the main entry point for counterfactual outputs on time-series experiments.

Test plan

• make test-patch-cov (176 tests)
• IV, ITS three-period, reporting, IPW, piecewise ITS, SDiD test modules
• prek run on changed files (ruff, mypy, notebook validation)
• Review feedback: drop # ponytail: comment leak; transform_y only rebuilds outcome design; runner.py --full IOPub timeout + include skip_notebooks.yml entries for local output refresh
• Re-execute affected notebooks with runner.py --full and commit refreshed outputs (sc_pymc, iv_pymc, iv_weak_instruments, sdid_pymc, its_placebo_in_time_analysis, geolift1, multi_cell_geolift)

Made with Cursor

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[drbenvincent]

@BugBot review

[cursor]

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

^{Reviewed by Cursor Bugbot for commit 9b0a715. Configure here.}

[read-the-docs-community]

Documentation build overview

📚 causalpy | 🛠️ Build #33145382 | 📁 Comparing 3478475 against latest (e287f23)

  🔍 Preview build  

23 files changed · ± 23 modified

± Modified

• 404.html
• notebooks/geolift1.html
• notebooks/its_lift_test.html
• notebooks/its_placebo_in_time_analysis.html
• notebooks/its_pymc.html
• notebooks/iv_pymc.html
• notebooks/iv_weak_instruments.html
• notebooks/multi_cell_geolift.html
• notebooks/sc_pymc.html
• notebooks/sdid_pymc.html
• and 13 more...

[codecov]

Codecov Report

❌ Patch coverage is 98.09725% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.44%. Comparing base (e287f23) to head (3478475).

Files with missing lines	Patch %	Lines
causalpy/experiments/diff_in_diff.py	87.50%	0 Missing and 3 partials ⚠️
causalpy/experiments/interrupted_time_series.py	97.26%	0 Missing and 2 partials ⚠️
causalpy/tests/test_three_period_its.py	98.19%	0 Missing and 2 partials ⚠️
causalpy/experiments/regression_kink.py	93.33%	0 Missing and 1 partial ⚠️
causalpy/reporting.py	75.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #966      +/-   ##
==========================================
+ Coverage   95.32%   95.44%   +0.11%     
==========================================
  Files          96       99       +3     
  Lines       15221    15352     +131     
  Branches      878      885       +7     
==========================================
+ Hits        14510    14652     +142     
+ Misses        502      494       -8     
+ Partials      209      206       -3     

☔ 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.

[daimon-pymclabs]

Review — bundling experiment attributes

Solid, well-structured refactor. All 16 checks are green (176 tests across 3.11/3.14, all notebook jobs execute, RTD builds, patch coverage holds). A cross-file caller trace confirms every internal read site — reporting.py, checks/placebo_in_time.py, the plotting methods — matches the new CausalResult / Scenario / PriorCalibration field names, so there's no internal breakage. The DesignTransform abstraction is a clean way to keep patsy off every formula experiment.

Per maintainer guidance, the public-API breakage (removing pre_pred/post_impact/ols_beta_params/etc. in favour of experiment.result.* and _prior_calibration.*) is acceptable pre-1.0 and is not a blocker. Worth a rename table in the release notes so users hit something searchable. Two concrete items below; I'm handling both in follow-up commits on this branch.

1. Leaked internal codename in shipped comment — causalpy/experiments/_design.py:118

# ponytail: eval_env=1 resolves custom transforms (step/ramp) from the caller's module.

# ponytail: is an agent/tool codename, not appropriate for library source. The comment content is correct and useful — just the prefix needs dropping. (Same class of accidental leak as #953.)

2. Affected notebooks have stale outputs

The code cells were updated to the new API, but most notebooks were not re-executed — outputs still reflect the old-API run. its_lift_test.ipynb was re-executed (outputs refreshed); these were not: sc_pymc, iv_pymc, iv_weak_instruments, sdid_pymc, its_placebo_in_time_analysis, geolift1, multi_cell_geolift (code edited, 0 output changes). CI executes notebooks but discards outputs, so the stale committed outputs slip through. Re-executing these with runner.py --full and committing the refreshed outputs.

Minor (non-blocking)

• _design.py:100 — transform_y calls build_design_matrices([y_design_info, x_design_info], …) and discards the X result, rebuilding the full predictor matrix on every outcome rebuild. Pass [self._y_design_info] only.
• _results.py:40 — impact_post_cumulative is a stored field always derived from impact_post; a @property would prevent the two drifting apart.

Not this PR (pre-existing)

• docs/dev/its_pymc copy.ipynb still references result.post_impact, but it's a stray "copy" scratch file already on main, outside the built docs. Worth deleting in a separate cleanup.

[daimon-pymclabs]

Correction to my note above — to be clear, the two action items below are requests for the PR author, not things I'm taking on:

1. Drop the # ponytail: prefix on _design.py:118 (keep the comment text).

2. Re-execute the affected notebooks locally with real sampling and commit the refreshed outputs in this PR. The code cells were migrated to the new API but most outputs are stale (only its_lift_test was re-run). Stale set — code edited, outputs unchanged:
sc_pymc, iv_pymc, iv_weak_instruments, sdid_pymc, its_placebo_in_time_analysis, geolift1, multi_cell_geolift.

Two things to watch when doing this:

• iv_pymc and iv_weak_instruments are in scripts/run_notebooks/skip_notebooks.yml, so CI never executes them — and they carry the _prior_calibration edits. They're the highest-risk notebooks and need a manual full run (JAX/numpyro env) to confirm the new IV API actually works end-to-end.
• CI's notebook job runs with mocked pm.sample (the runner injects a mock unless --full), so a green CI does not mean the new-API cells executed against real output. Use runner.py --full (note maketables must be installed for several notebooks).

[drbenvincent]

I like the direction of pulling the Patsy rebuild logic into _design.py, but I think the current DesignTransform abstraction is a bit too low-level and a bit too early.

If we really may support multiple formula engines later, I’d suggest aiming for a compiled-formula boundary rather than a generic transform protocol. Right now there is only one implementation, and the protocol is still mostly exposing a thin wrapper around Patsy-specific compiled state. That gives us extra abstraction today without yet giving us a clean swap point tomorrow.

Concretely, I think the simpler shape would be something like:

compiled = compile_formula(formula, data)
X, y = compiled.fit_matrices()
X_new = compiled.transform_x(new_data)
y_new = compiled.transform_y(new_data)
labels = compiled.labels
outcome_name = compiled.outcome_name

Benefits of that shape:

• experiments depend on one stable, experiment-facing contract instead of a prematurely generic protocol
• backend choice stays behind a single factory / module boundary
• Patsy-specific compilation details stay encapsulated in one object
• if we add a second backend later, we can align behavior at the compile_formula(...) boundary instead of having backend differences leak into experiments
• we avoid standardizing too early on an abstraction that may not actually be the right one once we hit real differences (categorical coding, custom transforms, mixed-model support, missing-data behavior, etc.)

So my suggestion would be: keep the extraction, but narrow/reframe it. _design.py looks worthwhile; I’m less convinced that DesignTransform as the central abstraction is the best long-term shape for multi-backend support.

[daimon-pymclabs]

Agreed — I think you've correctly identified the smell. As it stands DesignTransform is a Protocol with exactly one implementation, and the contract is split awkwardly: build_patsy_design/build_patsy_formula_sides hand back the fit arrays (X, y) as free-function returns, while the re-build path (transform_x/transform_y) lives on the object. So "compile this formula" is half a function call and half an object, which is the part that actually leaks Patsy-shaped assumptions into the experiments.

Your compile_formula(...) -> compiled shape fixes exactly that by making one object own compilation and the fit matrices and the transforms:

compiled = compile_formula(formula, data)
X, y = compiled.fit_matrices()

That's a single experiment-facing contract, and the backend choice collapses to the factory — which is the real swap point, not a generic transform protocol. +1 to "keep the extraction, narrow/reframe it."

One wrinkle to design against up front: build_patsy_formula_sides returns two transforms (treatment + covariates) from a single formula for the IV / two-sided case, which doesn't map cleanly onto one compiled.fit_matrices(). Worth deciding whether that's a second factory entry point (compile_formula_sides) or whether CompiledFormula exposes named sides, so the two-sided path doesn't get bolted on later and reintroduce the same split you're removing.

(To be clear, this is just review input — happy for you to drive the reshape.)