add integration tests for the pymc examples in the docs

Author: drbenvincent

causalpy/tests/test_integration_pymc_examples.py

@@ -0,0 +1,155 @@
+import pandas as pd
+import pytest
+
+import causalpy as cp
+
+
+@pytest.mark.integration
+def test_did():
+    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(),
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.pymc_experiments.DifferenceInDifferences)
+
+
+@pytest.mark.integration
+def test_did_banks():
+    df = (
+        cp.load_data("banks")
+        .filter(items=["bib6", "bib8", "year"])
+        .rename(columns={"bib6": "Sixth District", "bib8": "Eighth District"})
+        .groupby("year")
+        .median()
+    )
+    df.reset_index(level=0, inplace=True)
+    df_long = pd.melt(
+        df,
+        id_vars=["year"],
+        value_vars=["Sixth District", "Eighth District"],
+        var_name="district",
+        value_name="bib",
+    ).sort_values("year")
+    df_long["district"] = df_long["district"].astype("category")
+    df_long["unit"] = df_long["district"]
+    df_long["treated"] = (df_long.year >= 1931) & (df_long.district == "Sixth District")
+    result = cp.pymc_experiments.DifferenceInDifferences(
+        df_long[df_long.year.isin([1930, 1931])],
+        formula="bib ~ 1 + district + year + district:treated",
+        time_variable_name="year",
+        group_variable_name="district",
+        treated="Sixth District",
+        untreated="Eighth District",
+        prediction_model=cp.pymc_models.LinearRegression(),
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.pymc_experiments.DifferenceInDifferences)
+
+
+@pytest.mark.integration
+def test_rd():
+    df = cp.load_data("rd")
+    result = cp.pymc_experiments.RegressionDiscontinuity(
+        df,
+        formula="y ~ 1 + bs(x, df=6) + treated",
+        prediction_model=cp.pymc_models.LinearRegression(),
+        treatment_threshold=0.5,
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.pymc_experiments.RegressionDiscontinuity)
+
+
+@pytest.mark.integration
+def test_rd_drinking():
+    df = (
+        cp.load_data("drinking")
+        .rename(columns={"agecell": "age"})
+        .assign(treated=lambda df_: df_.age > 21)
+        .dropna(axis=0)
+    )
+    result = cp.pymc_experiments.RegressionDiscontinuity(
+        df,
+        formula="all ~ 1 + age + treated",
+        running_variable_name="age",
+        prediction_model=cp.pymc_models.LinearRegression(),
+        treatment_threshold=21,
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.pymc_experiments.RegressionDiscontinuity)
+
+
+@pytest.mark.integration
+def test_its():
+    df = cp.load_data("its")
+    df["date"] = pd.to_datetime(df["date"])
+    df.set_index("date", inplace=True)
+    treatment_time = pd.to_datetime("2017-01-01")
+    result = cp.pymc_experiments.SyntheticControl(
+        df,
+        treatment_time,
+        formula="y ~ 1 + t + C(month)",
+        prediction_model=cp.pymc_models.LinearRegression(),
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.pymc_experiments.SyntheticControl)
+
+
+@pytest.mark.integration
+def test_its_covid():
+    df = cp.load_data("covid")
+    df["date"] = pd.to_datetime(df["date"])
+    df = df.set_index("date")
+    treatment_time = pd.to_datetime("2020-01-01")
+    result = cp.pymc_experiments.SyntheticControl(
+        df,
+        treatment_time,
+        formula="standardize(deaths) ~ 0 + standardize(t) + C(month) + standardize(temp)",  # noqa E501
+        prediction_model=cp.pymc_models.LinearRegression(),
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.pymc_experiments.SyntheticControl)
+
+
+@pytest.mark.integration
+def test_sc():
+    df = cp.load_data("sc")
+    treatment_time = 70
+    result = cp.pymc_experiments.SyntheticControl(
+        df,
+        treatment_time,
+        formula="actual ~ 0 + a + b + c + d + e + f + g",
+        prediction_model=cp.pymc_models.WeightedSumFitter(),
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.pymc_experiments.SyntheticControl)
+
+
+@pytest.mark.integration
+def test_sc_brexit():
+    df = cp.load_data("brexit")
+    df["Time"] = pd.to_datetime(df["Time"])
+    df.set_index("Time", inplace=True)
+    df = df.iloc[df.index > "2009", :]
+    treatment_time = pd.to_datetime("2016 June 24")
+    df = df.drop(["Japan", "Italy", "US", "Spain"], axis=1)
+    target_country = "UK"
+    all_countries = df.columns
+    other_countries = all_countries.difference({target_country})
+    all_countries = list(all_countries)
+    other_countries = list(other_countries)
+    formula = target_country + " ~ " + "0 + " + " + ".join(other_countries)
+    result = cp.pymc_experiments.SyntheticControl(
+        df,
+        treatment_time,
+        formula=formula,
+        prediction_model=cp.pymc_models.WeightedSumFitter(),
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.pymc_experiments.SyntheticControl)