feat: variable selection priors for SSTS covariates

causalpy/pymc_models.py

@@ -2446,6 +2446,57 @@ class StateSpaceTimeSeries(PyMCModel):
         `default_priors`. The `P0` covariance is parameterized through its
         diagonal under the key `"P0_diag"`. Dims are resolved from the built
         state-space model, so priors do not need to declare them.
+    vs_prior_type : {"spike_and_slab", "horseshoe", "normal"}, optional
+        Variable selection prior for the exogenous regression coefficients.
+        Requires covariates. Takes precedence over a `beta_exog` entry in
+        `priors`.
+    vs_hyperparams : dict, optional
+        Hyperparameters for the variable selection prior. See
+        :class:`causalpy.variable_selection_priors.VariableSelectionPrior`.
+        The defaults work without hand-tuning on roughly unit-scale data:
+        the horseshoe scales its global shrinkage from the data with an
+        expected model size of ``min(5, p / 2)`` (Piironen & Vehtari, 2017),
+        while spike-and-slab uses a ``Beta(2, 2)`` inclusion prior (prior
+        inclusion probability centered on 0.5, no expected-model-size knob).
+
+    Examples
+    --------
+    Covariate selection through :class:`causalpy.InterruptedTimeSeries`:
+    pass many candidate covariates in the formula and let the model select.
+
+    >>> import numpy as np
+    >>> import pandas as pd
+    >>> import causalpy as cp
+    >>> rng = np.random.default_rng(7)
+    >>> n = 60
+    >>> dates = pd.date_range(start="2023-01-01", periods=n, freq="D")
+    >>> X = rng.normal(size=(n, 3))
+    >>> y = 5 + 2.0 * X[:, 0] + rng.normal(0, 0.3, size=n)
+    >>> df = pd.DataFrame(
+    ...     {"y": y, "x1": X[:, 0], "x2": X[:, 1], "x3": X[:, 2]}, index=dates
+    ... )
+    >>> model = cp.pymc_models.StateSpaceTimeSeries(
+    ...     level_order=1,
+    ...     seasonal_length=7,
+    ...     sample_kwargs={
+    ...         "chains": 1,
+    ...         "draws": 10,
+    ...         "tune": 10,
+    ...         "progressbar": False,
+    ...     },
+    ...     vs_prior_type="spike_and_slab",
+    ... )
+    >>> import io
+    >>> from contextlib import redirect_stdout
+    >>> with redirect_stdout(io.StringIO()):  # silence the model-build table
+    ...     result = cp.InterruptedTimeSeries(
+    ...         data=df,
+    ...         treatment_time=dates[45],
+    ...         formula="y ~ 0 + x1 + x2 + x3",
+    ...         model=model,
+    ...     )
+    >>> result.model.get_inclusion_probabilities().columns.tolist()
+    ['prob', 'selected', 'gamma_mean']
     """
 
     default_priors = {
@@ -2466,6 +2517,8 @@ def __init__(
         sample_kwargs: dict[str, Any] | None = None,
         mode: str | None = None,
         priors: dict[str, Prior] | None = None,
+        vs_prior_type: Literal["spike_and_slab", "horseshoe", "normal"] | None = None,
+        vs_hyperparams: dict[str, Any] | None = None,
     ):
         super().__init__(sample_kwargs=sample_kwargs, priors=priors)
 
@@ -2485,6 +2538,20 @@ def __init__(
         self._treated_units = ["unit_0"]
         self.ss_mod: Any = None
         self._exog_names: list[str] = []
+        self.vs_prior_type = vs_prior_type
+        self.vs_hyperparams = vs_hyperparams
+        self.vs_prior: VariableSelectionPrior | None = None
+        if vs_prior_type is not None:
+            # Validates the prior type eagerly
+            self.vs_prior = VariableSelectionPrior(vs_prior_type, vs_hyperparams or {})
+            if priors and "beta_exog" in priors:
+                warnings.warn(
+                    "Both vs_prior_type and a beta_exog entry in priors were "
+                    "given. The variable selection prior takes precedence for "
+                    "beta_exog.",
+                    UserWarning,
+                    stacklevel=2,
+                )
         self._validate_and_initialize_components()
 
     def _clone(self, priors: dict[str, Any] | None = None) -> "PyMCModel":
@@ -2501,6 +2568,8 @@ def _clone(self, priors: dict[str, Any] | None = None) -> "PyMCModel":
             sample_kwargs=dict(self.sample_kwargs),
             mode=self.mode,
             priors=self._user_priors if priors is None else priors,
+            vs_prior_type=self.vs_prior_type,
+            vs_hyperparams=self.vs_hyperparams,
         )
 
     def _validate_and_initialize_components(self):
@@ -2593,6 +2662,10 @@ def _extract_exog_names(self, X: xr.DataArray | None) -> list[str]:
                 )
         return names
 
+    def _exog_values(self, X: xr.DataArray) -> np.ndarray:
+        """Exogenous regressor values from X, in fit-time column order."""
+        return X.sel(coeffs=self._exog_names).values
+
     def build_model(
         self,
         X: xr.DataArray | None = None,
@@ -2672,6 +2745,12 @@ def build_model(
         season = self._get_seasonality_component()
         combined = trend + season
         self._exog_names = self._extract_exog_names(X)
+        if self.vs_prior is not None and not self._exog_names:
+            raise ValueError(
+                "vs_prior_type was set but the model has no exogenous "
+                "covariates. Pass covariates via X, e.g. with a "
+                "'y ~ 0 + x1 + x2' formula."
+            )
         if self._exog_names:
             from pymc_extras.statespace import structural as st
 
@@ -2724,6 +2803,13 @@ def build_model(
                     prior.dims = dims[0]
                     P0_diag = prior.create_variable("P0_diag")
                     pm.Deterministic("P0", pt.diag(P0_diag), dims=dims)
+                elif name == "beta_exog" and self.vs_prior is not None:
+                    self.vs_prior.create_prior(
+                        "beta_exog",
+                        n_params=len(self._exog_names),
+                        dims=dims,
+                        X=self._exog_values(X) if X is not None else None,
+                    )
                 else:
                     prior = deepcopy(self.priors[name])
                     prior.dims = dims
@@ -2739,7 +2825,7 @@ def build_model(
             df = pd.DataFrame({"y": y_values.flatten()}, index=datetime_index)
             if self._exog_names and X is not None:
                 # The state-space graph looks this variable up by name
-                pm.Data("data_exog", X.sel(coeffs=self._exog_names).values)
+                pm.Data("data_exog", self._exog_values(X))
             self.ss_mod.build_statespace_graph(df[["y"]])
 
     def fit(
@@ -2826,6 +2912,73 @@ def _smooth(self) -> xr.Dataset:
             else conditional_idata
         )
 
+    def _require_vs_diagnostics(self, what: str) -> tuple[VariableSelectionPrior, Any]:
+        """Guard the variable-selection accessors.
+
+        Returns the prior and the fitted idata, raising the same errors both
+        accessors documented: ValueError when the model was not configured
+        with `vs_prior_type`, RuntimeError when it has not been fit.
+        """
+        if self.vs_prior is None:
+            raise ValueError(
+                "Model was not configured with vs_prior_type; there are no "
+                f"{what} to report."
+            )
+        if self.idata is None:
+            raise RuntimeError("Model must be fit first.")
+        return self.vs_prior, self.idata
+
+    def get_inclusion_probabilities(
+        self, param_name: str = "beta_exog"
+    ) -> pd.DataFrame:
+        """
+        Posterior inclusion probabilities of the exogenous regressors.
+
+        Only available when the model was configured with
+        `vs_prior_type="spike_and_slab"` and has been fit.
+
+        Interpret the probabilities as a relative ranking of the candidate
+        regressors. The `beta_exog` point estimates shrink toward zero
+        under this prior (the state-space `P0` lets the regression states
+        drift from the parameter), but counterfactual forecasts use the
+        smoothed states and are not affected by that attenuation.
+
+        Parameters
+        ----------
+        param_name : str, optional
+            Name of the coefficient parameter. Defaults to "beta_exog".
+
+        Returns
+        -------
+        pd.DataFrame
+            One row per regressor with columns "prob" (inclusion
+            probability), "selected" (probability above 0.5), and
+            "gamma_mean" (mean of the selection indicator).
+        """
+        vs_prior, idata = self._require_vs_diagnostics("inclusion probabilities")
+        return vs_prior.get_inclusion_probabilities(idata, param_name)
+
+    def get_shrinkage_factors(self, param_name: str = "beta_exog") -> pd.DataFrame:
+        """
+        Shrinkage factors of the exogenous regressors.
+
+        Only available when the model was configured with
+        `vs_prior_type="horseshoe"` and has been fit.
+
+        Parameters
+        ----------
+        param_name : str, optional
+            Name of the coefficient parameter. Defaults to "beta_exog".
+
+        Returns
+        -------
+        pd.DataFrame
+            One row per regressor with the effective shrinkage applied to
+            its coefficient.
+        """
+        vs_prior, idata = self._require_vs_diagnostics("shrinkage factors")
+        return vs_prior.get_shrinkage_factors(idata, param_name)
+
     def _forecast(
         self,
         start: pd.Timestamp,
@@ -2907,7 +3060,7 @@ def predict(
                     raise ValueError(
                         f"X is missing exogenous columns used at fit time: {missing}."
                     )
-                scenario = X.sel(coeffs=self._exog_names).values
+                scenario = self._exog_values(X)
             last = self._train_index[-1]  # start forecasting after the last observed
             forecast_data = self._forecast(
                 start=last, periods=len(idx), scenario=scenario

causalpy/tests/test_integration_its_new_timeseries.py

@@ -288,3 +288,116 @@ def test_its_with_state_space_covariates():
     n_post = n - 80
     assert result.post_impact.sizes["obs_ind"] == n_post
     assert np.isfinite(result.post_impact.values).all()
+
+
+@pytest.mark.integration
+def test_its_with_state_space_variable_selection(mock_pymc_sample):
+    """ITS + StateSpaceTimeSeries with spike-and-slab covariate selection.
+
+    Structure-only assertions: the suite mocks pm.sample session-wide,
+    so posterior values come from the prior.
+    """
+    try:
+        from pymc_extras.statespace import structural  # noqa: F401
+    except ImportError:
+        pytest.skip("pymc-extras is required for StateSpaceTimeSeries tests")
+
+    rng = np.random.default_rng(seed=42)
+    n = 90
+    dates = pd.date_range(start="2020-01-01", periods=n, freq="D")
+    x1 = rng.normal(size=n)
+    x2 = rng.normal(size=n)
+    x3 = rng.normal(size=n)
+    y = 5 + 2.0 * x1 + rng.normal(0, 0.3, n)
+    df = pd.DataFrame({"y": y, "x1": x1, "x2": x2, "x3": x3}, index=dates)
+
+    model = cp.pymc_models.StateSpaceTimeSeries(
+        level_order=1,
+        seasonal_length=7,
+        sample_kwargs={
+            "chains": 1,
+            "draws": 50,
+            "tune": 50,
+            "progressbar": False,
+            "random_seed": 7,
+        },
+        vs_prior_type="spike_and_slab",
+    )
+
+    result = cp.InterruptedTimeSeries(
+        data=df,
+        treatment_time=dates[70],
+        formula="y ~ 0 + x1 + x2 + x3",
+        model=model,
+    )
+
+    assert "beta_exog" in result.idata.posterior
+    assert "gamma_beta_exog" in result.idata.posterior
+
+    incl = model.get_inclusion_probabilities()
+    assert isinstance(incl, pd.DataFrame)
+    assert len(incl) == 3
+    assert ((incl["prob"] >= 0) & (incl["prob"] <= 1)).all()
+
+    assert np.isfinite(result.post_impact.values).all()
+
+
+@pytest.mark.integration
+@pytest.mark.slow
+@pytest.mark.correctness
+def test_its_state_space_variable_selection_recovery():
+    """Irrelevant covariates shrink out under the spike-and-slab prior.
+
+    Real NUTS, no pm.sample mock: correctness-marked tests run in their own
+    lane (`make test-correctness`), where the session mock is never
+    instantiated. Selection is asserted through the inclusion-probability
+    ranking, not `beta_exog` point estimates, because the state-space `P0`
+    lets the regression states drift from the parameter, which attenuates
+    the point estimates without affecting the ranking.
+    """
+    try:
+        from pymc_extras.statespace import structural  # noqa: F401
+    except ImportError:
+        pytest.skip("pymc-extras is required for StateSpaceTimeSeries tests")
+
+    rng = np.random.default_rng(seed=157)
+    n = 120
+    dates = pd.date_range(start="2022-01-01", periods=n, freq="D")
+    X = rng.normal(size=(n, 6))
+    y = 3.0 + 2.0 * X[:, 0] - 1.5 * X[:, 1] + rng.normal(0, 0.3, size=n)
+    df = pd.DataFrame({"y": y, **{f"x{i + 1}": X[:, i] for i in range(6)}}, index=dates)
+
+    model = cp.pymc_models.StateSpaceTimeSeries(
+        level_order=1,
+        seasonal_length=7,
+        sample_kwargs={
+            "chains": 2,
+            "draws": 400,
+            "tune": 400,
+            "cores": 1,
+            "target_accept": 0.9,
+            "progressbar": False,
+            "random_seed": 157,
+        },
+        vs_prior_type="spike_and_slab",
+    )
+    cp.InterruptedTimeSeries(
+        data=df,
+        treatment_time=dates[100],
+        formula="y ~ 0 + x1 + x2 + x3 + x4 + x5 + x6",
+        model=model,
+    )
+
+    # Rows follow the formula column order: x1, x2 are in the DGP, x3-x6
+    # are noise. The attenuation documented in the class docstring pulls
+    # every inclusion probability toward the Beta(2, 2) prior mean of 0.5,
+    # so the gates are ranking and separation, not absolute levels.
+    # Calibration (seed 157, 2 chains, 400 draws/tune, target_accept 0.9):
+    # relevant probs ~0.42-0.45, irrelevant ~0.23-0.25, worst-pair gap
+    # 0.163, mean gap 0.187; the limits below keep roughly 2-3x headroom.
+    incl = model.get_inclusion_probabilities()
+    relevant = incl["prob"].iloc[:2]
+    irrelevant = incl["prob"].iloc[2:]
+    assert relevant.min() > irrelevant.max()
+    assert relevant.min() - irrelevant.max() >= 0.05
+    assert relevant.mean() - irrelevant.mean() >= 0.10