[niklashoelter]

Adds AIMNet2 as a fourth backend: model_type: "aimnet" (alias "aimnet2") with aimnet2, aimnet2-nse (open-shell) and aimnet2-2025 (B97-3c), ensemble member 0.

What changes

• Dependency: aimnet[ase,torchsim]==0.2.5. The lock only adds aimnet; nothing else moves (warp stays 1.17.0, torch-sim 0.6.2).
• models/aimnet.py follows the existing loader pattern: AIMNet2ASE for the ASE path, AIMNet2TorchSim wrapped in sevenn's generic Float64Wrapper for float64 batch states (AIMNet2 is float32-only, like SevenNet).
• Input checks before any model loads (API and CLI, exit 3), via check_structure_supported:
  • only the 14 trained elements (H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I);
  • aimnet2 / aimnet2-2025 are closed-shell: multiplicity > 1 points to aimnet2-nse;
  • aimnet2 / aimnet2-2025 reject H–H distances < 0.8 Å with heavy atoms present (spurious geminal-H collapse, aimnetcentral#139); free H2 is fine.
• d3_correction and model_path are rejected for AIMNet2 (D3 is built in; only registry names load). model_cache_dir works via AIMNET_CACHE_DIR.
• memory_scales_with: auto → n_atoms for AIMNet2 (see below).
• Docs, skill, examples/config_aimnet.json, AIMNet examples in the two example scripts, workflow-benchmark entries (aimnet2 joins the default set).

Memory metric: n_atoms

AIMNet2's gas-phase Coulomb builds an int32 neighbour matrix of (total batch atoms)². Forward peak at equal total atoms is identical across the 4k, amine and kraken sets although their edges/atom differ (22 vs 30): 14.4–14.5 GB at 60k atoms, ≈ 4 B × N² above ~20k atoms. So memory follows atoms, not edges. With n_edges, the amine set (few edges per atom) packed batches up to 23.9 of 24 GB; with n_atoms it peaked at 2.6 GB. Wall times are similar (53–73 s for ~1.6k–4.3k structures, all converged, both models).

The N² sizing is upstream (calculator.py: max_neighbors = N for infinite Coulomb cutoff); it could be sized by the largest molecule. Not addressed here.

Optimizer memory: warnings for BFGS and AIMNet2 + L-BFGS

torch-sim's autobatcher sizes batches from forward passes only; optimizer state on top is not counted. Overnight stress tests (9.7k and 29k structures from the bundled amine/kraken sets, plus 140–423-atom molecules, without gpuma's per-structure retry):

• FIRE and gradient descent: every set passes for all three AIMNet2 models (29k structures in ~7 min).
• BFGS (dense padded Hessian per system) runs out of memory at scale with every backend (AIMNet2, ORB, UMA).
• L-BFGS with AIMNet2 runs out of memory on the 29k set and the set with large molecules, because AIMNet2 batches fill the GPU (the N² Coulomb buffer above). With aimnet's Coulomb buffer sized by the largest molecule (tested locally, upstream fix in preparation), every L-BFGS case passes. ORB and UMA pass L-BFGS on the same sets.

gpuma now logs a warning once per batch job for bfgs (any backend) and for lbfgs with AIMNet2; docs and skill recommend fire for AIMNet2. The AIMNet2 warning can go once a released aimnet has the fix.

Testing

• Full suite: 832 passed, 1 skipped (nvMolKit usable).
• New tests/models/test_aimnet.py: registry vs. model metadata, NSE multiplicity, float64 torch-sim output matching ASE, per-system charge in one batch, cache dir, batch vs. sequential on mixed charges, all input checks; CLI dry-run and API pre-load checks.
• Real CLI runs: optimize (SMILES, XYZ), batch (multi-XYZ, dir), ensemble, NSE triplet; triplet on aimnet2 → exit 3.
• All three example scripts exit 0; mkdocs build --strict passes.
• Workflow benchmark (small/medium so far): aimnet2 optimizes the medium set (244 conformers) in 2.9 s vs. 13 s ORB-direct, 22 s SevenNet, 340 s UMA d4rev.