AIMNet2 v0.2 support: new flags, dynamic element validation

[isayev]

Add AIMNet2 v0.2 support: new flags, dynamic element validation

Brings OET's AIMNet2 calculator up to the AIMNet2Calculator API
introduced in aimnet v0.2 (verified against aimnet 0.2.0 on PyPI).
Salvages and adapts the AIMNet2 work originally proposed for OPI in
faccts/opi#228 (closed; cplett asked to retarget at OET). A separate
small PR (feature/server-vram-eviction) adds matching server-mode
VRAM safety; that PR is order-independent.

• New flag surface (9 flags total) covering performance (--compile,
  --nb-threshold, --ensemble-member), long-range Coulomb (--coulomb /
  --coulomb-method / --coulomb-cutoff), and DFT-D3 dispersion
  (--dispersion / --dftd3-cutoff / --dftd3-smoothing-fraction). The
  HuggingFace --revision/--token flags from PR #228 are intentionally
  omitted; per-method tuning of --dsf-alpha/--ewald-accuracy is left
  to the upstream Python API. --coulomb-cutoff without --coulomb-method
  raises SystemExit rather than silently ignoring.

• The hardcoded 15-element table in Aimnet2Calc is replaced with
  a full Z=1..118 periodic-table translator. Per-model element
  rejection is delegated to AIMNet2Calculator.eval (validate_species
  defaults to True; OET never overrides). Model-specific error
  messages bubble up unmodified.

• The CPU-forcing monkey-patch on torch.cuda.is_available is
  removed; device= is now passed directly to the constructor.
  --device auto maps to None (upstream auto-detect). Early
  cuda-not-available raise sits ABOVE the cache short-circuit so
  a cached CPU calc does not silently service a --device cuda
  request on a no-CUDA box.

• The .jpt extension append in get_model_file is removed: the v0.2
  registry uses .pt and the upstream resolver knows the right
  filename without OET's intervention.

• The mult key is now passed only when the loaded model reports
  is_nse == True (forward-safety against v0.2's stricter input
  validation).

• setup() rewritten as a single method with args-match idempotency
  (frozenset of items, not a positional tuple), partial-state
  rollback on post-ctor failure, validation guards on tristate /
  coulomb-method inputs, "simple" coulomb method skips cutoff
  forwarding, --coulomb-method=simple + non-default --coulomb-cutoff
  now warns instead of silently ignoring, torch.set_num_threads moves
  from per-call into setup() (one-shot per worker). Post-ctor failure
  now rolls back via self.release() so the model + external_coulomb +
  external_dftd3 are walked back to CPU and torch.cuda.empty_cache()
  is called before references drop — Python GC alone does not reclaim
  CUDA caching-allocator reservations promptly, so the previous
  null-only rollback would silently accumulate VRAM under a server
  loop with a repeating bad config.

• forces.detach() before listify, defending against any future
  upstream change to create_graph=True default.

• BaseCalc gains a release() no-op default method. Calculators that
  hold device-resident state (currently only Aimnet2Calc) override
  to move models to CPU, drop references, and call empty_cache().
  This makes the salvage / sibling-PR boundary genuinely order-
  independent: any future polymorphic calc.release() call from
  shared eviction code is safe on every BaseCalc subclass without
  AttributeError. Three lines, no behavior change for existing
  calculators (xtb, mopac, mlatom, gxtb, uma).

• Aimnet2Calc gains a release() override that drops the cached
  calculator, moves model + external_coulomb + external_dftd3 off
  GPU (logging to stderr on partial failure), and calls
  torch.cuda.empty_cache(). The server-side worker-cache eviction
  that calls release() lives in the sibling PR; this PR only adds
  the override (order-independent).

• Aimnet2-pd loads with a UserWarning that CPCM/THF solvation is
  baked into the model and energies are NOT gas-phase.

• README gains an "AIMNet2 options" subsection with a model-
  selection guide that distinguishes aimnet2 (default, wB97M-D3)
  from aimnet2-2025 (B97-3c, opt-in for non-covalent / screening;
  3-5 kcal/mol off vs default for reaction barriers), aimnet2-nse
  (charged/open-shell, 14-element coverage, single-reference
  applicability), aimnet2-rxn (NEUTRAL H/C/N/O only, 4.6 A Coulomb
  cutoff lock), and aimnet2-pd (CPCM/THF baked in, no As). A
  cross-family energy-incomparability WARNING ("tens of kcal/mol")
  sits ABOVE the model-selection table; a bold callout above the
  long-range Coulomb section emphasizes the rxn-family 4.6 A cutoff
  lock. Reproducibility note pins to canonical model keys. Two
  open-shell warnings are added under "AIMNet2 options" itself:
  (1) the default aimnet2 (and aimnet2-2025, aimnet2-b973c-d3*,
  aimnet2-rxn*) are closed-shell-trained — passing mult != 1 is
  silently accepted but yields a spin-restricted energy, not a
  true UKS-equivalent open-shell value, so users wanting genuine
  open-shell energetics must pass -m aimnet2-nse; and (2) no
  AIMNet2 family does broken-symmetry, so closed-shell singlet
  biradicals (mult=1 with two unpaired electrons coupled
  antiferromagnetically: carbenes, nitrenes, stretched singlet
  sigma bonds) are not modelable correctly by any model in the
  suite, NSE included, and need a multireference treatment.

