make lots of tests pass

Author: drbenvincent

CONTRIBUTING.md

@@ -128,12 +128,12 @@ We recommend that your contribution complies with the following guidelines befor
     make doctest
     ```
 
-- Doctest can also be run directly via pytest, which can be helpful to run only specific tests during development. The following commands run all doctests, only doctests in the pymc_models module, and only the doctests for the `ModelBuilder` class in pymc_models:
+- Doctest can also be run directly via pytest, which can be helpful to run only specific tests during development. The following commands run all doctests, only doctests in the pymc_models module, and only the doctests for the `PyMCModel` class in pymc_models:
 
     ```bash
     pytest --doctest-modules causalpy/
     pytest --doctest-modules causalpy/pymc_models.py
-    pytest --doctest-modules causalpy/pmyc_models.py::causalpy.pymc_models.ModelBuilder
+    pytest --doctest-modules causalpy/pmyc_models.py::causalpy.pymc_models.PyMCModel
     ```
 
 - To indicate a work in progress please mark the PR as `draft`. Drafts may be useful to (1) indicate you are working on something to avoid duplicated work, (2) request broad review of functionality or API, or (3) seek collaborators.

README.md

@@ -44,7 +44,7 @@ df = (
     )
 
 # Run the analysis
-result = cp.pymc_experiments.RegressionDiscontinuity(
+result = cp.RegressionDiscontinuity(
     df,
     formula="all ~ 1 + age + treated",
     running_variable_name="age",

causalpy/pymc_experiments.py

@@ -55,7 +55,7 @@
 #         >>> import causalpy as cp
 #         >>> df = cp.load_data("did")
 #         >>> seed = 42
-#         >>> result = cp.pymc_experiments.DifferenceInDifferences(
+#         >>> result = cp.DifferenceInDifferences(
 #         ...     df,
 #         ...     formula="y ~ 1 + group*post_treatment",
 #         ...     time_variable_name="t",
@@ -304,7 +304,7 @@
 #     >>> import causalpy as cp
 #     >>> df = cp.load_data("did")
 #     >>> seed = 42
-#     >>> result = cp.pymc_experiments.DifferenceInDifferences(
+#     >>> result = cp.DifferenceInDifferences(
 #     ...     df,
 #     ...     formula="y ~ 1 + group*post_treatment",
 #     ...     time_variable_name="t",
@@ -595,7 +595,7 @@
 #     >>> import causalpy as cp
 #     >>> df = cp.load_data("rd")
 #     >>> seed = 42
-#     >>> result = cp.pymc_experiments.RegressionDiscontinuity(
+#     >>> result = cp.RegressionDiscontinuity(
 #     ...     df,
 #     ...     formula="y ~ 1 + x + treated + x:treated",
 #     ...     model=cp.pymc_models.LinearRegression(
@@ -1004,7 +1004,7 @@
 #     >>> import causalpy as cp
 #     >>> df = cp.load_data("anova1")
 #     >>> seed = 42
-#     >>> result = cp.pymc_experiments.PrePostNEGD(
+#     >>> result = cp.PrePostNEGD(
 #     ...     df,
 #     ...     formula="post ~ 1 + C(group) + pre",
 #     ...     group_variable_name="group",
@@ -1358,7 +1358,7 @@
 #     >>> import causalpy as cp
 #     >>> df = cp.load_data("nhefs")
 #     >>> seed = 42
-#     >>> result = cp.pymc_experiments.InversePropensityWeighting(
+#     >>> result = cp.InversePropensityWeighting(
 #     ...     df,
 #     ...     formula="trt ~ 1 + age + race",
 #     ...     outcome_variable ="outcome",

causalpy/skl_experiments.py

@@ -67,7 +67,7 @@
 #     >>> import causalpy as cp
 #     >>> df = cp.load_data("sc")
 #     >>> treatment_time = 70
-#     >>> result = cp.skl_experiments.PrePostFit(
+#     >>> result = cp.PrePostFit(
 #     ...     df,
 #     ...     treatment_time,
 #     ...     formula="actual ~ 0 + a + b + c + d + e + f + g",
@@ -255,7 +255,7 @@
 #     ...     .set_index("date")
 #     ... )
 #     >>> treatment_time = pd.to_datetime("2017-01-01")
-#     >>> result = cp.skl_experiments.InterruptedTimeSeries(
+#     >>> result = cp.InterruptedTimeSeries(
 #     ...     df,
 #     ...     treatment_time,
 #     ...     formula="y ~ 1 + t + C(month)",
@@ -285,7 +285,7 @@
 #     >>> import causalpy as cp
 #     >>> df = cp.load_data("sc")
 #     >>> treatment_time = 70
-#     >>> result = cp.skl_experiments.SyntheticControl(
+#     >>> result = cp.SyntheticControl(
 #     ...     df,
 #     ...     treatment_time,
 #     ...     formula="actual ~ 0 + a + b + c + d + e + f + g",
@@ -334,7 +334,7 @@
 #     >>> import causalpy as cp
 #     >>> from sklearn.linear_model import LinearRegression
 #     >>> df = cp.load_data("did")
-#     >>> result = cp.skl_experiments.DifferenceInDifferences(
+#     >>> result = cp.DifferenceInDifferences(
 #     ...     df,
 #     ...     formula="y ~ 1 + group*post_treatment",
 #     ...     time_variable_name="t",
@@ -558,7 +558,7 @@
 #     >>> import causalpy as cp
 #     >>> from sklearn.linear_model import LinearRegression
 #     >>> data = cp.load_data("rd")
-#     >>> result = cp.skl_experiments.RegressionDiscontinuity(
+#     >>> result = cp.RegressionDiscontinuity(
 #     ...     data,
 #     ...     formula="y ~ 1 + x + treated",
 #     ...     model=LinearRegression(),

causalpy/tests/test_input_validation.py

@@ -40,7 +40,7 @@ def test_did_validation_post_treatment_formula():
     )
 
     with pytest.raises(FormulaException):
-        _ = cp.pymc_experiments.DifferenceInDifferences(
+        _ = cp.DifferenceInDifferences(
             df,
             formula="y ~ 1 + group*post_SOMETHING",
             time_variable_name="t",
@@ -49,7 +49,7 @@ def test_did_validation_post_treatment_formula():
         )
 
     with pytest.raises(FormulaException):
-        _ = cp.skl_experiments.DifferenceInDifferences(
+        _ = cp.DifferenceInDifferences(
             df,
             formula="y ~ 1 + group*post_SOMETHING",
             time_variable_name="t",
@@ -73,7 +73,7 @@ def test_did_validation_post_treatment_data():
     )
 
     with pytest.raises(DataException):
-        _ = cp.pymc_experiments.DifferenceInDifferences(
+        _ = cp.DifferenceInDifferences(
             df,
             formula="y ~ 1 + group*post_treatment",
             time_variable_name="t",
@@ -82,7 +82,7 @@ def test_did_validation_post_treatment_data():
         )
 
     with pytest.raises(DataException):
-        _ = cp.skl_experiments.DifferenceInDifferences(
+        _ = cp.DifferenceInDifferences(
             df,
             formula="y ~ 1 + group*post_treatment",
             time_variable_name="t",
@@ -106,7 +106,7 @@ def test_did_validation_unit_data():
     )
 
     with pytest.raises(DataException):
-        _ = cp.pymc_experiments.DifferenceInDifferences(
+        _ = cp.DifferenceInDifferences(
             df,
             formula="y ~ 1 + group*post_treatment",
             time_variable_name="t",
@@ -115,7 +115,7 @@ def test_did_validation_unit_data():
         )
 
     with pytest.raises(DataException):
-        _ = cp.skl_experiments.DifferenceInDifferences(
+        _ = cp.DifferenceInDifferences(
             df,
             formula="y ~ 1 + group*post_treatment",
             time_variable_name="t",
@@ -139,7 +139,7 @@ def test_did_validation_group_dummy_coded():
     )
 
     with pytest.raises(DataException):
-        _ = cp.pymc_experiments.DifferenceInDifferences(
+        _ = cp.DifferenceInDifferences(
             df,
             formula="y ~ 1 + group*post_treatment",
             time_variable_name="t",
@@ -148,7 +148,7 @@ def test_did_validation_group_dummy_coded():
         )
 
     with pytest.raises(DataException):
-        _ = cp.skl_experiments.DifferenceInDifferences(
+        _ = cp.DifferenceInDifferences(
             df,
             formula="y ~ 1 + group*post_treatment",
             time_variable_name="t",
@@ -168,7 +168,7 @@ def test_sc_input_error():
     with pytest.raises(BadIndexException):
         df = cp.load_data("sc")
         treatment_time = pd.to_datetime("2016 June 24")
-        _ = cp.pymc_experiments.SyntheticControl(
+        _ = cp.SyntheticControl(
             df,
             treatment_time,
             formula="actual ~ 0 + a + b + c + d + e + f + g",
@@ -178,7 +178,7 @@ def test_sc_input_error():
     with pytest.raises(BadIndexException):
         df = cp.load_data("sc")
         treatment_time = pd.to_datetime("2016 June 24")
-        _ = cp.skl_experiments.SyntheticControl(
+        _ = cp.SyntheticControl(
             df,
             treatment_time,
             formula="actual ~ 0 + a + b + c + d + e + f + g",
@@ -202,7 +202,7 @@ def test_sc_brexit_input_error():
         all_countries = list(all_countries)
         other_countries = list(other_countries)
         formula = target_country + " ~ " + "0 + " + " + ".join(other_countries)
-        _ = cp.pymc_experiments.SyntheticControl(
+        _ = cp.SyntheticControl(
             df,
             treatment_time,
             formula=formula,
@@ -224,7 +224,7 @@ def test_ancova_validation_2_levels():
     )
 
     with pytest.raises(DataException):
-        _ = cp.pymc_experiments.PrePostNEGD(
+        _ = cp.PrePostNEGD(
             df,
             formula="post ~ 1 + C(group) + pre",
             group_variable_name="group",
@@ -247,7 +247,7 @@ def test_rd_validation_treated_in_formula():
     )
 
     with pytest.raises(FormulaException):
-        _ = cp.pymc_experiments.RegressionDiscontinuity(
+        _ = cp.RegressionDiscontinuity(
             df,
             formula="y ~ 1 + x",
             model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
@@ -257,7 +257,7 @@ def test_rd_validation_treated_in_formula():
     with pytest.raises(FormulaException):
         from sklearn.linear_model import LinearRegression
 
-        _ = cp.skl_experiments.RegressionDiscontinuity(
+        _ = cp.RegressionDiscontinuity(
             df,
             formula="y ~ 1 + x",
             model=LinearRegression(),
@@ -276,7 +276,7 @@ def test_rd_validation_treated_is_dummy():
     )
 
     with pytest.raises(DataException):
-        _ = cp.pymc_experiments.RegressionDiscontinuity(
+        _ = cp.RegressionDiscontinuity(
             df,
             formula="y ~ 1 + x + treated",
             model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
@@ -286,7 +286,7 @@ def test_rd_validation_treated_is_dummy():
     from sklearn.linear_model import LinearRegression
 
     with pytest.raises(DataException):
-        _ = cp.skl_experiments.RegressionDiscontinuity(
+        _ = cp.RegressionDiscontinuity(
             df,
             formula="y ~ 1 + x + treated",
             model=LinearRegression(),
@@ -302,7 +302,7 @@ def test_iv_treatment_var_is_present():
     instruments_data = pd.DataFrame({"z": [1, 3, 4], "w": [2, 3, 4]})
 
     with pytest.raises(DataException):
-        _ = cp.pymc_experiments.InstrumentalVariable(
+        _ = cp.InstrumentalVariable(
             instruments_data=instruments_data,
             data=data,
             instruments_formula=instruments_formula,
@@ -347,7 +347,7 @@ def test_rkink_bandwidth_check():
     with pytest.raises(ValueError):
         kink = 0.5
         df = setup_regression_kink_data(kink)
-        _ = cp.pymc_experiments.RegressionKink(
+        _ = cp.RegressionKink(
             df,
             formula=f"y ~ 1 + x + I((x-{kink})*treated)",
             model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
@@ -358,7 +358,7 @@ def test_rkink_bandwidth_check():
     with pytest.raises(ValueError):
         kink = 0.5
         df = setup_regression_kink_data(kink)
-        _ = cp.pymc_experiments.RegressionKink(
+        _ = cp.RegressionKink(
             df,
             formula=f"y ~ 1 + x + I((x-{kink})*treated)",
             model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
@@ -372,7 +372,7 @@ def test_rkink_epsilon_check():
     with pytest.raises(ValueError):
         kink = 0.5
         df = setup_regression_kink_data(kink)
-        _ = cp.pymc_experiments.RegressionKink(
+        _ = cp.RegressionKink(
             df,
             formula=f"y ~ 1 + x + I((x-{kink})*treated)",
             model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
@@ -383,7 +383,7 @@ def test_rkink_epsilon_check():
     with pytest.raises(ValueError):
         kink = 0.5
         df = setup_regression_kink_data(kink)
-        _ = cp.pymc_experiments.RegressionKink(
+        _ = cp.RegressionKink(
             df,
             formula=f"y ~ 1 + x + I((x-{kink})*treated)",
             model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),