Adding the abstract class

Author: JeanVanDyk

.github/workflows/ci.yml

@@ -24,7 +24,7 @@ jobs:
       - name: Setup environment
         run: pip install -e .[test]
       - name: Run doctests
-        run: pytest --doctest-modules --ignore=causalpy/tests/ causalpy/
+        run: pytest --doctest-modules --ignore=causalpy/tests/ causalpy/ --config-file=causalpy/tests/conftest.py
       - name: Run extra tests
         run: pytest docs/source/.codespell/test_notebook_to_markdown.py
       - name: Run tests

.pre-commit-config.yaml

@@ -25,7 +25,7 @@ repos:
         exclude: &exclude_pattern 'iv_weak_instruments.ipynb'
         args: ["--maxkb=1500"]
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.13
+    rev: v0.12.1
     hooks:
       # Run the linter
       - id: ruff

Makefile

@@ -1,4 +1,4 @@
-.PHONY: init lint check_lint test uml html cleandocs
+.PHONY: init lint check_lint test uml html cleandocs doctest
 
 init:
 	python -m pip install -e . --no-deps
@@ -13,7 +13,7 @@ check_lint:
 	interrogate .
 
 doctest:
-	pytest --doctest-modules --ignore=causalpy/tests/ causalpy/
+	pytest --doctest-modules --ignore=causalpy/tests/ causalpy/ --config-file=causalpy/tests/conftest.py
 
 test:
 	pytest

causalpy/experiments/interrupted_time_series.py

@@ -15,6 +15,7 @@
 Interrupted Time Series Analysis
 """
 
+from abc import ABC, abstractmethod
 from typing import List, Union
 
 import arviz as az
@@ -34,7 +35,35 @@
 LEGEND_FONT_SIZE = 12
 
 
-class UnknownTreatmentTimeHandler:
+class TreatmentTimeHandler(ABC):
+    @abstractmethod
+    def data_preprocessing(self, data, treatment_time, model):
+        pass
+
+    @abstractmethod
+    def data_postprocessing(
+        self, model, data, idata, treatment_time, y, X, pre_y, pre_X
+    ):
+        pass
+
+    @abstractmethod
+    def plot_intervention_line(
+        self, ax, handles, labels, datapre, datapost, pre_pred, post_pred
+    ):
+        pass
+
+    @abstractmethod
+    def plot_impact_cumulative(self, ax, datapre, datapost, post_impact_cumulative):
+        pass
+
+    def plot_treated_counterfactual(
+        self, ax, handles, labels, datapre, datapost, pre_pred, post_pred
+    ):
+        """Optional: override if needed"""
+        pass
+
+
+class UnknownTreatmentTimeHandler(TreatmentTimeHandler):
     """
     A utility class for managing data preprocessing, postprocessing,
     and plotting steps for models that infer unknown treatment times.
@@ -185,7 +214,7 @@ def plot_intervention_line(
             )
 
 
-class KnownTreatmentTimeHandler:
+class KnownTreatmentTimeHandler(TreatmentTimeHandler):
     """
     Handles data preprocessing, postprocessing, and plotting logic for models
     where the treatment time is known in advance.
@@ -247,14 +276,6 @@ def data_postprocessing(
 
         return res
 
-    def plot_treated_counterfactual(
-        self, sax, handles, labels, datapre, datapost, pre_pred, post_pred
-    ):
-        """
-        Placeholder method to maintain interface compatibility with UnknownTreatmentTimeHandler.
-        """
-        pass
-
     def plot_impact_cumulative(self, ax, datapre, datapost, post_impact_cumulative):
         """
         Plot the cumulative causal impact for the post-intervention period.

causalpy/experiments/prepostnegd.py

@@ -72,7 +72,7 @@ class PrePostNEGD(BaseExperiment):
     ...         }
     ...     ),
     ... )
-    >>> result.summary(round_to=1)  # doctest: +NUMBER
+    >>> result.summary(round_to=1)  # doctest: +SKIP
     ==================Pretest/posttest Nonequivalent Group Design===================
     Formula: post ~ 1 + C(group) + pre
     <BLANKLINE>

causalpy/tests/conftest.py

@@ -20,10 +20,24 @@
 
 import numpy as np
 import pytest
+from pymc.testing import mock_sample, mock_sample_setup_and_teardown
 
 
 @pytest.fixture(scope="session")
 def rng() -> np.random.Generator:
     """Random number generator that can persist through a pytest session"""
     seed: int = sum(map(ord, "causalpy"))
     return np.random.default_rng(seed=seed)
+
+
+mock_pymc_sample = pytest.fixture(mock_sample_setup_and_teardown, scope="session")
+
+
+@pytest.fixture(autouse=True)
+def mock_sample_for_doctest(request):
+    if not request.config.getoption("--doctest-modules", default=False):
+        return
+
+    import pymc as pm
+
+    pm.sample = mock_sample

causalpy/tests/test_integration_pymc_examples.py

@@ -23,7 +23,7 @@
 
 
 @pytest.mark.integration
-def test_did():
+def test_did(mock_pymc_sample):
     """
     Test Difference in Differences (DID) PyMC experiment.
 
@@ -57,7 +57,7 @@ def test_did():
 
 
 @pytest.mark.integration
-def test_did_banks_simple():
+def test_did_banks_simple(mock_pymc_sample):
     """
     Test simple Differences In Differences Experiment on the 'banks' data set.
 
