adding notebook tests

.github/workflows/test_notebooks.yml

@@ -0,0 +1,46 @@
+name: Test notebooks
+
+on:
+    schedule:
+      - cron: '0 0 * * 6'  # Run at midnight on Saturdays
+    push:
+      branches:
+        - master
+    pull_request:
+      branches:
+        - master
+    workflow_dispatch:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.8]
+
+    steps:
+        - uses: actions/checkout@v2
+
+        - name: pip cache
+          uses: actions/cache@v2
+          with:
+            path: ~/.cache/pip
+            key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements.txt') }}
+
+
+        - name: Set up Python ${{ matrix.python-version }}
+          uses: actions/setup-python@v1
+          with:
+            python-version: ${{ matrix.python-version }}
+
+        - name: Setup Graphviz
+          uses: ts-graphviz/setup-graphviz@v1
+
+        - name: Install Python packages from requirements.txt
+          run: |
+            pip install --upgrade pip
+            pip install -e .[test]
+
+        - name: Run Notebook Test
+          run: | 
+           ./scripts/test_notebooks.sh

.gitignore

@@ -81,7 +81,9 @@ venv/
 /.vscode
 
 # Automatically generated files
+docs/source/lightning_logs
 docs/preconvert
 docs/build
 site/
 out
+*.matrix.gz

docs/source/tutorial_i.ipynb

@@ -88,9 +88,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import os\n",
     "import pyro\n",
     "from pyro.distributions import Bernoulli\n",
     "\n",
+    "smoke_test = ('CI' in os.environ)\n",
+    "num_samples = 1000 if not smoke_test else 10\n",
+    "num_iterations = 1000 if not smoke_test else 10\n",
+    "n_individuals = 10000 if not smoke_test else 10\n",
+    "\n",
     "def causal_model(stress_pt, smokes_cpt, cancer_cpt):\n",
     "    stress = pyro.sample(\"stress\", Bernoulli(stress_pt))\n",
     "    smokes = pyro.sample(\"smokes\", Bernoulli(smokes_cpt[stress.long()]))\n",
@@ -223,7 +229,7 @@
     "    with pyro.plate(\"individuals\", n_individuals, dim=-1):\n",
     "        return causal_model(stress_pt, smokes_cpt, cancer_cpt)\n",
     "\n",
-    "n_individuals = 10000\n",
+    "\n",
     "\n",
     "smokes_obs = population_causal_model(n_individuals, stress_pt, smokes_cpt, cancer_cpt)\n",
     "# stress_obs, smokes_obs, cancer_obs = population_causal_model(n_individuals, stress_pt, smokes_cpt, cancer_cpt)\n",
@@ -455,7 +461,7 @@
     }
    ],
    "source": [
-    "predictive_int_fixed = pyro.infer.Predictive(population_forced_smokes_model, num_samples=1000, return_sites=(\"cancer\", \"_RETURN\"))\n",
+    "predictive_int_fixed = pyro.infer.Predictive(population_forced_smokes_model, num_samples=num_samples, return_sites=(\"cancer\", \"_RETURN\"))\n",
     "prediction_int_fixed = predictive_int_fixed(n_individuals, stress_pt, smokes_cpt, cancer_cpt)\n",
     "\n",
     "visualize_proportions_single(prediction_int_fixed, \"Intervened Data - Fixed Parameters\")"
@@ -846,7 +852,7 @@
     }
    ],
    "source": [
-    "predictive_bayesian = pyro.infer.Predictive(bayesian_population_causal_model, num_samples=1000, return_sites=(\"cancer\", \"_RETURN\"))\n",
+    "predictive_bayesian = pyro.infer.Predictive(bayesian_population_causal_model, num_samples=num_samples, return_sites=(\"cancer\", \"_RETURN\"))\n",
     "prediction_bayesian = predictive_bayesian(n_individuals)\n",
     "visualize_proportions_multiple(prediction_bayesian, \"Observational Data - Uncertain Parameters\")"
    ]
@@ -914,7 +920,7 @@
     }
    ],
    "source": [
-    "predictive_int_bayesian = pyro.infer.Predictive(bayesian_population_forced_smokes_model, num_samples=1000, return_sites=(\"cancer\", \"_RETURN\"))\n",
+    "predictive_int_bayesian = pyro.infer.Predictive(bayesian_population_forced_smokes_model, num_samples=num_samples, return_sites=(\"cancer\", \"_RETURN\"))\n",
     "prediction_int_bayesian = predictive_int_bayesian(n_individuals)\n",
     "visualize_proportions_multiple(prediction_int_bayesian, \"Interventional Data - Uncertain Parameters\")"
    ]
@@ -1710,8 +1716,6 @@
     "\n",
     "svi = SVI(twin_model_conditioned, guide, adam, loss=elbo)\n",
     "\n",
-    "num_iterations = 1000\n",
-    "\n",
     "for j in range(num_iterations):\n",
     "    # calculate the loss and take a gradient step\n",
     "    loss = svi.step(n_individuals)\n",

scripts/test_notebooks.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+INCLUDED_NOTEBOOKS="docs/source/tutorial_i.ipynb docs/source/backdoor.ipynb docs/source/dr_learner.ipynb docs/source/mediation.ipynb docs/source/sdid.ipynb docs/source/slc.ipynb"
+
+CI=1 pytest --nbval-lax --dist loadscope -n auto $INCLUDED_NOTEBOOKS

setup.py

@@ -59,6 +59,7 @@
             "sphinx_rtd_theme==1.3.0",
             "myst_parser",
             "nbsphinx",
+            "nbval",
         ],
     },
     python_requires=">=3.8",