Merge pull request #130 from pymc-labs/idata

access idata through new idata property + show property in autogenerated docs

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -32,7 +32,7 @@ def idata(self):
     def print_coefficients(self):
         """Prints the model coefficients"""
         print("Model coefficients:")
-        coeffs = az.extract(self.prediction_model.idata.posterior, var_names="beta")
+        coeffs = az.extract(self.idata.posterior, var_names="beta")
         # Note: f"{name: <30}" pads the name with spaces so that we have alignment of
         # the stats despite variable names of different lengths
         for name in self.labels:
@@ -697,7 +697,7 @@ def __init__(
         self.pred_treated = self.prediction_model.predict(X=np.asarray(new_x))
 
         # Evaluate causal impact as equal to the trestment effect
-        self.causal_impact = self.prediction_model.idata.posterior["beta"].sel(
+        self.causal_impact = self.idata.posterior["beta"].sel(
             {"coeffs": self._get_treatment_effect_coeff()}
         )
 

causalpy/tests/test_integration_pymc_examples.py

@@ -20,14 +20,8 @@ def test_did():
     )
     assert isinstance(df, pd.DataFrame)
     assert isinstance(result, cp.pymc_experiments.DifferenceInDifferences)
-    assert (
-        len(result.prediction_model.idata.posterior.coords["chain"])
-        == sample_kwargs["chains"]
-    )
-    assert (
-        len(result.prediction_model.idata.posterior.coords["draw"])
-        == sample_kwargs["draws"]
-    )
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
 @pytest.mark.integration
@@ -61,14 +55,8 @@ def test_did_banks():
     )
     assert isinstance(df, pd.DataFrame)
     assert isinstance(result, cp.pymc_experiments.DifferenceInDifferences)
-    assert (
-        len(result.prediction_model.idata.posterior.coords["chain"])
-        == sample_kwargs["chains"]
-    )
-    assert (
-        len(result.prediction_model.idata.posterior.coords["draw"])
-        == sample_kwargs["draws"]
-    )
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
 @pytest.mark.integration
@@ -82,14 +70,8 @@ def test_rd():
     )
     assert isinstance(df, pd.DataFrame)
     assert isinstance(result, cp.pymc_experiments.RegressionDiscontinuity)
-    assert (
-        len(result.prediction_model.idata.posterior.coords["chain"])
-        == sample_kwargs["chains"]
-    )
-    assert (
-        len(result.prediction_model.idata.posterior.coords["draw"])
-        == sample_kwargs["draws"]
-    )
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
 @pytest.mark.integration
@@ -108,14 +90,8 @@ def test_rd_drinking():
     )
     assert isinstance(df, pd.DataFrame)
     assert isinstance(result, cp.pymc_experiments.RegressionDiscontinuity)
-    assert (
-        len(result.prediction_model.idata.posterior.coords["chain"])
-        == sample_kwargs["chains"]
-    )
-    assert (
-        len(result.prediction_model.idata.posterior.coords["draw"])
-        == sample_kwargs["draws"]
-    )
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
 @pytest.mark.integration
@@ -132,14 +108,8 @@ def test_its():
     )
     assert isinstance(df, pd.DataFrame)
     assert isinstance(result, cp.pymc_experiments.SyntheticControl)
-    assert (
-        len(result.prediction_model.idata.posterior.coords["chain"])
-        == sample_kwargs["chains"]
-    )
-    assert (
-        len(result.prediction_model.idata.posterior.coords["draw"])
-        == sample_kwargs["draws"]
-    )
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
 @pytest.mark.integration
@@ -156,14 +126,8 @@ def test_its_covid():
     )
     assert isinstance(df, pd.DataFrame)
     assert isinstance(result, cp.pymc_experiments.SyntheticControl)
-    assert (
-        len(result.prediction_model.idata.posterior.coords["chain"])
-        == sample_kwargs["chains"]
-    )
-    assert (
-        len(result.prediction_model.idata.posterior.coords["draw"])
-        == sample_kwargs["draws"]
-    )
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
 @pytest.mark.integration
@@ -178,14 +142,8 @@ def test_sc():
     )
     assert isinstance(df, pd.DataFrame)
     assert isinstance(result, cp.pymc_experiments.SyntheticControl)
-    assert (
-        len(result.prediction_model.idata.posterior.coords["chain"])
-        == sample_kwargs["chains"]
-    )
-    assert (
-        len(result.prediction_model.idata.posterior.coords["draw"])
-        == sample_kwargs["draws"]
-    )
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
 @pytest.mark.integration
@@ -210,14 +168,8 @@ def test_sc_brexit():
     )
     assert isinstance(df, pd.DataFrame)
     assert isinstance(result, cp.pymc_experiments.SyntheticControl)
-    assert (
-        len(result.prediction_model.idata.posterior.coords["chain"])
-        == sample_kwargs["chains"]
-    )
-    assert (
-        len(result.prediction_model.idata.posterior.coords["draw"])
-        == sample_kwargs["draws"]
-    )
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
 @pytest.mark.integration

docs/api_pymc_experiments.rst

@@ -7,3 +7,4 @@
 .. automodule:: causalpy.pymc_experiments
    :members:
    :undoc-members:
+   :inherited-members:

docs/notebooks/rd_pymc_drinking.ipynb

@@ -446,9 +446,9 @@
    "source": [
     "fig, ax = plt.subplots(1, 2, figsize=(10, 3))\n",
     "\n",
-    "az.plot_forest(result.prediction_model.idata.posterior, var_names=\"beta\", ax=ax[0])\n",
+    "az.plot_forest(result.idata.posterior, var_names=\"beta\", ax=ax[0])\n",
     "az.plot_posterior(\n",
-    "    result.prediction_model.idata.posterior.beta.sel(coeffs=\"treated[T.True]\"),\n",
+    "    result.idata.posterior.beta.sel(coeffs=\"treated[T.True]\"),\n",
     "    round_to=3,\n",
     "    ax=ax[1],\n",
     ")\n",
@@ -792,9 +792,7 @@
     }
    ],
    "source": [
-    "az.plot_forest(\n",
-    "    result2.prediction_model.idata.posterior, var_names=\"beta\", figsize=(10, 3)\n",
-    ");"
+    "az.plot_forest(result2.idata.posterior, var_names=\"beta\", figsize=(10, 3));"
    ]
   },
   {
@@ -1129,9 +1127,9 @@
    "source": [
     "fig, ax = plt.subplots(1, 2, figsize=(10, 3))\n",
     "\n",
-    "az.plot_forest(result3.prediction_model.idata.posterior, var_names=\"beta\", ax=ax[0])\n",
+    "az.plot_forest(result3.idata.posterior, var_names=\"beta\", ax=ax[0])\n",
     "az.plot_posterior(\n",
-    "    result3.prediction_model.idata.posterior.beta.sel(coeffs=\"treated[T.True]\"),\n",
+    "    result3.idata.posterior.beta.sel(coeffs=\"treated[T.True]\"),\n",
     "    round_to=3,\n",
     "    ax=ax[1],\n",
     ")\n",