@@ -113,7 +113,7 @@ def test_did_banks_simple():
 
 
 @pytest.mark.integration
-def test_did_banks_multi():
+def test_did_banks_multi(mock_pymc_sample):
     """
     Test multiple regression Differences In Differences Experiment on the 'banks'
     data set.
@@ -168,7 +168,7 @@ def test_did_banks_multi():
 
 
 @pytest.mark.integration
-def test_rd():
+def test_rd(mock_pymc_sample):
     """
     Test Regression Discontinuity experiment.
 
@@ -199,7 +199,7 @@ def test_rd():
 
 
 @pytest.mark.integration
-def test_rd_bandwidth():
+def test_rd_bandwidth(mock_pymc_sample):
     """
     Test Regression Discontinuity experiment with bandwidth parameter.
 
@@ -229,7 +229,7 @@ def test_rd_bandwidth():
 
 
 @pytest.mark.integration
-def test_rd_drinking():
+def test_rd_drinking(mock_pymc_sample):
     """
     Test Regression Discontinuity experiment on drinking age data.
 
@@ -289,7 +289,7 @@ def reg_kink_function(x, beta, kink):
 
 
 @pytest.mark.integration
-def test_rkink():
+def test_rkink(mock_pymc_sample):
     """
     Test Regression Kink design.
 
@@ -320,7 +320,7 @@ def test_rkink():
 
 
 @pytest.mark.integration
-def test_rkink_bandwidth():
+def test_rkink_bandwidth(mock_pymc_sample):
     """
     Test Regression Kink experiment with bandwidth parameter.
 
@@ -350,7 +350,7 @@ def test_rkink_bandwidth():
 
 
 @pytest.mark.integration
-def test_its():
+def test_its(mock_pymc_sample):
     """
     Test Interrupted Time-Series experiment.
 
@@ -461,7 +461,7 @@ def test_its_no_treatment_time():
 
 
 @pytest.mark.integration
-def test_its_covid():
+def test_its_covid(mock_pymc_sample):
     """
     Test Interrupted Time-Series experiment on COVID data.
 
@@ -515,7 +515,7 @@ def test_its_covid():
 
 
 @pytest.mark.integration
-def test_sc():
+def test_sc(mock_pymc_sample):
     """
     Test Synthetic Control experiment.
 
@@ -574,7 +574,7 @@ def test_sc():
 
 
 @pytest.mark.integration
-def test_sc_brexit():
+def test_sc_brexit(mock_pymc_sample):
     """
     Test Synthetic Control experiment on Brexit data.
 
@@ -637,7 +637,7 @@ def test_sc_brexit():
 
 
 @pytest.mark.integration
-def test_ancova():
+def test_ancova(mock_pymc_sample):
     """
     Test Pre-PostNEGD experiment on anova1 data.
 
@@ -669,7 +669,7 @@ def test_ancova():
 
 
 @pytest.mark.integration
-def test_geolift1():
+def test_geolift1(mock_pymc_sample):
     """
     Test Synthetic Control experiment on geo lift data.
 
@@ -706,7 +706,7 @@ def test_geolift1():
 
 
 @pytest.mark.integration
-def test_iv_reg():
+def test_iv_reg(mock_pymc_sample):
     df = cp.load_data("risk")
     instruments_formula = "risk  ~ 1 + logmort0"
     formula = "loggdp ~  1 + risk"
@@ -734,7 +734,7 @@ def test_iv_reg():
 
 
 @pytest.mark.integration
-def test_inverse_prop():
+def test_inverse_prop(mock_pymc_sample):
     """Test the InversePropensityWeighting class."""
     df = cp.load_data("nhefs")
     sample_kwargs = {

causalpy/tests/test_pymc_models.py

@@ -93,7 +93,7 @@ def test_fit_build_not_implemented(self):
         argvalues=[None, {"a": 1}],
         ids=["None-coords", "dict-coords"],
     )
-    def test_fit_predict(self, coords, rng) -> None:
+    def test_fit_predict(self, coords, rng, mock_pymc_sample) -> None:
         """
         Test fit and predict methods on MyToyModel.
 
@@ -122,7 +122,7 @@ def test_fit_predict(self, coords, rng) -> None:
         assert isinstance(predictions, az.InferenceData)
 
 
-def test_idata_property():
+def test_idata_property(mock_pymc_sample):
     """Test that we can access the idata property of the model"""
     df = cp.load_data("did")
     result = cp.DifferenceInDifferences(
@@ -140,7 +140,7 @@ def test_idata_property():
 
 
 @pytest.mark.parametrize("seed", seeds)
-def test_result_reproducibility(seed):
+def test_result_reproducibility(seed, mock_pymc_sample):
     """Test that we can reproduce the results from the model. We could in theory test
     this with all the model and experiment types, but what is being targeted is
     the PyMCModel.fit method, so we should be safe testing with just one model. Here

codecov.yml

@@ -1,4 +1,7 @@
 coverage:
+  ignore:
+    - "*/conftest.py"
+    - "*/tests/conftest.py"
   status:
     project:
       default: