Merge pull request #128 from pymc-labs/ancova2

add ANCOVA example to integration tests + README

Author: drbenvincent

README.md

@@ -90,9 +90,26 @@ This is appropriate when you have multiple units, one of which is treated. You b
 
 > The data (treated and untreated units), pre-treatment model fit, and counterfactual (i.e. the synthetic control) are plotted (top). The causal impact is shown as a blue shaded region. The Bayesian analysis shows shaded Bayesian credible regions of the model fit and counterfactual. Also shown is the causal impact (middle) and cumulative causal impact (bottom).
 
+### ANCOVA
+
+This is appropriate for non-equivalent group designs when you have a single pre and post intervention measurement and have a treament and a control group.
+
+| Group | pre | post |
+|------|---|-------|
+| 0    | $x_1$ | $y_1$ |
+| 0    | $x_2$ | $y_2$ |
+| 1    | $x_3$ | $y_3$ |
+| 1    | $x_4$ | $y_4$ |
+
+| Frequentist | Bayesian |
+|--|--|
+| coming soon | ![](img/anova_pymc.svg) |
+
+> The data from the control and treatment group are plotted, along with posterior predictive 94% credible intervals. The lower panel shows the estimated treatment effect.
+
 ### Difference in Differences
 
-This is appropriate when you have a single pre and post intervention measurement and have a treament and a control group.
+This is appropriate for non-equivalent group designs when you have pre and post intervention measurement and have a treament and a control group. Unlike the ANCOVA approach, difference in differences is appropriate when there are multiple pre and/or post treatment measurements.
 
 Data is expected to be in the following form. Shown are just two units - one in the treated group (`group=1`) and one in the untreated group (`group=0`), but there can of course be multiple units per group. This is panel data (also known as repeated measures) where each unit is measured at 2 time points.
 
@@ -107,7 +124,7 @@ Data is expected to be in the following form. Shown are just two units - one in
 |--|--|
 | ![](img/difference_in_differences_skl.svg) | ![](img/difference_in_differences_pymc.svg) |
 
-The data, model fit, and counterfactual are plotted. Frequentist model fits result in points estimates, but the Bayesian analysis results in posterior distributions, represented by the violin plots. The causal impact is the difference between the counterfactual prediction (treated group, post treatment) and the observed values for the treated group, post treatment.
+>The data, model fit, and counterfactual are plotted. Frequentist model fits result in points estimates, but the Bayesian analysis results in posterior distributions, represented by the violin plots. The causal impact is the difference between the counterfactual prediction (treated group, post treatment) and the observed values for the treated group, post treatment.
 
 ### Regression discontinuity designs
 

causalpy/tests/test_integration_pymc_examples.py

@@ -218,3 +218,19 @@ def test_sc_brexit():
         len(result.prediction_model.idata.posterior.coords["draw"])
         == sample_kwargs["draws"]
     )
+
+
+@pytest.mark.integration
+def test_ancova():
+    df = cp.load_data("anova1")
+    result = cp.pymc_experiments.PrePostNEGD(
+        df,
+        formula="post ~ 1 + C(group) + pre",
+        group_variable_name="group",
+        pretreatment_variable_name="pre",
+        prediction_model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.pymc_experiments.PrePostNEGD)
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]

docs/index.rst

@@ -69,10 +69,17 @@ This is appropriate when you have multiple units, one of which is treated. You b
 
 .. image:: ../img/synthetic_control_pymc.svg
 
+ANCOVA
+""""""
+
+This is appropriate when you have a single pre and post intervention measurement and have a treament and a control group.
+
+.. image:: ../img/anova_pymc.svg
+
 Difference in differences
 """""""""""""""""""""""""
 
-This is appropriate when you have a single pre and post intervention measurement and have a treament and a control group.
+This is appropriate when you have pre and post intervention measurement(s) and have a treament and a control group.
 
 .. image:: ../img/difference_in_differences_pymc.svg