• New pytest test file tests/aimnet2/test_aiment2_v02_contract.py
  covers: argparse contract (flag presence, defaults, choices,
  --coulomb-cutoff cross-flag rejection), NSE-mult gating with
  mocks, periodic-table coverage, setup() idempotency / args-match
  / device-cuda-unavailable / validation guards / partial-state
  rollback (asserting components are walked back to CPU before
  _calc is dropped), release() lifecycle including drop-and-resetup,
  BaseCalc.release() default no-op idempotency, plus one network-
  marked end-to-end smoke on the default model with a deliberately
  stretched water (gradient assertion 0.005 < |g| < 0.5 Eh/Bohr to
  actually catch unit/sign bugs). 30 contract tests + 1 network
  smoke. Existing test_aiment2_standalone.py and test_aiment2_client.py
  reference numerics regenerated against aimnet 0.2.0 PyPI: v0.2
  ships retrained model files (storage path includes /aimnet2v2/)
  whose energies differ from v0.1.x by ~1e-6 Eh on the H2O / OH-/
  OH. fixtures. Both wrappers and the server path produce bit-exact
  identical numerics, so assertAlmostEqual tolerance is held at
  places=9.

• Dependency pin: aimnet>=0.2,<0.3.

The new flags propagate to oet_server aimnet2 automatically via the
existing extend_parser plumbing in src/oet/server_client/server.py.
The cross-flag validation deliberately does NOT stash a parser
reference inside the parsed-args namespace: doing so was found to
poison server.py's calculator cache key (frozenset of namespace
items where ArgumentParser is identity-hashed), silently nullifying
calculator reuse for AIMNet2 in oet_server. The validation raises
SystemExit directly with an argparse-style message instead.

The default model is intentionally unchanged (aimnet2 =
aimnet2-wb97m-d3_0). aimnet2-2025 is documented as opt-in via -m.

Out of scope (separate PRs):

• Server-mode VRAM-aware eviction (sibling PR feature/server-vram-eviction).
• Analytical Hessian via .hess sidecar (pending design discussion
  with cplett on extending OET's output protocol; ORCA does not
  currently support Hessian read-in from external programs, so
  this is a longer-tail follow-up).
• Conceptual DFT / Fukui descriptors (best fit is a companion tool
  outside OET).
• Per-atom Hirshfeld charges (engrad bridge has no place to surface
  them; users who need them call aimnet directly).
• Thermochemistry post-processing (G, H, S, ZPE — provided by ORCA's
  own Freq block).

Related upstream work:

• refactor(registry): adopt aimnet2-<family>_<member> naming convention isayevlab/aimnetcentral#70 — registry rename to canonical
  aimnet2-_ form; all legacy aliases preserved.
• isayevlab/aimnetcentral v0.2 — adds the constructor / methods this
  PR consumes; ships an OET-surface integration test as part of
  release verification (tests/oet_compat/test_v0_2_surface.py).
• Add AIMNetCentral ORCA integration for maintained AIMNet2 models opi#228 — original (closed) OPI integration that this PR
  salvages and adapts.

Pre-PR coordination: #37 RFC; cplett
greenlit all four design questions (element-validation rewrite,
sibling-PR scope, default-model decision, Hessian deferral).

---

RFC: #37 (cplett's reply: #37 (comment))

[CLAassistant]

All committers have signed the CLA.

[cplett]

Thank you very much for the contribution. I really like the improvements!
There are just a few minor comments and suggestions, and then it should be good to go.

[isayev]

Thanks for the thoughtful review! All 12 comments are addressed in 1b1662f. Quick summary mapping each comment to the change:

Structure

• README structure → moved the long "AIMNet2 options" subsection out of the main README.md into a dedicated readmes/aimnet2.md. Main README now carries a brief pointer.
• get_model_file → generalised into oet.core.misc.resolve_model_file (takes alias-resolver / fetch / fetch-fallback as callables, so core/ stays calculator-agnostic). Aimnet2Calc.get_model_file becomes a thin wrapper. Future MLIP wrappers (UMA etc.) can reuse the helper directly.

Behavioural improvements

• Open-shell + non-NSE warning → serialize_input now emits a UserWarning when mult ≠ 1 is requested with a closed-shell-trained model, pointing the user at -m aimnet2-nse.
• aimnet2-rxn cutoff validation → setup() now warns when an aimnet2-rxn family model is combined with --coulomb-method and a --coulomb-cutoff other than the trained 4.6 Å. I went with UserWarning rather than a hard quit so legitimate testing / parameter sweeps remain possible — happy to flip to a hard error if you'd prefer.

Docstrings & help text

• Restored full Numpy-style Parameters/Returns/Raises sections on setup(), extend_parser(), serialize_input() (matching mopac.py / xtb.py house style); serialize_input documents the keys of the returned kwargs dict explicitly.
• --model help now points open-shell users at aimnet2-nse. --ensemble-member help reworded so it's explicit OET runs ONE model per call. README mirrors both reformulations.

Tooling

• pytest>=8.0 added under [project.optional-dependencies] test.

Five new contract tests cover the two new warning surfaces (test_open_shell_with_non_nse_model_warns, test_open_shell_with_nse_model_no_warn, test_rxn_with_non_trained_cutoff_warns, test_rxn_with_trained_cutoff_no_warn, test_rxn_without_coulomb_method_no_warn). Local: 36/36 contract pytest + 1 network smoke + 3/3 standalone + 3/3 client all green against aimnet 0.2.0 PyPI; all 7 nox -t static_check sessions green; CI shows 5/5 nox jobs green on this commit. Ready for another look whenever you have time.

[cplett]

Looks good to me. Thanks again for the great work!
Looking forward to seeing the models in action together with ORCA.