Increase prior predictive draws to 500 in notebook runner

Updated the mock for pm.sample to use 500 draws instead of 100 to ensure compatibility with notebook code that iterates over posterior samples, such as plot_ate which defaults to 500 draws. Adjusted documentation and injected.py accordingly.

Author: drbenvincent

scripts/run_notebooks/README.md

@@ -4,7 +4,7 @@ This script runs Jupyter notebooks from `docs/source/notebooks/` to validate the
 
 ## How It Works
 
-1. **Mocks `pm.sample()`** — Replaces MCMC sampling with prior predictive (1 chain × 100 draws) for speed
+1. **Mocks `pm.sample()`** — Replaces MCMC sampling with prior predictive (1 chain × 500 draws) for speed
 2. **Uses Papermill** — Executes notebooks programmatically
 3. **Discards outputs** — Only checks for errors, doesn't save results
 

scripts/run_notebooks/injected.py

@@ -4,6 +4,10 @@
 import pymc as pm
 import xarray as xr
 
+# Minimum draws needed to satisfy notebook code that iterates over posterior samples
+# (e.g., plot_ate uses ate_draws=500 by default)
+MIN_DRAWS = 500
+
 
 def mock_sample(*args, **kwargs):
     """Mock pm.sample using prior predictive sampling for speed."""
@@ -15,7 +19,10 @@ def mock_sample(*args, **kwargs):
         first_arg = args[0]
         if isinstance(first_arg, pm.Model):
             model = first_arg
-    n_draws = 100
+
+    # Always use MIN_DRAWS to ensure compatibility with code that iterates
+    # over posterior samples (ignoring the requested draws for speed)
+    n_draws = MIN_DRAWS
 
     idata = pm.sample_prior_predictive(
         model=model,