Lazy experiment lifecycle: configure → optional prior checks → fit() - #1175
Lazy experiment lifecycle: configure → optional prior checks → fit()#1175drbenvincent wants to merge 21 commits into
Conversation
Implements #1092. Experiment constructors no longer sample: __init__ validates and builds design matrices only; posterior inference runs through an explicit fit() (returns Self), and sample_prior_predictive() enables prior predictive checks before MCMC. Core: - _results.py: per-experiment result bundles (CausalResult family), exposed through raising result/prior_result properties backing the is_fitted/has_prior_predictive predicates; nothing draw-derived is assigned to self. - base.py: build() (public, idempotent, auto-called), fit(), sample_prior_predictive(), group guard in _render_plot, model setter resetting all state (the documented prior-revision path). - pymc_models.py: sampling split with build-time prior merge, prior_sample_kwargs (default draws=500) carried by _clone(), predict(group=) conditioning on idata.prior, re-arm of fit-time data nodes before posterior sampling so refits stay correct after forward sampling mutated them. - adapter surface: build/sample_prior_predictive/sample_posterior/ predict(group=)/coefficients(group=) plus capability flags; IV and StateSpaceTimeSeries declare no prior phase. - GroupNotSampleedException and PriorPredictiveNotSupportedException exported from causalpy. All twelve experiments migrate off algorithm(); plot/get_plot_data/ effect_summary take keyword-only group= with reduced prior panels and prior-plausibility prose. EstimateEffect and every re-fitting check call fit() explicitly. reporting/maketables helpers consume bundles.
Chain .fit() onto experiment construction, move flat-attribute assertions to bundle reads, update guard expectations to GroupNotSampleedException/PriorPredictiveNotSupportedException, accommodate keyword-only group= on plot/effect_summary/get_plot_data, and add the real_pymc_sampling undo fixture to modules asserting posterior-vs-prior relationships. Numeric expectations are unchanged. New test_experiment_lifecycle.py covers guards and their messages, Self returns, build() idempotency with graphviz inspection, model-assignment reset, second-fit overwrite semantics, sklearn capability errors, canonical prior containers, independent draw sizes, reduced prior plot layout, and property-based invariants (posterior HDI narrower than prior; prior tail probability near 0.5 under a neutral prior).
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## pymc6_and_pymcmarketing1_migration #1175 +/- ##
======================================================================
+ Coverage 97.15% 97.25% +0.09%
======================================================================
Files 123 126 +3
Lines 21744 23026 +1282
Branches 1177 1271 +94
======================================================================
+ Hits 21125 22393 +1268
- Misses 411 425 +14
Partials 208 208 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
drbenvincent
left a comment
drbenvincent
left a comment
There was a problem hiding this comment.
Adversarial review — PR #1175 (lazy experiment lifecycle)
I reviewed the full diff (~14.4k lines) and verified the suspicious paths against the working tree, including targeted runtime reproductions in the CausalPy env (no full-suite runs). The lifecycle scaffolding itself (build → optional prior → fit, bundle objects, guards, exception types) is solid, and the migration of every constructor call site (EstimateEffect.run(), all six checks, docs/doctests) to .fit() is complete. However, I found one blocker and three major issues, all rooted in the same mechanism: forward-sampling through predict() mutates the graph's shared data nodes, and only sample_posterior() re-arms them.
Findings
1. [BLOCKER] BBETS time-data nodes are never re-armed — documented prior-check→fit workflow crashes; refit silently corrupts or crashes
File: causalpy/pymc_models.py:445-447 (also 450-470, 2535-2596)
PyMCModel.build() records only {"X", "y"} in _build_data_nodes, and _rearm_fit_data() restores only those. But BayesianBasisExpansionTimeSeries.build_model() creates two more mutable nodes (t_trend_data, t_season_data, lines 2466-2476) and its _data_setter() mutates them on every predict(). Nothing records or restores them. The comment at line 444 says subclasses with different node names should override build() and populate the mapping themselves — BBETS does not, and no residual risk is documented anywhere (release notes / ARCHITECTURE are silent).
Reproduced end-to-end (CausalPy env):
exp = cp.InterruptedTimeSeries(df, treatment_time=idx[20], formula="y ~ 1",
model=BayesianBasisExpansionTimeSeries(...))
exp.sample_prior_predictive() # documented workflow: prior check first
exp.fit()
# AssertionError: Could not broadcast dimensions. Broadcasting is only allowed
# along axes that have a statically known length 1.
Mechanism: _finalize("prior") ends with predict(post_X, out_of_sample=True), which leaves t_trend_data/t_season_data pointing at the post window. The subsequent fit() → sample_posterior() → _rearm_fit_data() restores y to the 20 training rows but leaves t_trend_data at 10 post-window rows → PyTensor broadcast failure when lengths differ, and silently wrong trend/seasonality timestamps when pre/post windows have equal length (e.g. cp.load_data("its"), 60/60). Verified node state directly:
after _rearm_fit_data(): y=(20, 1), t_trend_data=(10,), t_season_data=(10,)
t_trend_data still spans [0.0548,0.0794] — build-time span was [0.0000,0.0520]
A plain refit (fit() twice) is corrupted identically. This breaks the exact workflow the release notes and SKILL.md recommend.
def build(
self,
X: xr.DataArray,
y: xr.DataArray,
coords: dict[str, Any] | None = None,
) -> None:
super().build(X=X, y=y, coords=coords)
# Record the time-feature nodes so _rearm_fit_data() can restore them;
# predict() permanently repoints them via _data_setter().
(
time_for_trend,
time_for_seasonality,
_X_for_pymc,
_num_obs,
) = self._prepare_time_and_exog_features(X)
self._build_data_nodes["t_trend_data"] = np.asarray(time_for_trend)
self._build_data_nodes["t_season_data"] = np.asarray(time_for_seasonality)
2. [MAJOR] Auto-populated prior groups are sampled against mutated data nodes (DiD, PrePostNEGD confirmed polluted)
File: causalpy/experiments/base.py:397-403 (root cause also causalpy/pymc_models.py:487-505)
fit() order is: sample_posterior() → _finalize("posterior") → sample_prior_predictive(). _finalize("posterior")'s last backend calls are predict()s whose _data_setter() permanently repoints the graph's X/y nodes at small scenario designs; sample_prior_predictive() does not re-arm them first, so pm.sample_prior_predictive evaluates deterministics/likelihood against whatever design was last predicted.
Reproduced (DiD, 160 training rows):
training rows: 160
posterior mu obs_ind: 160
prior mu obs_ind: 4 # x_pred_counterfactual scenario grid
prior_predictive y_hat obs_ind: 4
PrePostNEGD is affected identically (last predict is the 200-point interpolation grid). ITS/SC/RD/RK escape only by accident: their trailing score(X=<build X>, ...) call happens to route through predict() again, resetting the nodes — remove or reorder that score call and they pollute too. The historical eager fit sampled the prior immediately after pm.sample against the training design, so this regresses the public idata["prior"]/idata["prior_predictive"] containers relative to both history and the docstring's claim ("fills the prior groups … so idata is as complete as the historical eager fit produced"). Parameter draws (beta) are unaffected, so the prior_result bundles themselves remain correct — the pollution is in exp.idata.
Fix: re-arm before prior sampling, which is safe and idempotent:
def sample_prior_predictive(self, **kwargs: Any) -> xr.DataTree:
self.require_built()
self._rearm_fit_data()
resolved = {**self.prior_sample_kwargs, **kwargs}
3. [MAJOR] StaggeredDiD.effect_summary(group="prior") ignores the group — silently returns the posterior summary
File: causalpy/experiments/staggered_did.py:1635-1643
The method resolves the requested bundle and then discards it:
self._resolve_group(group) # result unused
...
return _effect_summary_staggered_did(self, ...)and _effect_summary_staggered_did reads experiment.result.att_event_time unconditionally (reporting.py:551). Consequences:
- After
sample_prior_predictive()andfit():effect_summary(group="prior")returns the posterior-based causal prose with no prior framing — contradicting its own docstring ("produces prior-appropriate prose — a prior plausibility statement, not a causal estimate"). - With only the prior phase sampled: it raises a misleading
GroupNotSampleedExceptiontelling the user to callfit()even though the requested group exists.
Every other draw-consuming reader of this PR threads the resolved bundle; this one doesn't.
bundle = self._resolve_group(group)
from causalpy.reporting import _effect_summary_staggered_did
summary = _effect_summary_staggered_did(
self,
direction=direction,
alpha=alpha,
min_effect=min_effect,
)
if group == "prior":
summary.text = (
"Prior predictive check (not a causal estimate):\n" + summary.text
)
return summary
(For a true prior-group table, change _effect_summary_staggered_did to accept the resolved bundle's att_event_time/hdi_prob instead of reaching back to experiment.result; the minimal patch above at least stops the mislabeling.)
4. [MAJOR] Prior-group effect summaries read as causal claims for ITS, SyntheticControl, and PiecewiseITS
Files: causalpy/experiments/interrupted_time_series.py:1363, causalpy/experiments/synthetic_control.py:1064, causalpy/experiments/piecewise_its.py:906
These three pass the prior bundle into _effect_summary_timeseries(...), which has no group parameter and applies no framing. DiD/RD/RK route prose through _apply_prior_grouping(...) and SDiD overrides prefix, so the framing contract is implemented everywhere except exactly the three experiments whose docstrings promise it most explicitly ("produces prior-appropriate prose — under a neutral prior, P(effect > 0) should sit near 0.5…"). A user running the recommended prior-check flow gets prose like "Post-period: the intervention had an effect of X (94% HDI …)" computed from prior draws, with no plausibility-check framing.
Fix: thread group into _effect_summary_timeseries and apply _apply_prior_grouping (or override prefix as SDiD does) before calling it from these three experiments.
5. [MINOR] ARCHITECTURE.md states the opposite phase ordering of fit()
File: ARCHITECTURE.md:73
The bullet says fit() "samples the absent prior phase first, then NUTS + posterior predictive". The code (and the BaseExperiment.fit docstring, which explains why posterior must run first) does the reverse: NUTS first, absent prior phase second. Same drift risk as finding 2 — if the ordering is later "fixed" per these docs, the data-node pollution becomes universal.
- **`fit(**kwargs)`** — posterior phase; samples NUTS + posterior predictive, then fills the absent prior phase so `idata` matches the historical eager output; populates `exp.result`. Returns `Self`, so call sites migrate with one appended token: `cp.InterruptedTimeSeries(...).fit()`.
6. [MINOR] RegressionDiscontinuity.plot() lost its keyword-only marker
File: causalpy/experiments/regression_discontinuity.py:364-366
The diff replaced * with group instead of inserting group after it, so plot(group=…, round_to=…, ci_prob=…, kind=…, ci_kind=…, num_samples=…, figsize=…, show=…, legend_kwargs=…) is now positional-or-keyword. Every other experiment kept *, group=…. Project convention (ARCHITECTURE.md, issue #886, test_public_signatures.py) mandates explicit keyword-only plot contracts, and this reintroduces positional ambiguity that a future parameter reorder would silently break.
def plot(
self,
*,
group: Literal["prior", "posterior"] = "posterior",
round_to: int | None = 2,
7. [NIT] Unreachable duplicate return in PiecewiseITS._plot_prior_checks
File: causalpy/experiments/piecewise_its.py:759-760
return fig, [ax]
return fig, ax # deadDelete the second return (its list-vs-single-axes inconsistency suggests a leftover from an earlier shape).
8. [NIT] IV summary() silently dropped the formula header lines
File: causalpy/experiments/instrumental_variable.py:378
Formula: / Instruments formula: prints were removed without a mention in the release notes' attribute-rename table. If intentional cleanup, fine — but it changes captured output for anyone diffing summaries, so it deserves a release-note line.
9. [NIT] Stale format_r2_score docstring reference
File: causalpy/plot_utils.py:51
Still says the score is "as stored on experiment.score"; scores now live on result.score (the signature correctly accepts None).
Checked and found sound
- PiecewiseITS bundle mapping (
predictions_pre= full-window fitted expectation,predictions_post/impact_post= post-first-interruption slices): consistent acrossget_plot_data,effect_summary, and_plot; full-window counterfactual recovery (predictions_pre − impact_pre) is algebraically exact.interruption_times[0]cannot IndexError because_validate_inputsrejects formulas without step/ramp terms. - SDID
_extract_weight_posteriors(group): readsrequire_idata()[group]; tau math identical for both groups since it operates on the extracted arrays. Mapping-input build means no data nodes to mutate. - StaggeredDiD hdi_prob semantics:
bundle.hdi_probpreserves oldhdi_prob_behavior (aggregation-time freeze, plot-time mismatchValueError,get_plot_datarecompute-on-mismatch, placebo helper reuse). - IV/IPW lazy overrides: warning + no-op
build()+supports_prior_predictive=False(flag verified atpymc_models.py:1497and state-space:2942, matching the new exception docstring); IPW's t-as-y flows correctly throughadapter.build→PropensityScore.build(y=treatment)with"t"recorded in_build_data_nodes. - PanelRegression
_supports_results=False: every draw reader (summary,_plot,_plot_coefficients_internal,get_plot_data,plot_unit_effects,plot_trajectories) is guarded.plot(group="prior")rendering posterior draws is documented in the docstring ("treated as posterior") — deliberate, though arguably surprising given PyMC-backed panels do acquire prior draws insidefit(). - Lifecycle plumbing:
EstimateEffect.run()and all checks call.fit();generate_reportno longer swallowsGroupNotSampleedException(onlyNotImplementedError); maketables readers guard all three failure modes; refit warning verified live; model reassignment reset verified by tests; harness_lazy_fitsentinel is scoped to the migration script.
drbenvincent
left a comment
drbenvincent
left a comment
There was a problem hiding this comment.
Adversarial review — PR #1175 (lazy experiment lifecycle)
I reviewed the full diff (~14.4k lines) and verified the suspicious paths against the working tree, including targeted runtime reproductions in the CausalPy env (no full-suite runs). The lifecycle scaffolding itself (build → optional prior → fit, bundle objects, guards, exception types) is solid, and the migration of every constructor call site (EstimateEffect.run(), all six checks, docs/doctests) to .fit() is complete. However, I found one blocker and three major issues, all rooted in the same mechanism: forward-sampling through predict() mutates the graph's shared data nodes, and only sample_posterior() re-arms them.
Findings
1. [BLOCKER] BBETS time-data nodes are never re-armed — documented prior-check→fit workflow crashes; refit silently corrupts or crashes
File: causalpy/pymc_models.py:445-447 (also 450-470, 2535-2596)
PyMCModel.build() records only {"X", "y"} in _build_data_nodes, and _rearm_fit_data() restores only those. But BayesianBasisExpansionTimeSeries.build_model() creates two more mutable nodes (t_trend_data, t_season_data, lines 2466-2476) and its _data_setter() mutates them on every predict(). Nothing records or restores them. The comment at line 444 says subclasses with different node names should override build() and populate the mapping themselves — BBETS does not, and no residual risk is documented anywhere (release notes / ARCHITECTURE are silent).
Reproduced end-to-end (CausalPy env):
exp = cp.InterruptedTimeSeries(df, treatment_time=idx[20], formula="y ~ 1",
model=BayesianBasisExpansionTimeSeries(...))
exp.sample_prior_predictive() # documented workflow: prior check first
exp.fit()
# AssertionError: Could not broadcast dimensions. Broadcasting is only allowed
# along axes that have a statically known length 1.
Mechanism: _finalize("prior") ends with predict(post_X, out_of_sample=True), which leaves t_trend_data/t_season_data pointing at the post window. The subsequent fit() → sample_posterior() → _rearm_fit_data() restores y to the 20 training rows but leaves t_trend_data at 10 post-window rows → PyTensor broadcast failure when lengths differ, and silently wrong trend/seasonality timestamps when pre/post windows have equal length (e.g. cp.load_data("its"), 60/60). Verified node state directly:
after _rearm_fit_data(): y=(20, 1), t_trend_data=(10,), t_season_data=(10,)
t_trend_data still spans [0.0548,0.0794] — build-time span was [0.0000,0.0520]
A plain refit (fit() twice) is corrupted identically. This breaks the exact workflow the release notes and SKILL.md recommend.
def build(
self,
X: xr.DataArray,
y: xr.DataArray,
coords: dict[str, Any] | None = None,
) -> None:
super().build(X=X, y=y, coords=coords)
# Record the time-feature nodes so _rearm_fit_data() can restore them;
# predict() permanently repoints them via _data_setter().
(
time_for_trend,
time_for_seasonality,
_X_for_pymc,
_num_obs,
) = self._prepare_time_and_exog_features(X)
self._build_data_nodes["t_trend_data"] = np.asarray(time_for_trend)
self._build_data_nodes["t_season_data"] = np.asarray(time_for_seasonality)
2. [MAJOR] Auto-populated prior groups are sampled against mutated data nodes (DiD, PrePostNEGD confirmed polluted)
File: causalpy/experiments/base.py:397-403 (root cause also causalpy/pymc_models.py:487-505)
fit() order is: sample_posterior() → _finalize("posterior") → sample_prior_predictive(). _finalize("posterior")'s last backend calls are predict()s whose _data_setter() permanently repoints the graph's X/y nodes at small scenario designs; sample_prior_predictive() does not re-arm them first, so pm.sample_prior_predictive evaluates deterministics/likelihood against whatever design was last predicted.
Reproduced (DiD, 160 training rows):
training rows: 160
posterior mu obs_ind: 160
prior mu obs_ind: 4 # x_pred_counterfactual scenario grid
prior_predictive y_hat obs_ind: 4
PrePostNEGD is affected identically (last predict is the 200-point interpolation grid). ITS/SC/RD/RK escape only by accident: their trailing score(X=<build X>, ...) call happens to route through predict() again, resetting the nodes — remove or reorder that score call and they pollute too. The historical eager fit sampled the prior immediately after pm.sample against the training design, so this regresses the public idata["prior"]/idata["prior_predictive"] containers relative to both history and the docstring's claim ("fills the prior groups … so idata is as complete as the historical eager fit produced"). Parameter draws (beta) are unaffected, so the prior_result bundles themselves remain correct — the pollution is in exp.idata.
Fix: re-arm before prior sampling, which is safe and idempotent:
def sample_prior_predictive(self, **kwargs: Any) -> xr.DataTree:
self.require_built()
self._rearm_fit_data()
resolved = {**self.prior_sample_kwargs, **kwargs}
3. [MAJOR] StaggeredDiD.effect_summary(group="prior") ignores the group — silently returns the posterior summary
File: causalpy/experiments/staggered_did.py:1635-1643
The method resolves the requested bundle and then discards it:
self._resolve_group(group) # result unused
...
return _effect_summary_staggered_did(self, ...)and _effect_summary_staggered_did reads experiment.result.att_event_time unconditionally (reporting.py:551). Consequences:
- After
sample_prior_predictive()andfit():effect_summary(group="prior")returns the posterior-based causal prose with no prior framing — contradicting its own docstring ("produces prior-appropriate prose — a prior plausibility statement, not a causal estimate"). - With only the prior phase sampled: it raises a misleading
GroupNotSampleedExceptiontelling the user to callfit()even though the requested group exists.
Every other draw-consuming reader of this PR threads the resolved bundle; this one doesn't.
bundle = self._resolve_group(group)
from causalpy.reporting import _effect_summary_staggered_did
summary = _effect_summary_staggered_did(
self,
direction=direction,
alpha=alpha,
min_effect=min_effect,
)
if group == "prior":
summary.text = (
"Prior predictive check (not a causal estimate):\n" + summary.text
)
return summary
(For a true prior-group table, change _effect_summary_staggered_did to accept the resolved bundle's att_event_time/hdi_prob instead of reaching back to experiment.result; the minimal patch above at least stops the mislabeling.)
4. [MAJOR] Prior-group effect summaries read as causal claims for ITS, SyntheticControl, and PiecewiseITS
Files: causalpy/experiments/interrupted_time_series.py:1363, causalpy/experiments/synthetic_control.py:1064, causalpy/experiments/piecewise_its.py:906
These three pass the prior bundle into _effect_summary_timeseries(...), which has no group parameter and applies no framing. DiD/RD/RK route prose through _apply_prior_grouping(...) and SDiD overrides prefix, so the framing contract is implemented everywhere except exactly the three experiments whose docstrings promise it most explicitly ("produces prior-appropriate prose — under a neutral prior, P(effect > 0) should sit near 0.5…"). A user running the recommended prior-check flow gets prose like "Post-period: the intervention had an effect of X (94% HDI …)" computed from prior draws, with no plausibility-check framing.
Fix: thread group into _effect_summary_timeseries and apply _apply_prior_grouping (or override prefix as SDiD does) before calling it from these three experiments.
5. [MINOR] ARCHITECTURE.md states the opposite phase ordering of fit()
File: ARCHITECTURE.md:73
The bullet says fit() "samples the absent prior phase first, then NUTS + posterior predictive". The code (and the BaseExperiment.fit docstring, which explains why posterior must run first) does the reverse: NUTS first, absent prior phase second. Same drift risk as finding 2 — if the ordering is later "fixed" per these docs, the data-node pollution becomes universal.
- **`fit(**kwargs)`** — posterior phase; samples NUTS + posterior predictive, then fills the absent prior phase so `idata` matches the historical eager output; populates `exp.result`. Returns `Self`, so call sites migrate with one appended token: `cp.InterruptedTimeSeries(...).fit()`.
6. [MINOR] RegressionDiscontinuity.plot() lost its keyword-only marker
File: causalpy/experiments/regression_discontinuity.py:364-366
The diff replaced * with group instead of inserting group after it, so plot(group=…, round_to=…, ci_prob=…, kind=…, ci_kind=…, num_samples=…, figsize=…, show=…, legend_kwargs=…) is now positional-or-keyword. Every other experiment kept *, group=…. Project convention (ARCHITECTURE.md, issue #886, test_public_signatures.py) mandates explicit keyword-only plot contracts, and this reintroduces positional ambiguity that a future parameter reorder would silently break.
def plot(
self,
*,
group: Literal["prior", "posterior"] = "posterior",
round_to: int | None = 2,
7. [NIT] Unreachable duplicate return in PiecewiseITS._plot_prior_checks
File: causalpy/experiments/piecewise_its.py:759-760
return fig, [ax]
return fig, ax # deadDelete the second return (its list-vs-single-axes inconsistency suggests a leftover from an earlier shape).
8. [NIT] IV summary() silently dropped the formula header lines
File: causalpy/experiments/instrumental_variable.py:378
Formula: / Instruments formula: prints were removed without a mention in the release notes' attribute-rename table. If intentional cleanup, fine — but it changes captured output for anyone diffing summaries, so it deserves a release-note line.
9. [NIT] Stale format_r2_score docstring reference
File: causalpy/plot_utils.py:51
Still says the score is "as stored on experiment.score"; scores now live on result.score (the signature correctly accepts None).
Checked and found sound
- PiecewiseITS bundle mapping (
predictions_pre= full-window fitted expectation,predictions_post/impact_post= post-first-interruption slices): consistent acrossget_plot_data,effect_summary, and_plot; full-window counterfactual recovery (predictions_pre − impact_pre) is algebraically exact.interruption_times[0]cannot IndexError because_validate_inputsrejects formulas without step/ramp terms. - SDID
_extract_weight_posteriors(group): readsrequire_idata()[group]; tau math identical for both groups since it operates on the extracted arrays. Mapping-input build means no data nodes to mutate. - StaggeredDiD hdi_prob semantics:
bundle.hdi_probpreserves oldhdi_prob_behavior (aggregation-time freeze, plot-time mismatchValueError,get_plot_datarecompute-on-mismatch, placebo helper reuse). - IV/IPW lazy overrides: warning + no-op
build()+supports_prior_predictive=False(flag verified atpymc_models.py:1497and state-space:2942, matching the new exception docstring); IPW's t-as-y flows correctly throughadapter.build→PropensityScore.build(y=treatment)with"t"recorded in_build_data_nodes. - PanelRegression
_supports_results=False: every draw reader (summary,_plot,_plot_coefficients_internal,get_plot_data,plot_unit_effects,plot_trajectories) is guarded.plot(group="prior")rendering posterior draws is documented in the docstring ("treated as posterior") — deliberate, though arguably surprising given PyMC-backed panels do acquire prior draws insidefit(). - Lifecycle plumbing:
EstimateEffect.run()and all checks call.fit();generate_reportno longer swallowsGroupNotSampleedException(onlyNotImplementedError); maketables readers guard all three failure modes; refit warning verified live; model reassignment reset verified by tests; harness_lazy_fitsentinel is scoped to the migration script.
- BBETS: record t_trend_data/t_season_data in _build_data_nodes at build so refits re-arm the trend/seasonality nodes after forecast-window conditioning (blocker: prior-check workflow crashed on unequal windows and silently corrupted equal-window refits). - sample_prior_predictive() re-arms fit-time data before sampling, so the auto-prior phase inside fit() draws against the training design instead of whatever scenario grid a prior finalize left on the graph. - StaggeredDiD.effect_summary consumes the resolved group bundle; reporting helper reads att_event_time/hdi_prob from it and applies prior-plausibility framing for group='prior'. - _effect_summary_timeseries gains group= with prior framing; ITS, SC, and PiecewiseITS pass it through. - RD plot() keyword-only marker restored; PiecewiseITS dead return removed; IV summary Formula headers restored for output parity; extract_r2_score docstring references result.score; ARCHITECTURE.md phase-ordering bullet corrected.
Fixes for round-1 review findingsThanks for the thorough review — every finding was reproduced where claimed and addressed. Summary of what changed and why: Blocker 1 — BBETS time-feature nodes never re-armedFixed. Major 2 — auto-prior phase sampled against mutated data nodesFixed as suggested. Major 3 — StaggeredDiD.effect_summary discarded the resolved bundleFixed. Major 4 — timeseries effect summaries lacked prior framingFixed at the helper level. Minor findings
Verification after fixes
|
drbenvincent
left a comment
drbenvincent
left a comment
There was a problem hiding this comment.
Round-2 review — verification of round-1 fixes + fresh sweep
All six round-1 findings were re-verified with targeted reproductions ($CONDA_EXE run -n CausalPy python); every one is genuinely fixed. The fix commit itself introduces one new edge-case defect (finding R2-1) and leaves the round-1 blocker without a pinned regression test (R2-2). No full-suite run was repeated per review instructions.
Verification of round-1 findings
Blocker 1 — BBETS t_trend_data/t_season_data never re-armed: FIXED ✅
Reproduced on an unequal-window setup (40 train / 15 forecast): build → sample_prior_predictive → sample_posterior → predict(15-row window) → sample_posterior keeps idata["posterior"]["mu"].sizes["obs_ind"] == 40 throughout; _build_data_nodes correctly records X, y, t_trend_data, t_season_data at build time and _rearm_fit_data() restores them after predict(). The no-exog case also works (the name in self.named_vars filter drops the absent "X" node from the payload).
Major 2 — prior phase sampled against mutated data nodes: FIXED ✅
With a LinearRegression conditioned onto a 4-row grid via predict(), a subsequent standalone sample_prior_predictive() now draws prior["mu"] with obs_ind == 20 (training design), not 4. Also confirmed: sample_prior_predictive() after fit() overwrites only prior/prior_predictive and preserves posterior sizes; refit-after-conditioning samples the training design. I traced PrePostNEGD's scenario-grid path specifically: _finalize(group) conditions through PyMCModel.predict(..., group=...), which calls pm.sample_posterior_predictive(conditioning_draws) directly and never routes through sample_prior_predictive() — so the new re-arm does not interfere with intentional grid conditioning. All internal second-build callers (BaseExperiment.build, legacy fused fit) are guarded by _built.
Major 3 — StaggeredDiD.effect_summary discarded resolved bundle: FIXED ✅
Verified end-to-end on tiny staggered draws: (a) prior-only phase — effect_summary(group="prior") returns the prior ATT table with "Prior predictive check (not a causal estimate):" framing instead of raising the misleading fit() error; (b) both phases — group="prior" reads att_event_time/hdi_prob from _prior_result while the default summary reads them from result (prose HDI % matches each bundle's stored hdi_prob); (c) back-compat — calling _effect_summary_staggered_did(exp) with the bundle omitted still falls back to experiment.result and produces identical output. No remaining caller or test mock uses the old positional layout.
Major 4 — timeseries effect summaries lacked prior framing: FIXED ✅
_effect_summary_timeseries(..., group="prior") prepends the plausibility framing on both the Bayesian-draws branch and the point-estimate branch; all three call sites (ITS interrupted_time_series.py:1374, PiecewiseITS piecewise_its.py:916, SC synthetic_control.py:1075) pass group=group. Posterior prose unchanged. (For OLS backends group="prior" is unreachable — the adapter raises PriorPredictiveNotSupportedException first — so framing the t-based branch is harmless.)
Minor findings: ALL FIXED ✅
- RD
plot():groupverifiedKEYWORD_ONLY. - IV
summary(): both Formula header lines restored (instrumental_variable.py:381-382);self.formula/self.instruments_formulaexist (set at construction, lines 161–162); targeted IV summary test passes. - ARCHITECTURE.md bullet now matches implementation (posterior first, then absent-prior fill, data-node rationale stated).
extract_r2_scoredocstring references the bundle'sscorefield, which exists on all result bundles.- PiecewiseITS dead return removed.
New findings
R2-1 (minor, leaning major) — BBETS build() override pollutes _build_data_nodes when invoked on an already-built graph after conditioning
BayesianBasisExpansionTimeSeries.build() (causalpy/pymc_models.py:2447-2458) runs its .update() unconditionally after super().build(). The base build() documents itself as a no-op once built ("a graph is built at most once per instance"), but the override still executes on that early-return path and reads back the current shared-node values. After predict() has re-purposed the graph for a 15-row forecast window, a direct second model.build(X_train, y_train) records t_trend_data/t_season_data at length 15 next to X/y at length 40; the next sample_posterior() then crashes in _rearm_fit_data():
pymc.exceptions.ShapeError: Length of new coordinate values does not match
the new dimension length. (actual 40 != expected 15)
Worse, if the prediction window happens to be the same length as training (different span), the override silently records forecast-window time features as the build-time design and the refit runs on the wrong trend/seasonality inputs with no error. The experiment-level lifecycle cannot trigger this (BaseExperiment.build() checks is_built; every internal caller is guarded or one-shot), but exp.model and model-level build()/fit() are public, documented surfaces, and this reintroduces exactly the stale-design class the round-1 blocker was about. Suggested fix — record only when the graph was actually constructed:
already_built = self._built
super().build(X=X, y=y, coords=coords)
if already_built:
return
# The graph's trend/seasonality inputs are derived from X at build
# time and stored under their own pm.Data node names; record them so
# _rearm_fit_data() can restore them after predict() re-purposes the
# nodes for forecast-window conditioning.
self._build_data_nodes.update(R2-2 (minor) — the round-1 blocker has no pinned regression test
No test in the suite touches _build_data_nodes, the BBETS re-arm, or a refit-after-predict sequence (grep across causalpy/tests confirms). The implementer's verification was manual only; R2-1 shows how easily this area regresses. A ~20-line integration test would pin it: unequal-window BBETS, fit → predict(forecast window) → fit, assert idata["posterior"]["mu"].sizes["obs_ind"] == n_train (and ideally assert _build_data_nodes shapes are untouched after a second build() call, which fails today per R2-1).
R2-3 (nit) — lifecycle suite never exercises sample_prior_predictive() after fit() directly
The auto-fill inside fit() covers the merge-with-existing-idata branch indirectly, but no test calls sample_prior_predictive() after a completed fit and asserts idata["posterior"] sizes are preserved while prior draws change. One assertion in test_experiment_lifecycle.py (e.g., extending test_second_fit_warns_and_preserves_prior) would close it.
Verdict
Round 1: all findings fixed, verified by reproduction. Round 2: one new patch-introduced defect (R2-1, narrow reachability, crash or silent wrong-design refit via documented public model API) plus two low-cost test-hardening suggestions. Nothing here reopens the round-1 blocker for the standard experiment lifecycle.
…regressions The build override now returns early when the graph already existed, so an idempotent second build() cannot record forecast-window values over the training design in _build_data_nodes (round-2 review R2-1). Adds lifecycle regression tests for the refit-after-predict re-arm path (round-1 blocker) and for prior-after-fit posterior preservation.
Fixes for round-2 review findingsAll three findings addressed: R2-1 — BBETS
|
drbenvincent
left a comment
drbenvincent
left a comment
There was a problem hiding this comment.
Round-3 review (final): APPROVE
Verified at HEAD d801d33c. Round-1/round-2 fixes confirmed with concrete reproductions (BBETS-backed ITS, 73 pre / 47 post, tiny draws). No ship-blocking defects found: nothing produces wrong numbers, nothing corrupts silently, and no guard fails open. One non-blocking coverage/claim finding below.
Round-2 fix verification (reproduced, not just read)
R2-1 guard (causalpy/pymc_models.py:2447-2453) — correct.
- HEAD:
fit()→plot()(forecast conditioning) → explicit secondmodel.build(X_pre, y_pre)leaves_build_data_nodesshape-neutral (all nodes stay at training shape). Refit-after-plot re-samples on the training window (posterior["mu"].obs_ind == 73), andsample_prior_predictive()→fit()yields both groups on the training design. All PASS. - Pre-fix simulation (guard removed via monkeypatch, no repo edits): the same second
build()records the 47-row forecast window over the 73-row training design int_trend_data/t_season_data, and the subsequentfit()raisesShapeError: Length of new coordinate values does not match... (actual 73 != expected 47)— mechanism confirmed, and the failure is loud, not silent.
R2-3 test (test_prior_after_fit_overwrites_prior_preserves_posterior, causalpy/tests/test_experiment_lifecycle.py:316) — pins what it claims. Asserts prior re-sampled to 40 draws while posterior chain/draw/obs sizes are byte-identical and the prior bundle object is recomputed.
R2-2 test (test_refit_after_predict_restores_training_design, line 273) — exists, passes, but does NOT pin the R2-1 fix (non-blocking, see finding). It constructs cp.pymc_models.LinearRegression, so the BBETS build() override never runs. I re-ran its exact assertion sequence with a deliberately broken (pre-fix) BBETS in-process and every assertion still passed — it cannot detect the R2-1 regression. The PR comment's claim that it "fails on the pre-fix code path R2-1 described" is inaccurate; it pins the base X/y re-arm contract (which worked since the initial PR commit — sample_posterior always called _rearm_fit_data). The BBETS round-1 re-arm and the R2-1 guard currently have no automated coverage (test_its_with_bsts_model fits and plots but never refits after plot). I verified both paths manually at HEAD; the guard is correct, so this is a coverage/claim-accuracy gap, not a runtime defect.
Final sweep (blockers only)
- All six re-fitting checks +
EstimateEffect.run()clone-then-.fit()on fresh instances (bandwidth.py:112,leave_one_out.py:107,placebo_in_time.py:521with per-fold seeds,placebo_in_space.py:123,outcome_falsification.py:195,prior_sensitivity.py:142,estimate_effect.py:128). No leftover eager call sites; remaining.algorithm()mentions are historical docstrings only. fit()ordering (posterior → finalize → absent-prior fill) matches ARCHITECTURE.md and the PR body; verified dynamically on BBETS (prior and posterior both on the training design).prior_sample_kwargsdefaults{"draws": 500, "random_seed": posterior seed}(pymc_models.py:274-281) match the PR body.- All five exceptions re-exported from
causalpy/__init__.pyand pinned bytest_exceptions_are_exported. - Prior-group prose framing applied at the helper level across scalar and timeseries summaries (
reporting.py:464-467and call sites). test_experiment_lifecycle.py: 16/16 pass at HEAD.
Commit story
Five commits read coherently: core lazy lifecycle → test-suite migration → docs → round-1 fixes → round-2 guard + regressions. PR body claims (fit ordering, prior_sample_kwargs defaults, exception exports, Self returns, bundle-backed state) all match final behavior.
Non-blocking finding
test_refit_after_predict_restores_training_design does not exercise the BBETS path it cites — causalpy/tests/test_experiment_lifecycle.py:273-314, P2, confidence 0.95. Docstring says it "guards the BBETS time-feature re-arm path (round-1 review blocker) and the base X/y re-arm alike", but with LinearRegression the BBETS override is never invoked, and its second-build() assertion passes even with the R2-1 bug present (demonstrated). Fix when convenient: parametrize the test over LinearRegression and a BayesianBasisExpansionTimeSeries(trend_component=MockComponent(), seasonality_component=MockComponent()) variant (the coverage-suite's MockComponent avoids the pymc-marketing dependency) — the BBETS variant fails pre-fix and passes at HEAD, per my reproduction. If regressed, the failure mode is a loud ShapeError on an experimental model, so this does not block merge.
Verdict: APPROVE — ship it. The finding above is follow-up test hygiene, not a merge condition.
Parametrize test_refit_after_predict_restores_training_design over LinearRegression and a mock-component BBETS variant, reorder the sequence so the second build() call happens while the shared nodes are still forecast-conditioned (before the healing refit re-arms them), and assert every recorded data node keeps training-length shapes. Verified: the bbets variant fails with the round-2 guard disabled and passes at HEAD (round-3 review P2).
|
Round-3 P2 addressed in c7c2e07: |
Patch coverage was 84% against the 96% gate: the prior-group branches added across nine experiments and the adapter guard paths had no direct tests. New test_prior_phase.py pins the observable contract end to end: plot(group='prior') renders exactly one axes per capable experiment, effect_summary(group='prior') prose is framed as a plausibility check, get_plot_data(group='prior') returns without caching, and every documented capability/guard error on the PyMC, sklearn, forecast, and base adapters raises as specified. Coverage work exposed two real defects, both fixed: - RegressionDiscontinuity.effect_summary(group='prior') double-prefixed the plausibility framing (helper-level grouping plus a leftover manual prefix block). - reporting._extract_window crashed with TypeError for datetime-sliced windows: pandas 3.x no longer accepts Timestamp/string-bounded slice objects on DatetimeIndex.__getitem__. Both branches now boolean-mask.
Remote
|
Append .fit() at all 80 experiment-construction call sites across the 37 gallery and knowledgebase notebooks, and remap the 40 reads of removed flat draw-derived attributes to their result-bundle locations (post_impact -> result.impact_post, att_event_time_ -> result.att_event_time, etc.). Regenerate all notebook outputs by executing the corpus in place with real MCMC (scripts/run_notebooks/runner.py --full): 33/34 gallery and 3/3 knowledgebase notebooks pass. Local-only environment notes: nutpie and CPU jax were installed to execute interrupted-time-series-placebo- in-time-analysis (nuts_sampler="nutpie"); staggered-difference-in- differences-pymc hit the known PyMC 6.0.x macOS fork-worker EOFError once and passed on retry. Skip gallery/instrumental-variables-variable-selection-priors.ipynb in the notebook CI lanes: like its two sibling IV notebooks, it requires JAX, which the test environment does not install.
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
Docs pass complete: all notebooks migrated and executedWhat changed (e8caf65)
Local execution notes (not repo changes)
Verification
|
test_production_skip_configuration_is_consistent pins the exact contents of skip_notebooks.yml; extend the expected set with gallery/instrumental-variables-variable-selection-priors.ipynb, which requires JAX like its two sibling IV notebooks.
The attribute-remap sweep was not idempotent: replacements contain their own search patterns, so re-running it after the variable-selection notebook restore doubled the bundle prefix (result.result.discontinuity_at_threshold). Collapse repeated .result prefixes to a single one in the five affected notebooks. Verified with a mock-mode run of the gallery-pymc lane (exit 0).
|
Fixed the gallery-pymc notebook lane (021a0af): the attribute-remap sweep used replacements containing their own search patterns, so re-running it doubled the bundle prefix (e.g. |
|
Boom! Big API change PR ready for review. So far zero attempt to document the new workflow other than compatibility / at the new Feedback welcome |
anevolbap
left a comment
anevolbap
left a comment
There was a problem hiding this comment.
Read the whole change against issue #1092 and the conventions in ARCHITECTURE.md. This is a first pass over the full diff; more detailed passes to follow.
Most of it checks out. Confirmed by reading the code: predict(group=) really does condition on idata[group], a second fit() warns and keeps the prior draws, _clone() carries prior_sample_kwargs, and the new tests do not pin seeds. The migration notes are already written at docs/source/release_notes.md:136, so the unchecked box in the description is out of date.
Two things look worth settling before merge: the misspelled GroupNotSampleedException, which is already exported, and the removed plot guard in causalpy/steps/report.py, which turns a partial report into a crash for two experiments. The rest are small: repeated code, a few missing type hints, and two doc lines that no longer match the code.
|
Third wave handled (4b064b8) — inline replies posted on each thread:
Full suite at HEAD: 2410 passed, 2 skipped; |
|
Think that's all comments addressed now @anevolbap. Thanks for this :) This is a big one, so could do with more critique - maybe @cetagostini or @juanitorduz could unleash an adversarial review :) |
|
Re-checked at 4b064b8. Three things left, all still reproducing there.
An IV refit leaves the predictive groups stale ( Two smaller ones. Nothing here blocks from my side; the second item is the only one I would want either fixed or consciously deferred before merge. |
…pc choice, capability-aware prior guard - PrePostNEGD._finalize reuses the control scenario's prediction for scenario_counterfactual: both sweep the same shared pretest grid with the group indicator at zero, so the extra predict pass computed an identical array twice per group (twice per fit with prior auto-fill). Bundle stays fully populated; one posterior-predictive pass saved. - InstrumentalVariable.fit keeps the previous ppc_sampler when a refit omits it: falling back to None left posterior_predictive (and any prior-predictive groups from ppc_sampler='pymc') stale against the freshly resampled posterior while the warning only spoke about the posterior. Explicit None still opts out; the refit test pins all four transitions including persistence and explicit opt-back-out. - BaseExperiment._resolve_group raises PriorPredictiveNotSupportedException directly when group='prior' on a backend without the capability, mirroring what sample_prior_predictive() would raise instead of pointing users at a call that cannot succeed (sklearn-backed PanelRegression hit this); guard tests updated accordingly. - diff_in_diff plot drops the stale Optional comment and dead None guard on scenario_counterfactual, required since the fully-populated-bundle invariant.
|
All three addressed in 7e1a63c — thanks for the careful re-check at HEAD. 1. PrePostNEGD 2. IV refit leaving predictive groups stale 3a. Stale Optional handling in DiD 3b. Capability-blind prior guard message Verification at 7e1a63c: full suite 2410 passed, 2 skipped, |
|
|
||
| summary = InterruptedTimeSeries._comparison_period_summary( | ||
| result, | ||
| InterruptedTimeSeries._comparison_period_summary( |
There was a problem hiding this comment.
This test doesn't check anything any more. The return value is thrown away, so it passes no matter what _comparison_period_summary gives back.
Before this PR it pinned two ArviZ 0.22 HDI baselines:
assert tuple(
summary.table.loc["intervention", ["hdi_lower", "hdi_upper"]]
) == pytest.approx((0.07791736162436769, 2.333331414052698), rel=1e-12, abs=1e-12)
assert tuple(
summary.table.loc["post_intervention", ["hdi_lower", "hdi_upper"]]
) == pytest.approx((0.05178753076996489, 2.3244471772455273), rel=1e-12, abs=1e-12)Both are gone now. grep -c assert over the body of the test returns 0.
What that breaks: the test name, the docstring ("retain the ArviZ 0.22 94% baseline") and the comment a few lines up ("the frozen expectations below") all still say the baseline is pinned, but nothing checks it. So anyone auditing acceptance criterion 7 by reading test names would think this is still covered when it isn't, and a real HDI change would go through quietly.
Worth saying: as far as I can tell this is the only test in the whole diff that lost expected values. I looked at every numeric literal inside an assert across all 37 changed test files and found no loosened tolerances, no changed seeds and no new xfail/skip. So the rest of the AC7 parity argument holds up fine, which is why this one stands out.
Suggested fix: keep the new bundle-shaped stub, put summary = back on the call, and re-add the two pytest.approx asserts. If the seeded draw order really did change (the comment about the intervention slice consuming the first draws suggests it did), just recompute the four numbers and pin the new ones instead. Either way, better than shipping it with no assertions.
| @@ -0,0 +1,721 @@ | |||
| # Copyright 2022 - 2026 The PyMC Labs Developers | |||
There was a problem hiding this comment.
This file never opts out of the session-wide mock_pymc_sample fixture from conftest.py:61, so in CI every .fit() here runs against a mocked pm.sample. test_experiment_lifecycle.py:34, test_parameter_recovery.py and test_pymc_forecast_adapter.py all add a fixture to undo it — this one doesn't.
Easy to see:
test_prior_phase.py alone -> 43 passed in 56.62s (real MCMC)
test_default_models.py + test_prior_phase.py -> 54 passed in 8.81s (mocked)
What that breaks: in the one file about prior vs posterior contracts, the "posterior" in CI is prior sampling. It also makes the file's behaviour depend on test order, and pytest-randomly changes that every run.
Fix: add the same opt-out fixture you used in test_experiment_lifecycle.py, or apply mock_sample per test and keep real sampling only for the two property tests that need it.
| try: | ||
| return experiment.result | ||
| except (GroupNotSampledException, NotImplementedError, AttributeError): | ||
| return None |
There was a problem hiding this comment.
Catching GroupNotSampledException here turns a clear error into a confusing one:
ETable(fitted) -> OK
ETable(unfitted) -> TypeError: No extractor available for model type: InterruptedTimeSeries
The user just forgot .fit(), but the message says the plugin doesn't support their experiment.
This is the thing #1092 wrote a whole rejected-design row about — it refused sklearn.NotFittedError so guards wouldn't get swallowed into a None in these hooks. The PR avoided the inheritance and then got the same outcome by catching the new exception on purpose.
Fix: let GroupNotSampledException through and keep AttributeError for the duck-typed stubs. test_maketables_plugin.py has no unfitted-state test, so nothing catches this today either.
| if window.stop is not None | ||
| else post_index.max() | ||
| ) | ||
| window_coords = post_index[(post_index >= start) & (post_index <= stop)] |
There was a problem hiding this comment.
This branch drops slice.step, while the integer branch at :773 keeps it (post_index[mask][::step]). On a 6-month DatetimeIndex:
slice(None, None, 2) before -> 3 timestamps
now -> all 6 <- silent wrong answer
slice(0, 3) before -> first 3
now -> ValueError: Window contains no time points
Base used result.datapost.index[window], which was positional and handled both.
Reachable from effect_summary(window=slice(...)) on ITS, SC and PiecewiseITS. No test covers a datetime slice — test_reporting.py:238 and :2311 both use integer indexes.
Fix: apply [::step] here too, and either keep positional behaviour for integer bounds on a DatetimeIndex or reject them clearly instead of turning 0 into 1970-01-01.
| # Point counterfactual predictions per observation; treatment | ||
| # effects are observed outcome minus prediction. | ||
| y_hat0 = y_pred.mean(dim=["chain", "draw"]).isel(treated_units=0).values | ||
| treated_positions = self.data.index.get_indexer(treated_data.index) |
There was a problem hiding this comment.
get_indexer needs a unique index. Base did self.data["y_hat0"] = ....values then a boolean .loc, which was fine with duplicates.
OLS, unique RangeIndex -> OK, att(e=0)=2.0334
OLS, duplicated index -> InvalidIndexError: Reindexing only valid with uniquely valued Index objects
Long panels built with pd.concat and no ignore_index=True hit this straight away, and to_pandas() doesn't reset the index.
Same pattern at :714 and :1329.
Fix: go back to positional masks, e.g. np.flatnonzero(~self.data["_is_untreated"].to_numpy()), and drop the label .loc lookups on _observed_outcome.
For context, the Bayesian path at :588 was already index-fragile, so the class always wanted a positional index — but this removes the one branch that coped.
| assert re.search(r"^Other Parameters\n-+$", doc, flags=re.MULTILINE) | ||
| assert re.search( | ||
| rf"^\*\*{re.escape(expected_parameter)}\s*$", doc, flags=re.MULTILINE | ||
| # Every forwarder documents its kwargs contract under a uniform |
There was a problem hiding this comment.
This swaps two structural checks for one substring check. The numpydoc header with its underline, and the actual **kwargs entry, are no longer verified — a docstring that just mentions the words in prose now passes.
I ran the old strict assertions against all 9 exemptions on this branch: 0 failures. So nothing needed loosening, and the bar dropped right as 7 new entries were added to this same test.
Fix: put the two regexes back.
| np.testing.assert_allclose(bundle.impact_post.to_numpy()[..., 0], expected_post) | ||
|
|
||
|
|
||
| class TestInputValidation: |
There was a problem hiding this comment.
test_datetime_index_requires_timestamp_treatment_time and test_int_index_rejects_timestamp_treatment_time were dropped here. I ran both against this branch and they still pass unchanged, and SyntheticDifferenceInDifferences.input_validation is still there (synthetic_difference_in_differences.py:180).
The class docstring still says "Both BadIndexException branches in input_validation" but the class now has neither.
Fix: restore both, or drop the docstring claim. The shared helper is covered at test_pandas_compat.py:356, so what's actually lost is the pin that SDID wires to it.
| #: a follow-up (issue #1092). | ||
| supports_prior_predictive = False | ||
|
|
||
| def __init__( |
There was a problem hiding this comment.
This __init__ never got prior_sample_kwargs, so the inherited _clone() passes an argument it can't take and raises TypeError.
PyMCModel._clone (:308), BBETS (:2295) and StateSpaceTimeSeries (:2768) all carry it — this one is the gap.
That breaks acceptance criterion 4 for IV, and it matters in practice because checks/base.py:clone_model is what prior_sensitivity.py uses, so any check cloning an IV model fails.
Fix: add prior_sample_kwargs to the signature and pass it through to super().__init__.
| min_effect=min_effect, | ||
| group=group, | ||
| ) | ||
| else: |
There was a problem hiding this comment.
On the has_posterior_draws() false branch this passes self instead of the resolved bundle, and the helper then reads self.result — so group is ignored.
sample_prior_predictive(draws=1) is a documented per-call override and it flips that condition on a real Bayesian model:
fitted + prior draws=1 -> ValueError: cannot reshape array of size 200 into shape ()
unfitted + prior draws=1 -> GroupNotSampledException: ...Call RegressionDiscontinuity.fit() first
The second is the worse one: you asked for the prior and it tells you to fit.
The group-aware version already exists — reporting.py:622-658 has a full fallback behind an experiment=None sentinel that nothing in the library ever passes.
Fix: route through it, e.g. _effect_summary_rd(bundle, ..., experiment=self, group=group). Same shape at diff_in_diff.py:793 and prepostnegd.py:571.
| figsize : tuple of (float, float), optional | ||
| Width and height of the figure in inches. Defaults to ``(7, 8)``. | ||
| """ | ||
| bundle = self._require_bundle(group) |
There was a problem hiding this comment.
The group == "prior" early return happens before the treated_unit resolve/validate block at :629-635, so the prior plot always draws unit 0 and never validates the name:
plot(group="prior", treated_unit="actual2") -> draws "actual" (range -0.08..29.81, not 29.97..41.93)
plot(group="prior", treated_unit="NOT_A_UNIT") -> no error
get_plot_data(group="prior") does honour treated_unit (:921), so the two prior readers disagree.
Fix: move the resolve/validate block above the early return and pass the unit into _plot_prior_checks, replacing the hardcoded isel(treated_units=0) at :833.
| post_y = self.post_design["y"].isel(treated_units=0) | ||
|
|
||
| fig, ax = plt.subplots(1, 1, figsize=(7, 4)) | ||
| style: _PosteriorPlotStyle = { |
There was a problem hiding this comment.
_plot_prior_checks hardcodes its style dict and figsize and takes no kwargs, so everything plot() accepted gets dropped:
plot(group="prior", ci_prob=0.10, kind="spaghetti", num_samples=3, figsize=(3,3))
-> same ribbon vertices as the bare call, figsize (7.0, 4.0)
ci_prob is the one that misleads — you ask for a 10% band and quietly get a 94% HDI.
Same in 8 other files: synthetic_control.py:839, piecewise_its.py:716, diff_in_diff.py:685, prepostnegd.py:475, regression_discontinuity.py:625, regression_kink.py:492, synthetic_difference_in_differences.py:886. #1092 asked for a reduced panel set, not for the knobs to stop working.
Fix: forward ci_prob/kind/ci_kind/num_samples/figsize, or raise on the ones the reduced layout can't honour.
| ) | ||
|
|
||
|
|
||
| def _apply_prior_grouping(text: str, group: Literal["prior", "posterior"]) -> str: |
There was a problem hiding this comment.
This only adds a prefix — the sentence inside still says "posterior". Real output for group="prior" on a DiD:
Prior predictive check (not a causal estimate): The average treatment effect was 16.20
(95% HDI [-88.60, 84.70]). The posterior probability of an increase is 0.520.
That last number is the sign-neutrality check #1092 calls the most valuable prior diagnostic, and it's labelled "posterior".
The string is hardcoded at :284 and :290 in _render_bayesian_decision, which never sees group.
Fix: pass group into _render_bayesian_decision and say "prior probability" for the prior group.
| dim="obs_ind" | ||
| ) | ||
|
|
||
| def _comparison_period_summary( |
There was a problem hiding this comment.
_comparison_period_summary has no group parameter and never calls _apply_prior_grouping, unlike every other summary path. On a three-period ITS after sample_prior_predictive() only:
period="intervention" -> Prior predictive check (not a causal estimate): During the ...
period="post" -> Prior predictive check (not a causal estimate): During the ...
period="comparison" -> Effect persistence: The post-intervention effect (-429.4, 95% HDI [-9435.7, 9209.5])
The third reads like a real persistence finding, built from prior draws, with a ±9400 interval.
Fix: add group to the signature, pass it from the call site at :1322, and wrap both prose branches (:518 and :525) in _apply_prior_grouping.
| "first.", | ||
| group="prior", | ||
| ) | ||
| return None |
There was a problem hiding this comment.
The PriorPredictiveNotSupportedException pre-check only lives in the not self._supports_results branch above (:471-489), so the nine bundle-backed experiments send sklearn users to a dead end:
RD + sklearn: plot(group="prior") -> GroupNotSampledException: ...Call RegressionDiscontinuity.sample_prior_predictive() first
sample_prior_predictive() -> PriorPredictiveNotSupportedException
PanelRegression (bundle-less, same backend) correctly raises the capability error straight away.
That undoes the reason #1092 gives for having two exception types: "you skipped a step" vs "this backend can never do this".
Fix: hoist the capability check above the _supports_results split. The prior_result property at :295 needs the same pre-check.
| for this call only. | ||
|
|
||
| """ | ||
| self.build() |
There was a problem hiding this comment.
If the auto prior phase at :419-423 raises, fit() propagates but the object is already fully fitted. Patching the prior phase to fail:
fit() raised: ValueError -> prior spec is broken
sample_posterior calls = 1 <- MCMC already paid for
exp.is_fitted -> True
exp.result -> CausalResult
exp.plot() -> works
So a caller's except sees failure while the object says success, and you only find out the prior is broken after the expensive part. Putting the prior first was #1092's way of avoiding the second half of that.
Fix: either make a failure in the auto prior fill warn instead of raising out of fit(), or roll _result back before re-raising.
| identity is the model instance, so stale draws would be incoherent. | ||
| This assignment is the documented prior-revision mechanism. | ||
| """ | ||
| self._model_backend = make_model_adapter( |
There was a problem hiding this comment.
This clears _result and _prior_result but not idata — it assumes the incoming model is fresh. With a used instance:
exp.model = <already-fitted model>
exp.is_fitted -> False
exp.idata -> still has all 7 groups
So is_fitted says no and .idata says yes. The documented path is fine — I checked, assigning a brand new model does set idata to None — this is only the reuse case.
Fix: null the backend's idata in the setter instead of trusting the argument.
| Builds the graph (idempotent), samples NUTS plus posterior predictive | ||
| draws, then — when the backend supports a prior phase and no prior | ||
| state exists yet — fills the prior groups and :attr:`prior_result` so | ||
| ``idata`` is as complete as the historical eager fit produced. The |
There was a problem hiding this comment.
No argument with the order itself — but this reason isn't the real one, and as written it will stop someone touching it later for the wrong cause.
_rearm_fit_data() is the first thing both sample_prior_predictive (pymc_models.py:506) and sample_posterior (:545) do, so node mutation can't reach either sampler and order can't affect correctness. Checked both ways on the shipped tree:
prior-first vs posterior-first, seeded -> max |diff| = 0.0 (identical)
The other reason you gave in the review thread — RNG-stream parity with the eager baseline — is the one that actually holds:
unseeded -> max |diff| = 14.44
Fix: say it keeps the RNG stream identical to the eager baseline so the unchanged integration suite stays a valid parity test, and state the real invariant (_rearm_fit_data() runs before every phase). Same wording in ARCHITECTURE.md and the #1092 comment.
| """Run the experiment algorithm: fit OLS, 2SLS, and Bayesian IV model.""" | ||
| self.get_naive_OLS_fit() | ||
| self.get_2SLS_fit() | ||
| def build(self) -> Self: |
There was a problem hiding this comment.
The reason in this docstring isn't true any more. This PR added a non-sampling InstrumentalVariableRegression.build() at pymc_models.py:1747, and :275-288 here already calls model.build(...) and model.sample_posterior(...) as two separate steps — so the graph isn't fused into fit.
Calling it by hand works:
iv.model.build(X=..., Z=..., y=..., t=..., coords=..., priors=..., ppc_sampler=None)
-> basic_RVs ['beta_t','beta_z','chol_cov','likelihood'], idata None, graphviz populated, is_built True
So the AC14 waiver looks like a leftover rather than a real limit, and right now exp.build() quietly gives an empty graph.
Fix: delegate to self.model.build(...) with ppc_sampler=None, then drop the waiver from the docstring and ARCHITECTURE.md.
| group: Literal["prior", "posterior"] = "posterior", | ||
| round_to: int | None = None, | ||
| ci_prob: float = HDI_PROB, | ||
| kind: Literal["ribbon", "histogram", "spaghetti"] = "ribbon", |
There was a problem hiding this comment.
This advice is copied from the coefficient experiments and doesn't hold here. SDID's summary runs on impact_post = observed minus prior synthetic, whose level comes from the omega0 prior and the control panel, not a symmetric contrast. On a correct model with default priors:
SDID: p_gt_0 = 0.383, P(tau>0) = 0.263
DiD 0.548 | PrePostNEGD 0.537 | RD 0.515 | RK 0.525
So users will read a perfectly normal SDID prior as a design-matrix bug.
Fix: keep the sign-neutrality wording on the experiments where it works, and say something here about whether the prior synthetic trajectory looks plausible instead.
| # path: store canonical in-sample predictions at fit time. | ||
| if self._model_backend.is_bayesian: | ||
| # PanelRegression is posterior-only; no group= phase. | ||
| mu = self._model_backend.require_idata().posterior["mu"] |
There was a problem hiding this comment.
This is the exact line #1092 names in its .posterior sweep, and it's unchanged from base (:656 there). get_plot_data still has no group=, same at plot_trajectories :898. The other four all took it: synthetic_control.py:879, interrupted_time_series.py:983, piecewise_its.py:758, staggered_did.py:1478.
The comment "PanelRegression is posterior-only; no group= phase" also argues with the plot() docstring ~40 lines up, which documents a working prior phase — and 57531220 made that prior plot real.
mu is present in the prior group for LinearRegression, so it's a small change.
Fix: add keyword-only group, read require_idata()[group], guard via _resolve_group(group). If it's staying posterior-only on purpose, worth saying so in the PR body rather than in a comment that contradicts its neighbour.
| | **HDI_PROB** | Project default is 0.94 (ArviZ default), not 0.95. | | ||
| | **create_causalpy_compatible_class** | Applied during `make_model_adapter()` for sklearn backends; clones the user instance before patching. | | ||
|
|
||
| ## Adding New Code |
There was a problem hiding this comment.
The HDI_PROB row got dropped from this table in the rewrite. Base had:
| **HDI_PROB** | Project default is 0.94 (ArviZ default), not 0.95. |
That's still true, still surprising (most people assume 0.95), and this was the only place it was written down. Looks like collateral from removing the "Eager fitting" row right next to it.
Fix: put the row back.
| ) | ||
|
|
||
| def effect_summary(self) -> NoReturn: | ||
| def effect_summary( |
There was a problem hiding this comment.
effect_summary took group= here "for base-contract parity", but plot() at :361 didn't:
iv.plot(group="prior") -> TypeError: InstrumentalVariable.plot() got an unexpected keyword argument 'group'
Same in inverse_propensity_weighting.py:545. #1092 says every experiment's read methods gain a keyword-only group, so anyone looping over all twelve gets a TypeError from two of them.
Fix: add group= to both plot() stubs, for the same parity reason already accepted here.
| return self._model_backend.is_built or self.is_fitted | ||
|
|
||
| @property | ||
| def is_fitted(self) -> bool: |
There was a problem hiding this comment.
For the three bundle-less experiments this reads idata off the model instance, not off the experiment, so fitted state travels with a shared model. A second experiment constructed with an already-fitted model never fits anything and still answers:
e2.is_fitted -> True
e1 y_fitted[:3] -> [-0.3208 -0.0317 0.3118]
e2 y_fitted[:3] -> [-0.3208 -0.0317 0.3118] identical
e2 residual magnitude -> 100.0 (d2 is d1 with y + 100)
get_plot_data() pairs the second dataset's y_actual with the first experiment's y_fitted, and summary() prints the first experiment's coefficients. No warning, no exception. At a2f1302 PanelRegression.__init__ refit the shared model, so the numbers were right; the lazy lifecycle removed the refit and left the predicate keyed on the shared object.
This is next to the model-setter case in the other thread but reaches through construction, so nulling idata in the setter would not close it.
Would a per-experiment flag set in BaseExperiment.fit() work here, with has_prior_predictive given the same treatment? The alternative is rejecting a model that already carries draws at construction, the way checks/ do with clone_model.
| Reserved for future non-centred parameterisations of the | ||
| coefficient prior. Currently informational only. | ||
| """ | ||
| if self._built: |
There was a problem hiding this comment.
This is the only build() in the file without the input-fingerprint check: bare if self._built: return, and _built_input_fingerprint is never assigned. PyMCModel.build, InstrumentalVariableRegression.build, StateSpaceTimeSeries.build and both record-only adapters all compare and raise.
BaseExperiment.build's docstring promises that RuntimeError, so the guarantee is a no-op for exactly one experiment. It bites harder here than elsewhere because the IPW estimand is computed outside the graph, from self.X, so the propensity draws come from one dataset and the weighting from the other:
e1 = InversePropensityWeighting(df1, ..., model=m).fit()
e2 = InversePropensityWeighting(df2, ..., model=m).fit() # df2 is a 783-row subset
e2 rows -> 783
e2 posterior -> p_dim_0: 1566 e1's design
same pattern on DiD -> RuntimeError: This model is already built with different inputs
at a2f13029 -> ValueError: Variable name X already exists
So a reused instance went from failing loudly to returning an ATE for the wrong dataset.
Would mirroring the sibling overrides be enough: compute the fingerprint before the early return, raise on mismatch, assign it next to self._built = True? test_propensity_score_build_is_idempotent covers only the same-inputs no-op, so a changed-input case would pin it.
|
Last pass from my side, with Opus 5: From-scratch pass over the whole diff at 7e1a63c. The inline comments already on it cover the RD prior dispatch, the dropped The
Release note names the wrong experiment ( Shipped text still teaches the eager API. Seven files under Prior-group draw values are asserted nowhere, only their presence, which is the other half of the mocked-sampler point. Three mutations leave all 2394 tests passing here: hardcoding Smaller things:
Checked and clean: no One aside, not this PR: |
3 participants
Lazy experiment lifecycle: configure → optional prior checks → fit()
Fixes #1092
Summary
Implements the lazy experiment lifecycle PRD: experiment construction no longer samples.
__init__validates and builds design matrices only; posterior inference happens through an explicitfit()(returningSelf), and an optionalsample_prior_predictive()phase enables prior predictive checks before MCMC. Draw-derived state moves from ~15 flat per-experiment attributes into per-group result bundles exposed through raisingresult/prior_resultproperties. Read methods (plot,get_plot_data,effect_summary) take keyword-onlygroup=and fail fast with actionableGroupNotSampleedExceptions. Two new exceptions are exported fromcausalpy.Because
fit()returnsSelf, every existing call site migrates with one appended token:cp.InterruptedTimeSeries(...).fit().This is a breaking API change landing on the
pymc6_and_pymcmarketing1_migrationline (1.0), per the issue's sequencing decision. The bundling prerequisite (#1093 / step 3) is carried by this same PR because it had not landed when this work started.Changes
Core lifecycle
causalpy/experiments/base.py:build()(public, idempotent, auto-called),fit(**kwargs),sample_prior_predictive(**kwargs)— all returnSelf. State predicatesis_configured/is_built/is_fitted/has_prior_predictivedefined over the two private bundle slots.result/prior_resultraise instead of returningNone.modelassignment resets all state (documented prior-revision path; noset_priors()). Secondfit()warns and overwrites posterior only._render_plotresolvesgroup=once for all subclasses.causalpy/experiments/_results.py(new):CausalResult,SyntheticDifferenceInDifferencesResult,CoefficientResult,GroupComparisonScenario,DiscontinuityResult,KinkResult,StaggeredDifferenceInDifferencesResult— fully-populated bundles,xr.DataArray-typed, one per experiment family.causalpy/experiments/model_adapter.py:build(X, y, coords)/sample_prior_predictive(**kwargs)/sample_posterior(**kwargs)/predict(..., group=)/coefficients(*, group=)/ capability properties (supports_prior_predictive,is_built,has_posterior,has_prior) across PyMC, sklearn, and pymc-forecast adapters.causalpy/pymc_models.py: sampling split —build()(prior merge at build time; idempotent),sample_prior_predictive(),sample_posterior()(overwrite semantics via DataTree group assignment;xr.DataTree.extenddoes not exist on xarray 2026.7).prior_sample_kwargs(default draws=500) added and carried by_clone().predict(group=)conditions forward sampling onidata.prior(verified on PyMC 6.0.1: forward samples land inposterior_predictiveregardless of conditioning group).IVandStateSpaceTimeSeriesdeclaresupports_prior_predictive=False;PropensityScoreandBayesianBasisExpansionTimeSeriessplit cleanly; SDID weight fitter gainsbuild_mapping.causalpy/custom_exceptions.py:GroupNotSampleedException(carriesgroup, message names the missing call) andPriorPredictiveNotSupportedException— following the repo's...Exception(Exception)convention; sklearn'sNotFittedErrorrejected per issue (itsAttributeErrorMRO would be swallowed bygetattrcall sites).causalpy/__init__.py: all five exceptions re-exported.All 12 experiments migrated
algorithm()removed everywhere (no alias). Each experiment implements_fit_inputs()+_finalize(group); nothing draw-derived is assigned toself.plot/get_plot_data/effect_summarytake keyword-onlygroupfirst;plot(group="prior")renders the reduced counterfactual-vs-observed panel set; prioreffect_summaryprose reads as a plausibility check. IV / IPW / PanelRegression store no bundle (_supports_results=False, fitted-state keys off the backend) per the issue's partition. IV's OLS/2SLS pre-step always runs in__init__(deterministic point estimates that also feedsummary(); never samples); only the default-prior derivation is guarded bypriors is None.Ecosystem
EstimateEffect.run()constructs then calls.fit()..fit()on their fresh instances.reporting.pyhelpers take containers/bundles;maketables_adaptersresolve score/hdi_probfrom bundles;placebo_in_time/pre_treatment_placeboread bundle fields.generate_report()no longer swallows guard exceptions (narrowed toNotImplementedError).Docs & tests
ARCHITECTURE.mdlifecycle section + conventions rewritten;running-causalpy-experimentsskill updated.causalpy/tests/test_experiment_lifecycle.py: guards and messages,Selfreturns,build()idempotency + graphviz inspection, model-assignment reset, second-fit overwrite semantics, sklearn capability error, canonical prior containers, independent prior/posterior draw sizes, reduced prior plot layout, and property-based invariants (posterior HDI narrower than prior; prior tail probability near 0.5 under a neutral prior). No seed-pinned baselines, per the issue's testing strategy..fit(); expected values unchanged (the parity argument from the issue).Testing
$CONDA_EXE run -n CausalPy python -m pytest causalpy/tests/test_experiment_lifecycle.py$CONDA_EXE run -n CausalPy python -m pytest— 2357 passed at final headprek run --all-filesbuild()+pm.model_to_graphviz→sample_prior_predictive()→plot(group="prior")/effect_summary(group="prior")→fit()→plot()→ refit warning, all verified on synthetic ITS data.Checklist
__init__for any experiment (incl. PiecewiseITS)EstimateEffectand all checks call.fit()prior_sample_kwargswith per-call overrides; carried by_clone()plot(group="prior")reduced panels for prior-capable backendsfit()/sample_prior_predictive()build()/fit()/sample_prior_predictive()returnSelf;build()idempotent + auto-calledis_fitted≡_result is not None(bundle-backed experiments); raisingresult/prior_resultselfin any experimentpm.model_to_graphviz(exp.model)populated afterbuild()with no draws inidataARCHITECTURE.md+ skill updated