adding vs module

Signed-off-by: Nathaniel <[email protected]>

Author: NathanielF

causalpy/__init__.py

@@ -41,4 +41,5 @@
     "RegressionKink",
     "skl_models",
     "SyntheticControl",
+    "variable_selection_priors",
 ]

causalpy/experiments/instrumental_variable.py

@@ -52,6 +52,12 @@ class InstrumentalVariable(BaseExperiment):
                                     "eta": 2,
                                     "lkj_sd": 2,
                                     }
+    :param vs_prior_type : str or None, default=None
+        Type of variable selection prior: 'spike_and_slab', 'horseshoe', or None.
+        If None, uses standard normal priors.
+    :param vs_hyperparams : dict, optional
+        Hyperparameters for variable selection priors. Only used if vs_prior_type
+        is not None.
 
     Example
     --------
@@ -99,6 +105,8 @@ def __init__(
         formula: str,
         model=None,
         priors=None,
+        vs_prior_type=None,
+        vs_hyperparams=None,
         **kwargs,
     ):
         super().__init__(model=model)
@@ -108,6 +116,8 @@ def __init__(
         self.formula = formula
         self.instruments_formula = instruments_formula
         self.model = model
+        self.vs_prior_type = (vs_prior_type,)
+        self.vs_hyperparams = vs_hyperparams or {}
         self.input_validation()
 
         y, X = dmatrices(formula, self.data)
@@ -139,7 +149,14 @@ def __init__(
             }
         self.priors = priors
         self.model.fit(
-            X=self.X, Z=self.Z, y=self.y, t=self.t, coords=COORDS, priors=self.priors
+            X=self.X,
+            Z=self.Z,
+            y=self.y,
+            t=self.t,
+            coords=COORDS,
+            priors=self.priors,
+            vs_prior_type=vs_prior_type,
+            vs_hyperparams=vs_hyperparams,
         )
 
     def input_validation(self):

causalpy/pymc_models.py

@@ -26,6 +26,7 @@
 from pymc_extras.prior import Prior
 
 from causalpy.utils import round_num
+from causalpy.variable_selection_priors import VariableSelectionPrior
 
 
 class PyMCModel(pm.Model):
@@ -604,7 +605,9 @@ class InstrumentalVariableRegression(PyMCModel):
     Inference data...
     """
 
-    def build_model(self, X, Z, y, t, coords, priors):
+    def build_model(
+        self, X, Z, y, t, coords, priors, vs_prior_type=None, vs_hyperparams=None
+    ):
         """Specify model with treatment regression and focal regression data and priors
 
         :param X: A pandas dataframe used to predict our outcome y
@@ -618,23 +621,47 @@ def build_model(self, X, Z, y, t, coords, priors):
                       sigmas of both regressions
                       :code:`priors = {"mus": [0, 0], "sigmas": [1, 1],
                       "eta": 2, "lkj_sd": 2}`
+        :param vs_prior_type: An optional string. Can be "spike_and_slab"
+                              or "horseshoe" or "normal
+        :param vs_hyperparams: An optional dictionary of priors for the
+                               variable selection hyperparameters
+
         """
 
         # --- Priors ---
         with self:
             self.add_coords(coords)
-            beta_t = pm.Normal(
-                name="beta_t",
-                mu=priors["mus"][0],
-                sigma=priors["sigmas"][0],
-                dims="instruments",
-            )
-            beta_z = pm.Normal(
-                name="beta_z",
-                mu=priors["mus"][1],
-                sigma=priors["sigmas"][1],
-                dims="covariates",
-            )
+
+            # Create coefficient priors
+            if vs_prior_type:
+                # Use variable selection priors
+                vs_prior_treatment = VariableSelectionPrior(
+                    vs_prior_type, vs_hyperparams
+                )
+                vs_prior_outcome = VariableSelectionPrior(vs_prior_type, vs_hyperparams)
+
+                beta_t = vs_prior_treatment.create_prior(
+                    name="beta_t", n_params=Z.shape[1], dims="instruments", X=Z
+                )
+
+                beta_z = vs_prior_outcome.create_prior(
+                    name="beta_z", n_params=X.shape[1], dims="covariates", X=X
+                )
+            else:
+                # Use standard normal priors
+                beta_t = pm.Normal(
+                    name="beta_t",
+                    mu=priors["mus"][0],
+                    sigma=priors["sigmas"][0],
+                    dims="instruments",
+                )
+                beta_z = pm.Normal(
+                    name="beta_z",
+                    mu=priors["mus"][1],
+                    sigma=priors["sigmas"][1],
+                    dims="covariates",
+                )
+
             sd_dist = pm.Exponential.dist(priors["lkj_sd"], shape=2)
             chol, corr, sigmas = pm.LKJCholeskyCov(
                 name="chol_cov",
@@ -689,7 +716,18 @@ def sample_predictive_distribution(self, ppc_sampler="jax"):
                     )
                 )
 
-    def fit(self, X, Z, y, t, coords, priors, ppc_sampler=None):
+    def fit(
+        self,
+        X,
+        Z,
+        y,
+        t,
+        coords,
+        priors,
+        ppc_sampler=None,
+        vs_prior_type=None,
+        vs_hyperparams=None,
+    ):
         """Draw samples from posterior distribution and potentially
         from the prior and posterior predictive distributions. The
         fit call can take values for the
@@ -703,7 +741,7 @@ def fit(self, X, Z, y, t, coords, priors, ppc_sampler=None):
         # sample_posterior_predictive() if provided in sample_kwargs.
         # Use JAX for ppc sampling of multivariate likelihood
 
-        self.build_model(X, Z, y, t, coords, priors)
+        self.build_model(X, Z, y, t, coords, priors, vs_prior_type, vs_hyperparams)
         with self:
             self.idata = pm.sample(**self.sample_kwargs)
         self.sample_predictive_distribution(ppc_sampler=ppc_sampler)

causalpy/tests/test_integration_pymc_examples.py

@@ -676,6 +676,38 @@ def test_iv_reg(mock_pymc_sample):
         result.get_plot_data()
 
 
+@pytest.mark.integration
+def test_iv_reg_vs_prior(mock_pymc_sample):
+    df = cp.load_data("risk")
+    instruments_formula = "risk  ~ 1 + logmort0"
+    formula = "loggdp ~  1 + risk"
+    instruments_data = df[["risk", "logmort0"]]
+    data = df[["loggdp", "risk"]]
+
+    result = cp.InstrumentalVariable(
+        instruments_data=instruments_data,
+        data=data,
+        instruments_formula=instruments_formula,
+        formula=formula,
+        model=cp.pymc_models.InstrumentalVariableRegression(
+            sample_kwargs=sample_kwargs
+        ),
+        vs_prior_type="spike_and_slab",
+        vs_hyperparams={"pi_alpha": 5},
+    )
+    result.model.sample_predictive_distribution(ppc_sampler="pymc")
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(data, pd.DataFrame)
+    assert isinstance(instruments_data, pd.DataFrame)
+    assert isinstance(result, cp.InstrumentalVariable)
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
+    with pytest.raises(NotImplementedError):
+        result.get_plot_data()
+    assert "gamma_beta_t" in result.model.named_vars
+    assert "pi_beta_t" in result.model.named_vars
+
+
 @pytest.mark.integration
 def test_inverse_prop(mock_pymc_sample):
     """Test the InversePropensityWeighting class."""