feat: variable selection priors for SSTS covariates - #1112
cetagostini merged 8 commits into
Conversation
853ffdf to
c58c9e3
Compare
e98a03b to
374ece7
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## pymc6_and_pymcmarketing1_migration #1112 +/- ##
====================================================================
Coverage 97.24% 97.25%
====================================================================
Files 129 129
Lines 23553 23670 +117
Branches 1318 1324 +6
====================================================================
+ Hits 22904 23020 +116
- Misses 440 443 +3
+ Partials 209 207 -2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
374ece7 to
f4eb592
Compare
c58c9e3 to
d158d16
Compare
f4eb592 to
0f3e88d
Compare
d158d16 to
b9da23f
Compare
0f3e88d to
fcee78a
Compare
b9da23f to
031c83d
Compare
fcee78a to
af7713f
Compare
e13d969 to
e855435
Compare
Automated triageRecommendation: Why:
Review focus:
Confidence: high |
drbenvincent
left a comment
There was a problem hiding this comment.
Agent-drafted review.
PR #1112 Review
Scope: declared base issue-758-ssts-covariates (merged #1111, b9761e78) .. head issue-758-ssts-variable-selection (29d9a109). Diff is +398/−2 across 3 files; draft descendant #1113 excluded. Isolated worktree /Users/benjamv/git/.worktrees/pr-1112. Primary tree left on issue-1092-lazy-experiment-lifecycle with its pre-existing dirty .agents/skills/review-pr/SKILL.md.
Executive Summary
Recommendation: approve.
Value delivered:
StateSpaceTimeSeriesgains a public, constructor-level variable-selection API (vs_prior_type/vs_hyperparams) that reusesVariableSelectionPriorunchanged and is reachable fromInterruptedTimeSerieswithout adapter changes — the intentional departure from IV’sbuild_modelkwargs is justified by the ITS fit path.- Structure tests pin API shape, dims/coords, precedence warning, build-time covariate guard, clone round-trip, and inclusion/shrinkage helpers; a correctness-marked real-NUTS gate (
test_its_state_space_variable_selection_recovery) passed in remote CI on all three Python/pandas matrix jobs; thebeta_exogattenuation caveat is documented at the helper and class surfaces and the recovery test correctly asserts ranking rather than point estimates.
Review focus:
- Highest-risk assumption was that selection would be wired with wrong dims/precedence, or that the correctness gate would be too thin or flaky. Probes falsified the wiring risks; remote correctness (4 passed including this recovery test on 3.12/3.14 × pandas 2.3/3.x) supports the ranking claim. Residual merge friction is the failing
gallery-othernotebook job, which is not caused by this PR’s diff.
Findings
No must-fix items.
-
[should-fix]
docs/requirements.txt(inherited from base, unchanged in this PR): remotenotebooks (gallery-other, …)fails withImportError: cannot import name 'extract_dims' from 'pymc.model.fgraph'inside pinnedpymc-marketing@d710dc0, breaking BBETS-using gallery notebooks (interrupted-time-series-lift-test,interrupted-time-series-placebo-in-time-analysis, and the IV VS notebook). This is the same dependency-drift failure already fixed onpymc6_and_pymcmarketing1_migrationby41b6416b(re-pin tof8b552c+pymc-extras>=0.14,<0.15). Parent #1111 had green gallery-other; this head is 0 behind its declared base but that base is behind the migration tip that carries the fix. Green the notebook lane by bringing the stack’s docs pin forward (rebase/sync or a tiny pin bump) before merge; do not treat it as a defect in the VS wiring. -
[should-fix]
causalpy/pymc_models.py(get_inclusion_probabilities/get_shrinkage_factors) and the shared factory: returned DataFrames use a bareRangeIndex([0, 1, 2, …]) even though posterior coords carrystate_exognames. Local probe confirmed labels are absent; correctness and structure tests rely on formula/ilocorder. For a public selection API aimed at P6 inclusion plots, attach_exog_names(or posterior coord labels) as the index in the SSTS wrappers — or document positional order as the contract. Prefer fixing at the thin SSTS pass-through so IV’s factory behavior stays untouched if desired. -
[nit] No structure test for
vs_prior_type="normal". Local real-sampling probe confirmedbeta_exogwithstate_exogcoords and nogamma_*. Worth a one-liner next to the horseshoe structure test. Note also that factorynormaldefaults tosigma=1, which is not the SSTS defaultbeta_exogprior (sigma=50); that is fine for a baseline comparator but easy to misread. -
[nit] Codecov patch is ~96.8% with a few uncovered lines;
codecov/patchstill passed. Not a merge blocker.
Merge Readiness
Verdict: approve.
Branch status: 6 commits ahead of issue-758-ssts-covariates, 0 behind that declared base; GitHub MERGEABLE / UNSTABLE. Stack is behind pymc6_and_pymcmarketing1_migration (docs-pin fix and other tip commits).
CI status: core matrix green (prek, three Test jobs including correctness lane, codecov/patch, Read the Docs, gallery-pymc, gallery-sklearn, knowledgebase, zizmor). Required notebook shard gallery-other failing for the inherited docs pin above — not PR-caused, but must be greened before a clean merge.
PR Summary
This PR closes #981 (P0 of #758) by wiring spike-and-slab, horseshoe, and normal variable-selection priors onto SSTS exogenous coefficients via constructor kwargs, with build-time refusal when VS is configured without covariates, VS-over-priors["beta_exog"] precedence with a warning, _clone preservation, and thin inclusion/shrinkage accessors. The important trade-off is documented and accepted: beta_exog point estimates attenuate because state-space P0 lets regression states drift, so selection is ranked by inclusion probability and counterfactuals ride on smoothed states; pinning exog P0 was rejected for NUTS runtime. Full notebook treatment is deferred to #1113.
Test Evidence
Remote (gh pr checks 1112):
- Pass: Test (3.12/pandas 2.3, 3.12/pandas 3.x, 3.14/pandas 3.x), prek, check-environment-yml, codecov/patch, docs/readthedocs, gallery-pymc, gallery-sklearn, knowledgebase, zizmor, sync.
- Fail: notebooks (gallery-other) —
extract_dimsImportError from stalepymc-marketingpin (see Findings). - Correctness lane on all three Test jobs:
pytest -m correctnessreported4 passed, includingtest_its_state_space_variable_selection_recovery(~4–6.5 min for the correctness step).
Local (isolated worktree /Users/benjamv/git/.worktrees/pr-1112, mamba run -n CausalPy after editable install into the shared env):
pytest …TestStateSpaceTimeSeriesCoverage -k vs_+test_its_with_state_space_variable_selection --no-cov: 7 + 1 passed.- Probes:
normalbuilds withstate_exogcoords; spike-and-slab inclusion table shape OK (index positional); horseshoe shrinkage columns OK; VS precedence warning +gamma_beta_exogpresent;predictunder VS returns a DataTree without error. doctest.testmod(causalpy.pymc_models): 0 failures / 88 attempted (includes the new SSTS selection doctest with tiny real NUTS).
Not re-run locally: full ~3–6 min correctness recovery (already green on remote matrix ×3).
Open Questions
None that block approval of the VS feature itself. Process only: whether to green gallery-other by syncing this stack onto the migration docs-pin fix versus a one-file pin bump on this PR — either is fine; the feature diff does not need to change.
Probe notes (review brief)
Private hypotheses and outcomes:
- Constructor/
vs_prior_typecontract wrong or invalid types accepted late — falsified;Literal+ eagerVariableSelectionPriorvalidation;test_vs_prior_invalid_type/ without-covariates / precedence warning all pass. - Prior factory not used or
priors["beta_exog"]wins silently — falsified; special-case in the param loop; warning + gamma present under dual config. - Exog dims/coords or slicing break selection or forecasts — falsified;
beta_exogcoords["x1","x2","x3"];_exog_valuesshared by build/pm.Data/predict; ITS structure test keeps finitepost_impact; localpredictunder VS succeeded. - Inclusion/shrinkage helpers mis-gated or wrong prior-type behavior — falsified; shared
_require_vs_diagnostics; horseshoe cannot call inclusion; spike-and-slab cannot call shrinkage (structure tests). normal/ horseshoe / spike-and-slab paths broken for SSTS — falsified by structure tests (HS/SS) and local real sampling (all three).- Correctness gate too weak, order-dependent, or flaky — weakened but not overturned; ranking+gap assertions with documented headroom; passed remote correctness on three matrix jobs; absolute inclusion levels correctly not asserted given attenuation.
- Documented attenuation caveat overstates forecast safety — not overturned; counterfactual path still uses smoothed states/
_prepare_idataunchanged; no evidence VS breaks ITS impact finiteness. Residual: no dedicated real-NUTS forecast-recovery assert in this PR (acceptable given documented scope). _clonedrops VS config (MF-1) — falsified;test_vs_clone_preserves_config.- Notebook CI failure is PR-caused — falsified; identical failure mode and notebooks as migration fix
41b6416b; this PR does not touch docs pins or those notebooks. - Callers/adapters need changes — falsified; ITS consumes constructor-configured model via existing adapter
fit(X,y,coords)only.
Draft GitHub review (not posted): approve on merits; ask for stack docs-pin sync before merge; optional follow-ups for named inclusion index and a normal structure test.
REVIEWED | /Users/benjamv/git/CausalPy/.scratch/pr-reviews/2026-09-08-author-anevolbap/pr-1112.md
29d9a10 to
bdff957
Compare
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
…able Rebuilding a DatetimeIndex from raw values drops its freq, so every fit emitted a pymc-extras 'No frequency was specific on the data's DateTimeIndex' warning and the forecast index had to re-infer it. Recover it when the observations are regularly spaced, and leave it unset when they are not. combined.build() also printed a rich 'Model Requirements' table to stdout on every construction, telling the reader to declare priors that this class declares itself a few lines later. The docstring example in #1112 had to wrap the call in redirect_stdout to hide it.
Re-run against the current stack, so the stored outputs no longer carry the pymc-extras requirements table or the frequency warning. That also drops three absolute site-packages paths from the stderr output. Three prose fixes. The plot title reports a Bayesian R2 of 1 with a standard deviation around 1e-9, because the Kalman smoother conditions on the observed outcome and interpolates the pre-period; the text now says so instead of leaving the reader with an apparently perfect fit. The two divergences in the covariates fit are acknowledged. The rhat paragraph told the reader to fall back on the inclusion probabilities, which are averaged over the very chains that disagree, so it now calls the ranking provisional for that reason. The hand-written index on the inclusion chart is gone: the labels come from the model since #1112.
…able Rebuilding a DatetimeIndex from raw values drops its freq, so every fit emitted a pymc-extras 'No frequency was specific on the data's DateTimeIndex' warning and the forecast index had to re-infer it. Recover it when the observations are regularly spaced, and leave it unset when they are not. combined.build() also printed a rich 'Model Requirements' table to stdout on every construction, telling the reader to declare priors that this class declares itself a few lines later. The docstring example in #1112 had to wrap the call in redirect_stdout to hide it.
Re-run against the current stack, so the stored outputs no longer carry the pymc-extras requirements table or the frequency warning. That also drops three absolute site-packages paths from the stderr output. Three prose fixes. The plot title reports a Bayesian R2 of 1 with a standard deviation around 1e-9, because the Kalman smoother conditions on the observed outcome and interpolates the pre-period; the text now says so instead of leaving the reader with an apparently perfect fit. The two divergences in the covariates fit are acknowledged. The rhat paragraph told the reader to fall back on the inclusion probabilities, which are averaged over the very chains that disagree, so it now calls the ranking provisional for that reason. The hand-written index on the inclusion chart is gone: the labels come from the model since #1112.
|
|
@anevolbap if you solve the issues (conflicts) then we are good for merge. |
Wire VariableSelectionPrior (spike-and-slab, horseshoe, normal) to the StateSpaceTimeSeries regression coefficients via vs_prior_type and vs_hyperparams, mirroring InstrumentalVariableRegression. The generic prior loop special-cases beta_exog when configured. Expose get_inclusion_probabilities() and get_shrinkage_factors() pass-throughs (groundwork for the inclusion plot in #758 P6). Prototype-verified: sampling is clean and inclusion probabilities rank true predictors first. beta_exog point estimates attenuate toward zero because P0 lets the regression states drift from the parameter; pinning the exog P0 entries fixes that on paper but makes NUTS pathologically slow, so the caveat is documented instead. Counterfactual forecasts use the smoothed states and are unaffected.
…stacklevel get_inclusion_probabilities() and get_shrinkage_factors() returned a bare RangeIndex, so the caller had to know the fit-time column order to read them. That is the wrong contract for a feature whose output is a claim about which covariates matter: the correctness test sliced positionally and the notebook in #1113 relabelled the index by hand. Rows are now indexed by the regressor names the model already tracks in _exog_names. The beta_exog precedence warning goes to stacklevel 3. pm.Model's metaclass calls __init__, so at 2 the warning was attributed to pymc/model/core.py instead of the line that passed both arguments. The vs_hyperparams docstring said the horseshoe scales its global shrinkage from the data. Only the sample size comes from the data; the residual scale in the Piironen and Vehtari rule is held at 1.
vs_prior_type="normal" had no structure test. It builds a plain Normal on beta_exog with state_exog coords, neither the spike-and-slab indicators nor the horseshoe scales, so both diagnostics must refuse. The vs_hyperparams docstring now says the factory's normal default is Normal(0, 1), much tighter than the Normal(0, 50) this class uses for beta_exog without a selection prior.
a897cec to
cce4b29
Compare
|
Rebased onto |
…able Rebuilding a DatetimeIndex from raw values drops its freq, so every fit emitted a pymc-extras 'No frequency was specific on the data's DateTimeIndex' warning and the forecast index had to re-infer it. Recover it when the observations are regularly spaced, and leave it unset when they are not. combined.build() also printed a rich 'Model Requirements' table to stdout on every construction, telling the reader to declare priors that this class declares itself a few lines later. The docstring example in #1112 had to wrap the call in redirect_stdout to hide it.
Re-run against the current stack, so the stored outputs no longer carry the pymc-extras requirements table or the frequency warning. That also drops three absolute site-packages paths from the stderr output. Three prose fixes. The plot title reports a Bayesian R2 of 1 with a standard deviation around 1e-9, because the Kalman smoother conditions on the observed outcome and interpolates the pre-period; the text now says so instead of leaving the reader with an apparently perfect fit. The two divergences in the covariates fit are acknowledged. The rhat paragraph told the reader to fall back on the inclusion probabilities, which are averaged over the very chains that disagree, so it now calls the ranking provisional for that reason. The hand-written index on the inclusion chart is gone: the labels come from the model since #1112.
143e823
into
pymc6_and_pymcmarketing1_migration
Closes #981. Part of #758 (P0). Targets
pymc6_and_pymcmarketing1_migration.Summary
Adds variable selection priors for the covariates of
StateSpaceTimeSeries, reusingcausalpy/variable_selection_priors.pyunchanged. This is the same factory the IV class uses.API
vs_prior_type:"spike_and_slab","horseshoe"or"normal", the same set the factory accepts. When set, the factory builds thebeta_exogprior instead of the default.vs_hyperparams: optional dict passed to the factory.get_inclusion_probabilities()andget_shrinkage_factors()pass through to the factory and index the rows by regressor name, since the factory's tables are positional. This prepares the P6 plotting item of Feature parity with Google's CausalImpact #758.Behavior
"normal"is a plainNormal(0, 1)with no selection. It is much tighter than the class defaultNormal(0, 50)forbeta_exog, and the docstring says so.priors={"beta_exog": ...}andvs_prior_typelets variable selection win with a warning, same as the IV class. The warning points at the calling line.vs_prior_typewithout covariates raises at build time. The check cannot run earlier, because covariates come from the experiment formula and only reach the model when the experiment builds it.Design
Unlike the IV class, the two arguments are constructor kwargs rather than
build_modelarguments.InterruptedTimeSeriesonly feedsX, y, coordsthrough the shared model adapter, so constructor configuration makes selection usable from the experiment with no adapter changes. It also validates the prior type before any data work starts.Known caveat
beta_exogpoint estimates shrink toward zero, becauseP0(the initial state covariance, unrelated to the P0 priority label of #758) lets the regression states drift away from the parameter. Pinning the exogP0entries fixes the statistics on paper but makes NUTS far slower (15+ minutes against about 2 in an earlier run), so it is not done here. Counterfactual forecasts use the smoothed states, which do recover the truth, and selection uses the inclusion-probability ranking, which is correct. The class docstring documents this.Tests
pm.sample.make test-correctness, real NUTS) covers the selection contract from Wire spike-and-slab covariate selection into the BSTS/ITS path #981: 6 candidate covariates with 2 in the data-generating process, gated on the inclusion-probability ranking and separation.InterruptedTimeSeries. The full notebook treatment is in feat!: graduate StateSpaceTimeSeries #1113.Verification
Rebased onto the lazy experiment lifecycle: the new tests and the docstring example call
.fit()and readresult.result.impact_post.Local run with pymc 6.2.0, arviz 1.3.0, pymc-extras 0.14.0:
--no-cov): 41 passed, 13 skipped.diff-coveragainstpymc6_and_pymcmarketing1_migration: all 32 changed executable lines covered.🤖 Generated with Claude Code