Skip to content

tech-debt: v3.5 docs refresh, ANI2xt fp64 energies, optimizer benchmark#97

Merged
isayev merged 5 commits into
mainfrom
tech-debt/batch-2
Jun 12, 2026
Merged

tech-debt: v3.5 docs refresh, ANI2xt fp64 energies, optimizer benchmark#97
isayev merged 5 commits into
mainfrom
tech-debt/batch-2

Conversation

@isayev

@isayev isayev commented Jun 12, 2026

Copy link
Copy Markdown
Contributor

Closes the remaining workstreams from the two post-hardening plans (W7–W9 of the tech-debt roadmap; both plans were otherwise already implemented across PRs #93/#94).

Changes

  • Docs refresh for v3.5 (W7) — removed all documentation of the deleted ensemble feature (AUTO3D_USE_ENSEMBLE, --use-ensemble, "8-model ensemble") across 10 doc pages; fixed version floors (Python 3.11, PyTorch 2.8); documented aimnet as a core dependency, the [ani]/[ase] pip extras, the ~/.cache/aimnet model download/cache (AIMNET_CACHE_DIR), and the AIMNet2 registry model names. Sphinx build is clean (0 warnings).
  • ANI2xt fp64 accumulation (W8) — accumulate the ANI2xt forward's per-atom and self-atomic energies in float64 so the float64 energy_shifts buffer is no longer silently truncated. Cosmetic: self-atomic shifts cancel in conformer energy differences, so ranking is unaffected; energy is returned as float64 and forces as float32, matching the AIMNet2 adapter contract. The fp32 network output still caps usable precision.
  • Optimizer benchmark harness (W9)scripts/bench_optimizer.py, an opt-in micro-benchmark (outside tests/, never gates CI) that times optimizing(...).run() on a fixed, deterministic drug-like conformer set and reports wall time, throughput, and peak memory. Documented in CONTRIBUTING.
  • Small: clarified the @ tautomer-separator docstring in results.py; recorded the two plan docs; bumped a stale "Python 3.10" mention in CONTRIBUTING to 3.11.

Notes for review

  • W8 deviates from the written plan, on purpose. The plan's "build atom_energies in float64" crashes at runtime (Index put requires source and destination dtypes match — fp64 destination, fp32 network output). This PR adds the explicit elem_energies.to(torch.float64) cast the plan omitted; verified via a standalone autograd reproduction (fp64 energy / fp32 forces, finite grads). ANI2xt cannot run without the optional torchani extra, so this path must be confirmed by the CI [ani] job — it is not exercised by the local suite.

Verification

  • ruff: clean (src/, scripts/)
  • Sphinx docs build: clean, 0 warnings
  • Test gate (pytest -m "not slow"): 631 passed, 8 skipped (torchani-gated), 45 slow deselected

@isayev isayev merged commit e7e46e6 into main Jun 12, 2026
7 of 8 checks passed
@isayev isayev deleted the tech-debt/batch-2 branch June 12, 2026 04:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant