sample_kwargs: speed up pymc integration tests + add related asserts

Author: drbenvincent

causalpy/tests/test_integration_pymc_examples.py

@@ -3,6 +3,8 @@
 
 import causalpy as cp
 
+sample_kwargs = {"tune": 20, "draws": 20, "chains": 2, "cores": 2}
+
 
 @pytest.mark.integration
 def test_did():
@@ -14,10 +16,18 @@ def test_did():
         group_variable_name="group",
         treated=1,
         untreated=0,
-        prediction_model=cp.pymc_models.LinearRegression(),
+        prediction_model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
     )
     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"]
+    )
 
 
 @pytest.mark.integration
@@ -47,10 +57,18 @@ def test_did_banks():
         group_variable_name="district",
         treated="Sixth District",
         untreated="Eighth District",
-        prediction_model=cp.pymc_models.LinearRegression(),
+        prediction_model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
     )
     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"]
+    )
 
 
 @pytest.mark.integration
@@ -59,11 +77,19 @@ def test_rd():
     result = cp.pymc_experiments.RegressionDiscontinuity(
         df,
         formula="y ~ 1 + bs(x, df=6) + treated",
-        prediction_model=cp.pymc_models.LinearRegression(),
+        prediction_model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
         treatment_threshold=0.5,
     )
     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"]
+    )
 
 
 @pytest.mark.integration
@@ -77,11 +103,19 @@ def test_rd_drinking():
         df,
         formula="all ~ 1 + age + treated",
         running_variable_name="age",
-        prediction_model=cp.pymc_models.LinearRegression(),
+        prediction_model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
         treatment_threshold=21,
     )
     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"]
+    )
 
 
 @pytest.mark.integration
@@ -94,10 +128,18 @@ def test_its():
         df,
         treatment_time,
         formula="y ~ 1 + t + C(month)",
-        prediction_model=cp.pymc_models.LinearRegression(),
+        prediction_model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
     )
     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"]
+    )
 
 
 @pytest.mark.integration
@@ -110,10 +152,18 @@ def test_its_covid():
         df,
         treatment_time,
         formula="standardize(deaths) ~ 0 + standardize(t) + C(month) + standardize(temp)",  # noqa E501
-        prediction_model=cp.pymc_models.LinearRegression(),
+        prediction_model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
     )
     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"]
+    )
 
 
 @pytest.mark.integration
@@ -124,10 +174,18 @@ def test_sc():
         df,
         treatment_time,
         formula="actual ~ 0 + a + b + c + d + e + f + g",
-        prediction_model=cp.pymc_models.WeightedSumFitter(),
+        prediction_model=cp.pymc_models.WeightedSumFitter(sample_kwargs=sample_kwargs),
     )
     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"]
+    )
 
 
 @pytest.mark.integration
@@ -148,7 +206,15 @@ def test_sc_brexit():
         df,
         treatment_time,
         formula=formula,
-        prediction_model=cp.pymc_models.WeightedSumFitter(),
+        prediction_model=cp.pymc_models.WeightedSumFitter(sample_kwargs=sample_kwargs),
     )
     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"]
+    )