fix(jax_engine): correct gradients for MM3 angle term at near-collinear geometries

[ericchansen]

Summary

Fixes the MM3 angle gradient correctness bug from #284 §1. The previously "no improvement available" verdicts for rh-conjugate and heck-relay were wrong — they were spurious stationary points caused by this bug. Post-fix, both systems show real improvements that the optimizer can now find.

Companion data PR: ericchansen/q2mm-data#10

Headline results

System	Pre-fix Δ% (q2mm-data#9)	Post-fix Δ%	Status
ch3f	99.83 % (det.)	99.83 % (det.)	unchanged ✅
rh-enamide	44.66 % ± 0.29 %	44.73 % ± 0.29 %	unchanged ✅
pd-allyl	−0.029 % ± 0.34 %	−0.010 % ± 0.40 %	still NOT SIG ❌
rh-conjugate	−0.080 % ± 1.18 %	18.00 % ± 4.17 %	🚀 NEWLY UNLOCKED
heck-relay	−0.59 % ± 3.26 %	52.82 % ± 1.54 %	🚀 NEWLY UNLOCKED

Root cause

The MM3 bond-angle energy used arccos(clip(cos θ, −1+ε, 1−ε)) to compute the angle from atomic positions. When a geometry minimization drives an angle to collinear (cos θ → ±1), the clip's gradient becomes zero in the boundary region — and the autodiff chain ∂E/∂atom = ∂E/∂θ · ∂θ/∂cos · ∂cos/∂atom underestimates the real gradient. The optimizer sees a spurious stationary point with angles wedged at ≈ 180°.

For rh-conjugate at the BFGS-converged geometry, the harmonic-only angle term produced JAX gradient norm 982 vs FD norm 575 (1.7× discrepancy). The discrepancy localized to four angles where cos θ saturated at the clip value −1+1e-7. Bond and torsion terms agreed with FD to ~1e-7.

Fix

Replace _safe_arccos(cos θ) with _angle_from_vectors(a, b):

• Forward: atan2(|a×b|, a·b) — well-conditioned at collinear
• Custom VJP: analytic gradient ∂θ/∂a = (cos θ · a/|a|² − b/(|a||b|)) / sin θ, with sin θ floored at 1e-12 to keep direction stable at exact collinearity

The decorator is applied inside an lru_cache(maxsize=1) factory because jax is None at module import time (populated lazily by _ensure_jax()).

Touches three call sites in q2mm/backends/mm/jax_engine.py:

• _harmonic_angle_energy (OPLSAA harmonic angle)
• _mm3_angle_energy (MM3 sextic anharmonic angle)
• _mm3_stretch_bend_energy (stretch-bend cross-term)

Deletes _safe_arccos and _safe_norm_keepdims (no remaining callers; AGENTS.md §2 alpha discipline — no shims).

Tests

• New regression: test/test_mm3_jax.py::TestMM3AngleEnergy::test_near_collinear_gradient_matches_fd
  Places three atoms at near-antiparallel geometry, checks ∂E/∂atom_k matches FD on the perpendicular axis to 1e-3 relative. Fails on pre-fix (JAX returns 0), passes on this fix.
• All 41 MM3-jax tests pass
• All 115 JAX-marked tests pass
• 680 unit tests pass
• ruff check + format clean

Doc updates (in same commit)

• docs/systems/rh-conjugate.md — verdict flipped from "NOT SIGNIFICANT, JaxLoss local min" to "SIGNIFICANT, 18.00 %", explanation of how the gradient bug had been hiding the real descent direction
• docs/systems/heck-relay.md — verdict flipped from "NOT SIGNIFICANT, --ratio-tol none doesn't help" to "SIGNIFICANT, 52.82 %"; the system no longer needs --ratio-tol none (ratio now 1.085, within default gate)
• docs/systems/pd-allyl.md — numbers refreshed (still NOT SIGNIFICANT, the FF truly is at a local min here)

Closes

• MM3 energy is non-smooth: JAX grad returns wrong subgradients at converged metal-TS geometries #284 §1 (the gradient correctness bug; the noise side §2 was separately mitigated by feat(regen): add --n-evals + median+CI reporting to distinguish signal from noise #286 + docs: confirm 'no real improvement' verdicts for metal-TS systems (n=10) #287)

[Copilot]

Pull request overview

Fixes a JAX gradient-correctness bug in MM3/harmonic angle energies where arccos(clip(cos θ, …)) zeroed gradients at near-collinear geometries, causing the optimizer to land on spurious stationary points. Replaces the clip-based angle with an atan2-based formulation backed by a custom VJP using analytic ∂θ/∂a expressions, with sin θ floored at 1e-12 to keep direction stable at collinearity. Unlocks real >18% (rh-conjugate) and >52% (heck-relay) ObjectiveFunction reductions previously hidden behind the bug.

Changes:

• New _angle_from_vectors(a, b) with jax.custom_vjp analytic gradients; applied to harmonic angle, MM3 angle, and (via inlined atan2) MM3 stretch-bend.
• Deletes superseded _safe_arccos and _safe_norm_keepdims per alpha-discipline.
• Adds near-collinear FD-vs-JAX regression test; updates rh-conjugate / heck-relay / pd-allyl docs with post-fix verdicts.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
q2mm/backends/mm/jax_engine.py	New atan2-based angle helper with custom VJP; replaces clip-arccos at three call sites; removes dead helpers.
test/test_mm3_jax.py	Adds test_near_collinear_gradient_matches_fd regression test for the gradient correctness fix.
docs/systems/rh-conjugate.md	Flips verdict to SIGNIFICANT (18.00 % ± 4.17 %) with explanation linking to #284.
docs/systems/heck-relay.md	Flips verdict to SIGNIFICANT (52.82 % ± 1.54 %); ratio improves 1.378 → 1.085.
docs/systems/pd-allyl.md	Refreshes numbers; verdict unchanged (NOT SIGNIFICANT, true local min).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.