Add automated notebook testing with Papermill

.github/workflows/test_notebook.yml

@@ -0,0 +1,48 @@
+name: Test Notebooks
+
+on:
+  pull_request:
+    branches: [main]
+    paths:
+      - "pyproject.toml"
+      - "causalpy/**"
+      - ".github/workflows/test_notebook.yml"
+      - "scripts/run_notebooks/**"
+      - "docs/source/notebooks/**"
+  push:
+    branches: [main]
+    paths:
+      - "pyproject.toml"
+      - "causalpy/**"
+      - ".github/workflows/test_notebook.yml"
+      - "scripts/run_notebooks/**"
+      - "docs/source/notebooks/**"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  notebooks:
+    runs-on: ubuntu-latest
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        split:
+          - "--pattern *_pymc*.ipynb"
+          - "--pattern *_skl*.ipynb"
+          - "--exclude-pattern _pymc --exclude-pattern _skl"
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install dependencies
+        run: |
+          pip install --upgrade pip
+          pip install -e ".[test,docs]"
+
+      - name: Run notebooks
+        run: python scripts/run_notebooks/runner.py ${{ matrix.split }}

pyproject.toml

@@ -91,7 +91,7 @@ docs = [
     "sphinx-togglebutton",
 ]
 lint = ["interrogate", "pre-commit", "ruff", "mypy"]
-test = ["pytest", "pytest-cov", "codespell", "nbformat", "nbconvert"]
+test = ["pytest", "pytest-cov", "codespell", "nbformat", "nbconvert", "papermill"]
 
 [project.urls]
 Homepage = "https://github.com/pymc-labs/CausalPy"

scripts/run_notebooks/README.md

@@ -0,0 +1,37 @@
+# Notebook Runner
+
+This script runs Jupyter notebooks from `docs/source/notebooks/` to validate they execute without errors.
+
+## How It Works
+
+1. **Mocks `pm.sample()`** — Replaces MCMC sampling with prior predictive (10 draws) for speed
+2. **Uses Papermill** — Executes notebooks programmatically
+3. **Discards outputs** — Only checks for errors, doesn't save results
+
+## Usage
+
+```bash
+# Run all notebooks
+python scripts/run_notebooks/runner.py
+
+# Run only PyMC notebooks
+python scripts/run_notebooks/runner.py --pattern "*_pymc*.ipynb"
+
+# Run only sklearn notebooks
+python scripts/run_notebooks/runner.py --pattern "*_skl*.ipynb"
+
+# Exclude PyMC and sklearn notebooks (run others)
+python scripts/run_notebooks/runner.py --exclude-pattern _pymc --exclude-pattern _skl
+```
+
+## CI Integration
+
+The GitHub Actions workflow (`.github/workflows/test_notebook.yml`) runs this script in parallel:
+- Job 1: PyMC notebooks
+- Job 2: Sklearn notebooks
+- Job 3: Other notebooks
+
+## Files
+
+- `runner.py` — Main script
+- `injected.py` — Code injected into notebooks to mock `pm.sample()`

scripts/run_notebooks/injected.py

@@ -0,0 +1,44 @@
+"""Injected code to mock pm.sample for faster notebook execution."""
+
+import numpy as np
+import pymc as pm
+import xarray as xr
+
+
+def mock_sample(*args, **kwargs):
+    """Mock pm.sample using prior predictive sampling for speed."""
+    random_seed = kwargs.get("random_seed")
+    model = kwargs.get("model")
+    samples = 10
+
+    idata = pm.sample_prior_predictive(
+        model=model,
+        random_seed=random_seed,
+        draws=samples,
+    )
+    idata.add_groups(posterior=idata.prior)
+
+    # Create mock sample stats with diverging data
+    if "sample_stats" not in idata:
+        n_chains = 1
+        n_draws = samples
+        sample_stats = xr.Dataset(
+            {
+                "diverging": xr.DataArray(
+                    np.zeros((n_chains, n_draws), dtype=int),
+                    dims=("chain", "draw"),
+                )
+            }
+        )
+        idata.add_groups(sample_stats=sample_stats)
+
+    del idata.prior
+    if "prior_predictive" in idata:
+        del idata.prior_predictive
+
+    return idata
+
+
+pm.sample = mock_sample
+pm.HalfFlat = pm.HalfNormal
+pm.Flat = pm.Normal

