Merge pull request #126 from pymc-labs/idata-property

allow direct access to idata via a property

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -24,6 +24,11 @@ def __init__(self, prediction_model=None, **kwargs):
         if self.prediction_model is None:
             raise ValueError("fitting_model not set or passed.")
 
+    @property
+    def idata(self):
+        """Access to the InferenceData object"""
+        return self.prediction_model.idata
+
     def print_coefficients(self):
         """Prints the model coefficients"""
         print("Model coefficients:")

causalpy/tests/test_pymc_models.py

@@ -4,8 +4,11 @@
 import pymc as pm
 import pytest
 
+import causalpy as cp
 from causalpy.pymc_models import ModelBuilder
 
+sample_kwargs = {"tune": 20, "draws": 20, "chains": 2, "cores": 2}
+
 
 class MyToyModel(ModelBuilder):
     def build_model(self, X, y, coords):
@@ -67,3 +70,19 @@ def test_fit_predict(self, coords, rng) -> None:
         assert isinstance(score, pd.Series)
         assert score.shape == (2,)
         assert isinstance(predictions, az.InferenceData)
+
+
+def test_idata_property():
+    """Test that we can access the idata property of the model"""
+    df = cp.load_data("did")
+    result = cp.pymc_experiments.DifferenceInDifferences(
+        df,
+        formula="y ~ 1 + group + t + treated:group",
+        time_variable_name="t",
+        group_variable_name="group",
+        treated=1,
+        untreated=0,
+        prediction_model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+    )
+    assert hasattr(result, "idata")
+    assert isinstance(result.idata, az.InferenceData)