update adding more tests

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

Author: NathanielF

causalpy/pymc_models.py

@@ -684,9 +684,9 @@ def build_model(  # type: ignore
             Dictionary of priors for the mus and sigmas of both
             regressions. Example: ``priors = {"mus": [0, 0],
             "sigmas": [1, 1], "eta": 2, "lkj_sd": 2}``.
-        :param vs_prior_type: An optional string. Can be "spike_and_slab"
+        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
+        vs_hyperparams: An optional dictionary of priors for the
                                variable selection hyperparameters
 
         """
@@ -705,16 +705,18 @@ def build_model(  # type: ignore
             # Create coefficient priors
             if vs_prior_type:
                 # Use variable selection priors
-                vs_prior_treatment = VariableSelectionPrior(
+                self.vs_prior_treatment = VariableSelectionPrior(
+                    vs_prior_type, vs_hyperparams
+                )
+                self.vs_prior_outcome = VariableSelectionPrior(
                     vs_prior_type, vs_hyperparams
                 )
-                vs_prior_outcome = VariableSelectionPrior(vs_prior_type, vs_hyperparams)
 
-                beta_t = vs_prior_treatment.create_prior(
+                beta_t = self.vs_prior_treatment.create_prior(
                     name="beta_t", n_params=Z.shape[1], dims="instruments", X=Z
                 )
 
-                beta_z = vs_prior_outcome.create_prior(
+                beta_z = self.vs_prior_outcome.create_prior(
                     name="beta_z", n_params=X.shape[1], dims="covariates", X=X
                 )
             else:
@@ -733,7 +735,7 @@ def build_model(  # type: ignore
                 )
 
             sd_dist = pm.Exponential.dist(priors["lkj_sd"], shape=2)
-            chol, corr, sigmas = pm.LKJCholeskyCov(
+            chol, _, _ = pm.LKJCholeskyCov(
                 name="chol_cov",
                 eta=priors["eta"],
                 n=2,

causalpy/tests/test_integration_pymc_examples.py

@@ -706,6 +706,45 @@ def test_iv_reg_vs_prior(mock_pymc_sample):
         result.get_plot_data()
     assert "gamma_beta_t" in result.model.named_vars
     assert "pi_beta_t" in result.model.named_vars
+    summary = result.model.vs_prior_outcome.get_inclusion_probabilities(
+        result.idata, "beta_z"
+    )
+    assert isinstance(summary, pd.DataFrame)
+
+
+@pytest.mark.integration
+def test_iv_reg_vs_prior_hs(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="horseshoe",
+    )
+    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 "tau_beta_t" in result.model.named_vars
+    assert "tau_beta_z" in result.model.named_vars
+    summary = result.model.vs_prior_outcome.get_shrinkage_factors(
+        result.idata, "beta_z"
+    )
+    assert isinstance(summary, pd.DataFrame)
 
 
 @pytest.mark.integration

causalpy/variable_selection_priors.py

@@ -23,6 +23,7 @@
 from typing import Any, Dict, Optional, Union
 
 import numpy as np
+import pandas as pd
 import pymc as pm
 import pytensor.tensor as pt
 from pymc_extras.prior import Prior
@@ -65,9 +66,10 @@ class SpikeAndSlabPrior:
     Creates a mixture prior with a point mass at zero (spike) and a diffuse
     normal distribution (slab), implemented as:
 
-    β_j = γ_j × β_j^raw
-
-    where γ_j ∈ [0,1] is a relaxed indicator and β_j^raw ~ N(0, σ_slab²).
+    .. math::
+        \beta_{j} = \gamma_{j} \cdot \beta_{j}^{\text{raw}} \\
+        \beta_{j}^{\text{raw}} \sim \mathcal{N}(0, \sigma_{\text{slab}}^{2}), \qquad
+        \gamma_{j} \in [0,1].
 
     Parameters
     ----------
@@ -145,9 +147,9 @@ class HorseshoePrior:
     Provides continuous shrinkage with heavy tails, allowing strong signals
     to escape shrinkage while weak signals are dampened:
 
-    β_j = τ · λ̃_j · β_j^raw
-
-    where λ̃_j = √(c²λ_j² / (c² + τ²λ_j²)) is the regularized local shrinkage.
+    .. math::
+        \beta_{j} & =  \tau \cdot \lambda_{j} \cdot \beta_{j}^{raw}  \\
+        \lambda_{j} & = \sqrt{ \dfrac{c^{2}\lambda_{j}^{2}}{c^{2} + \tau^{2}\lambda_{j}^{2}} }
 
     Parameters
     ----------
@@ -423,7 +425,7 @@ def create_prior(
 
     def get_inclusion_probabilities(
         self, idata, param_name: str, threshold: float = 0.5
-    ) -> Dict[str, np.ndarray]:
+    ) -> pd.DataFrame:
         """
         Extract variable inclusion probabilities from fitted model.
 
@@ -472,17 +474,24 @@ def get_inclusion_probabilities(
         gamma = az.extract(idata.posterior[gamma_name])
 
         # Compute inclusion probabilities
-        probabilities = (gamma > threshold).mean(dim="sample").values
-        gamma_mean = gamma.mean(dim="sample").values
+        probabilities = (gamma > threshold).mean(dim="sample").to_array()
+        gamma_mean = gamma.mean(dim="sample").to_array()
         selected = probabilities > threshold
 
-        return {
+        summary = {
             "probabilities": probabilities,
             "selected": selected,
             "gamma_mean": gamma_mean,
         }
+        probs = summary["probabilities"].T
+        df = pd.DataFrame(index=list(range(len(probs))))
+
+        df["prob"] = probs
+        df["selected"] = summary["selected"].T
+        df["gamma_mean"] = summary["gamma_mean"].T
+        return df
 
-    def get_shrinkage_factors(self, idata, param_name: str) -> Dict[str, np.ndarray]:
+    def get_shrinkage_factors(self, idata, param_name: str) -> pd.DataFrame:
         """
         Extract shrinkage factors from horseshoe prior.
 
@@ -524,17 +533,26 @@ def get_shrinkage_factors(self, idata, param_name: str) -> Dict[str, np.ndarray]
             raise ValueError(f"Could not find '{lambda_tilde_name}' in posterior")
 
         # Extract components
-        tau = az.extract(idata.posterior[tau_name])
-        lambda_tilde = az.extract(idata.posterior[lambda_tilde_name])
+        tau = az.extract(idata.posterior[tau_name]).to_array()
+        lambda_tilde = az.extract(idata.posterior[lambda_tilde_name]).to_array()
 
-        # Compute shrinkage factors
-        shrinkage_factors = (tau * lambda_tilde).mean(dim="sample").values
+        shrinkage_factor = np.array(
+            [tau[0, i] * lambda_tilde[0, :, :] for i in range(len(tau))]
+        )
+        shrinkage_factor = shrinkage_factor.mean(axis=2)
 
-        return {
-            "shrinkage_factors": shrinkage_factors,
-            "tau": tau.mean().values,
-            "lambda_tilde": lambda_tilde.mean(dim="sample").values,
+        summary = {
+            "shrinkage_factors": shrinkage_factor,
+            "tau": tau.mean(),
+            "lambda_tilde": lambda_tilde.mean(dim=("sample")),
         }
+        probs = summary["shrinkage_factors"].T
+        df = pd.DataFrame(index=list(range(len(probs))))
+        df["shrinkage_factor"] = probs
+
+        df["lambda_tilde"] = summary["lambda_tilde"].T
+        df["tau"] = np.mean(tau).item()
+        return df
 
 
 def create_variable_selection_prior(