Adjust notebooks and test code to allow for pytest testing

examples/Linear_Regression_Starter.ipynb

@@ -52,6 +52,7 @@
    ],
    "source": [
     "import numpy as np\n",
+    "from pathlib import Path\n",
     "\n",
     "import keras\n",
     "import bayesflow as bf"
@@ -987,7 +988,9 @@
    ],
    "source": [
     "# Recommended - full serialization (checkpoints folder must exist)\n",
-    "workflow.approximator.save(filepath=\"checkpoints/regression.keras\")\n",
+    "filepath = Path(\"checkpoints\") / \"regression.keras\"\n",
+    "filepath.parent.mkdir(exist_ok=True)\n",
+    "workflow.approximator.save(filepath=filepath)\n",
     "\n",
     "# Not recommended due to adapter mismatches - weights only\n",
     "# approximator.save_weights(filepath=\"checkpoints/regression.h5\")"
@@ -1016,7 +1019,7 @@
    ],
    "source": [
     "# Load approximator\n",
-    "approximator = keras.saving.load_model(\"checkpoints/regression.keras\")"
+    "approximator = keras.saving.load_model(filepath)"
    ]
   },
   {

examples/Lotka_Volterra_Point_Estimation_and_Expert_Stats.ipynb

@@ -37,6 +37,7 @@
    "source": [
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "from pathlib import Path\n",
     "import seaborn as sns\n",
     "\n",
     "import scipy\n",
@@ -748,7 +749,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "checkpoint_path = \"checkpoints/model.keras\"\n",
+    "checkpoint_path = Path(\"checkpoints\") / \"model.keras\"\n",
+    "checkpoint_path.parent.mkdir(exist_ok=True)\n",
     "keras.saving.save_model(point_inference_workflow.approximator, checkpoint_path)"
    ]
   },

tests/test_examples/test_examples.py

@@ -3,6 +3,7 @@
 from tests.utils import run_notebook
 
 
+@pytest.mark.skip(reason="requires setting up Stan")
 @pytest.mark.slow
 def test_bayesian_experimental_design(examples_path):
     run_notebook(examples_path / "Bayesian_Experimental_Design.ipynb")
@@ -30,7 +31,7 @@ def test_one_sample_ttest(examples_path):
 
 @pytest.mark.slow
 def test_sir_posterior_estimation(examples_path):
-    run_notebook(examples_path / "SIR_Posterior_estimation.ipynb")
+    run_notebook(examples_path / "SIR_Posterior_Estimation.ipynb")
 
 
 @pytest.mark.slow

tests/utils/jupyter.py

@@ -1,11 +1,27 @@
 import nbformat
 from nbconvert.preprocessors import ExecutePreprocessor
 
+from pathlib import Path
+import shutil
+
 
 def run_notebook(path):
+    path = Path(path)
+    checkpoint_path = path.parent / "checkpoints"
+    # only clean up if the directory did not exist before the test
+    cleanup_checkpoints = not checkpoint_path.exists()
     with open(str(path)) as f:
         nb = nbformat.read(f, nbformat.NO_CONVERT)
 
-    kernel = ExecutePreprocessor(timeout=600, kernel_name="python3")
+    kernel = ExecutePreprocessor(timeout=600, kernel_name="python3", resources={"metadata": {"path": path.parent}})
+
+    try:
+        result = kernel.preprocess(nb)
+    except Exception as e:
+        raise e
+    finally:
+        if cleanup_checkpoints and checkpoint_path.exists():
+            # clean up if the directory was created by the test
+            shutil.rmtree(checkpoint_path)
 
-    return kernel.preprocess(nb)
+    return result