scripts/run_notebooks/runner.py

@@ -0,0 +1,135 @@
+"""Script to run notebooks in docs/source/notebooks directory.
+
+Examples
+--------
+Run all notebooks:
+
+    python scripts/run_notebooks/runner.py
+
+Run only PyMC notebooks:
+
+    python scripts/run_notebooks/runner.py --pattern "*_pymc*.ipynb"
+
+Run only sklearn notebooks:
+
+    python scripts/run_notebooks/runner.py --pattern "*_skl*.ipynb"
+
+Exclude PyMC and sklearn notebooks (run others):
+
+    python scripts/run_notebooks/runner.py --exclude-pattern _pymc --exclude-pattern _skl
+
+"""
+
+import argparse
+import logging
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+
+import papermill
+from nbformat.notebooknode import NotebookNode
+from papermill.iorw import load_notebook_node, write_ipynb
+
+HERE = Path(__file__).parent
+NOTEBOOKS_PATH = Path("docs/source/notebooks")
+KERNEL_NAME = "python3"
+
+INJECTED_CODE_FILE = HERE / "injected.py"
+INJECTED_CODE = INJECTED_CODE_FILE.read_text()
+
+
+def setup_logging() -> None:
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s",
+    )
+
+
+def inject_mock_code(cells: list) -> None:
+    """Inject mock pm.sample code at the start of the notebook."""
+    cells.insert(
+        0,
+        NotebookNode(
+            id="mock-injection",
+            execution_count=0,
+            cell_type="code",
+            metadata={"tags": []},
+            outputs=[],
+            source=INJECTED_CODE,
+        ),
+    )
+
+
+def run_notebook(notebook_path: Path) -> None:
+    """Run a notebook with mocked pm.sample."""
+    logging.info(f"Running notebook: {notebook_path.name}")
+
+    nb = load_notebook_node(str(notebook_path))
+    inject_mock_code(nb.cells)
+
+    with NamedTemporaryFile(suffix=".ipynb", delete=False) as f:
+        write_ipynb(nb, f.name)
+        try:
+            papermill.execute_notebook(
+                input_path=f.name,
+                output_path=None,  # Discard output
+                kernel_name=KERNEL_NAME,
+                progress_bar=True,
+                cwd=notebook_path.parent,
+            )
+        except Exception as e:
+            logging.error(f"Error running notebook: {notebook_path.name}")
+            raise e
+
+
+def get_notebooks(
+    pattern: str | None = None,
+    exclude_patterns: list[str] | None = None,
+) -> list[Path]:
+    """Get list of notebooks to run, optionally filtered."""
+    notebooks = list(NOTEBOOKS_PATH.glob("*.ipynb"))
+
+    if pattern:
+        notebooks = [nb for nb in notebooks if Path(nb).match(pattern)]
+
+    if exclude_patterns:
+        for exc in exclude_patterns:
+            notebooks = [nb for nb in notebooks if exc not in nb.name]
+
+    return sorted(notebooks)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run CausalPy notebooks.")
+    parser.add_argument(
+        "--pattern",
+        type=str,
+        default=None,
+        help="Glob pattern to filter notebooks (e.g., '*_pymc*.ipynb')",
+    )
+    parser.add_argument(
+        "--exclude-pattern",
+        type=str,
+        action="append",
+        dest="exclude_patterns",
+        help="Pattern to exclude from notebook names (can be used multiple times)",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    setup_logging()
+    args = parse_args()
+
+    notebooks = get_notebooks(
+        pattern=args.pattern,
+        exclude_patterns=args.exclude_patterns,
+    )
+
+    logging.info(f"Found {len(notebooks)} notebooks to run")
+    for nb in notebooks:
+        logging.info(f"  - {nb.name}")
+
+    for notebook in notebooks:
+        run_notebook(notebook)
+
+    logging.info("All notebooks completed successfully!")