jejjohnson / jejjohnson/pyrox

docs(notebooks): add bayesian_neural_fields.ipynb flagship demo with calibration plot

Open
#73 0 comments 0 reactions 0 assignees View on GitHub
area:docs area:nn area:testing type:docs wave:4-structured
Dominant language
Python
Stars
1
Forks
0
Avg merge
19h 2m
Merged PRs (30d)
16

Description

## Problem / Request

Ship the canonical **Bayesian Neural Field demo notebook** — `docs/notebooks/bayesian_neural_fields.ipynb` — that exercises the full `BNFEstimator` stack end-to-end on a synthetic spatiotemporal problem and serves as the package's flagship "Bayesian NN with calibrated uncertainty" example. This is the last piece of the BNF port; it ties together [#70 (ensemble runner)](https://github.com/jejjohnson/pyrox/issues/70), [#71 (Estimator facade)](https://github.com/jejjohnson/pyrox/issues/71), [#72 (BNF layers)](https://github.com/jejjohnson/pyrox/issues/72), and [gaussx#121 (mixture quantile)](https://github.com/jejjohnson/gaussx/issues/121) into a single pedagogical artifact.

## User Story

As a prospective user evaluating `pyrox`, I want to open a single docs page — "Bayesian Neural Fields" — and see:

1. A short motivation + math framing (why BNF, where it sits between GPs and vanilla MLPs).
2. A self-contained executable cell chain: generate synthetic `(t, x, y)` data from a known GP truth → `BNFEstimator(...).fit(df, seed=0)` → `.predict(df_test, quantiles=(0.025, 0.5, 0.975))` → plot posterior mean + 95% band against the held-out truth grid.
3. A calibration plot that shows the 95% interval covers ~95% of held-out points.
4. A reproducibility line (`%watermark -p pyrox,gaussx,numpyro,jax,matplotlib`).

And I want to be able to click "Open in Colab" and run the whole thing with zero local setup.

## Motivation

- **Docs debt.** The BNF port (#70–#72) adds a lot of public surface; without a flagship notebook the package looks like a loose collection of primitives.
- **Integration test.** The notebook doubles as a "does the full stack still compose?" integration gate — every PR that touches the inference runner, the estimator facade, or the BNF layers breaks or passes this notebook's execution under `jupyter nbconvert --execute`.
- **Reference for future notebooks.** The BNF demo establishes the calibration-plot / held-out-band / Colab-badge pattern that every subsequent `pyrox` notebook should follow.

## Proposed Structure

Following [`.github/instructions/docs-examples.instructions.md`](https://github.com/jejjohnson/pyrox/blob/main/.github/instructions/docs-examples.instructions.md):

1. **Title + Colab badge** (markdown).
2. **Background** — short motivation: BNF as a non-stationary feature-gated Bayesian MLP, equivalence to a composite GP at the prior-predictive level, the three bayesnf observation families.
3. **Setup** — Colab detection + `pyrox[colab]` install via `subprocess`.
4. **Imports + `jax_enable_x64` + watermark**.
5. **Problem setup** — draw a 3D synthetic truth from a `pyrox.gp` ground-truth GP on `(t, x, y)` with a known RBF kernel; sample ~500 training points on a random jittered grid, hold out a dense 50×50 test grid.
6. **Core demonstration** — `BNFEstimator(width=64, depth=4, seasonality_periods=(), num_seasonal_harmonics=())` → `.fit(df, seed=0)` → `.predict(df_test, quantiles=(0.025, 0.5, 0.975))`.
7. **Visualizations** (matplotlib defaults; repo convention):
- 1×3 subplot `(18, 5)`: truth / posterior-mean / per-point posterior-std on a heatmap.
- 1D slice at fixed `t`: scatter of training points, `k--` truth, `C0` mean, `C0 alpha=0.2` 95% band.
- Calibration plot: empirical coverage vs. nominal coverage at quantiles `(0.1, 0.2, ..., 0.95)`.
8. **Summary** — takeaways + links to the API docs for each piece.

## Mathematical Notes (for the Background cell)

### BNF as a conditional Gaussian process

The BNF forward pass is, schematically,

$$f(\mathbf{x}) \;=\; w_{\text{out}} \cdot \phi_L \!\Bigl(\, W_L \, \psi_{L-1}\!\bigl(\cdots \psi_1\bigl(W_1 \, h_0(\mathbf{x})\bigr)\cdots\bigr) / \sqrt{w_L}\,\Bigr)$$

where $h_0(\mathbf{x}) = [\tilde x,\, \phi_{\text{Fourier}}(\tilde x),\, \phi_{\text{seasonal}}(t),\, \phi_{\text{interaction}}(\tilde x)]$ is the concatenated feature block and every $W_\ell$, gain, and scale has a Logistic$(0, 1)$ prior. In the infinite-width limit, with the $1/\sqrt{\text{fan-in}}$ pre-normalization, this converges to a composite GP whose kernel is a product of the feature-space RBF-like kernel and the activation-induced kernel (NTK-style analysis, cf. [Lee et al. 2018](https://arxiv.org/abs/1711.00165) and [Tancik et al. 2020](https://arxiv.org/abs/2006.10739)).

### Ensemble posterior

With `ensemble_size=E` MAP fits, the predictive distribution at a new point $\mathbf{x}^*$ is the mixture

$$p(y^* \mid \mathbf{x}^*) \;\approx\; \frac{1}{E} \sum_{e=1}^{E} p\!\bigl(y^* \mid \mathbf{x}^*,\; \hat\theta_e\bigr).$$

For Gaussian observation models the mixture has closed-form mean and variance; for NB / ZINB, quantiles require CDF inversion via `gaussx.mixture_quantile` ([gaussx#121](https://github.com/jejjohnson/gaussx/issues/121)).

### Calibration

If the model is well-calibrated, the fraction of held-out targets falling below the $q$-th predictive quantile should equal $q$. The calibration plot bins the $q$'s and compares empirical coverage to nominal — a 45° line means perfect calibration; a curve above/below indicates over/under-confidence.

## Reference Code (bayesnf tutorial)

The closest equivalent in bayesnf is their [Chickenpox tutorial](https://github.com/google/bayesnf/blob/main/notebooks/chickenpox.ipynb) (weekly cases, multiple regions, strong seasonality). We deliberately pick a synthetic GP-drawn target instead, for three reasons:

1. **Reproducible without a data download.** Runs under `nbconvert --execute` in CI with zero external I/O.
2. **Known ground truth.** Calibration assertions become exact rather than "it looks reasonable."
3. **Pedagogy.** The target being a GP makes the "BNF approximates a GP" framing concrete.

## References

1. Saad, F. A., et al. (2024). *Scalable spatiotemporal prediction with Bayesian neural fields.* Nat. Commun. 15, 7942. [DOI:10.1038/s41467-024-51477-5](https://doi.org/10.1038/s41467-024-51477-5). Tutorials: [google/bayesnf/notebooks](https://github.com/google/bayesnf/tree/main/notebooks).
2. Lakshminarayanan, B., Pritzel, A., & Blundell, C. (2017). *Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles.* NeurIPS. [arXiv:1612.01474](https://arxiv.org/abs/1612.01474) — the deep-ensembles framing used in the Summary cell.
3. Kuleshov, V., Fenner, N., & Ermon, S. (2018). *Accurate Uncertainties for Deep Learning Using Calibrated Regression.* ICML. [arXiv:1807.00263](https://arxiv.org/abs/1807.00263) — source of the calibration plot.
4. Tancik, M., et al. (2020). *Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains.* NeurIPS. [arXiv:2006.10739](https://arxiv.org/abs/2006.10739) — the feature-embedding framing.
5. [mkdocs-jupyter docs](https://github.com/danielfrg/mkdocs-jupyter) — the renderer; `execute: false` is set in `mkdocs.yml` so the committed `.ipynb` outputs are what the user sees.
6. pyrox repo `.github/instructions/docs-examples.instructions.md` — the authoring checklist we follow exactly.

## Implementation Steps

- [ ] Author `docs/notebooks/bayesian_neural_fields.py` in jupytext percent format (dev source).
- [ ] Header + first markdown cell with title + Colab badge to `jejjohnson/pyrox@main`.
- [ ] Setup cell: Colab detection + `pip install -q "pyrox[colab] @ git+https://github.com/jejjohnson/pyrox@main"`.
- [ ] Imports, `jax.config.update("jax_enable_x64", True)`, `%watermark -p jax,equinox,numpyro,gaussx,pyrox,matplotlib`.
- [ ] Generate synthetic 3D ground-truth via `pyrox.gp.GPPrior` with a fixed RBF kernel; sample ~500 training points + a 50×50 hold-out grid.
- [ ] Fit `BNFEstimator(observation_model="NORMAL", width=64, depth=4, ensemble_size=16, num_epochs=5_000)`.
- [ ] Predict with `quantiles=(0.025, 0.5, 0.975)`; plot the three visualizations (truth/mean/std heatmap, 1D slice with band, calibration plot).
- [ ] Summary cell with takeaways + cross-links to the `BNFEstimator`, `ensemble_map`, and `mixture_quantile` API pages.
- [ ] Convert: `uv run --group docs jupytext --to notebook docs/notebooks/bayesian_neural_fields.py`.
- [ ] Execute in place: `uv run --group docs jupyter nbconvert --to notebook --execute docs/notebooks/bayesian_neural_fields.ipynb --inplace --ExecutePreprocessor.timeout=300`.
- [ ] Delete the `.py` (per workflow); commit the executed `.ipynb`.
- [ ] Add `mkdocs.yml` nav entry under "Tutorials" or equivalent section.

## Definition of Done

- `docs/notebooks/bayesian_neural_fields.ipynb` is committed with embedded cell outputs (heatmap + 1D slice + calibration plot).
- `uv run --group docs mkdocs build --strict` succeeds with the notebook rendering cleanly.
- Running `uv run --group docs jupyter nbconvert --execute docs/notebooks/bayesian_neural_fields.ipynb --inplace --ExecutePreprocessor.timeout=300` on a clean env reproduces the outputs end-to-end.
- The calibration plot shows ≥90% coverage at the 95% nominal level and ≥45% coverage at the 50% nominal level (sanity bounds; exact values printed in a cell assertion).
- Posterior-mean RMSE < target-std on the synthetic problem (printed in a cell assertion).
- `mkdocs.yml` nav includes the notebook.

## Testing

- The notebook's own `assert` cells are the primary test: calibration bounds + RMSE bound.
- CI hook: add a docs-smoke job (if not already present) that runs `jupyter nbconvert --execute --inplace` on all `docs/notebooks/*.ipynb` to catch regressions introduced by #70 / #71 / #72 changes.

## Documentation

- This *is* the docs deliverable. No separate mkdocstrings page is needed for the notebook itself; the API pages created in #72 are the reference-level companion.

## Relationships

- **Blocked by:** #72 (BNF layers), #71 (Estimator facade), #70 (ensemble runner), [gaussx#121](https://github.com/jejjohnson/gaussx/issues/121) (mixture quantile).
- Parent (theme epic): #36
- **Related:** `.github/instructions/docs-examples.instructions.md` — the authoring standard this notebook must satisfy.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.