feat: add notebook quality control pipeline with nbmake, auto-save hooks, and plot extraction

DEVELOPMENT.md

@@ -67,6 +67,7 @@ All of the following MUST pass before merging to `main`:
 
 ```bash
 pytest                           # all tests pass
+nbmake notebooks/                # execute and validate notebooks
 black --check .                  # code formatting
 ruff check .                     # linting
 mypy .                           # type checking (strict mode)
@@ -74,6 +75,38 @@ mypy .                           # type checking (strict mode)
 
 These are configured in `pyproject.toml`. Run them individually to identify failures.
 
+### Notebook Quality Control
+
+Notebooks in `notebooks/` are executed as pytest tests using `nbmake`. This ensures:
+
+- All code cells execute without errors
+- Full error traces are available for debugging
+- Plots are extracted for LLM-based visual analysis
+
+**Workflow:**
+
+1. Write notebooks with explicit figure saving:
+   ```python
+   fig.savefig('_plots/notebook_name_cell{N}.png', dpi=100, bbox_inches='tight')
+   from IPython.display import display, Image
+   display(Image(filename='_plots/notebook_name_cell{N}.png'))
+   ```
+
+2. Execute and extract plots:
+   ```bash
+   python scripts/run_notebooks.py notebooks/
+   ```
+
+3. Send extracted plots to LLM for visual review:
+   ```
+   @notebooks/_plots/notebook_name_cell4.png
+   ```
+
+4. Run full quality gates:
+   ```bash
+   pytest notebooks/ --nbmake  # or just `pytest` (includes notebooks/)
+   ```
+
 ## 5. Agent Instructions
 
 When working on this project, agents must:

notebooks/conftest.py

@@ -0,0 +1,75 @@
+# conftest.py — shared fixtures for notebooks
+import matplotlib
+matplotlib.use("Agg")  # non-interactive backend for CI
+
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+import pytest
+
+# Auto-save figures directory
+PLOTS_DIR = os.path.join(os.path.dirname(__file__), "_plots")
+os.makedirs(PLOTS_DIR, exist_ok=True)
+_plot_counter = 0
+
+
+def _save_figure_on_close(fig):
+    """Save figure to PNG before closing."""
+    global _plot_counter
+    plot_path = os.path.join(PLOTS_DIR, f"notebook_{_plot_counter:03d}.png")
+    try:
+        fig.savefig(plot_path, dpi=100, bbox_inches="tight")
+        _plot_counter += 1
+    except Exception:
+        pass
+
+
+# Patch plt.close to auto-save before closing
+_original_close = plt.close
+
+
+def _patched_close(fig=None):
+    if fig is not None:
+        _save_figure_on_close(fig)
+    return _original_close(fig)
+
+
+plt.close = _patched_close
+plt.rcParams.update({"figure.autolayout": True})
+
+
+@pytest.fixture(autouse=True)
+def close_figures():
+    """Close all matplotlib figures after each notebook cell execution."""
+    yield
+    plt.close("all")
+
+
+@pytest.fixture()
+def W_low():
+    """Low-energy total energy: T ≈ 10 keV → W = 1 + 0.01/0.511 ≈ 1.02."""
+    return np.array([1.02])
+
+
+@pytest.fixture()
+def W_mid():
+    """Mid-range energies (bulk of spectrum)."""
+    return np.linspace(1.1, 2.0, 5)
+
+
+@pytest.fixture()
+def W_high():
+    """High-energy total energy near endpoint: T ≈ 3 MeV → W ≈ 7."""
+    return np.array([6.8])
+
+
+@pytest.fixture()
+def W_full():
+    """Full spectrum grid from ~1 keV to just below endpoint."""
+    return np.linspace(1.003, 5.8, 100)
+
+
+@pytest.fixture()
+def large_W():
+    """Large array for performance/shape checks."""
+    return np.linspace(1.002, 4.9, 1000)

pyproject.toml

@@ -24,6 +24,7 @@ where = ["."]
 [project.optional-dependencies]
 dev = [
     "pytest",
+    "nbmake",
     "black",
     "ruff",
     "mypy",
@@ -41,7 +42,7 @@ strict = true
 line-length = 88
 
 [tool.pytest.ini_options]
-testpaths = ["tests"]
+testpaths = ["tests", "notebooks"]
 
 [tool.ruff]
 line-length = 88

scripts/notebook_plots.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+"""Extract plots from executed notebooks and save as PNGs for LLM review.
+
+Usage:
+    python scripts/notebook_plots.py notebooks/smoke_test.ipynb
+    python scripts/notebook_plots.py notebooks/          # all notebooks
+
+Outputs PNG files to notebooks/_plots/ directory.
+Each plot gets a descriptive filename based on notebook name and cell index.
+"""
+
+import json
+import os
+import sys
+from pathlib import Path
+
+
+def extract_plots(notebook_path: str, output_dir: str = None) -> list:
+    """Extract all matplotlib figures from an executed notebook.
+
+    Looks for PNG image data in cell outputs (standard Jupyter output format).
+    """
+    nb_path = Path(notebook_path)
+    with open(nb_path) as f:
+        nb = json.load(f)
+
+    if output_dir is None:
+        output_dir = nb_path.parent / "_plots"
+    output_dir = Path(output_dir)
+    output_dir.mkdir(exist_ok=True)
+
+    plots = []
+    nb_name = nb_path.stem
+
+    for cell_idx, cell in enumerate(nb["cells"]):
+        if cell["cell_type"] != "code":
+            continue
+        for output in cell.get("outputs", []):
+            if output.get("output_type") != "display_data":
+                continue
+            # Jupyter stores images as base64 in data["image/png"]
+            data = output.get("data", {})
+            png_b64 = data.get("image/png")
+            if png_b64:
+                plot_path = output_dir / f"{nb_name}_cell{cell_idx}.png"
+                import base64
+
+                with open(plot_path, "wb") as f:
+                    f.write(base64.b64decode(png_b64))
+                plots.append(str(plot_path))
+
+    return plots
+
+
+def main():
+    if len(sys.argv) < 2:
+        print("Usage: python scripts/notebook_plots.py <notebook_or_dir> ...")
+        print("  Pass a .ipynb file or a directory (processes all .ipynb files)")
+        sys.exit(1)
+
+    all_plots = []
+    for arg in sys.argv[1:]:
+        p = Path(arg)
+        if p.suffix == ".ipynb":
+            all_plots.extend(extract_plots(str(p)))
+        elif p.is_dir():
+            for nb in sorted(p.glob("*.ipynb")):
+                if nb.name == "conftest.py":
+                    continue
+                all_plots.extend(extract_plots(str(nb)))
+
+    if not all_plots:
+        print("No plots found.")
+        sys.exit(0)
+
+    print(f"Extracted {len(all_plots)} plot(s):")
+    for p in all_plots:
+        print(f"  {p}")
+    print()
+    print("To review with LLM, send the PNG paths using @ notation:")
+    print("  @ " + " @ ".join(all_plots))
+
+
+if __name__ == "__main__":
+    main()

scripts/run_notebooks.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python3
+"""Execute notebooks and extract plots for LLM review.
+
+This script:
+1. Executes all notebooks (or specified ones) using nbconvert
+2. Saves executed notebooks with outputs
+3. Extracts PNG figures from cell outputs
+4. Reports plot locations for LLM analysis
+
+Usage:
+    python scripts/run_notebooks.py              # all notebooks
+    python scripts/run_notebooks.py nb1.ipynb    # specific notebook
+    python scripts/run_notebooks.py notebooks/   # directory
+"""
+
+import json
+import os
+import sys
+import base64
+from pathlib import Path
+
+
+def execute_notebook(nb_path: str) -> bool:
+    """Execute a notebook using nbconvert and save outputs."""
+    import nbformat
+    from nbclient import NotebookClient
+
+    nb = nbformat.read(str(nb_path), as_version=4)
+    client = NotebookClient(nb, timeout=600, kernel_name="python3")
+    try:
+        client.execute()
+        nbformat.write(nb, str(nb_path))
+        return True
+    except Exception as e:
+        print(f"  ERROR: {e}")
+        return False
+
+
+def extract_plots(nb_path: str, output_dir: str = None) -> list:
+    """Extract PNG figures from an executed notebook's cell outputs."""
+    nb_path = Path(nb_path)
+    with open(nb_path) as f:
+        nb = json.load(f)
+
+    if output_dir is None:
+        output_dir = nb_path.parent / "_plots"
+    output_dir = Path(output_dir)
+    output_dir.mkdir(exist_ok=True)
+
+    plots = []
+    nb_name = nb_path.stem
+
+    for cell_idx, cell in enumerate(nb["cells"]):
+        if cell["cell_type"] != "code":
+            continue
+        for output in cell.get("outputs", []):
+            if output.get("output_type") != "display_data":
+                continue
+            data = output.get("data", {})
+            png_b64 = data.get("image/png")
+            if png_b64:
+                plot_path = output_dir / f"{nb_name}_cell{cell_idx}.png"
+                with open(plot_path, "wb") as f:
+                    f.write(base64.b64decode(png_b64))
+                plots.append(str(plot_path))
+
+    return plots
+
+
+def main():
+    if len(sys.argv) < 2:
+        print("Usage: python scripts/run_notebooks.py <notebook_or_dir> ...")
+        sys.exit(1)
+
+    all_plots = []
+    failed = []
+
+    for arg in sys.argv[1:]:
+        p = Path(arg)
+        if p.suffix == ".ipynb":
+            print(f"Executing {p.name}...")
+            if execute_notebook(str(p)):
+                plots = extract_plots(str(p))
+                all_plots.extend(plots)
+                print(f"  Extracted {len(plots)} plot(s)")
+            else:
+                failed.append(p.name)
+        elif p.is_dir():
+            for nb in sorted(p.glob("*.ipynb")):
+                if nb.name.startswith("_") or nb.name == "conftest.py":
+                    continue
+                print(f"Executing {nb.name}...")
+                if execute_notebook(str(nb)):
+                    plots = extract_plots(str(nb))
+                    all_plots.extend(plots)
+                    print(f"  Extracted {len(plots)} plot(s)")
+                else:
+                    failed.append(nb.name)
+
+    print()
+    print(f"Done. {len(all_plots)} plot(s) extracted, {len(failed)} notebook(s) failed.")
+    if all_plots:
+        print()
+        print("Plots saved to: notebooks/_plots/")
+        print("To review with LLM, send the PNG paths using @ notation:")
+        for p in all_plots:
+            print(f"  @{p}")
+
+
+if __name__ == "__main__":
+    main()