Skip to content

feat: variable selection priors for SSTS covariates - #1112

Merged
cetagostini merged 8 commits into
pymc6_and_pymcmarketing1_migrationfrom
issue-758-ssts-variable-selection
Sep 24, 2026
Merged

cetagostini merged 8 commits into
pymc6_and_pymcmarketing1_migrationfrom
issue-758-ssts-variable-selection

Conversation

@anevolbap

@anevolbap anevolbap commented Jul 30, 2026 •

Copy link
Copy Markdown
Collaborator

Closes #981. Part of #758 (P0). Targets pymc6_and_pymcmarketing1_migration.

Summary

Adds variable selection priors for the covariates of StateSpaceTimeSeries, reusing causalpy/variable_selection_priors.py unchanged. 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 the beta_exog prior instead of the default.
  • vs_hyperparams: optional dict passed to the factory.
  • get_inclusion_probabilities() and get_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 plain Normal(0, 1) with no selection. It is much tighter than the class default Normal(0, 50) for beta_exog, and the docstring says so.
  • Passing both priors={"beta_exog": ...} and vs_prior_type lets variable selection win with a warning, same as the IV class. The warning points at the calling line.
  • Setting vs_prior_type without 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_model arguments. InterruptedTimeSeries only feeds X, y, coords through 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_exog point estimates shrink toward zero, because P0 (the initial state covariance, unrelated to the P0 priority label of #758) lets the regression states drift away from the parameter. Pinning the exog P0 entries 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

  • Structure tests cover all three prior types. They assert API shape, dims and coords only, since the suite mocks pm.sample.
  • A correctness-marked test (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.
  • The class docstring documents the default hyperparameters and carries a doctest example of selection through 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 read result.result.impact_post.

Local run with pymc 6.2.0, arviz 1.3.0, pymc-extras 0.14.0:

  • Full suite: 2508 passed, 18 skipped, 4 deselected.
  • Correctness lane: 4 passed.
  • Doctests (the CI invocation, with --no-cov): 41 passed, 13 skipped.
  • diff-cover against pymc6_and_pymcmarketing1_migration: all 32 changed executable lines covered.

🤖 Generated with Claude Code

@read-the-docs-community

read-the-docs-community Bot commented Jul 30, 2026 •

Copy link
Copy Markdown

Documentation build overview

📚 causalpy | 🛠️ Build #34634115 | 📁 Comparing cce4b29 against latest (7e23946)

  🔍 Preview build  

831 files changed · + 125 added · ± 675 modified · - 31 deleted

+ Added

± Modified

- Deleted

@codecov

codecov Bot commented Jul 31, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.27891% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 97.25%. Comparing base (0c17458) to head (cce4b29).
⚠️ Report is 10 commits behind head on pymc6_and_pymcmarketing1_migration.

Files with missing lines Patch % Lines
...salpy/tests/test_integration_its_new_timeseries.py 91.11% 4 Missing ⚠️
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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@anevolbap
anevolbap force-pushed the issue-758-ssts-covariates branch from 374ece7 to f4eb592 Compare July 31, 2026 13:06
@anevolbap
anevolbap force-pushed the issue-758-ssts-variable-selection branch from c58c9e3 to d158d16 Compare July 31, 2026 13:06
@anevolbap
anevolbap force-pushed the issue-758-ssts-covariates branch from f4eb592 to 0f3e88d Compare July 31, 2026 17:28
@anevolbap
anevolbap force-pushed the issue-758-ssts-variable-selection branch from d158d16 to b9da23f Compare July 31, 2026 17:28
@anevolbap
anevolbap marked this pull request as ready for review July 31, 2026 17:32
@anevolbap
anevolbap force-pushed the issue-758-ssts-covariates branch from 0f3e88d to fcee78a Compare August 3, 2026 20:19
@anevolbap
anevolbap force-pushed the issue-758-ssts-variable-selection branch from b9da23f to 031c83d Compare August 3, 2026 20:19
@anevolbap
anevolbap force-pushed the issue-758-ssts-covariates branch from fcee78a to af7713f Compare August 4, 2026 02:36
@anevolbap
anevolbap force-pushed the issue-758-ssts-variable-selection branch 2 times, most recently from e13d969 to e855435 Compare August 4, 2026 11:23
@drbenvincent drbenvincent added the review:high High-impact change requiring thorough human review label Aug 10, 2026
@drbenvincent

Copy link
Copy Markdown
Collaborator

Automated triage

Recommendation: review:high — no decision gate identified.

Why:

  • Adds variable-selection prior support (spike-and-slab, horseshoe, normal) to StateSpaceTimeSeries — a public API and behavioural change for a causal model class.
  • Introduces new public methods get_inclusion_probabilities() and get_shrinkage_factors().
  • 103 new lines in causalpy/pymc_models.py and 182 new test lines. Targets pymc6_and_pymcmarketing1_migration release branch.
  • All CI checks pass cleanly; patch coverage is verified. Stacked on feat: covariate support for StateSpaceTimeSeries #1111.

Review focus:

  1. Confirm the variable-selection prior factory integration matches the pattern established by InstrumentalVariableRegression.
  2. Verify the documented P0 covariance caveat (NUTS slowdown when pinning exog P0 entries) is acceptable at this stage.
  3. Check that mocked-sample tests sufficiently cover API shape/dims/coords until posterior-behaviour tests land in the follow-up notebook PR.

Confidence: high

@drbenvincent drbenvincent left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agent-drafted review.

PR #1112 Review

#1112

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:

  • StateSpaceTimeSeries gains a public, constructor-level variable-selection API (vs_prior_type / vs_hyperparams) that reuses VariableSelectionPrior unchanged and is reachable from InterruptedTimeSeries without adapter changes — the intentional departure from IV’s build_model kwargs 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; the beta_exog attenuation 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-other notebook 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): remote notebooks (gallery-other, …) fails with ImportError: cannot import name 'extract_dims' from 'pymc.model.fgraph' inside pinned pymc-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 on pymc6_and_pymcmarketing1_migration by 41b6416b (re-pin to f8b552c + 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 bare RangeIndex ([0, 1, 2, …]) even though posterior coords carry state_exog names. Local probe confirmed labels are absent; correctness and structure tests rely on formula/iloc order. 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 confirmed beta_exog with state_exog coords and no gamma_*. Worth a one-liner next to the horseshoe structure test. Note also that factory normal defaults to sigma=1, which is not the SSTS default beta_exog prior (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/patch still 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_dims ImportError from stale pymc-marketing pin (see Findings).
  • Correctness lane on all three Test jobs: pytest -m correctness reported 4 passed, including test_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: normal builds with state_exog coords; spike-and-slab inclusion table shape OK (index positional); horseshoe shrinkage columns OK; VS precedence warning + gamma_beta_exog present; predict under 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:

  1. Constructor/vs_prior_type contract wrong or invalid types accepted late — falsified; Literal + eager VariableSelectionPrior validation; test_vs_prior_invalid_type / without-covariates / precedence warning all pass.
  2. Prior factory not used or priors["beta_exog"] wins silently — falsified; special-case in the param loop; warning + gamma present under dual config.
  3. Exog dims/coords or slicing break selection or forecasts — falsified; beta_exog coords ["x1","x2","x3"]; _exog_values shared by build/pm.Data/predict; ITS structure test keeps finite post_impact; local predict under VS succeeded.
  4. 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).
  5. normal / horseshoe / spike-and-slab paths broken for SSTS — falsified by structure tests (HS/SS) and local real sampling (all three).
  6. 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.
  7. Documented attenuation caveat overstates forecast safety — not overturned; counterfactual path still uses smoothed states/_prepare_idata unchanged; no evidence VS breaks ITS impact finiteness. Residual: no dedicated real-NUTS forecast-recovery assert in this PR (acceptable given documented scope).
  8. _clone drops VS config (MF-1) — falsified; test_vs_clone_preserves_config.
  9. 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.
  10. 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

@anevolbap
anevolbap force-pushed the issue-758-ssts-variable-selection branch from 29d9a10 to bdff957 Compare September 8, 2026 23:01
@review-notebook-app

Copy link
Copy Markdown

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@anevolbap
anevolbap deleted the branch pymc6_and_pymcmarketing1_migration September 8, 2026 23:04
@anevolbap anevolbap closed this Sep 8, 2026
@anevolbap anevolbap reopened this Sep 8, 2026
@anevolbap
anevolbap removed this pull request from stack #1114 September 9, 2026 17:17
@anevolbap
anevolbap changed the base branch from issue-758-ssts-covariates to pymc6_and_pymcmarketing1_migration September 9, 2026 17:19
anevolbap added a commit that referenced this pull request Sep 11, 2026
…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.
anevolbap added a commit that referenced this pull request Sep 11, 2026
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 added a commit that referenced this pull request Sep 11, 2026
…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.
anevolbap added a commit that referenced this pull request Sep 11, 2026
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

anevolbap commented Sep 11, 2026 •

Copy link
Copy Markdown
Collaborator Author
  • Docs pin: the stack now carries 41b6416b; docs/requirements.txt pins pymc-marketing f8b552c and pymc-extras>=0.14,<0.15, and gallery-other passes.
  • Inclusion and shrinkage tables: the SSTS pass-throughs now index rows by regressor name; the factory and the IV path are unchanged. The correctness test selects rows by name instead of position.
  • normal: structure test added next to the horseshoe one, and the vs_hyperparams docstring notes the factory default is Normal(0, 1) against the class default Normal(0, 50) for beta_exog.
  • Patch coverage: make test-patch-cov against the migration branch covers all 39 changed executable lines.

@cetagostini

Copy link
Copy Markdown
Collaborator

@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.
@anevolbap
anevolbap force-pushed the issue-758-ssts-variable-selection branch from a897cec to cce4b29 Compare September 18, 2026 11:56
@anevolbap

Copy link
Copy Markdown
Collaborator Author

Rebased onto pymc6_and_pymcmarketing1_migration, conflicts resolved. The new tests and the docstring example now call .fit() and read result.result.impact_post, to match the lazy experiment lifecycle. Nothing else changed.

anevolbap added a commit that referenced this pull request Sep 18, 2026
…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.
anevolbap added a commit that referenced this pull request Sep 18, 2026
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.
@cetagostini
cetagostini merged commit 143e823 into pymc6_and_pymcmarketing1_migration Sep 24, 2026
14 of 16 checks passed
@cetagostini
cetagostini deleted the issue-758-ssts-variable-selection branch September 24, 2026 07:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

review:high High-impact change requiring thorough human review

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants