EDU-2 — Marimo theory notebooks + forward-demo with a CI'd numerics cross-check

[aacostadiaz]

Depends on: EDU-1 (#1607), ARCH-4 (#1563) · Blocks: —

Context: tutorials/ is the educational implementation of MACE — a top-level directory, a sibling of packages/ and the frozen legacy mace/. It is not a package and not the "reference backend" (that is the plain-torch/jax kernel oracle living inside the packages, mace_torch/backends/reference/ and mace_jax/kernels/reference.py). It exists so chemists and physicists can read and inspect the model blocks — pure-function JAX, readable over fast, evaluation/inspection-only: not a training path (there is no JAX training anywhere) and not a mirror of the latest architecture. EDU-1 (#1607) lays down the tutorials/ scaffold and mace_forward.py, a pure-function nanoGPT-style JAX MACE forward that computes energies with no closures. This ticket adds the pedagogy on top of it: marimo notebooks that walk through the equivariant math with inline LaTeX and a runnable forward/eval demo, plus a CI job that executes everything and pins a numerics cross-check against a real MACE forward so the tutorial can never silently drift from the model it teaches. This supersedes the hand-written theory Colab ("MACE theory and code (advanced)", the third detailed tutorial linked from the README). Nothing in packages/ depends on tutorials/; the dependency arrow runs one way (the cross-check test imports an oracle, never the reverse).

Interface & constraints:

1. Directory layout (all under the repo-root tutorials/, established by EDU-1 (EDU-1 — tutorials/ scaffold + pure-function JAX MACE forward (energies) #1607)):

tutorials/
├── mace_forward.py        # EDU-1 (#1607): pure-function JAX MACE forward (energies), no closures
├── notebooks/             # EDU-2 (#1608): marimo notebooks (this ticket)
│   ├── 01_math_walkthrough.py     # SH, Clebsch-Gordan, message passing, many-body — inline LaTeX
│   └── 02_forward_demo.py         # forward/eval demo on a tiny model
└── tests/                 # EDU-2 (#1608): notebook execution + numerics cross-check

2. Notebooks are marimo, i.e. plain executable .py files (reactive notebooks), not .ipynb. They must run top-to-bottom headless in CI. Keep them import-light: they read tutorials/mace_forward.py and standard scientific-python (jax/numpy/matplotlib); they do not import packages/.

3. Math walkthrough (01_math_walkthrough.py) covers, in order, with derivations rendered as inline LaTeX (marimo mo.md(r"...")):

   • Spherical harmonics — the real (Y_\ell^m) basis used for edge directions, in the component normalisation the model uses.
   • Clebsch-Gordan coupling — how two irreps combine into an output irrep (prose says Clebsch-Gordan, never bare "CG"); tie this to the tensor-product convolution.
   • Message passing — one interaction layer as a scatter over the graph.
   • Many-body — the symmetric contraction up to correlation order as the many-body term.
     Each section is small, self-contained, and executable (short numeric examples the reader can perturb).

4. Forward/eval demo (02_forward_demo.py): load a tiny model + a small structure, run mace_forward.py to get the energy, and display intermediate tensors so the reader can see the blocks light up. Eval/inspection only — no optimizer, no training loop, no forces (EDU-1 (EDU-1 — tutorials/ scaffold + pure-function JAX MACE forward (energies) #1607)'s forward exposes none).

5. Golden fixture and the edge list. The fixture is a tiny model + structure, committed (or committed as a reproducible recipe) so the cross-check is deterministic. The edge list is part of the fixture, not recomputed by the tutorial, and it is built the way mace.data.get_neighborhood builds it (mace/data/neighborhood.py:7-60), whose cell handling has three regimes the fixture must not straddle: the neighbour search always runs in an extended cell whose non-periodic rows are resized to the atom extent plus 2 * cutoff + 1; a fully aperiodic system returns that extended cell, deliberately, because long-range models need a non-degenerate box; a system with any periodic axis returns the physical cell, with only all-zero rows patched from the extended one, because stress and the electrostatic volume corrections divide by det(cell). Pin the fixture to a fully aperiodic molecule: the tutorial has no cell-dependent term, so the returned cell must not be able to influence the comparison at all.

6. Numerics cross-check (tutorials/tests/): assert the JAX educational forward reproduces a real MACE forward on the fixture to fp64 tolerance on energy. The oracle is versioned:

   • Before ARCH-4 (ARCH-4 — Two-phase forward and derivative engine (forces via autograd, stress via strain) #1563) lands, the oracle is the frozen legacy mace — ScaleShiftMACE, default_dtype=float64, compute_force=False. The notebooks therefore land pinned, not on an inert assertion that activates two phases later.
   • When ARCH-4 (ARCH-4 — Two-phase forward and derivative engine (forces via autograd, stress via strain) #1563) lands, the cross-check is repointed at mace_torch's two-phase forward and the legacy import is deleted. v1 reproduces develop's energies to the golden tolerances, so this is an oracle swap, not a tolerance change; the ticket closes on the v1 oracle.
     The tolerance is a strict fp64 row, not an eyeballed epsilon. This test is the only file in tutorials/ allowed to import an oracle; mace_forward.py and both notebooks import none of them.

7. CI execution: a dedicated tutorials job in .github/workflows/ci-core.yaml (CPU-only, needs: lint, unconditional — GOV-1 (GOV-1 — CONTRIBUTING.md and the project governance document #1609)'s examples-must-not-break rule makes this a merge-time gate, and the job is cheap enough not to need a paths filter). It installs through .github/actions/setup-mace (extras: dev, plus pip-packages: -r requirements/tutorials.txt, a new pinned requirements file for jax[cpu] and marimo, following the requirements/les.txt / requirements/polar.txt convention for deps that are not extras) and runs .github/actions/run-tests with tests: tutorials/tests. Notebook execution is itself a pytest case, so one action invocation covers both. Two properties of that action apply here: tests/conftest.py governs only the tests/ subtree, so the capability contract is inert for tutorials/tests and require-caps stays empty; and --cov=mace is hardcoded in the coverage branch, so coverage stays "false".

Task:

1. Write tutorials/notebooks/01_math_walkthrough.py (marimo): spherical harmonics → Clebsch-Gordan coupling → message passing → many-body, each with inline-LaTeX derivation and a short runnable numeric example built on tutorials/mace_forward.py.
2. Write tutorials/notebooks/02_forward_demo.py (marimo): tiny-model forward/eval demo, surfacing intermediate block outputs; energy only, no training.
3. Commit the tiny golden fixture (fully aperiodic molecule + model + its edge list, or a reproducible recipe) under tutorials/tests/.
4. Add tutorials/tests/test_notebooks.py, which executes each notebook as a script and fails on any raise, and tutorials/tests/test_cross_check.py, which compares the JAX forward against the oracle at fp64 — legacy mace today, mace_torch once ARCH-4 (ARCH-4 — Two-phase forward and derivative engine (forces via autograd, stress via strain) #1563) lands.
5. Add requirements/tutorials.txt (pinned jax[cpu], marimo) and the tutorials job in .github/workflows/ci-core.yaml wired as in §7.

Out of scope: the tutorials/ scaffold and mace_forward.py itself (EDU-1 (#1607)); any JAX training or run_train (there is no JAX training anywhere); mirroring the latest architecture (the tutorials are inspection-only, not a spec); the eval CLI and goldens inside packages/; every row marked DROP in EDU-1 (#1607)'s scope table.

Acceptance criteria:

• tutorials/notebooks/01_math_walkthrough.py and 02_forward_demo.py exist as marimo notebooks and execute top-to-bottom headless with no error.
• The math walkthrough renders inline-LaTeX derivations for spherical harmonics, Clebsch-Gordan coupling, message passing, and many-body, each with a runnable numeric example.
• The forward demo runs the JAX forward on a tiny model and displays intermediate block tensors; no training/optimizer code is present.
• A committed golden fixture drives a cross-check asserting the JAX forward matches a real MACE forward to fp64 tolerance on energy; the fixture is fully aperiodic and carries its own edge list.
• The cross-check is green against the frozen legacy mace from the day it lands, and against mace_torch once ARCH-4 (ARCH-4 — Two-phase forward and derivative engine (forces via autograd, stress via strain) #1563) exists — the ticket does not close on the legacy oracle.
• A CPU-only tutorials job in ci-core.yaml executes both notebooks and runs the cross-check; a notebook that raises, or a cross-check outside tolerance, fails the job.
• tutorials/ imports nothing from packages/; only the cross-check test imports an oracle (one-way arrow), and neither notebook does.

Verify:

# Marimo notebooks are executable Python; `marimo run` serves them and never exits,
# so headless execution is a plain interpreter run.
python tutorials/notebooks/01_math_walkthrough.py
python tutorials/notebooks/02_forward_demo.py
python -m pytest tutorials/tests -v                       # notebook execution + fp64 cross-check
python - <<'PY'                                           # only the cross-check may import an oracle
import ast, pathlib
roots = {"mace", "mace_core", "mace_torch", "mace_jax"}
offenders = []
for path in pathlib.Path("tutorials").rglob("*.py"):
    if path.name == "test_cross_check.py":
        continue
    tree = ast.parse(path.read_text())
    mods = [n.module or "" for n in ast.walk(tree) if isinstance(n, ast.ImportFrom)]
    mods += [a.name for n in ast.walk(tree) if isinstance(n, ast.Import) for a in n.names]
    offenders += [(str(path), m) for m in mods if m.split(".")[0] in roots]
assert not offenders, offenders
print("one-way arrow: OK")
PY

Review focus: pedagogical correctness of the math (a physics reader should recognise the derivations) and that the cross-check tolerance is a strict fp64 row, not loosened to hide drift; that the fixture is fully aperiodic so no cell convention can leak into the comparison; confirm the one-way dependency (packages/ never imports tutorials/, and only the cross-check imports an oracle).

---