From b2c4ebb89f580b045a6608835959317e078e3d05 Mon Sep 17 00:00:00 2001 From: Olexandr Isayev Date: Sun, 3 May 2026 17:59:17 -0400 Subject: [PATCH 1/4] Add AIMNet2 v0.2 support: new flags, dynamic element validation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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: - 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). - faccts/opi#228 — original (closed) OPI integration that this PR salvages and adapts. Pre-PR coordination: faccts/orca-external-tools#37 RFC; cplett greenlit all four design questions (element-validation rewrite, sibling-PR scope, default-model decision, Hessian deferral). --- README.md | 140 +++++ pyproject.toml | 7 + requirements/aimnet2.txt | 2 +- src/oet/calculator/aimnet2.py | 581 +++++++++++++++------ src/oet/core/base_calc.py | 12 + tests/aimnet2/test_aiment2_client.py | 52 +- tests/aimnet2/test_aiment2_standalone.py | 58 +- tests/aimnet2/test_aiment2_v02_contract.py | 490 +++++++++++++++++ 8 files changed, 1118 insertions(+), 224 deletions(-) create mode 100644 tests/aimnet2/test_aiment2_v02_contract.py diff --git a/README.md b/README.md index 42aeae8..c77a43a 100644 --- a/README.md +++ b/README.md @@ -65,6 +65,146 @@ If you want to keep multiple servers running for different types of calculations you have to specify different ports for the server and clients with the `-b :` keyword. Provide the keyword to the client via the ORCA input line `Ext_Params "-b :"`. +### AIMNet2 options + +Default model is `aimnet2` (= `aimnet2-wb97m-d3_0`, ωB97M-D3 trained). +Other models are selectable via `-m`. For non-covalent or screening +work, `aimnet2-2025` (B97-3c trained, faster) is a good alternative; +for general thermochemistry and reaction barriers, `aimnet2` remains +the recommended default per upstream guidance. + +> **Open-shell users — read this.** The default `aimnet2` (and +> `aimnet2-2025`, `aimnet2-b973c-d3*`, `aimnet2-rxn*`) are +> closed-shell-trained. Passing `mult ≠ 1` (e.g. an ORCA input with +> `* xyzfile 0 2 OH.xyz`) is **silently accepted** and produces a +> spin-restricted energy, not a true UKS-equivalent open-shell value. +> For genuine open-shell or charged-species energetics, use +> `-m aimnet2-nse`. +> +> No AIMNet2 family does broken-symmetry, so a closed-shell-singlet +> biradical (`mult = 1` with two unpaired electrons coupled +> antiferromagnetically — carbenes, nitrenes, stretched singlet σ +> bonds) is not modelable correctly by **any** model in the suite, +> NSE included. Treat such systems with a multireference method. + +#### Choosing a model + +**WARNING — energies from different model families are NOT comparable.** +Different families were trained on different reference data, with +different functionals, with or without solvation. Mixing models within +a single workflow (e.g. optimize with `aimnet2-rxn`, compute SP energies +with `aimnet2-2025`) produces wrong reaction energies of order tens of +kcal/mol — comparable to or larger than typical reaction barriers, with +no failure indication. The aimnet runtime emits a one-time +`UserWarning` when two families are constructed in the same Python +process; this is informational, not a bug. + +Recommended pattern: TS search in `aimnet2-nse` → SP energies in +`aimnet2-nse`. Never optimize in one family and compute SP in another. +The exception is `aimnet2-pd`, which can be used end-to-end for +Pd-catalyzed THF-solvated workflows but must not be mixed with any +other family at all. + +| Workflow | Recommended model | Notes | +|----------|-------------------|-------| +| Closed-shell, neutral, equilibrium geometry, gas-phase | `aimnet2` | ωB97M-D3 trained | +| Closed-shell, neutral, non-covalent / screening / large systems | `aimnet2-2025` | B97-3c trained; faster but barriers typically 3-5 kcal/mol off vs the ωB97M-D3 default. Fine for relative ranking and screening, not absolute kinetics. | +| Open-shell or charged | `aimnet2-nse` | NSE: handles arbitrary charge / multiplicity. 14-element coverage (H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I). Covers single-reference DFT regimes. NOT reliable for biradicals (two unpaired electrons on different atoms with weak coupling), near-degenerate spin states, or stretched singlet bonds (e.g. homolytic dissociation on a singlet surface). | +| Reactive (TS, NEB, IRC), broad element set | `aimnet2-nse` | Same caveats as above; covers reactive open-shell trajectories. | +| Reactive, neutral H/C/N/O only, faster | `aimnet2-rxn` | NEUTRAL only (charge ≠ 0 raises `ValueError`). H/C/N/O only (other elements raise `ValueError`). Coulomb cutoff locked at 4.6 Å (other values fire an upstream `UserWarning`). | +| Pd-containing (catalysis, organometallic) | `aimnet2-pd` | B97-3c with **implicit CPCM/THF solvation BAKED IN** — energies are NOT gas-phase. Replaces As with Pd vs default; AsR3 ligands (arsine ligands on Pd, e.g. Pd(AsPh3)4-style complexes) are not supported. | + +#### Reproducibility + +Aliases (`aimnet2`, `aimnet2-2025`, …) may be repointed in future +aimnet releases. For bit-stable reproducibility across versions, pin +to canonical keys: `-m aimnet2-wb97m-d3_0` rather than `-m aimnet2`. + +#### Performance flags + +| Flag | Default | Effect | +|------|---------|--------| +| `--compile` | False | torch.compile JIT. **Server mode only** — standalone re-pays JIT every ORCA call. First-call latency 10–60s. Recompiles on shape change (catastrophic in NEB). Incompatible with Hessian. Do not use with NEB / OptTS / IRC. For server mode, set `TORCHINDUCTOR_CACHE_DIR=/persistent/path` to keep the warm cache across worker restarts. | +| `--nb-threshold N` | 120 | Adaptive neighbor-list batch size. | +| `--ensemble-member {0,1,2,3}` | 0 | Use a single ensemble member (default 0). OET runs ONE model per call. To get an ensemble mean for production accuracy, run with each `--ensemble-member` value (0,1,2,3) and average outside ORCA. Upstream recommends ensemble averaging for production calculations. | + +#### Numerical precision + +AIMNet2 internally evaluates in float32. For large systems (energy magnitudes +above ~1000 eV), absolute energies are precise to ~1e-4 eV ≈ 4e-6 Eh. ORCA's +default `TolE=5e-6 Eh` is at this noise floor; if optimization stalls on +`Energy convergence not reached`, loosen with `! TightOpt` → `%scf TolE 1e-5 end` +or accept slightly looser convergence. + +#### GPU server deployment + +Each `oet_server aimnet2 -d cuda` worker holds the model resident on GPU +until evicted. `torch.cuda.empty_cache()` returns memory to PyTorch's allocator, +not to the OS — `nvidia-smi` shows steady-state usage equal to the sum of +resident workers' models. Plan worker count by VRAM/model-size, not per-call +peak. For multi-worker GPU deployments, set `CUDA_VISIBLE_DEVICES` per worker +or use `--workers 1`. (Server-side worker-cache eviction with `release()` is +in a sibling PR.) + +> **`aimnet2-rxn` family CAUTION**: the Coulomb cutoff is locked at 4.6 Å +> in training. Always pass `--coulomb-cutoff 4.6` when using `-m aimnet2-rxn*`, +> otherwise the upstream `UserWarning` fires and your electrostatics are +> physically suspect. + +#### Long-range Coulomb flags + +| Flag | Default | Effect | +|------|---------|--------| +| `--coulomb {auto,on,off}` | `auto` | Force on/off the long-range Coulomb module. `auto` defers to model. | +| `--coulomb-method {simple,dsf,ewald}` | unset | Override the model's method. | +| `--coulomb-cutoff R` | 15.0 | Cutoff in Å (used by dsf/ewald). Rejected without `--coulomb-method`. **For `aimnet2-rxn` family, pass `--coulomb-cutoff 4.6`** (training-frozen value). | + +#### Dispersion (DFT-D3) flags + +| Flag | Default | Effect | +|------|---------|--------| +| `--dispersion {auto,on,off}` | `auto` | Force on/off the D3 module. | +| `--dftd3-cutoff R` | unset | Override D3 cutoff in Å. | +| `--dftd3-smoothing-fraction f` | unset | Override D3 smoothing fraction. | + +#### Examples + +CLI smoke (default model, water gradient): + +```bash +oet_aimnet2 input.extinp.tmp +``` + +Condensed phase, DSF Coulomb at 12 Å: + +```bash +oet_aimnet2 input.extinp.tmp --coulomb-method dsf --coulomb-cutoff 12.0 +``` + +ORCA OptTS with the NSE model on a charged radical (no `--compile`): + +```text +! OptTS +%method + ProgExt "/path/to/oet_aimnet2" + Ext_Params "-m aimnet2-nse -d cuda" +end +* xyzfile -1 2 ts_radical.xyz +``` + +ORCA NEB-CI with the RXN model (neutral H/C/N/O only; no `--compile`, +`--coulomb-cutoff 4.6`): + +```text +! NEB-CI +%neb Product "product.xyz" NImages 8 end +%method + ProgExt "/path/to/oet_aimnet2" + Ext_Params "-m aimnet2-rxn --coulomb-method dsf --coulomb-cutoff 4.6" +end +* xyzfile 0 1 reactant.xyz +``` + ## Interface All scripts must be executable as: diff --git a/pyproject.toml b/pyproject.toml index 6943fcc..3334a1f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -115,6 +115,13 @@ line-length = 100 ignore = ["E203"] +# //// PYTEST //// +[tool.pytest.ini_options] +markers = [ + "network: tests that download models from the network (run via OET_RUN_NETWORK_TESTS=1)", +] + + # //// CODESPELL //// [tool.codespell] skip = '*.pyc,*.densities*,*.gbw,./build/,./_venv' diff --git a/requirements/aimnet2.txt b/requirements/aimnet2.txt index 278698e..062c4a3 100644 --- a/requirements/aimnet2.txt +++ b/requirements/aimnet2.txt @@ -1 +1 @@ -aimnet>=0.1.0 +aimnet>=0.2,<0.3 diff --git a/src/oet/calculator/aimnet2.py b/src/oet/calculator/aimnet2.py index c6100b4..3da95b0 100755 --- a/src/oet/calculator/aimnet2.py +++ b/src/oet/calculator/aimnet2.py @@ -13,6 +13,7 @@ import shutil import sys +import warnings from argparse import ArgumentParser from pathlib import Path from typing import Any @@ -43,32 +44,41 @@ DEFAULT_MODEL_PATH = ASSETS_DIR / "aimnet2" +# Periodic table for symbol→Z conversion. +# Per-model element rejection is delegated to AIMNet2Calculator.eval() +# against model.metadata["implemented_species"]. +_PERIODIC_TABLE: tuple[str, ...] = ( + "H", "He", + "Li", "Be", "B", "C", "N", "O", "F", "Ne", + "Na", "Mg", "Al", "Si", "P", "S", "Cl", "Ar", + "K", "Ca", "Sc", "Ti", "V", "Cr", "Mn", "Fe", "Co", "Ni", + "Cu", "Zn", "Ga", "Ge", "As", "Se", "Br", "Kr", + "Rb", "Sr", "Y", "Zr", "Nb", "Mo", "Tc", "Ru", "Rh", "Pd", + "Ag", "Cd", "In", "Sn", "Sb", "Te", "I", "Xe", + "Cs", "Ba", + "La", "Ce", "Pr", "Nd", "Pm", "Sm", "Eu", "Gd", "Tb", "Dy", "Ho", "Er", "Tm", "Yb", "Lu", + "Hf", "Ta", "W", "Re", "Os", "Ir", "Pt", "Au", "Hg", "Tl", "Pb", "Bi", "Po", "At", "Rn", + "Fr", "Ra", + "Ac", "Th", "Pa", "U", "Np", "Pu", "Am", "Cm", "Bk", "Cf", "Es", "Fm", "Md", "No", "Lr", + "Rf", "Db", "Sg", "Bh", "Hs", "Mt", "Ds", "Rg", "Cn", "Nh", "Fl", "Mc", "Lv", "Ts", "Og", +) +assert len(_PERIODIC_TABLE) == 118, f"expected 118 elements, got {len(_PERIODIC_TABLE)}" +_SYMBOL_TO_Z: dict[str, int] = {sym: i + 1 for i, sym in enumerate(_PERIODIC_TABLE)} + class Aimnet2Calc(BaseCalc): - # Elements covered by AIMNet2 - ELEMENT_TO_ATOMIC_NUMBER = { - "H": 1, - "B": 5, - "C": 6, - "N": 7, - "O": 8, - "F": 9, - "Si": 14, - "P": 15, - "S": 16, - "Cl": 17, - "As": 33, - "Se": 34, - "Br": 35, - "Pd": 46, - "I": 53, + """ORCA ExtTool calculator backed by AIMNet2 (aimnet>=0.2,<0.3).""" + + # Wired into extend_parser's choices=, single source of truth. + _SUPPORTED_DEVICES: tuple[str, ...] = ("cpu", "cuda", "auto") + _COULOMB_METHODS: tuple[str, ...] = ("simple", "dsf", "ewald") + # Single source of truth for tri-state choices AND value translation. + _TRISTATE_MAP: dict[str, bool | None] = { + "auto": None, "on": True, "off": False, } - # Supported devices - supported_devices = ("cpu", "cuda", "auto") - - # AIMNet2 calculator used to compute energy and grad _calc: AIMNet2Calculator | None = None + _setup_args: frozenset | None = None def get_calculator(self) -> AIMNet2Calculator: """ @@ -81,26 +91,6 @@ def get_calculator(self) -> AIMNet2Calculator: """ return self._calc - def set_calculator(self, model: str, device: str) -> None: - """ - Set the calculator - - Parameters - ---------- - model: str - Model of the calculator - device: str - device to use - """ - if device == "cpu": - # Monkey-patch torch.cuda.is_available() to return False - # Another way to do it would be to set CUDA_VISIBLE_DEVICES="" *before* the initial import of torch - # Ideally, AIMNet2Calculator would just have an extra argument for this. - torch.cuda.is_available = lambda: False - elif device == "cuda" and not torch.cuda.is_available(): - raise RuntimeError("CUDA requested but not available") - self._calc = AIMNet2Calculator(model=model) - @staticmethod def get_model_file(model: str, model_dir: str) -> Path: """ @@ -111,7 +101,7 @@ def get_model_file(model: str, model_dir: str) -> Path: Parameters ---------- model - Model name, e.g. "aimnet2_wb97m", or filename, e.g. "aimnet2_wb97m_0.jpt", or absolute path + Model name, e.g. "aimnet2", or filename, e.g. "aimnet2-wb97m-d3_0.pt", or absolute path model_dir directory to look for or store model file @@ -125,7 +115,7 @@ def get_model_file(model: str, model_dir: str) -> Path: FileNotFoundError If the model file is given by absolute path and does not exist FileExistsError - If `model_dir` exists but is not a directory or `model_path` exists but is not a file + If `model_dir` exists but is not a directory or `cached_path` exists but is not a file """ # Check if `model` is already an absolute path if (model_path := Path(model)).is_absolute(): @@ -136,136 +126,371 @@ def get_model_file(model: str, model_dir: str) -> Path: else: # Check aliases # First, check if the model is available in the registry. If not, assume it is a - # filename and treat it as is. + # filename and treat it as is. The upstream resolver knows the canonical filename + # (including extension), so we do not append one here. model_registry = load_model_registry() if model in model_registry["aliases"]: model_file = model_registry["aliases"][model] else: model_file = model - # add jpt extension if not already present - if not model_file.endswith(".jpt"): - model_file += ".jpt" - # strip any directories - model_file = Path(model_file).name + # strip any directories for the local cache lookup + model_basename = Path(model_file).name # make sure the directory exists - model_path = Path(model_dir) - if model_path.exists() and not model_path.is_dir(): + model_dir_path = Path(model_dir) + if model_dir_path.exists() and not model_dir_path.is_dir(): raise FileExistsError(f'Path "{model_dir}" exists but is not a directory') - model_path.mkdir(parents=True, exist_ok=True) - # construct the full path - model_path = model_path / model_file - # if the file exists, pass the path to `get_model_path` - if model_path.exists(): - if model_path.is_file(): - model = str(model_path) + model_dir_path.mkdir(parents=True, exist_ok=True) + # if a cached file with the same basename already exists, hand it to `get_model_path` + cached_path = model_dir_path / model_basename + if cached_path.exists(): + if cached_path.is_file(): + model = str(cached_path) else: - raise FileExistsError(f'Path "{model_path}" exists but is not a file') + raise FileExistsError(f'Path "{cached_path}" exists but is not a file') # obtain the file from AIMNet2 try: actual_path = Path(get_model_path(model)) except HTTPError as e: - # If the URL is not found, it's possible the user requested, e.g. "aimnet2_wb97m_1.jpt" - # This is actually under "aimnet2/aimnet2_wb97m_1.jpt" and also not in the `model_registry_aliases` + # If the URL is not found, it's possible the user requested, e.g. "aimnet2_wb97m_1" + # This is actually under "aimnet2/aimnet2_wb97m_1" and also not in the `model_registry_aliases` if "/" not in model: # look for "aimnet2_..." under "aimnet2/aimnet2_..." model_subdir = model.split("_")[0] + "/" + model print( f'Failed to find model "{model}" at URL: {e.response.url}\n' - f'Trying again with model name "model_subdir"', + f'Trying again with model name "{model_subdir}"', file=sys.stderr, ) actual_path = Path(get_model_path(model_subdir)) else: raise e - # move it to the correct destination for subsequent runs - if not (model_path.exists() and model_path.samefile(actual_path)): - shutil.move(actual_path, model_path) + # The resolver returns the canonical filename (including its extension). + # Place it under `model_dir` using that filename for subsequent runs. + final_path = model_dir_path / actual_path.name + if not (final_path.exists() and final_path.samefile(actual_path)): + shutil.move(actual_path, final_path) # finally return the path - return model_path + return final_path - def setup(self, model: str, model_dir: str, device: str) -> None: + def setup( + self, + model: str, + model_dir: str, + device: str | None, + ncores: int, + *, + compile_model: bool = False, + nb_threshold: int = 120, + ensemble_member: int = 0, + coulomb: str = "auto", + dispersion: str = "auto", + coulomb_method: str | None = None, + coulomb_cutoff: float = 15.0, + dftd3_cutoff: float | None = None, + dftd3_smoothing_fraction: float | None = None, + ) -> None: + """Construct the calculator + apply post-ctor configuration. + + First call wins: subsequent calls with the same args are a no-op; + calls with different args raise (server cache is responsible for + keying on args, so a single Aimnet2Calc instance only ever sees + one configuration). Per-call work belongs in run_aimnet2. """ - Sets the calculator. Does nothing, if it is already set. + # Validate device availability BEFORE the cache short-circuit so + # a cached CPU calc does not silently service a --device cuda + # request on a no-CUDA box. + if device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA requested but not available") - Parameters - ---------- - model: str - Model that the calculator should have - model_dir: str - Path to the model files - device: str - device to use + # Normalize device for storage so "auto" and None compare equal + # in the args-match check below (both resolve to None upstream). + device_arg = None if device == "auto" else device + + # Args-match short-circuit (uses normalized device). + # frozenset of items (not tuple) so adding a new kwarg doesn't + # require editing positional indices everywhere. + new_args = frozenset({ + "model": model, + "model_dir": model_dir, + "device": device_arg, + "ncores": ncores, + "compile_model": compile_model, + "nb_threshold": nb_threshold, + "ensemble_member": ensemble_member, + "coulomb": coulomb, + "dispersion": dispersion, + "coulomb_method": coulomb_method, + "coulomb_cutoff": coulomb_cutoff, + "dftd3_cutoff": dftd3_cutoff, + "dftd3_smoothing_fraction": dftd3_smoothing_fraction, + }.items()) + if self._calc is not None: + if new_args == self._setup_args: + return + raise RuntimeError( + "Aimnet2Calc.setup() called with different args than the " + "cached calculator. Server-mode callers must key the cache " + "on setup args." + ) + + # Validate string-choice args with helpful error messages + # (argparse normally enforces this for CLI users, but direct-API + # callers benefit from the early check too). + for arg_name, arg_val, allowed in ( + ("coulomb", coulomb, tuple(self._TRISTATE_MAP)), + ("dispersion", dispersion, tuple(self._TRISTATE_MAP)), + ): + if arg_val not in allowed: + raise ValueError( + f"setup(): {arg_name}={arg_val!r} not in {allowed}" + ) + if coulomb_method is not None and coulomb_method not in self._COULOMB_METHODS: + raise ValueError( + f"setup(): coulomb_method={coulomb_method!r} not in " + f"{self._COULOMB_METHODS}" + ) + + # When coulomb_method=='simple', the cutoff is meaningless (simple has + # no cutoff in the upstream LR module). Warn loudly so users don't + # think a custom value applies. + if coulomb_method == "simple" and coulomb_cutoff != 15.0: + warnings.warn( + f"--coulomb-cutoff={coulomb_cutoff} is ignored when " + "--coulomb-method=simple (simple Coulomb has no cutoff; " + "cutoff applies only to dsf/ewald methods).", + UserWarning, + stacklevel=2, + ) + + model_path = str(self.get_model_file(model, model_dir)) + + # Warn if user picked aimnet2-pd: it carries baked-in CPCM/THF + # solvation, so energies are NOT gas-phase. Easy to miss from the + # README alone. + model_path_lower = Path(model_path).stem.lower() + if ( + "aimnet2-pd" in model_path_lower + or "aimnet2_pd" in model_path_lower + ): + warnings.warn( + "aimnet2-pd embeds B97-3c with implicit CPCM/THF solvation; " + "energies are NOT gas-phase. Do not mix aimnet2-pd energies " + "with any other family.", + UserWarning, + stacklevel=2, + ) + + # Translate tri-state values for upstream ctor. + nc = self._TRISTATE_MAP[coulomb] + nd = self._TRISTATE_MAP[dispersion] + + self._calc = AIMNet2Calculator( + model_path, + nb_threshold=nb_threshold, + needs_coulomb=nc, + needs_dispersion=nd, + device=device_arg, + compile_model=compile_model, + ensemble_member=ensemble_member, + ) + + # Post-ctor configuration. If any of these raise, roll back via + # release() so device-resident tensors that the ctor moved to CUDA + # are walked back to CPU before the reference is dropped (otherwise + # a repeat-bad-config server loop accumulates VRAM). + try: + # Skip cutoff for "simple" mode (no upstream cutoff parameter). + if coulomb_method is not None: + if coulomb_method == "simple": + self._calc.set_lrcoulomb_method(coulomb_method) + else: + self._calc.set_lrcoulomb_method(coulomb_method, cutoff=coulomb_cutoff) + + if dftd3_cutoff is not None or dftd3_smoothing_fraction is not None: + self._calc.set_dftd3_cutoff(dftd3_cutoff, dftd3_smoothing_fraction) + + # Process-wide thread setting; one-shot here, never per-call. + torch.set_num_threads(ncores) + except Exception: + self.release() + raise + + self._setup_args = new_args + + def release(self) -> None: + """Release device-side resources held by the cached calculator. + + Server-mode callers (sibling PR + 2026-04-26-oet-server-vram-eviction-design.md) invoke this before + evicting a cached calculator from the worker cache to reclaim GPU + memory. Until that sibling PR lands, this method has no live caller. + + NOTE: torch.cuda.empty_cache() returns memory to PyTorch's caching + allocator, not to the OS — nvidia-smi reservation does not drop. + Compiled models (compile_model=True) keep Inductor kernel artifacts + in process-global state that .to('cpu') + empty_cache() cannot + reclaim; only torch._dynamo.reset() can, and that would invalidate + every other compiled module in the process. """ - if not self._calc: - # Sanitize the model path and fetch the actual file - model_path = str(self.get_model_file(model, model_dir)) - self.set_calculator(model=model_path, device=device) + if self._calc is not None: + for component_name in ("model", "external_coulomb", "external_dftd3"): + component = getattr(self._calc, component_name, None) + if component is None: + continue + try: + component.to("cpu") + except Exception as e: + # Don't let a partial-cleanup failure cascade into the + # server's eviction logic, but log so silent VRAM leaks + # are debuggable. + print( + f"Aimnet2Calc.release(): failed to move {component_name} to cpu: {e!r}", + file=sys.stderr, + ) + self._calc = None + self._setup_args = None + if torch.cuda.is_available(): + torch.cuda.empty_cache() @classmethod def extend_parser(cls, parser: ArgumentParser) -> None: - """Add AimNet2 parsing options. + """Add AIMNet2 v0.2 options to an argument parser. - Parameters - ---------- - parser: ArgumentParser - Parser that should be extended + The same parser is used by both the standalone `oet_aimnet2` CLI + and by `oet_server aimnet2` (server.py:278 calls extend_parser to + build its CLI), so flags propagate to server mode automatically. """ + # --- model selection --------------------------------------------- parser.add_argument( - "-m", - "--model", + "-m", "--model", type=str, dest="model", default="aimnet2", - help="The AIMNet2 model name, file name, or absolute path. " - "If an absolute path is given, the file must exist locally. " - "Otherwise, the model will be downloaded to DIR if needed. " - "Note that different models are available. For example, `aimnet2nse` " - "is designed for open-shell/radical systems, such as those occurring in " - "transition states during bond breaking or formation. " - 'Default: "aimnet2".', + help=( + "Model name (registry alias e.g. aimnet2, aimnet2-2025, " + "aimnet2-nse, aimnet2-rxn, aimnet2-pd), canonical key " + "(e.g. aimnet2-wb97m-d3_0), HuggingFace repo id, or " + "absolute local path to a .pt file. " + "Default: aimnet2 (= aimnet2-wb97m-d3_0). " + "For non-covalent / screening: try aimnet2-2025." + ), ) parser.add_argument( - "-p", - "--model-path", + "-p", "--model-path", metavar="DIR", dest="model_dir", type=str, default=str(DEFAULT_MODEL_PATH), - help=f'The directory to look for and store AIMNet2 model files. Default: "{DEFAULT_MODEL_PATH}". ', + help=f"Local cache directory for downloaded model files. Default: {DEFAULT_MODEL_PATH}.", ) parser.add_argument( - "-d", - "--device", + "-d", "--device", metavar="DEVICE", dest="device", type=str, - choices=["cpu", "cuda", "auto"], + choices=cls._SUPPORTED_DEVICES, default="cpu", - help="Device to perform the calculation on. " - "Options: cpu, cuda, or auto (i.e. use cuda if available, otherwise cpu). " - "Default: cpu. ", + help="Compute device. 'auto' lets upstream auto-detect (None). Default: cpu.", ) - def atomic_symbol_to_number(self, symbol: str) -> int: - """Convert an element symbol to an atomic number. + # --- performance ------------------------------------------------- + parser.add_argument( + "--compile", + dest="compile", + action="store_true", + help=( + "Enable torch.compile JIT. SERVER MODE ONLY - standalone " + "oet_aimnet2 is a fresh process per ORCA call and re-pays " + "JIT cost every step. First-call latency 10-60s. Recompiles " + "on shape change. Incompatible with Hessian. Do NOT use with " + "NEB / OptTS / IRC." + ), + ) + parser.add_argument( + "--nb-threshold", + dest="nb_threshold", + type=int, + default=120, + help="Adaptive neighbor-list batch size. Default: 120.", + ) + parser.add_argument( + "--ensemble-member", + dest="ensemble_member", + type=int, + choices=(0, 1, 2, 3), + default=0, + help="Use a single ensemble member instead of the mean. Default: 0.", + ) - Parameters - ---------- - symbol: str - Element symbol, e.g. "Cl" + # --- long-range Coulomb ------------------------------------------ + parser.add_argument( + "--coulomb", + dest="coulomb", + type=str, + choices=tuple(cls._TRISTATE_MAP), + default="auto", + help="Force on/off the model's long-range Coulomb module. Default: auto (model decides).", + ) + parser.add_argument( + "--coulomb-method", + dest="coulomb_method", + type=str, + choices=cls._COULOMB_METHODS, + default=None, + help=( + "Override the model's long-range Coulomb method. " + "If unset, the model default is used (no post-ctor call)." + ), + ) + parser.add_argument( + "--coulomb-cutoff", + dest="coulomb_cutoff", + type=float, + default=15.0, + help=( + "Cutoff in Angstrom for dsf/ewald methods. Default: 15.0. " + "Requires --coulomb-method (rejected otherwise). " + "NOTE: aimnet2-rxn family was trained with cutoff frozen at " + "4.6 A; passing other values fires an upstream UserWarning." + ), + ) - Returns - ------- - int - atomic number of the element + # --- dispersion (DFT-D3) ----------------------------------------- + parser.add_argument( + "--dispersion", + dest="dispersion", + type=str, + choices=tuple(cls._TRISTATE_MAP), + default="auto", + help="Force on/off the model's D3 dispersion module. Default: auto (model decides).", + ) + parser.add_argument( + "--dftd3-cutoff", + dest="dftd3_cutoff", + type=float, + default=None, + help="Override D3 dispersion cutoff in Angstrom. Default: model's value.", + ) + parser.add_argument( + "--dftd3-smoothing-fraction", + dest="dftd3_smoothing_fraction", + type=float, + default=None, + help="Override D3 cutoff smoothing fraction. Default: model's value.", + ) - Raises - ------ - ValueError - if the element is not in `ELEMENT_TO_ATOMIC_NUMBER` + def atomic_symbol_to_number(self, symbol: str) -> int: + """Convert an element symbol to atomic number. + + Per-model element rejection happens upstream in + AIMNet2Calculator.eval(validate_species=True). For models that + do NOT populate metadata["implemented_species"] (e.g. legacy raw + nn.Module .pt files), upstream silently accepts ANY atomic number + and may produce undefined output for unsupported elements; OET + relies on the upstream check and does not second-guess. """ try: - return self.ELEMENT_TO_ATOMIC_NUMBER[symbol.title()] + return _SYMBOL_TO_Z[symbol.title()] except KeyError: raise ValueError(f"Unknown element symbol: {symbol}") @@ -277,7 +502,10 @@ def serialize_input( mult: int, dograd: bool, ) -> dict[str, Any]: - """Serialize the input data into kwargs for AIMNet2Calculator + """Build kwargs for AIMNet2Calculator.eval(). + + `mult` is only included for NSE-class models; non-NSE models reject + the key in v0.2. Parameters ---------- @@ -298,13 +526,17 @@ def serialize_input( kwargs for AIMNet2Calculator.eval() """ numbers = [self.atomic_symbol_to_number(sym) for sym in atom_types] + data: dict[str, Any] = { + "coord": [coordinates], + "numbers": [numbers], + "charge": [charge], + } + # run_aimnet2 guarantees _calc is set before this is reached; + # the previous defensive `_calc is not None` was unreachable. + if self._calc.is_nse: + data["mult"] = [mult] return { - "data": { - "coord": [coordinates], - "numbers": [numbers], - "charge": [charge], - "mult": [mult], - }, + "data": data, "forces": dograd, "stress": False, "hessian": False, @@ -336,9 +568,6 @@ def run_aimnet2( Flattened gradient vector (Eh/Bohr), if computed, otherwise empty. """ - # set the number of threads - torch.set_num_threads(calc_data.ncores) - # make ase atoms object for calculation aimnet2_input = self.serialize_input( atom_types=atom_types, @@ -355,6 +584,8 @@ def run_aimnet2( energy = float(results["energy"].detach()) / ENERGY_CONVERSION["eV"] gradient = [] if (forces := results.get("forces", None)) is not None: + # detach() defends against any future create_graph=True default. + forces = forces.detach() # unit conversion & factor of -1 to convert from forces to gradient fac = -LENGTH_CONVERSION["Ang"] / ENERGY_CONVERSION["eV"] gradient = (forces * fac).flatten().tolist() @@ -367,60 +598,68 @@ def calc( args_parsed: dict[str, Any], args_not_parsed: list[str], ) -> tuple[float, list[float]]: - """ - Routine for calculating energy and optional gradient. - Writes ORCA output + """Routine for calculating energy + optional gradient. - - Parameters - ---------- - calc_data: CalculationData - Object with calculation data for the run - args_parsed: dict[str, Any] - Arguments parsed as defined in extend_parser - args_not_parsed: list[str] - Arguments not parsed so far - - Returns - ------- - float - The computed energy (Eh) - list[float] - Flattened gradient vector (Eh/Bohr), if computed, otherwise empty + Validates cross-flag constraints, then sets up the calculator and + runs run_aimnet2. """ - # Get the arguments parsed as defined in extend_parser - model = args_parsed.get("model") + # --- cross-flag validation ---------------------------------------- + # --coulomb-cutoff is only meaningful with --coulomb-method. + # The 15.0 != check has one corner: a user who passes --coulomb-cutoff 15.0 + # explicitly without --coulomb-method silently gets a default-cutoff + # Coulomb-disabled run rather than a rejection. They should specify + # --coulomb-method anyway, so the false-negative is acceptable. + # + # We raise SystemExit(msg) directly rather than holding a parser ref: + # argparse's parser.error() is itself just print + SystemExit(2), and + # stashing the parser in args_parsed used to poison server.py's cache + # key (frozenset of args_parsed.items()) because ArgumentParser is + # hashed by identity. SystemExit's message goes to stderr automatically. + user_set_cutoff = args_parsed.get("coulomb_cutoff", 15.0) != 15.0 + if user_set_cutoff and args_parsed.get("coulomb_method") is None: + raise SystemExit( + "oet_aimnet2: error: --coulomb-cutoff requires --coulomb-method" + ) + + # --- read parsed args (defaults match extend_parser) ------------- + model = args_parsed.get("model", "aimnet2") model_dir = args_parsed.get("model_dir") - device = str(args_parsed.get("device")) - # Check if device is allowed - if device not in tuple(self.supported_devices): + device = str(args_parsed.get("device", "cpu")) + if device not in self._SUPPORTED_DEVICES: raise RuntimeError( - f"Device {device} not applicable for AIMNet2 calculation." - "Please use `--device` to choose between {supported_devices}." + f"Device {device} not supported. Use one of {self._SUPPORTED_DEVICES}." ) - if not isinstance(model, str) or not isinstance(model_dir, str): - raise RuntimeError("Problems detecting model parameters.") - # setup calculator if not already set - # this is important as usage on a server would otherwise cause - # initialization with every call so that nothing is gained - self.setup(model=model, model_dir=model_dir, device=device) - # process the XYZ file - atom_types, coordinates = xyzfile_to_at_coord(calc_data.xyzfile) - # run uma - energy, gradient = self.run_aimnet2( - atom_types=atom_types, coordinates=coordinates, calc_data=calc_data + # --- set up calculator (idempotent) ------------------------------ + self.setup( + model=model, + model_dir=model_dir, + device=device, + ncores=calc_data.ncores, + compile_model=args_parsed.get("compile", False), + nb_threshold=args_parsed.get("nb_threshold", 120), + ensemble_member=args_parsed.get("ensemble_member", 0), + coulomb=args_parsed.get("coulomb", "auto"), + dispersion=args_parsed.get("dispersion", "auto"), + coulomb_method=args_parsed.get("coulomb_method"), + coulomb_cutoff=args_parsed.get("coulomb_cutoff", 15.0), + dftd3_cutoff=args_parsed.get("dftd3_cutoff"), + dftd3_smoothing_fraction=args_parsed.get("dftd3_smoothing_fraction"), ) - return energy, gradient + # --- read XYZ and run -------------------------------------------- + atom_types, coordinates = xyzfile_to_at_coord(calc_data.xyzfile) + return self.run_aimnet2( + atom_types=atom_types, + coordinates=coordinates, + calc_data=calc_data, + ) -def main() -> None: - """ - Main routine for execution - """ +def main(argv: list[str] | None = None) -> None: + """Main routine for execution.""" calculator = Aimnet2Calc() - inputfile, args, args_not_parsed = calculator.parse_args() + inputfile, args, args_not_parsed = calculator.parse_args(argv) calculator.run(inputfile=inputfile, args_parsed=args, args_not_parsed=args_not_parsed) diff --git a/src/oet/core/base_calc.py b/src/oet/core/base_calc.py index 09ac426..373cd01 100644 --- a/src/oet/core/base_calc.py +++ b/src/oet/core/base_calc.py @@ -340,6 +340,18 @@ def extend_parser(self, parser: ArgumentParser) -> None: """ pass + def release(self) -> None: + """ + Release device-side resources held by this calculator. + + No-op by default. Calculators that hold device-resident state + (GPU models, compiled artifacts, etc.) override to move models + to CPU, drop references, and call torch.cuda.empty_cache(). + Server-mode worker cache eviction calls this before evicting + a cached calculator. Must be idempotent. + """ + pass + def _check_python_version(self) -> None: """ Checks whether the Python version matches the minimum requirement diff --git a/tests/aimnet2/test_aiment2_client.py b/tests/aimnet2/test_aiment2_client.py index c83c4cc..e72a69b 100644 --- a/tests/aimnet2/test_aiment2_client.py +++ b/tests/aimnet2/test_aiment2_client.py @@ -71,17 +71,17 @@ def test_H2O_engrad(self): ) run_aimnet2(input_file, output_file) expected_num_atoms = 3 - expected_energy = -7.647682644970e+01 + expected_energy = -7.647682538153e+01 expected_gradients = [ - -1.020941790193e-02, - -7.558942306787e-03, - 5.339907016605e-03, - 3.577796975151e-03, - 9.023884311318e-03, - 1.832915237173e-03, - 6.631625350565e-03, - -1.464942586608e-03, - -7.172828074545e-03 + -1.020942814648e-02, + -7.558954879642e-03, + 5.339907482266e-03, + 3.577803261578e-03, + 9.023892693222e-03, + 1.832913840190e-03, + 6.631619296968e-03, + -1.464935485274e-03, + -7.172822486609e-03 ] try: @@ -109,14 +109,14 @@ def test_OH_anion_eng_grad(self): ) run_aimnet2(input_file, output_file) expected_num_atoms = 2 - expected_energy = -7.582629740302e+01 + expected_energy = -7.582629635076e+01 expected_gradients = [ - -4.858186002821e-04, - -1.563774305396e-03, - -4.455401503947e-04, - 4.858186002821e-04, - 1.563771977089e-03, - 4.455401503947e-04 + -4.858376923949e-04, + -1.563820987940e-03, + -4.455552552827e-04, + 4.858376923949e-04, + 1.563823316246e-03, + 4.455552552827e-04 ] try: @@ -144,14 +144,14 @@ def test_OH_rad_eng_grad(self): ) run_aimnet2(input_file, output_file) expected_num_atoms = 2 - expected_energy = -7.568258800204e+01 + expected_energy = -7.568258700191e+01 expected_gradients = [ - -3.783911699429e-03, - -1.217970903963e-02, - -3.470176830888e-03, - 3.783911699429e-03, - 1.217970997095e-02, - 3.470176830888e-03 + -3.783945925534e-03, + -1.217983383685e-02, + -3.470211755484e-03, + 3.783945692703e-03, + 1.217983569950e-02, + 3.470211755484e-03 ] try: @@ -162,9 +162,9 @@ def test_OH_rad_eng_grad(self): ) from e self.assertEqual(num_atoms, expected_num_atoms) - self.assertAlmostEqual(energy, expected_energy, places=7) + self.assertAlmostEqual(energy, expected_energy, places=9) for g1, g2 in zip(gradients, expected_gradients): - self.assertAlmostEqual(g1, g2, places=7) + self.assertAlmostEqual(g1, g2, places=9) if __name__ == "__main__": diff --git a/tests/aimnet2/test_aiment2_standalone.py b/tests/aimnet2/test_aiment2_standalone.py index 88baed1..4e45716 100644 --- a/tests/aimnet2/test_aiment2_standalone.py +++ b/tests/aimnet2/test_aiment2_standalone.py @@ -11,6 +11,12 @@ write_xyz_file, ) +# Reference values regenerated against aimnet v0.2 (PyPI). The v0.2 +# release ships retrained model files (storage path .../aimnet2v2/...) +# whose energies differ from v0.1.x by ~1e-6 Eh at this geometry. +# v0.2 is bit-exact deterministic across runs and between standalone +# wrapper and server paths, so places=9 holds. + # Path to the script, adjust if needed. aimnet2_script_path = ROOT_DIR / "../../bin/oet_aimnet2" @@ -34,17 +40,17 @@ def test_H2O_engrad(self): ) run_aimnet2(input_file, output_file) expected_num_atoms = 3 - expected_energy = -7.647682644970e+01 + expected_energy = -7.647682538153e+01 expected_gradients = [ - -1.020941790193e-02, - -7.558942306787e-03, - 5.339907016605e-03, - 3.577796975151e-03, - 9.023884311318e-03, - 1.832915237173e-03, - 6.631625350565e-03, - -1.464942586608e-03, - -7.172828074545e-03 + -1.020942814648e-02, + -7.558954879642e-03, + 5.339907482266e-03, + 3.577803261578e-03, + 9.023892693222e-03, + 1.832913840190e-03, + 6.631619296968e-03, + -1.464935485274e-03, + -7.172822486609e-03 ] try: @@ -72,14 +78,14 @@ def test_OH_anion_eng_grad(self): ) run_aimnet2(input_file, output_file) expected_num_atoms = 2 - expected_energy = -7.582629740302e+01 + expected_energy = -7.582629635076e+01 expected_gradients = [ - -4.858186002821e-04, - -1.563774305396e-03, - -4.455401503947e-04, - 4.858186002821e-04, - 1.563771977089e-03, - 4.455401503947e-04 + -4.858376923949e-04, + -1.563820987940e-03, + -4.455552552827e-04, + 4.858376923949e-04, + 1.563823316246e-03, + 4.455552552827e-04 ] try: @@ -107,14 +113,14 @@ def test_OH_rad_eng_grad(self): ) run_aimnet2(input_file, output_file) expected_num_atoms = 2 - expected_energy = -7.568258800204e+01 + expected_energy = -7.568258700191e+01 expected_gradients = [ - -3.783911699429e-03, - -1.217970903963e-02, - -3.470176830888e-03, - 3.783911699429e-03, - 1.217970997095e-02, - 3.470176830888e-03 + -3.783945925534e-03, + -1.217983383685e-02, + -3.470211755484e-03, + 3.783945692703e-03, + 1.217983569950e-02, + 3.470211755484e-03 ] try: @@ -125,9 +131,9 @@ def test_OH_rad_eng_grad(self): ) from e self.assertEqual(num_atoms, expected_num_atoms) - self.assertAlmostEqual(energy, expected_energy, places=7) + self.assertAlmostEqual(energy, expected_energy, places=9) for g1, g2 in zip(gradients, expected_gradients): - self.assertAlmostEqual(g1, g2, places=7) + self.assertAlmostEqual(g1, g2, places=9) if __name__ == "__main__": diff --git a/tests/aimnet2/test_aiment2_v02_contract.py b/tests/aimnet2/test_aiment2_v02_contract.py new file mode 100644 index 0000000..60f8fbe --- /dev/null +++ b/tests/aimnet2/test_aiment2_v02_contract.py @@ -0,0 +1,490 @@ +"""Contract tests for the AIMNet2 v0.2 salvage. + +Lives alongside test_aiment2_standalone.py / test_aiment2_client.py +(unittest-style; preserved untouched). New behaviors get pytest-style +tests here. +""" +import argparse +import os +from unittest.mock import MagicMock + +import pytest + +from oet.calculator.aimnet2 import Aimnet2Calc + + +@pytest.fixture +def extinp_water_fixture(tmp_path): + """Build a minimal ORCA extinp + xyz for water (charge=0, mult=1, ncores=1). + + O-H stretched to 1.06 A so the gradient is non-trivial — useful for + smoke tests that want to exercise the gradient pathway, not just + assert finiteness near equilibrium. + """ + xyz = tmp_path / "h2o_EXT.xyz" + xyz.write_text( + "3\n\n" + "O 0.000000 0.000000 0.000000\n" + "H 0.000000 0.000000 1.060000\n" + "H 0.930000 0.000000 -0.240000\n" + ) + extinp = tmp_path / "h2o_EXT.extinp.tmp" + # extinp format: xyz_filename, charge, mult, ncores, dograd, [pointcharges] + extinp.write_text(f"{xyz}\n0\n1\n1\n1\n") + return extinp + + +@pytest.fixture +def extinp_h2_no_grad(tmp_path): + """Minimal H2 extinp with dograd=0 — used by argparse-rejection tests + that need a valid extinp but don't actually run the calculator. + """ + xyz = tmp_path / "h2.xyz" + xyz.write_text("2\n\nH 0 0 0\nH 0 0 0.74\n") + extinp = tmp_path / "h2_EXT.extinp.tmp" + extinp.write_text(f"{xyz}\n0\n1\n1\n0\n") + return extinp + + +class TestPeriodicTableCoverage: + """Test C from the design spec — periodic-table translator coverage.""" + + def setup_method(self): + self.calc = Aimnet2Calc() + + def test_first_row(self): + assert self.calc.atomic_symbol_to_number("H") == 1 + assert self.calc.atomic_symbol_to_number("He") == 2 + + def test_main_group(self): + assert self.calc.atomic_symbol_to_number("C") == 6 + assert self.calc.atomic_symbol_to_number("O") == 8 + assert self.calc.atomic_symbol_to_number("Cl") == 17 + + def test_transition_metal(self): + assert self.calc.atomic_symbol_to_number("Pd") == 46 + + def test_actinide(self): + assert self.calc.atomic_symbol_to_number("U") == 92 + + def test_transactinide(self): + assert self.calc.atomic_symbol_to_number("Og") == 118 + + def test_case_insensitive(self): + assert self.calc.atomic_symbol_to_number("h") == 1 + assert self.calc.atomic_symbol_to_number("cl") == 17 + assert self.calc.atomic_symbol_to_number("CL") == 17 + + def test_unknown_raises(self): + with pytest.raises(ValueError, match="Unknown element symbol"): + self.calc.atomic_symbol_to_number("Xx") + with pytest.raises(ValueError, match="Unknown element symbol"): + self.calc.atomic_symbol_to_number("Bq") # ghost atoms not supported + + +class TestNseMultGating: + """Test B from the design spec - mult key only included for NSE models.""" + + def setup_method(self): + self.calc = Aimnet2Calc() + + def _make_input_kwargs(self, is_nse: bool): + self.calc._calc = MagicMock() + self.calc._calc.is_nse = is_nse + return self.calc.serialize_input( + atom_types=["H", "H"], + coordinates=[(0.0, 0.0, 0.0), (0.0, 0.0, 0.74)], + charge=0, + mult=1, + dograd=True, + ) + + def test_mult_omitted_for_non_nse_model(self): + out = self._make_input_kwargs(is_nse=False) + assert "mult" not in out["data"] + assert "coord" in out["data"] + assert "numbers" in out["data"] + assert "charge" in out["data"] + + def test_mult_included_for_nse_model(self): + out = self._make_input_kwargs(is_nse=True) + assert out["data"]["mult"] == [1] + + def test_forces_passes_through_dograd(self): + out = self._make_input_kwargs(is_nse=False) + assert out["forces"] is True + assert out["stress"] is False + assert out["hessian"] is False + + +class TestSetupSignature: + """Sanity checks on the new setup() method (no model load).""" + + def test_cuda_raises_when_unavailable(self, monkeypatch): + monkeypatch.setattr("torch.cuda.is_available", lambda: False) + calc = Aimnet2Calc() + with pytest.raises(RuntimeError, match="CUDA requested but not available"): + calc.setup( + model="aimnet2", model_dir="/tmp", device="cuda", ncores=1 + ) + + def test_args_match_short_circuit(self): + """Second setup() call with same args is a no-op.""" + calc = Aimnet2Calc() + # Pre-populate as if first setup ran. MagicMock() is sufficient + # here because the args-match branch returns before touching + # any method on _calc. + calc._calc = MagicMock() + calc._setup_args = frozenset({ + "model": "aimnet2", "model_dir": "/tmp", "device": "cpu", + "ncores": 1, "compile_model": False, "nb_threshold": 120, + "ensemble_member": 0, "coulomb": "auto", "dispersion": "auto", + "coulomb_method": None, "coulomb_cutoff": 15.0, + "dftd3_cutoff": None, "dftd3_smoothing_fraction": None, + }.items()) + # Same args -> no-op (does not raise, does not touch _calc). + calc.setup( + model="aimnet2", model_dir="/tmp", device="cpu", ncores=1, + ) + assert calc._calc is not None # still the original mock + + def test_args_mismatch_raises(self): + """Second setup() call with different args raises.""" + calc = Aimnet2Calc() + # MagicMock() is sufficient here because the args-mismatch branch + # raises before touching any method on _calc. + calc._calc = MagicMock() + calc._setup_args = frozenset({ + "model": "aimnet2", "model_dir": "/tmp", "device": "cpu", + "ncores": 1, "compile_model": False, "nb_threshold": 120, + "ensemble_member": 0, "coulomb": "auto", "dispersion": "auto", + "coulomb_method": None, "coulomb_cutoff": 15.0, + "dftd3_cutoff": None, "dftd3_smoothing_fraction": None, + }.items()) + with pytest.raises(RuntimeError, match="different args"): + calc.setup( + model="aimnet2-2025", model_dir="/tmp", device="cpu", ncores=1, + ) + + def test_auto_and_none_device_compare_equal(self): + """device='auto' and device=None should both normalize to None + in _setup_args, so back-to-back setup() calls swapping between + them are idempotent (not a different-args raise).""" + calc = Aimnet2Calc() + calc._calc = MagicMock() + # Stash as if first call was device="auto" (normalized to None) + calc._setup_args = frozenset({ + "model": "aimnet2", "model_dir": "/tmp", "device": None, + "ncores": 1, "compile_model": False, "nb_threshold": 120, + "ensemble_member": 0, "coulomb": "auto", "dispersion": "auto", + "coulomb_method": None, "coulomb_cutoff": 15.0, + "dftd3_cutoff": None, "dftd3_smoothing_fraction": None, + }.items()) + # Second call with device=None — same after normalization → no-op + calc.setup( + model="aimnet2", model_dir="/tmp", device=None, ncores=1, + ) + assert calc._calc is not None + + def test_invalid_coulomb_choice_raises(self, monkeypatch): + monkeypatch.setattr("torch.cuda.is_available", lambda: False) + calc = Aimnet2Calc() + with pytest.raises(ValueError, match="coulomb='yes'"): + calc.setup( + model="aimnet2", model_dir="/tmp", device="cpu", ncores=1, + coulomb="yes", + ) + + def test_invalid_coulomb_method_raises(self, monkeypatch): + monkeypatch.setattr("torch.cuda.is_available", lambda: False) + calc = Aimnet2Calc() + with pytest.raises(ValueError, match="coulomb_method='wolf'"): + calc.setup( + model="aimnet2", model_dir="/tmp", device="cpu", ncores=1, + coulomb_method="wolf", + ) + + +class TestReleaseHook: + """Server-side eviction calls release(); verify it clears state.""" + + def test_release_drops_calc_and_args(self, monkeypatch): + monkeypatch.setattr("torch.cuda.is_available", lambda: False) + calc = Aimnet2Calc() + calc._calc = MagicMock() + calc._setup_args = frozenset({"sentinel": "dummy"}.items()) + calc.release() + assert calc._calc is None + assert calc._setup_args is None + + def test_release_no_op_on_uninitialized(self, monkeypatch): + monkeypatch.setattr("torch.cuda.is_available", lambda: False) + calc = Aimnet2Calc() + # Must not raise even though _calc was never set. + calc.release() + assert calc._calc is None + + def test_release_then_setup_with_different_args(self, monkeypatch): + """After release(), a fresh setup() with different args should work + (release() clears _setup_args so the args-mismatch check passes). + Verifies the lifecycle: setup() -> release() -> setup() -> ... + """ + from unittest.mock import patch + monkeypatch.setattr("torch.cuda.is_available", lambda: False) + calc = Aimnet2Calc() + # Pre-populate with first-setup state + calc._calc = MagicMock() + calc._setup_args = frozenset({"sentinel": "first"}.items()) + + calc.release() + assert calc._calc is None + assert calc._setup_args is None + + # Second setup with completely different args should construct fresh + # (not raise the "different args" RuntimeError). + with patch("oet.calculator.aimnet2.AIMNet2Calculator") as MockCalc: + MockCalc.return_value = MagicMock() + calc.setup( + model="aimnet2-2025", model_dir="/tmp", device="cpu", ncores=2, + ) + assert calc._calc is not None + assert calc._setup_args is not None + + def test_setup_post_ctor_failure_clears_calc(self, monkeypatch): + """If set_lrcoulomb_method raises after the ctor moved the model to + device, setup() must release device-resident state before re-raising + (otherwise a repeat-bad-config server loop accumulates VRAM). + """ + from unittest.mock import patch + monkeypatch.setattr("torch.cuda.is_available", lambda: False) + calc = Aimnet2Calc() + + bad_calc = MagicMock() + bad_calc.set_lrcoulomb_method.side_effect = RuntimeError("bad config") + bad_calc.model = MagicMock() + bad_calc.external_coulomb = MagicMock() + bad_calc.external_dftd3 = None + + with patch("oet.calculator.aimnet2.AIMNet2Calculator", return_value=bad_calc): + with pytest.raises(RuntimeError, match="bad config"): + calc.setup( + model="aimnet2", model_dir="/tmp", device="cpu", ncores=2, + coulomb_method="dsf", coulomb_cutoff=12.0, + ) + + # Rollback contract: components walked back to cpu BEFORE _calc dropped. + bad_calc.model.to.assert_called_with("cpu") + bad_calc.external_coulomb.to.assert_called_with("cpu") + assert calc._calc is None + assert calc._setup_args is None + + +class TestBaseCalcReleaseDefault: + """BaseCalc.release() exists, is callable, and is a no-op by default + (M4 — makes the sibling server-VRAM PR genuinely order-independent + by ensuring polymorphic `calc.release()` calls are safe on any + BaseCalc subclass that hasn't overridden it). + """ + + def test_release_default_is_callable_and_no_op(self): + from argparse import ArgumentParser + + from oet.core.base_calc import BaseCalc + + class _StubCalc(BaseCalc): + def calc(self, calc_data, args_parsed, args_not_parsed): # pragma: no cover + return 0.0, [] + + stub = _StubCalc() + # Must not raise; must return None. + assert stub.release() is None + # Idempotent. + assert stub.release() is None + + +class TestArgparseContract: + """Test A from the design spec - flag presence, defaults, choices, mappings. + + Catches argparse-side regressions (rename, accidental drop) at near-zero + cost. No network, no model load. + """ + + def _make_parser(self): + parser = argparse.ArgumentParser() + Aimnet2Calc.extend_parser(parser) + # extend_parser doesn't add the inputfile positional (BaseCalc does); + # add a stub so parse_args() works for tests. + if not any(a.dest == "inputfile" for a in parser._actions): + parser.add_argument("inputfile", nargs="?", default="dummy.tmp") + return parser + + def test_all_v02_flags_present(self): + parser = self._make_parser() + all_opts = {opt for a in parser._actions for opt in a.option_strings} + for flag in [ + "-m", "--model", + "-p", "--model-path", + "-d", "--device", + "--compile", + "--nb-threshold", + "--ensemble-member", + "--coulomb", + "--coulomb-method", + "--coulomb-cutoff", + "--dispersion", + "--dftd3-cutoff", + "--dftd3-smoothing-fraction", + ]: + assert flag in all_opts, f"{flag!r} missing from parser options" + + def test_dropped_flags_absent(self): + parser = self._make_parser() + all_opts = {opt for a in parser._actions for opt in a.option_strings} + for flag in [ + "--revision", + "--token", + "--dsf-alpha", + "--ewald-accuracy", + "--needs-coulomb", + "--no-needs-coulomb", + "--needs-dispersion", + "--no-needs-dispersion", + ]: + assert flag not in all_opts, f"{flag!r} should have been dropped" + + def test_defaults(self): + parser = self._make_parser() + ns = parser.parse_args(["dummy.tmp"]) + assert ns.model == "aimnet2" + assert ns.device == "cpu" + assert ns.compile is False + assert ns.nb_threshold == 120 + assert ns.ensemble_member == 0 + assert ns.coulomb == "auto" + assert ns.dispersion == "auto" + assert ns.coulomb_method is None + assert ns.coulomb_cutoff == 15.0 + assert ns.dftd3_cutoff is None + assert ns.dftd3_smoothing_fraction is None + + def test_device_choices(self): + parser = self._make_parser() + ns = parser.parse_args(["--device", "auto", "dummy.tmp"]) + assert ns.device == "auto" + with pytest.raises(SystemExit): + parser.parse_args(["--device", "tpu", "dummy.tmp"]) + + def test_coulomb_choices(self): + parser = self._make_parser() + ns = parser.parse_args(["--coulomb", "on", "dummy.tmp"]) + assert ns.coulomb == "on" + with pytest.raises(SystemExit): + parser.parse_args(["--coulomb", "yes", "dummy.tmp"]) + + def test_coulomb_method_choices(self): + parser = self._make_parser() + for m in ("simple", "dsf", "ewald"): + ns = parser.parse_args(["--coulomb-method", m, "dummy.tmp"]) + assert ns.coulomb_method == m + with pytest.raises(SystemExit): + parser.parse_args(["--coulomb-method", "wolf", "dummy.tmp"]) + + def test_ensemble_member_choices(self): + parser = self._make_parser() + # Valid choices 0..3 + for n in (0, 1, 2, 3): + ns = parser.parse_args(["--ensemble-member", str(n), "dummy.tmp"]) + assert ns.ensemble_member == n + # Out of range rejects + with pytest.raises(SystemExit): + parser.parse_args(["--ensemble-member", "4", "dummy.tmp"]) + + def test_dispersion_choices(self): + parser = self._make_parser() + ns = parser.parse_args(["--dispersion", "off", "dummy.tmp"]) + assert ns.dispersion == "off" + with pytest.raises(SystemExit): + parser.parse_args(["--dispersion", "yes", "dummy.tmp"]) + + def test_coulomb_cutoff_without_method_rejects(self, extinp_h2_no_grad, monkeypatch): + """--coulomb-cutoff without --coulomb-method must raise SystemExit. + + The validation is post-parse (in calc()), so we must invoke the + full main() entry point — parser.parse_args() alone wouldn't + trigger it. + """ + # Force CPU so the test doesn't try to use CUDA in CI + monkeypatch.setattr("torch.cuda.is_available", lambda: False) + + from oet.calculator.aimnet2 import main + + with pytest.raises(SystemExit): + main([str(extinp_h2_no_grad), "--coulomb-cutoff", "12.0"]) + + +@pytest.mark.network +@pytest.mark.skipif( + os.environ.get("OET_RUN_NETWORK_TESTS") != "1", + reason="set OET_RUN_NETWORK_TESTS=1 to run network tests", +) +class TestSmokeDefaultModel: + """Test D from the design spec — end-to-end on the default model. + + Downloads aimnet2 (= aimnet2-wb97m-d3_0) on first run; subsequent + runs use the local cache. + """ + + def test_default_model_round_trip(self, extinp_water_fixture, tmp_path): + """Default model loads, runs SP+grad, writes engrad parseable by ORCA. + + Uses the perturbed-water fixture (one O-H stretched to 1.06 A) so the + gradient is non-trivial; the bound below catches sign-flip / unit + bugs that an equilibrium geometry would silently mask. + """ + from oet.calculator.aimnet2 import main + + main([str(extinp_water_fixture)]) + + # Assert the engrad file exists and parses. + engrad = tmp_path / "h2o_EXT.engrad" + assert engrad.exists(), ( + f"engrad file not written; tmp_path contains: {list(tmp_path.iterdir())}" + ) + text = engrad.read_text() + + # Parse the engrad: skip comments and the natoms line, then + # the next float is the energy and the following 9 are gradient. + floats: list[float] = [] + for line in text.splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue + try: + floats.append(float(line)) + except ValueError: + continue + # First float is the natoms-line value (3); skip it. + # Then: energy + 9 gradient components. + assert floats[0] == 3.0, ( + f"first non-comment line should be natoms=3, got {floats[0]}\n" + f"engrad text:\n{text}" + ) + energy = floats[1] + grad = floats[2:11] + assert -100.0 < energy < 0.0, ( + f"water energy {energy} Eh implausible\nengrad text:\n{text}" + ) + assert len(grad) == 9, ( + f"expected 9 grad components, got {len(grad)}\nengrad text:\n{text}" + ) + # Stretched O-H at 1.06 A gives a non-trivial gradient on the + # stretched H (~0.03-0.10 Eh/Bohr). Assert at least one component + # is in the order-of-magnitude band for a real gradient (catches + # sign-flip / unit bugs that would either zero the gradient or + # blow it up). + max_grad = max(abs(g) for g in grad) + assert 0.005 < max_grad < 0.5, ( + f"max grad component {max_grad} outside expected (0.005, 0.5) Eh/Bohr " + f"— possible unit/sign bug\nengrad text:\n{text}" + ) + From 6183ea65e6b16d418c2501bea69347cb1d9cd43c Mon Sep 17 00:00:00 2001 From: Olexandr Isayev Date: Sun, 3 May 2026 18:27:11 -0400 Subject: [PATCH 2/4] Fix mypy errors in Aimnet2Calc reported by OET CI The squashed v0.2 support commit passed pytest locally but tripped mypy's strict mode in OET's nox -t static_check on all five OS x Python combinations. Three errors, all type-annotation drift: - _setup_args was annotated `frozenset | None` (missing type args for the generic). Tightened to `frozenset[tuple[str, Any]] | None` to match what setup() actually stores. - serialize_input dereferences self._calc.is_nse without an explicit None-check; the runtime invariant is that run_aimnet2 (the only caller) has called setup() first, but mypy can't see method ordering. Added an `assert self._calc is not None` to make the invariant explicit for the type-checker. - args_parsed.get("model_dir") had no default, so its type became `Any | None`, which is incompatible with setup()'s `model_dir: str` parameter. Mirrored the argparse default `str(DEFAULT_MODEL_PATH)` at the .get() call site so the inferred type is `Any` (compatible). Local mypy now reports "Success: no issues found in 18 source files". Pytest contract suite remains 31/31 green (incl. 1 network smoke). No runtime behavior change. --- src/oet/calculator/aimnet2.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/oet/calculator/aimnet2.py b/src/oet/calculator/aimnet2.py index 3da95b0..6e92c98 100755 --- a/src/oet/calculator/aimnet2.py +++ b/src/oet/calculator/aimnet2.py @@ -78,7 +78,7 @@ class Aimnet2Calc(BaseCalc): } _calc: AIMNet2Calculator | None = None - _setup_args: frozenset | None = None + _setup_args: frozenset[tuple[str, Any]] | None = None def get_calculator(self) -> AIMNet2Calculator: """ @@ -533,6 +533,7 @@ def serialize_input( } # run_aimnet2 guarantees _calc is set before this is reached; # the previous defensive `_calc is not None` was unreachable. + assert self._calc is not None # for mypy; runtime invariant per caller if self._calc.is_nse: data["mult"] = [mult] return { @@ -623,7 +624,9 @@ def calc( # --- read parsed args (defaults match extend_parser) ------------- model = args_parsed.get("model", "aimnet2") - model_dir = args_parsed.get("model_dir") + # Mirror the argparse default so mypy sees a guaranteed-str (the + # parser sets default=str(DEFAULT_MODEL_PATH)). + model_dir = args_parsed.get("model_dir", str(DEFAULT_MODEL_PATH)) device = str(args_parsed.get("device", "cpu")) if device not in self._SUPPORTED_DEVICES: raise RuntimeError( From c51d245848962a5d1571a312149dc607cd061905 Mon Sep 17 00:00:00 2001 From: Olexandr Isayev Date: Sun, 3 May 2026 18:34:09 -0400 Subject: [PATCH 3/4] Pass OET nox -t static_check (codespell ignore + ruff format) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI's static_check tag failed after the mypy fix on a second issue — codespell flagged element symbols Te (tellurium) and Nd (neodymium) in the new 118-element _PERIODIC_TABLE, plus a local variable `nd` in setup() (short for needs_dispersion). And ruff format wanted the periodic table tuple, _TRISTATE_MAP, and a few argparse blocks re-laid out one-element-per-line. pyproject.toml Add ignore-words-list = "te,nd" to [tool.codespell]. These are chemical element symbols; any future chemistry contribution to OET would hit the same false positive. Comment in the toml documents the rationale. src/oet/calculator/aimnet2.py Inline the local `nc` / `nd` shorthand for needs_coulomb / needs_dispersion straight into the AIMNet2Calculator constructor call (each was used once). Removes the codespell hit and improves readability (no jump from declaration to use). Also: ruff format reflow of _PERIODIC_TABLE, _TRISTATE_MAP, and the argparse block to OET's house style (one element per line, trailing commas). tests/aimnet2/test_aiment2_*.py ruff format reflow only — trailing commas, blank lines after module docstrings, scientific-notation literal style (e+01 -> e01; both parse to the same float, no numerics change). Reference values unchanged. assertAlmostEqual tolerance unchanged. All 7 static_check nox sessions (type_check, remove_unused_imports, sort_imports, lint, format_code, spell_check, dead_code) green locally. 31/31 contract tests still pass against aimnet 0.2.0 PyPI. No runtime behavior change. --- pyproject.toml | 4 + src/oet/calculator/aimnet2.py | 200 +++++++++++++++------ tests/aimnet2/test_aiment2_client.py | 12 +- tests/aimnet2/test_aiment2_standalone.py | 12 +- tests/aimnet2/test_aiment2_v02_contract.py | 139 +++++++++----- 5 files changed, 260 insertions(+), 107 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3334a1f..bd03dac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -125,6 +125,10 @@ markers = [ # //// CODESPELL //// [tool.codespell] skip = '*.pyc,*.densities*,*.gbw,./build/,./_venv' +# Whitelist chemical element symbols that codespell mistakes for typos. +# Te = tellurium (Z=52), Nd = neodymium (Z=60) — both appear in +# src/oet/calculator/aimnet2.py's _PERIODIC_TABLE. +ignore-words-list = "te,nd" count = true summary = true quiet-level = 3 diff --git a/src/oet/calculator/aimnet2.py b/src/oet/calculator/aimnet2.py index 6e92c98..fdcdff0 100755 --- a/src/oet/calculator/aimnet2.py +++ b/src/oet/calculator/aimnet2.py @@ -48,19 +48,124 @@ # Per-model element rejection is delegated to AIMNet2Calculator.eval() # against model.metadata["implemented_species"]. _PERIODIC_TABLE: tuple[str, ...] = ( - "H", "He", - "Li", "Be", "B", "C", "N", "O", "F", "Ne", - "Na", "Mg", "Al", "Si", "P", "S", "Cl", "Ar", - "K", "Ca", "Sc", "Ti", "V", "Cr", "Mn", "Fe", "Co", "Ni", - "Cu", "Zn", "Ga", "Ge", "As", "Se", "Br", "Kr", - "Rb", "Sr", "Y", "Zr", "Nb", "Mo", "Tc", "Ru", "Rh", "Pd", - "Ag", "Cd", "In", "Sn", "Sb", "Te", "I", "Xe", - "Cs", "Ba", - "La", "Ce", "Pr", "Nd", "Pm", "Sm", "Eu", "Gd", "Tb", "Dy", "Ho", "Er", "Tm", "Yb", "Lu", - "Hf", "Ta", "W", "Re", "Os", "Ir", "Pt", "Au", "Hg", "Tl", "Pb", "Bi", "Po", "At", "Rn", - "Fr", "Ra", - "Ac", "Th", "Pa", "U", "Np", "Pu", "Am", "Cm", "Bk", "Cf", "Es", "Fm", "Md", "No", "Lr", - "Rf", "Db", "Sg", "Bh", "Hs", "Mt", "Ds", "Rg", "Cn", "Nh", "Fl", "Mc", "Lv", "Ts", "Og", + "H", + "He", + "Li", + "Be", + "B", + "C", + "N", + "O", + "F", + "Ne", + "Na", + "Mg", + "Al", + "Si", + "P", + "S", + "Cl", + "Ar", + "K", + "Ca", + "Sc", + "Ti", + "V", + "Cr", + "Mn", + "Fe", + "Co", + "Ni", + "Cu", + "Zn", + "Ga", + "Ge", + "As", + "Se", + "Br", + "Kr", + "Rb", + "Sr", + "Y", + "Zr", + "Nb", + "Mo", + "Tc", + "Ru", + "Rh", + "Pd", + "Ag", + "Cd", + "In", + "Sn", + "Sb", + "Te", + "I", + "Xe", + "Cs", + "Ba", + "La", + "Ce", + "Pr", + "Nd", + "Pm", + "Sm", + "Eu", + "Gd", + "Tb", + "Dy", + "Ho", + "Er", + "Tm", + "Yb", + "Lu", + "Hf", + "Ta", + "W", + "Re", + "Os", + "Ir", + "Pt", + "Au", + "Hg", + "Tl", + "Pb", + "Bi", + "Po", + "At", + "Rn", + "Fr", + "Ra", + "Ac", + "Th", + "Pa", + "U", + "Np", + "Pu", + "Am", + "Cm", + "Bk", + "Cf", + "Es", + "Fm", + "Md", + "No", + "Lr", + "Rf", + "Db", + "Sg", + "Bh", + "Hs", + "Mt", + "Ds", + "Rg", + "Cn", + "Nh", + "Fl", + "Mc", + "Lv", + "Ts", + "Og", ) assert len(_PERIODIC_TABLE) == 118, f"expected 118 elements, got {len(_PERIODIC_TABLE)}" _SYMBOL_TO_Z: dict[str, int] = {sym: i + 1 for i, sym in enumerate(_PERIODIC_TABLE)} @@ -74,7 +179,9 @@ class Aimnet2Calc(BaseCalc): _COULOMB_METHODS: tuple[str, ...] = ("simple", "dsf", "ewald") # Single source of truth for tri-state choices AND value translation. _TRISTATE_MAP: dict[str, bool | None] = { - "auto": None, "on": True, "off": False, + "auto": None, + "on": True, + "off": False, } _calc: AIMNet2Calculator | None = None @@ -209,21 +316,23 @@ def setup( # Args-match short-circuit (uses normalized device). # frozenset of items (not tuple) so adding a new kwarg doesn't # require editing positional indices everywhere. - new_args = frozenset({ - "model": model, - "model_dir": model_dir, - "device": device_arg, - "ncores": ncores, - "compile_model": compile_model, - "nb_threshold": nb_threshold, - "ensemble_member": ensemble_member, - "coulomb": coulomb, - "dispersion": dispersion, - "coulomb_method": coulomb_method, - "coulomb_cutoff": coulomb_cutoff, - "dftd3_cutoff": dftd3_cutoff, - "dftd3_smoothing_fraction": dftd3_smoothing_fraction, - }.items()) + new_args = frozenset( + { + "model": model, + "model_dir": model_dir, + "device": device_arg, + "ncores": ncores, + "compile_model": compile_model, + "nb_threshold": nb_threshold, + "ensemble_member": ensemble_member, + "coulomb": coulomb, + "dispersion": dispersion, + "coulomb_method": coulomb_method, + "coulomb_cutoff": coulomb_cutoff, + "dftd3_cutoff": dftd3_cutoff, + "dftd3_smoothing_fraction": dftd3_smoothing_fraction, + }.items() + ) if self._calc is not None: if new_args == self._setup_args: return @@ -241,13 +350,10 @@ def setup( ("dispersion", dispersion, tuple(self._TRISTATE_MAP)), ): if arg_val not in allowed: - raise ValueError( - f"setup(): {arg_name}={arg_val!r} not in {allowed}" - ) + raise ValueError(f"setup(): {arg_name}={arg_val!r} not in {allowed}") if coulomb_method is not None and coulomb_method not in self._COULOMB_METHODS: raise ValueError( - f"setup(): coulomb_method={coulomb_method!r} not in " - f"{self._COULOMB_METHODS}" + f"setup(): coulomb_method={coulomb_method!r} not in {self._COULOMB_METHODS}" ) # When coulomb_method=='simple', the cutoff is meaningless (simple has @@ -268,10 +374,7 @@ def setup( # solvation, so energies are NOT gas-phase. Easy to miss from the # README alone. model_path_lower = Path(model_path).stem.lower() - if ( - "aimnet2-pd" in model_path_lower - or "aimnet2_pd" in model_path_lower - ): + if "aimnet2-pd" in model_path_lower or "aimnet2_pd" in model_path_lower: warnings.warn( "aimnet2-pd embeds B97-3c with implicit CPCM/THF solvation; " "energies are NOT gas-phase. Do not mix aimnet2-pd energies " @@ -280,15 +383,11 @@ def setup( stacklevel=2, ) - # Translate tri-state values for upstream ctor. - nc = self._TRISTATE_MAP[coulomb] - nd = self._TRISTATE_MAP[dispersion] - self._calc = AIMNet2Calculator( model_path, nb_threshold=nb_threshold, - needs_coulomb=nc, - needs_dispersion=nd, + needs_coulomb=self._TRISTATE_MAP[coulomb], + needs_dispersion=self._TRISTATE_MAP[dispersion], device=device_arg, compile_model=compile_model, ensemble_member=ensemble_member, @@ -362,7 +461,8 @@ def extend_parser(cls, parser: ArgumentParser) -> None: """ # --- model selection --------------------------------------------- parser.add_argument( - "-m", "--model", + "-m", + "--model", type=str, dest="model", default="aimnet2", @@ -376,7 +476,8 @@ def extend_parser(cls, parser: ArgumentParser) -> None: ), ) parser.add_argument( - "-p", "--model-path", + "-p", + "--model-path", metavar="DIR", dest="model_dir", type=str, @@ -384,7 +485,8 @@ def extend_parser(cls, parser: ArgumentParser) -> None: help=f"Local cache directory for downloaded model files. Default: {DEFAULT_MODEL_PATH}.", ) parser.add_argument( - "-d", "--device", + "-d", + "--device", metavar="DEVICE", dest="device", type=str, @@ -618,9 +720,7 @@ def calc( # hashed by identity. SystemExit's message goes to stderr automatically. user_set_cutoff = args_parsed.get("coulomb_cutoff", 15.0) != 15.0 if user_set_cutoff and args_parsed.get("coulomb_method") is None: - raise SystemExit( - "oet_aimnet2: error: --coulomb-cutoff requires --coulomb-method" - ) + raise SystemExit("oet_aimnet2: error: --coulomb-cutoff requires --coulomb-method") # --- read parsed args (defaults match extend_parser) ------------- model = args_parsed.get("model", "aimnet2") diff --git a/tests/aimnet2/test_aiment2_client.py b/tests/aimnet2/test_aiment2_client.py index e72a69b..6fda0c0 100644 --- a/tests/aimnet2/test_aiment2_client.py +++ b/tests/aimnet2/test_aiment2_client.py @@ -71,7 +71,7 @@ def test_H2O_engrad(self): ) run_aimnet2(input_file, output_file) expected_num_atoms = 3 - expected_energy = -7.647682538153e+01 + expected_energy = -7.647682538153e01 expected_gradients = [ -1.020942814648e-02, -7.558954879642e-03, @@ -81,7 +81,7 @@ def test_H2O_engrad(self): 1.832913840190e-03, 6.631619296968e-03, -1.464935485274e-03, - -7.172822486609e-03 + -7.172822486609e-03, ] try: @@ -109,14 +109,14 @@ def test_OH_anion_eng_grad(self): ) run_aimnet2(input_file, output_file) expected_num_atoms = 2 - expected_energy = -7.582629635076e+01 + expected_energy = -7.582629635076e01 expected_gradients = [ -4.858376923949e-04, -1.563820987940e-03, -4.455552552827e-04, 4.858376923949e-04, 1.563823316246e-03, - 4.455552552827e-04 + 4.455552552827e-04, ] try: @@ -144,14 +144,14 @@ def test_OH_rad_eng_grad(self): ) run_aimnet2(input_file, output_file) expected_num_atoms = 2 - expected_energy = -7.568258700191e+01 + expected_energy = -7.568258700191e01 expected_gradients = [ -3.783945925534e-03, -1.217983383685e-02, -3.470211755484e-03, 3.783945692703e-03, 1.217983569950e-02, - 3.470211755484e-03 + 3.470211755484e-03, ] try: diff --git a/tests/aimnet2/test_aiment2_standalone.py b/tests/aimnet2/test_aiment2_standalone.py index 4e45716..519d3c2 100644 --- a/tests/aimnet2/test_aiment2_standalone.py +++ b/tests/aimnet2/test_aiment2_standalone.py @@ -40,7 +40,7 @@ def test_H2O_engrad(self): ) run_aimnet2(input_file, output_file) expected_num_atoms = 3 - expected_energy = -7.647682538153e+01 + expected_energy = -7.647682538153e01 expected_gradients = [ -1.020942814648e-02, -7.558954879642e-03, @@ -50,7 +50,7 @@ def test_H2O_engrad(self): 1.832913840190e-03, 6.631619296968e-03, -1.464935485274e-03, - -7.172822486609e-03 + -7.172822486609e-03, ] try: @@ -78,14 +78,14 @@ def test_OH_anion_eng_grad(self): ) run_aimnet2(input_file, output_file) expected_num_atoms = 2 - expected_energy = -7.582629635076e+01 + expected_energy = -7.582629635076e01 expected_gradients = [ -4.858376923949e-04, -1.563820987940e-03, -4.455552552827e-04, 4.858376923949e-04, 1.563823316246e-03, - 4.455552552827e-04 + 4.455552552827e-04, ] try: @@ -113,14 +113,14 @@ def test_OH_rad_eng_grad(self): ) run_aimnet2(input_file, output_file) expected_num_atoms = 2 - expected_energy = -7.568258700191e+01 + expected_energy = -7.568258700191e01 expected_gradients = [ -3.783945925534e-03, -1.217983383685e-02, -3.470211755484e-03, 3.783945692703e-03, 1.217983569950e-02, - 3.470211755484e-03 + 3.470211755484e-03, ] try: diff --git a/tests/aimnet2/test_aiment2_v02_contract.py b/tests/aimnet2/test_aiment2_v02_contract.py index 60f8fbe..e1f2bfc 100644 --- a/tests/aimnet2/test_aiment2_v02_contract.py +++ b/tests/aimnet2/test_aiment2_v02_contract.py @@ -4,6 +4,7 @@ (unittest-style; preserved untouched). New behaviors get pytest-style tests here. """ + import argparse import os from unittest.mock import MagicMock @@ -124,9 +125,7 @@ def test_cuda_raises_when_unavailable(self, monkeypatch): monkeypatch.setattr("torch.cuda.is_available", lambda: False) calc = Aimnet2Calc() with pytest.raises(RuntimeError, match="CUDA requested but not available"): - calc.setup( - model="aimnet2", model_dir="/tmp", device="cuda", ncores=1 - ) + calc.setup(model="aimnet2", model_dir="/tmp", device="cuda", ncores=1) def test_args_match_short_circuit(self): """Second setup() call with same args is a no-op.""" @@ -135,16 +134,29 @@ def test_args_match_short_circuit(self): # here because the args-match branch returns before touching # any method on _calc. calc._calc = MagicMock() - calc._setup_args = frozenset({ - "model": "aimnet2", "model_dir": "/tmp", "device": "cpu", - "ncores": 1, "compile_model": False, "nb_threshold": 120, - "ensemble_member": 0, "coulomb": "auto", "dispersion": "auto", - "coulomb_method": None, "coulomb_cutoff": 15.0, - "dftd3_cutoff": None, "dftd3_smoothing_fraction": None, - }.items()) + calc._setup_args = frozenset( + { + "model": "aimnet2", + "model_dir": "/tmp", + "device": "cpu", + "ncores": 1, + "compile_model": False, + "nb_threshold": 120, + "ensemble_member": 0, + "coulomb": "auto", + "dispersion": "auto", + "coulomb_method": None, + "coulomb_cutoff": 15.0, + "dftd3_cutoff": None, + "dftd3_smoothing_fraction": None, + }.items() + ) # Same args -> no-op (does not raise, does not touch _calc). calc.setup( - model="aimnet2", model_dir="/tmp", device="cpu", ncores=1, + model="aimnet2", + model_dir="/tmp", + device="cpu", + ncores=1, ) assert calc._calc is not None # still the original mock @@ -154,16 +166,29 @@ def test_args_mismatch_raises(self): # MagicMock() is sufficient here because the args-mismatch branch # raises before touching any method on _calc. calc._calc = MagicMock() - calc._setup_args = frozenset({ - "model": "aimnet2", "model_dir": "/tmp", "device": "cpu", - "ncores": 1, "compile_model": False, "nb_threshold": 120, - "ensemble_member": 0, "coulomb": "auto", "dispersion": "auto", - "coulomb_method": None, "coulomb_cutoff": 15.0, - "dftd3_cutoff": None, "dftd3_smoothing_fraction": None, - }.items()) + calc._setup_args = frozenset( + { + "model": "aimnet2", + "model_dir": "/tmp", + "device": "cpu", + "ncores": 1, + "compile_model": False, + "nb_threshold": 120, + "ensemble_member": 0, + "coulomb": "auto", + "dispersion": "auto", + "coulomb_method": None, + "coulomb_cutoff": 15.0, + "dftd3_cutoff": None, + "dftd3_smoothing_fraction": None, + }.items() + ) with pytest.raises(RuntimeError, match="different args"): calc.setup( - model="aimnet2-2025", model_dir="/tmp", device="cpu", ncores=1, + model="aimnet2-2025", + model_dir="/tmp", + device="cpu", + ncores=1, ) def test_auto_and_none_device_compare_equal(self): @@ -173,16 +198,29 @@ def test_auto_and_none_device_compare_equal(self): calc = Aimnet2Calc() calc._calc = MagicMock() # Stash as if first call was device="auto" (normalized to None) - calc._setup_args = frozenset({ - "model": "aimnet2", "model_dir": "/tmp", "device": None, - "ncores": 1, "compile_model": False, "nb_threshold": 120, - "ensemble_member": 0, "coulomb": "auto", "dispersion": "auto", - "coulomb_method": None, "coulomb_cutoff": 15.0, - "dftd3_cutoff": None, "dftd3_smoothing_fraction": None, - }.items()) + calc._setup_args = frozenset( + { + "model": "aimnet2", + "model_dir": "/tmp", + "device": None, + "ncores": 1, + "compile_model": False, + "nb_threshold": 120, + "ensemble_member": 0, + "coulomb": "auto", + "dispersion": "auto", + "coulomb_method": None, + "coulomb_cutoff": 15.0, + "dftd3_cutoff": None, + "dftd3_smoothing_fraction": None, + }.items() + ) # Second call with device=None — same after normalization → no-op calc.setup( - model="aimnet2", model_dir="/tmp", device=None, ncores=1, + model="aimnet2", + model_dir="/tmp", + device=None, + ncores=1, ) assert calc._calc is not None @@ -191,7 +229,10 @@ def test_invalid_coulomb_choice_raises(self, monkeypatch): calc = Aimnet2Calc() with pytest.raises(ValueError, match="coulomb='yes'"): calc.setup( - model="aimnet2", model_dir="/tmp", device="cpu", ncores=1, + model="aimnet2", + model_dir="/tmp", + device="cpu", + ncores=1, coulomb="yes", ) @@ -200,7 +241,10 @@ def test_invalid_coulomb_method_raises(self, monkeypatch): calc = Aimnet2Calc() with pytest.raises(ValueError, match="coulomb_method='wolf'"): calc.setup( - model="aimnet2", model_dir="/tmp", device="cpu", ncores=1, + model="aimnet2", + model_dir="/tmp", + device="cpu", + ncores=1, coulomb_method="wolf", ) @@ -230,6 +274,7 @@ def test_release_then_setup_with_different_args(self, monkeypatch): Verifies the lifecycle: setup() -> release() -> setup() -> ... """ from unittest.mock import patch + monkeypatch.setattr("torch.cuda.is_available", lambda: False) calc = Aimnet2Calc() # Pre-populate with first-setup state @@ -245,7 +290,10 @@ def test_release_then_setup_with_different_args(self, monkeypatch): with patch("oet.calculator.aimnet2.AIMNet2Calculator") as MockCalc: MockCalc.return_value = MagicMock() calc.setup( - model="aimnet2-2025", model_dir="/tmp", device="cpu", ncores=2, + model="aimnet2-2025", + model_dir="/tmp", + device="cpu", + ncores=2, ) assert calc._calc is not None assert calc._setup_args is not None @@ -256,6 +304,7 @@ def test_setup_post_ctor_failure_clears_calc(self, monkeypatch): (otherwise a repeat-bad-config server loop accumulates VRAM). """ from unittest.mock import patch + monkeypatch.setattr("torch.cuda.is_available", lambda: False) calc = Aimnet2Calc() @@ -268,8 +317,12 @@ def test_setup_post_ctor_failure_clears_calc(self, monkeypatch): with patch("oet.calculator.aimnet2.AIMNet2Calculator", return_value=bad_calc): with pytest.raises(RuntimeError, match="bad config"): calc.setup( - model="aimnet2", model_dir="/tmp", device="cpu", ncores=2, - coulomb_method="dsf", coulomb_cutoff=12.0, + model="aimnet2", + model_dir="/tmp", + device="cpu", + ncores=2, + coulomb_method="dsf", + coulomb_cutoff=12.0, ) # Rollback contract: components walked back to cpu BEFORE _calc dropped. @@ -287,7 +340,6 @@ class TestBaseCalcReleaseDefault: """ def test_release_default_is_callable_and_no_op(self): - from argparse import ArgumentParser from oet.core.base_calc import BaseCalc @@ -322,9 +374,12 @@ def test_all_v02_flags_present(self): parser = self._make_parser() all_opts = {opt for a in parser._actions for opt in a.option_strings} for flag in [ - "-m", "--model", - "-p", "--model-path", - "-d", "--device", + "-m", + "--model", + "-p", + "--model-path", + "-d", + "--device", "--compile", "--nb-threshold", "--ensemble-member", @@ -466,17 +521,12 @@ def test_default_model_round_trip(self, extinp_water_fixture, tmp_path): # First float is the natoms-line value (3); skip it. # Then: energy + 9 gradient components. assert floats[0] == 3.0, ( - f"first non-comment line should be natoms=3, got {floats[0]}\n" - f"engrad text:\n{text}" + f"first non-comment line should be natoms=3, got {floats[0]}\nengrad text:\n{text}" ) energy = floats[1] grad = floats[2:11] - assert -100.0 < energy < 0.0, ( - f"water energy {energy} Eh implausible\nengrad text:\n{text}" - ) - assert len(grad) == 9, ( - f"expected 9 grad components, got {len(grad)}\nengrad text:\n{text}" - ) + assert -100.0 < energy < 0.0, f"water energy {energy} Eh implausible\nengrad text:\n{text}" + assert len(grad) == 9, f"expected 9 grad components, got {len(grad)}\nengrad text:\n{text}" # Stretched O-H at 1.06 A gives a non-trivial gradient on the # stretched H (~0.03-0.10 Eh/Bohr). Assert at least one component # is in the order-of-magnitude band for a real gradient (catches @@ -487,4 +537,3 @@ def test_default_model_round_trip(self, extinp_water_fixture, tmp_path): f"max grad component {max_grad} outside expected (0.005, 0.5) Eh/Bohr " f"— possible unit/sign bug\nengrad text:\n{text}" ) - From 1b1662f5263a0a16b779135af4b0ee35496f875e Mon Sep 17 00:00:00 2001 From: Olexandr Isayev Date: Mon, 4 May 2026 22:45:50 -0400 Subject: [PATCH 4/4] Address PR #38 review feedback from cplett Twelve review comments, all minor / suggestion-grade. Folded into a single follow-up commit on the same branch so cplett's diff is one focused readback. PR auto-updates on push; no force-push, no rebase. README structure (#10): Move the long "AIMNet2 options" subsection out of the main README into a dedicated readmes/aimnet2.md. Main README now carries a brief pointer paragraph + link, keeping the top-level focused on installation and generic ORCA wiring as cplett requested. The AIMNet2-specific content (model-selection table, open-shell warnings, performance flags, GPU server deployment, examples) is preserved verbatim in readmes/aimnet2.md plus the reformulations below. get_model_file generalised + moved (#11): Pull the alias-lookup / cache / fetch / fallback skeleton out of Aimnet2Calc.get_model_file and into oet.core.misc.resolve_model_file. The new helper takes alias_resolver / fetch / fetch_fallback as callables, so core/ stays calculator-agnostic (no aimnet import). Aimnet2Calc.get_model_file becomes a thin wrapper that supplies aimnet's load_model_registry and get_model_path; the HTTPError-fallback "look under //" subdirectory is now an aimnet-specific lambda living next to the wrapper. Open-shell + non-NSE runtime warning (#8): serialize_input now emits a UserWarning when mult != 1 is requested with a non-NSE (closed-shell-trained) model. Tells the user the returned energy is spin-restricted and points them at -m aimnet2-nse for genuine UKS-equivalent open-shell. Two new tests verify the warning fires for non-NSE+mult=2 and is silent for NSE+mult=2 and for non-NSE+mult=1. aimnet2-rxn cutoff validation (#9): setup() now emits a UserWarning when an aimnet2-rxn family model is combined with --coulomb-method and a --coulomb-cutoff that is not the trained 4.6 A. Chose a warning over a hard quit so legitimate testing / parameter sweeps remain possible; the warning is loud enough that production users will not miss it. Three new tests cover the warn / no-warn / no-coulomb-method matrix. Numpy-style docstrings restored (#5/#6/#7): setup(), extend_parser(), and serialize_input() now carry full Parameters / Returns / Raises sections in the OET house style (matching mopac.py / xtb.py). serialize_input documents the keys of the returned kwargs dict explicitly. Help-text reformulations (#1/#2/#3/#4): --model help now points open-shell users at aimnet2-nse. --ensemble-member help is reworded so it is explicit that OET runs ONE model per call and the user must average outside ORCA. README --compile and --ensemble-member rows mirror these reformulations (in readmes/aimnet2.md). pytest dependency (#12): Add a [project.optional-dependencies] test = ["pytest>=8.0"] group so future contributors do not have to guess where pytest comes from. Existing nox sessions are unaffected. mypy: _aimnet_alias_resolver explicitly returns str | None (load_model_registry is typed as Any, which made dict.get() leak Any to mypy strict mode). All 7 nox -t static_check sessions green: type_check, remove_unused_imports, sort_imports, lint, format_code, spell_check, dead_code. 36/36 contract tests green (was 31; +5 for the new warnings) + 1 network smoke. 3/3 standalone + 3/3 client tests still pass against aimnet 0.2.0 PyPI. --- README.md | 144 +------------ pyproject.toml | 3 + readmes/aimnet2.md | 148 ++++++++++++++ src/oet/calculator/aimnet2.py | 226 ++++++++++++++------- src/oet/core/misc.py | 110 +++++++++- tests/aimnet2/test_aiment2_v02_contract.py | 105 +++++++++- 6 files changed, 521 insertions(+), 215 deletions(-) create mode 100644 readmes/aimnet2.md diff --git a/README.md b/README.md index c77a43a..83b3135 100644 --- a/README.md +++ b/README.md @@ -67,143 +67,13 @@ Provide the keyword to the client via the ORCA input line `Ext_Params "-b **Open-shell users — read this.** The default `aimnet2` (and -> `aimnet2-2025`, `aimnet2-b973c-d3*`, `aimnet2-rxn*`) are -> closed-shell-trained. Passing `mult ≠ 1` (e.g. an ORCA input with -> `* xyzfile 0 2 OH.xyz`) is **silently accepted** and produces a -> spin-restricted energy, not a true UKS-equivalent open-shell value. -> For genuine open-shell or charged-species energetics, use -> `-m aimnet2-nse`. -> -> No AIMNet2 family does broken-symmetry, so a closed-shell-singlet -> biradical (`mult = 1` with two unpaired electrons coupled -> antiferromagnetically — carbenes, nitrenes, stretched singlet σ -> bonds) is not modelable correctly by **any** model in the suite, -> NSE included. Treat such systems with a multireference method. - -#### Choosing a model - -**WARNING — energies from different model families are NOT comparable.** -Different families were trained on different reference data, with -different functionals, with or without solvation. Mixing models within -a single workflow (e.g. optimize with `aimnet2-rxn`, compute SP energies -with `aimnet2-2025`) produces wrong reaction energies of order tens of -kcal/mol — comparable to or larger than typical reaction barriers, with -no failure indication. The aimnet runtime emits a one-time -`UserWarning` when two families are constructed in the same Python -process; this is informational, not a bug. - -Recommended pattern: TS search in `aimnet2-nse` → SP energies in -`aimnet2-nse`. Never optimize in one family and compute SP in another. -The exception is `aimnet2-pd`, which can be used end-to-end for -Pd-catalyzed THF-solvated workflows but must not be mixed with any -other family at all. - -| Workflow | Recommended model | Notes | -|----------|-------------------|-------| -| Closed-shell, neutral, equilibrium geometry, gas-phase | `aimnet2` | ωB97M-D3 trained | -| Closed-shell, neutral, non-covalent / screening / large systems | `aimnet2-2025` | B97-3c trained; faster but barriers typically 3-5 kcal/mol off vs the ωB97M-D3 default. Fine for relative ranking and screening, not absolute kinetics. | -| Open-shell or charged | `aimnet2-nse` | NSE: handles arbitrary charge / multiplicity. 14-element coverage (H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I). Covers single-reference DFT regimes. NOT reliable for biradicals (two unpaired electrons on different atoms with weak coupling), near-degenerate spin states, or stretched singlet bonds (e.g. homolytic dissociation on a singlet surface). | -| Reactive (TS, NEB, IRC), broad element set | `aimnet2-nse` | Same caveats as above; covers reactive open-shell trajectories. | -| Reactive, neutral H/C/N/O only, faster | `aimnet2-rxn` | NEUTRAL only (charge ≠ 0 raises `ValueError`). H/C/N/O only (other elements raise `ValueError`). Coulomb cutoff locked at 4.6 Å (other values fire an upstream `UserWarning`). | -| Pd-containing (catalysis, organometallic) | `aimnet2-pd` | B97-3c with **implicit CPCM/THF solvation BAKED IN** — energies are NOT gas-phase. Replaces As with Pd vs default; AsR3 ligands (arsine ligands on Pd, e.g. Pd(AsPh3)4-style complexes) are not supported. | - -#### Reproducibility - -Aliases (`aimnet2`, `aimnet2-2025`, …) may be repointed in future -aimnet releases. For bit-stable reproducibility across versions, pin -to canonical keys: `-m aimnet2-wb97m-d3_0` rather than `-m aimnet2`. - -#### Performance flags - -| Flag | Default | Effect | -|------|---------|--------| -| `--compile` | False | torch.compile JIT. **Server mode only** — standalone re-pays JIT every ORCA call. First-call latency 10–60s. Recompiles on shape change (catastrophic in NEB). Incompatible with Hessian. Do not use with NEB / OptTS / IRC. For server mode, set `TORCHINDUCTOR_CACHE_DIR=/persistent/path` to keep the warm cache across worker restarts. | -| `--nb-threshold N` | 120 | Adaptive neighbor-list batch size. | -| `--ensemble-member {0,1,2,3}` | 0 | Use a single ensemble member (default 0). OET runs ONE model per call. To get an ensemble mean for production accuracy, run with each `--ensemble-member` value (0,1,2,3) and average outside ORCA. Upstream recommends ensemble averaging for production calculations. | - -#### Numerical precision - -AIMNet2 internally evaluates in float32. For large systems (energy magnitudes -above ~1000 eV), absolute energies are precise to ~1e-4 eV ≈ 4e-6 Eh. ORCA's -default `TolE=5e-6 Eh` is at this noise floor; if optimization stalls on -`Energy convergence not reached`, loosen with `! TightOpt` → `%scf TolE 1e-5 end` -or accept slightly looser convergence. - -#### GPU server deployment - -Each `oet_server aimnet2 -d cuda` worker holds the model resident on GPU -until evicted. `torch.cuda.empty_cache()` returns memory to PyTorch's allocator, -not to the OS — `nvidia-smi` shows steady-state usage equal to the sum of -resident workers' models. Plan worker count by VRAM/model-size, not per-call -peak. For multi-worker GPU deployments, set `CUDA_VISIBLE_DEVICES` per worker -or use `--workers 1`. (Server-side worker-cache eviction with `release()` is -in a sibling PR.) - -> **`aimnet2-rxn` family CAUTION**: the Coulomb cutoff is locked at 4.6 Å -> in training. Always pass `--coulomb-cutoff 4.6` when using `-m aimnet2-rxn*`, -> otherwise the upstream `UserWarning` fires and your electrostatics are -> physically suspect. - -#### Long-range Coulomb flags - -| Flag | Default | Effect | -|------|---------|--------| -| `--coulomb {auto,on,off}` | `auto` | Force on/off the long-range Coulomb module. `auto` defers to model. | -| `--coulomb-method {simple,dsf,ewald}` | unset | Override the model's method. | -| `--coulomb-cutoff R` | 15.0 | Cutoff in Å (used by dsf/ewald). Rejected without `--coulomb-method`. **For `aimnet2-rxn` family, pass `--coulomb-cutoff 4.6`** (training-frozen value). | - -#### Dispersion (DFT-D3) flags - -| Flag | Default | Effect | -|------|---------|--------| -| `--dispersion {auto,on,off}` | `auto` | Force on/off the D3 module. | -| `--dftd3-cutoff R` | unset | Override D3 cutoff in Å. | -| `--dftd3-smoothing-fraction f` | unset | Override D3 smoothing fraction. | - -#### Examples - -CLI smoke (default model, water gradient): - -```bash -oet_aimnet2 input.extinp.tmp -``` - -Condensed phase, DSF Coulomb at 12 Å: - -```bash -oet_aimnet2 input.extinp.tmp --coulomb-method dsf --coulomb-cutoff 12.0 -``` - -ORCA OptTS with the NSE model on a charged radical (no `--compile`): - -```text -! OptTS -%method - ProgExt "/path/to/oet_aimnet2" - Ext_Params "-m aimnet2-nse -d cuda" -end -* xyzfile -1 2 ts_radical.xyz -``` - -ORCA NEB-CI with the RXN model (neutral H/C/N/O only; no `--compile`, -`--coulomb-cutoff 4.6`): - -```text -! NEB-CI -%neb Product "product.xyz" NImages 8 end -%method - ProgExt "/path/to/oet_aimnet2" - Ext_Params "-m aimnet2-rxn --coulomb-method dsf --coulomb-cutoff 4.6" -end -* xyzfile 0 1 reactant.xyz -``` +`oet_aimnet2` and the `oet_server aimnet2` server expose model-selection +and performance flags specific to AIMNet2. See +[`readmes/aimnet2.md`](readmes/aimnet2.md) for the full reference: model +families and when to use which, open-shell / biradical caveats, +performance flags (`--compile`, `--nb-threshold`, `--ensemble-member`), +long-range Coulomb (`--coulomb-method`, `--coulomb-cutoff`), DFT-D3 +flags, GPU server deployment notes, and complete ORCA-input examples. ## Interface diff --git a/pyproject.toml b/pyproject.toml index bd03dac..4052799 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,6 +81,9 @@ spell-check = [ dead-code = [ "vulture>=2.14", ] +test = [ + "pytest>=8.0", +] # //// MYPY //// [tool.mypy] diff --git a/readmes/aimnet2.md b/readmes/aimnet2.md new file mode 100644 index 0000000..fd269e0 --- /dev/null +++ b/readmes/aimnet2.md @@ -0,0 +1,148 @@ +# AIMNet2 calculator — options and usage + +This document covers AIMNet2-specific options for OET's `oet_aimnet2` +wrapper. For installation, generic ORCA wiring, and server / client +basics see the main [README](../README.md). + +## Default model + +Default model is `aimnet2` (= `aimnet2-wb97m-d3_0`, ωB97M-D3 trained). +Other models are selectable via `-m`. For non-covalent or screening +work, `aimnet2-2025` (B97-3c trained, faster) is a good alternative; +for general thermochemistry and reaction barriers, `aimnet2` remains +the recommended default per upstream guidance. For open-shell or +charged systems, switch to `aimnet2-nse` (see warning below). + +> **Open-shell users — read this.** The default `aimnet2` (and +> `aimnet2-2025`, `aimnet2-b973c-d3*`, `aimnet2-rxn*`) are +> closed-shell-trained. Passing `mult ≠ 1` (e.g. an ORCA input with +> `* xyzfile 0 2 OH.xyz`) is **silently accepted** and produces a +> spin-restricted energy, not a true UKS-equivalent open-shell value. +> For genuine open-shell or charged-species energetics, use +> `-m aimnet2-nse`. OET emits a one-time `UserWarning` at calculator +> setup when this combination is detected. +> +> No AIMNet2 family does broken-symmetry, so a closed-shell-singlet +> biradical (`mult = 1` with two unpaired electrons coupled +> antiferromagnetically — carbenes, nitrenes, stretched singlet σ +> bonds) is not modelable correctly by **any** model in the suite, +> NSE included. Treat such systems with a multireference method. + +## Choosing a model + +**WARNING — energies from different model families are NOT comparable.** +Different families were trained on different reference data, with +different functionals, with or without solvation. Mixing models within +a single workflow (e.g. optimize with `aimnet2-rxn`, compute SP energies +with `aimnet2-2025`) produces wrong reaction energies of order tens of +kcal/mol — comparable to or larger than typical reaction barriers, with +no failure indication. The aimnet runtime emits a one-time +`UserWarning` when two families are constructed in the same Python +process; this is informational, not a bug. + +Recommended pattern: TS search in `aimnet2-nse` → SP energies in +`aimnet2-nse`. Never optimize in one family and compute SP in another. +The exception is `aimnet2-pd`, which can be used end-to-end for +Pd-catalyzed THF-solvated workflows but must not be mixed with any +other family at all. + +| Workflow | Recommended model | Notes | +|----------|-------------------|-------| +| Closed-shell, neutral, equilibrium geometry, gas-phase | `aimnet2` | ωB97M-D3 trained | +| Closed-shell, neutral, non-covalent / screening / large systems | `aimnet2-2025` | B97-3c trained; faster but barriers typically 3-5 kcal/mol off vs the ωB97M-D3 default. Fine for relative ranking and screening, not absolute kinetics. | +| Open-shell or charged | `aimnet2-nse` | NSE: handles arbitrary charge / multiplicity. 14-element coverage (H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I). Covers single-reference DFT regimes. NOT reliable for biradicals (two unpaired electrons on different atoms with weak coupling), near-degenerate spin states, or stretched singlet bonds (e.g. homolytic dissociation on a singlet surface). | +| Reactive (TS, NEB, IRC), broad element set | `aimnet2-nse` | Same caveats as above; covers reactive open-shell trajectories. | +| Reactive, neutral H/C/N/O only, faster | `aimnet2-rxn` | NEUTRAL only (charge ≠ 0 raises `ValueError`). H/C/N/O only (other elements raise `ValueError`). Coulomb cutoff locked at 4.6 Å (other values fire an upstream `UserWarning`; OET also warns at setup time). | +| Pd-containing (catalysis, organometallic) | `aimnet2-pd` | B97-3c with **implicit CPCM/THF solvation BAKED IN** — energies are NOT gas-phase. Replaces As with Pd vs default; AsR3 ligands (arsine ligands on Pd, e.g. Pd(AsPh3)4-style complexes) are not supported. | + +## Reproducibility + +Aliases (`aimnet2`, `aimnet2-2025`, …) may be repointed in future +aimnet releases. For bit-stable reproducibility across versions, pin +to canonical keys: `-m aimnet2-wb97m-d3_0` rather than `-m aimnet2`. + +## Performance flags + +| Flag | Default | Effect | +|------|---------|--------| +| `--compile` | False | **Server mode only.** Compile the model with `torch.compile` for faster subsequent runs after the first call (introduces an additional latency of 10–60 s on the first call). Recompiles on shape changes (catastrophic in NEB) and is incompatible with Hessians, thus **do not use with NEB / OptTS / IRC**. For server mode, set `TORCHINDUCTOR_CACHE_DIR=/persistent/path` to keep the warm cache across worker restarts. | +| `--nb-threshold N` | 120 | Adaptive neighbor-list batch size. | +| `--ensemble-member {0,1,2,3}` | 0 | Define the member of the method ensemble (default 0). OET runs ONE model per call. For production-quality accuracy, run each member (`0`,`1`,`2`,`3`) separately and average the energies / gradients outside ORCA; upstream recommends ensemble averaging for production calculations. | + +## Numerical precision + +AIMNet2 internally evaluates in float32. For large systems (energy magnitudes +above ~1000 eV), absolute energies are precise to ~1e-4 eV ≈ 4e-6 Eh. ORCA's +default `TolE=5e-6 Eh` is at this noise floor; if optimization stalls on +`Energy convergence not reached`, loosen with `! TightOpt` → `%scf TolE 1e-5 end` +or accept slightly looser convergence. + +## GPU server deployment + +Each `oet_server aimnet2 -d cuda` worker holds the model resident on GPU +until evicted. `torch.cuda.empty_cache()` returns memory to PyTorch's allocator, +not to the OS — `nvidia-smi` shows steady-state usage equal to the sum of +resident workers' models. Plan worker count by VRAM/model-size, not per-call +peak. For multi-worker GPU deployments, set `CUDA_VISIBLE_DEVICES` per worker +or use `--workers 1`. (Server-side worker-cache eviction with `release()` is +in a sibling PR.) + +> **`aimnet2-rxn` family CAUTION**: the Coulomb cutoff is locked at 4.6 Å +> in training. Always pass `--coulomb-cutoff 4.6` when using `-m aimnet2-rxn*`. +> OET emits a `UserWarning` at setup time if a different value is used; the +> upstream `UserWarning` from `aimnet` also fires, and your electrostatics +> are physically suspect outside the trained cutoff. + +## Long-range Coulomb flags + +| Flag | Default | Effect | +|------|---------|--------| +| `--coulomb {auto,on,off}` | `auto` | Force on/off the long-range Coulomb module. `auto` defers to model. | +| `--coulomb-method {simple,dsf,ewald}` | unset | Override the model's method. | +| `--coulomb-cutoff R` | 15.0 | Cutoff in Å (used by dsf/ewald). Rejected without `--coulomb-method`. **For `aimnet2-rxn` family, pass `--coulomb-cutoff 4.6`** (training-frozen value). | + +## Dispersion (DFT-D3) flags + +| Flag | Default | Effect | +|------|---------|--------| +| `--dispersion {auto,on,off}` | `auto` | Force on/off the D3 module. | +| `--dftd3-cutoff R` | unset | Override D3 cutoff in Å. | +| `--dftd3-smoothing-fraction f` | unset | Override D3 smoothing fraction. | + +## Examples + +CLI smoke (default model, water gradient): + +```bash +oet_aimnet2 input.extinp.tmp +``` + +Condensed phase, DSF Coulomb at 12 Å: + +```bash +oet_aimnet2 input.extinp.tmp --coulomb-method dsf --coulomb-cutoff 12.0 +``` + +ORCA OptTS with the NSE model on a charged radical (no `--compile`): + +```text +! OptTS +%method + ProgExt "/path/to/oet_aimnet2" + Ext_Params "-m aimnet2-nse -d cuda" +end +* xyzfile -1 2 ts_radical.xyz +``` + +ORCA NEB-CI with the RXN model (neutral H/C/N/O only; no `--compile`, +`--coulomb-cutoff 4.6`): + +```text +! NEB-CI +%neb Product "product.xyz" NImages 8 end +%method + ProgExt "/path/to/oet_aimnet2" + Ext_Params "-m aimnet2-rxn --coulomb-method dsf --coulomb-cutoff 4.6" +end +* xyzfile 0 1 reactant.xyz +``` diff --git a/src/oet/calculator/aimnet2.py b/src/oet/calculator/aimnet2.py index fdcdff0..b167e9f 100755 --- a/src/oet/calculator/aimnet2.py +++ b/src/oet/calculator/aimnet2.py @@ -11,7 +11,6 @@ Main function """ -import shutil import sys import warnings from argparse import ArgumentParser @@ -22,7 +21,12 @@ from oet import ASSETS_DIR from oet.core.base_calc import BaseCalc, CalculationData -from oet.core.misc import ENERGY_CONVERSION, LENGTH_CONVERSION, xyzfile_to_at_coord +from oet.core.misc import ( + ENERGY_CONVERSION, + LENGTH_CONVERSION, + resolve_model_file, + xyzfile_to_at_coord, +) try: from aimnet.calculators import AIMNet2Calculator @@ -199,85 +203,66 @@ def get_calculator(self) -> AIMNet2Calculator: return self._calc @staticmethod - def get_model_file(model: str, model_dir: str) -> Path: + def _aimnet_alias_resolver(name: str) -> str | None: + """Look up `name` in aimnet's model registry; return canonical + name if it is an alias, else None.""" + canonical = load_model_registry()["aliases"].get(name) + return canonical if canonical is None else str(canonical) + + @staticmethod + def _aimnet_subdir_fallback(name: str, exc: Exception) -> str | None: + """If aimnet's URL fetch 404s on a flat name, retry once under + the registry's `/` subdirectory layout (e.g. + `aimnet2_wb97m_1` is actually stored at `aimnet2/aimnet2_wb97m_1`).""" + if isinstance(exc, HTTPError) and "/" not in name: + subdir_name = name.split("_")[0] + "/" + name + print( + f'Failed to find model "{name}" at URL: {exc.response.url}\n' + f'Trying again with model name "{subdir_name}"', + file=sys.stderr, + ) + return subdir_name + return None + + @classmethod + def get_model_file(cls, model: str, model_dir: str) -> Path: """ Make sure model file exists in the correct location. If `model` is an absolute path, it must already exist. Otherwise, let AIMNet2 download it, then move it to `model_dir`. + Thin wrapper that supplies aimnet-specific alias / fetch / fallback + callables to the calculator-agnostic + :func:`oet.core.misc.resolve_model_file`. + Parameters ---------- - model - Model name, e.g. "aimnet2", or filename, e.g. "aimnet2-wb97m-d3_0.pt", or absolute path - model_dir - directory to look for or store model file + model : str + Model name, e.g. "aimnet2", or filename, e.g. + "aimnet2-wb97m-d3_0.pt", or absolute path. + model_dir : str + Directory to look for or store the model file. Returns ------- Path - Full path to the model file + Full path to the model file. Raises ------ FileNotFoundError - If the model file is given by absolute path and does not exist + If the model file is given by absolute path and does not exist. FileExistsError - If `model_dir` exists but is not a directory or `cached_path` exists but is not a file + If `model_dir` exists but is not a directory, or `cached_path` + exists but is not a file. """ - # Check if `model` is already an absolute path - if (model_path := Path(model)).is_absolute(): - if not model_path.exists(): - raise FileNotFoundError(f'Model file "{model_path}" not found') - return model_path - # `model` must be the name of a model - else: - # Check aliases - # First, check if the model is available in the registry. If not, assume it is a - # filename and treat it as is. The upstream resolver knows the canonical filename - # (including extension), so we do not append one here. - model_registry = load_model_registry() - if model in model_registry["aliases"]: - model_file = model_registry["aliases"][model] - else: - model_file = model - # strip any directories for the local cache lookup - model_basename = Path(model_file).name - # make sure the directory exists - model_dir_path = Path(model_dir) - if model_dir_path.exists() and not model_dir_path.is_dir(): - raise FileExistsError(f'Path "{model_dir}" exists but is not a directory') - model_dir_path.mkdir(parents=True, exist_ok=True) - # if a cached file with the same basename already exists, hand it to `get_model_path` - cached_path = model_dir_path / model_basename - if cached_path.exists(): - if cached_path.is_file(): - model = str(cached_path) - else: - raise FileExistsError(f'Path "{cached_path}" exists but is not a file') - # obtain the file from AIMNet2 - try: - actual_path = Path(get_model_path(model)) - except HTTPError as e: - # If the URL is not found, it's possible the user requested, e.g. "aimnet2_wb97m_1" - # This is actually under "aimnet2/aimnet2_wb97m_1" and also not in the `model_registry_aliases` - if "/" not in model: - # look for "aimnet2_..." under "aimnet2/aimnet2_..." - model_subdir = model.split("_")[0] + "/" + model - print( - f'Failed to find model "{model}" at URL: {e.response.url}\n' - f'Trying again with model name "{model_subdir}"', - file=sys.stderr, - ) - actual_path = Path(get_model_path(model_subdir)) - else: - raise e - # The resolver returns the canonical filename (including its extension). - # Place it under `model_dir` using that filename for subsequent runs. - final_path = model_dir_path / actual_path.name - if not (final_path.exists() and final_path.samefile(actual_path)): - shutil.move(actual_path, final_path) - # finally return the path - return final_path + return resolve_model_file( + model, + model_dir, + fetch=get_model_path, + alias_resolver=cls._aimnet_alias_resolver, + fetch_fallback=cls._aimnet_subdir_fallback, + ) def setup( self, @@ -302,6 +287,52 @@ def setup( calls with different args raise (server cache is responsible for keying on args, so a single Aimnet2Calc instance only ever sees one configuration). Per-call work belongs in run_aimnet2. + + Parameters + ---------- + model : str + Model name (e.g. "aimnet2"), filename, or absolute path. Forwarded + to :meth:`get_model_file` for resolution and caching. + model_dir : str + Local cache directory for downloaded model files. + device : str | None + Compute device. One of {"cpu", "cuda", "auto"}; "auto" maps to + None upstream (auto-detect). When "cuda" is requested, this + method raises if torch.cuda.is_available() is False. + ncores : int + Number of CPU threads. Sets torch.set_num_threads process-wide + once per worker; does not change per call. + compile_model : bool, default: False + If True, wrap the model with torch.compile. Server mode only; + see readmes/aimnet2.md. + nb_threshold : int, default: 120 + Adaptive neighbor-list batch size forwarded to the constructor. + ensemble_member : int, default: 0 + Which ensemble member (0..3) to load. OET runs ONE model per call. + coulomb : str, default: "auto" + Tri-state ("auto" / "on" / "off") forcing the long-range Coulomb + module on or off. "auto" defers to the model's default. + dispersion : str, default: "auto" + Tri-state for the DFT-D3 module, same semantics as `coulomb`. + coulomb_method : str | None, default: None + One of {"simple", "dsf", "ewald"} or None. If set, applied via + `set_lrcoulomb_method` after construction. + coulomb_cutoff : float, default: 15.0 + Coulomb cutoff in Ångström (used by dsf/ewald). For + `aimnet2-rxn*` models the training-frozen value is 4.6; + other values trigger a UserWarning. + dftd3_cutoff : float | None, default: None + Override the D3 cutoff via `set_dftd3_cutoff` if non-None. + dftd3_smoothing_fraction : float | None, default: None + Override the D3 smoothing fraction via `set_dftd3_cutoff` if + non-None. + + Raises + ------ + RuntimeError + If `device` is "cuda" but no CUDA device is available, or if + this method has already been called with a different argument + set on this instance. """ # Validate device availability BEFORE the cache short-circuit so # a cached CPU calc does not silently service a --device cuda @@ -383,6 +414,21 @@ def setup( stacklevel=2, ) + # Warn if user picked an aimnet2-rxn family model with a + # non-trained --coulomb-cutoff override. The rxn family was + # trained with a 4.6 A long-range cutoff; other values produce + # physically-suspect electrostatics. Only relevant when the user + # explicitly enabled the long-range method. + is_rxn_family = "aimnet2-rxn" in model_path_lower or "aimnet2_rxn" in model_path_lower + if is_rxn_family and coulomb_method is not None and coulomb_cutoff != 4.6: + warnings.warn( + f"--coulomb-cutoff={coulomb_cutoff} with -m {model} disagrees with " + "the aimnet2-rxn training cutoff (4.6 A). Pass --coulomb-cutoff 4.6 " + "for physically meaningful long-range electrostatics.", + UserWarning, + stacklevel=2, + ) + self._calc = AIMNet2Calculator( model_path, nb_threshold=nb_threshold, @@ -458,6 +504,11 @@ def extend_parser(cls, parser: ArgumentParser) -> None: The same parser is used by both the standalone `oet_aimnet2` CLI and by `oet_server aimnet2` (server.py:278 calls extend_parser to build its CLI), so flags propagate to server mode automatically. + + Parameters + ---------- + parser : ArgumentParser + Parser that should be extended. """ # --- model selection --------------------------------------------- parser.add_argument( @@ -472,7 +523,9 @@ def extend_parser(cls, parser: ArgumentParser) -> None: "(e.g. aimnet2-wb97m-d3_0), HuggingFace repo id, or " "absolute local path to a .pt file. " "Default: aimnet2 (= aimnet2-wb97m-d3_0). " - "For non-covalent / screening: try aimnet2-2025." + "For non-covalent / screening: try aimnet2-2025. " + "For open-shell or charged systems: use aimnet2-nse " + "(other families are closed-shell-trained)." ), ) parser.add_argument( @@ -521,7 +574,11 @@ def extend_parser(cls, parser: ArgumentParser) -> None: type=int, choices=(0, 1, 2, 3), default=0, - help="Use a single ensemble member instead of the mean. Default: 0.", + help=( + "Define the member of the method ensemble to use (default 0). " + "OET runs ONE model per call; for production-quality accuracy, " + "run each member (0,1,2,3) separately and average outside ORCA." + ), ) # --- long-range Coulomb ------------------------------------------ @@ -612,20 +669,31 @@ def serialize_input( Parameters ---------- atom_types : list[str] - List of element symbols (e.g., ["O", "H", "H"]) + List of element symbols (e.g. ["O", "H", "H"]). coordinates : list[tuple[float, float, float]] - List of (x, y, z) coordinates + List of (x, y, z) coordinates in Ångström. charge : int - Molecular charge + Molecular charge. mult : int - Spin multiplicity + Spin multiplicity (2S + 1). dograd : bool - Whether to compute the gradient + Whether to compute the gradient. Returns ------- dict[str, Any] - kwargs for AIMNet2Calculator.eval() + kwargs ready to splat into ``AIMNet2Calculator.eval(**...)``. + Keys: + + - ``data`` : dict with ``coord`` (list of [N, 3] coordinates), + ``numbers`` (list of [N] atomic numbers), ``charge`` + (list of [B] total charges), and — only when the loaded + model is NSE — ``mult`` (list of [B] multiplicities). + - ``forces`` : bool, mirrors `dograd`. + - ``stress`` : bool, always False (unsupported by the + ORCA external-tool protocol). + - ``hessian`` : bool, always False (Hessian sidecar is + deferred per cplett's RFC reply). """ numbers = [self.atomic_symbol_to_number(sym) for sym in atom_types] data: dict[str, Any] = { @@ -638,6 +706,20 @@ def serialize_input( assert self._calc is not None # for mypy; runtime invariant per caller if self._calc.is_nse: data["mult"] = [mult] + elif mult != 1: + # Non-NSE models are closed-shell-trained; passing mult != 1 + # silently yields a spin-restricted energy, not a true UKS- + # equivalent open-shell value. Warn loudly (once per process + # per call site, via the default warnings filter) and direct + # the user to aimnet2-nse. + warnings.warn( + f"mult={mult} requested but the loaded model is closed-shell-" + "trained (non-NSE). The returned energy is spin-restricted, NOT " + "a UKS-equivalent open-shell value. For genuine open-shell " + "energetics use -m aimnet2-nse.", + UserWarning, + stacklevel=2, + ) return { "data": data, "forces": dograd, diff --git a/src/oet/core/misc.py b/src/oet/core/misc.py index 5336858..17f8c9b 100644 --- a/src/oet/core/misc.py +++ b/src/oet/core/misc.py @@ -6,7 +6,7 @@ import shutil import subprocess import sys -from collections.abc import Iterable, Sequence +from collections.abc import Callable, Iterable, Sequence from pathlib import Path from shutil import which @@ -125,6 +125,114 @@ def check_multi_progs(keys: Sequence[str]) -> Path | None: return None +def resolve_model_file( + model: str, + model_dir: str | Path, + *, + fetch: Callable[[str], str | Path], + alias_resolver: Callable[[str], str | None] | None = None, + fetch_fallback: Callable[[str, Exception], str | None] | None = None, +) -> Path: + """ + Resolve a model name / alias / absolute path to a local model file, + downloading via `fetch` when necessary and caching the result under + `model_dir`. + + Calculator-agnostic: callers supply the alias lookup and the fetch + function so this helper does not depend on any specific MLIP package. + + Parameters + ---------- + model : str + Either an absolute path to an existing model file, or a model + name / alias / relative filename that the calculator's registry + understands. + model_dir : str | Path + Local directory used as the cache for downloaded model files. + Created if missing; must be a directory if it exists. + fetch : Callable[[str], str | Path] + Function that takes a (possibly canonicalised) model name and + returns a local filesystem path to the file (downloading on + cache miss). Used for the actual model retrieval. + alias_resolver : Callable[[str], str | None] | None + Optional alias-to-canonical-name map function. Returns the + canonical name for an alias, or None if `model` is not an alias. + When None, `model` is passed unchanged to `fetch`. + fetch_fallback : Callable[[str, Exception], str | None] | None + Optional retry hook. Called with the original name and the + exception raised by `fetch`; returns a different name to retry + with, or None to re-raise the original exception. + + Returns + ------- + Path + Full path to the (possibly newly cached) model file. + + Raises + ------ + FileNotFoundError + If `model` is an absolute path that does not exist. + FileExistsError + If `model_dir` exists but is not a directory, or the cached path + exists but is not a file. + """ + # Absolute path: must already exist; return as-is. + if (model_path := Path(model)).is_absolute(): + if not model_path.exists(): + raise FileNotFoundError(f'Model file "{model_path}" not found') + return model_path + + # Resolve alias to a canonical name (if applicable). + if alias_resolver is not None: + canonical = alias_resolver(model) + if canonical is not None: + model_file = canonical + else: + model_file = model + else: + model_file = model + + # Strip any directories for the local-cache lookup. + model_basename = Path(model_file).name + + # Make sure the cache directory exists. + model_dir_path = Path(model_dir) + if model_dir_path.exists() and not model_dir_path.is_dir(): + raise FileExistsError(f'Path "{model_dir}" exists but is not a directory') + model_dir_path.mkdir(parents=True, exist_ok=True) + + # If a cached file with the same basename already exists, hand its path + # to `fetch` so that fetchers which short-circuit on absolute paths + # return immediately. + cached_path = model_dir_path / model_basename + if cached_path.exists(): + if cached_path.is_file(): + model = str(cached_path) + else: + raise FileExistsError(f'Path "{cached_path}" exists but is not a file') + + # Obtain the file from the fetcher; allow caller to retry with a + # different name on failure (e.g. registry-subdirectory fallback). + try: + actual_path = Path(fetch(model)) + except Exception as e: + if fetch_fallback is not None: + retry_name = fetch_fallback(model, e) + if retry_name is not None: + actual_path = Path(fetch(retry_name)) + else: + raise + else: + raise + + # Move the fetched file into the cache directory under its canonical + # filename for subsequent runs. + final_path = model_dir_path / actual_path.name + if not (final_path.exists() and final_path.samefile(actual_path)): + shutil.move(actual_path, final_path) + return final_path + + def print_filecontent(outfile: str | Path) -> None: """ Print the file content, e.g. the output file, to STDOUT diff --git a/tests/aimnet2/test_aiment2_v02_contract.py b/tests/aimnet2/test_aiment2_v02_contract.py index e1f2bfc..293cfc6 100644 --- a/tests/aimnet2/test_aiment2_v02_contract.py +++ b/tests/aimnet2/test_aiment2_v02_contract.py @@ -7,7 +7,8 @@ import argparse import os -from unittest.mock import MagicMock +import warnings +from unittest.mock import MagicMock, patch import pytest @@ -117,6 +118,35 @@ def test_forces_passes_through_dograd(self): assert out["stress"] is False assert out["hessian"] is False + def test_open_shell_with_non_nse_model_warns(self): + """mult != 1 + non-NSE model emits a UserWarning that points to + aimnet2-nse (open-shell guidance).""" + self.calc._calc = MagicMock() + self.calc._calc.is_nse = False + with pytest.warns(UserWarning, match="aimnet2-nse"): + self.calc.serialize_input( + atom_types=["O", "H"], + coordinates=[(0.0, 0.0, 0.0), (0.0, 0.0, 0.97)], + charge=0, + mult=2, + dograd=True, + ) + + def test_open_shell_with_nse_model_no_warn(self): + """mult != 1 + NSE model passes mult through and does not warn.""" + self.calc._calc = MagicMock() + self.calc._calc.is_nse = True + with warnings.catch_warnings(): + warnings.simplefilter("error") # turn warnings into errors + out = self.calc.serialize_input( + atom_types=["O", "H"], + coordinates=[(0.0, 0.0, 0.0), (0.0, 0.0, 0.97)], + charge=0, + mult=2, + dograd=True, + ) + assert out["data"]["mult"] == [2] + class TestSetupSignature: """Sanity checks on the new setup() method (no model load).""" @@ -248,6 +278,75 @@ def test_invalid_coulomb_method_raises(self, monkeypatch): coulomb_method="wolf", ) + def test_rxn_with_non_trained_cutoff_warns(self, monkeypatch, tmp_path): + """aimnet2-rxn family + --coulomb-method set + cutoff != 4.6 must + emit a UserWarning pointing the user at the trained cutoff.""" + monkeypatch.setattr("torch.cuda.is_available", lambda: False) + # Stub get_model_file so we don't hit the network; produce a path + # whose stem encodes the rxn family. + rxn_path = tmp_path / "aimnet2-rxn_0.pt" + rxn_path.write_bytes(b"") + calc = Aimnet2Calc() + with patch.object(Aimnet2Calc, "get_model_file", return_value=rxn_path): + with patch("oet.calculator.aimnet2.AIMNet2Calculator") as MockCalc: + MockCalc.return_value = MagicMock() + with pytest.warns(UserWarning, match="aimnet2-rxn training cutoff"): + calc.setup( + model="aimnet2-rxn", + model_dir=str(tmp_path), + device="cpu", + ncores=1, + coulomb_method="dsf", + coulomb_cutoff=12.0, + ) + + def test_rxn_with_trained_cutoff_no_warn(self, monkeypatch, tmp_path): + """aimnet2-rxn + cutoff = 4.6 must NOT emit the rxn-cutoff warning.""" + monkeypatch.setattr("torch.cuda.is_available", lambda: False) + rxn_path = tmp_path / "aimnet2-rxn_0.pt" + rxn_path.write_bytes(b"") + calc = Aimnet2Calc() + with patch.object(Aimnet2Calc, "get_model_file", return_value=rxn_path): + with patch("oet.calculator.aimnet2.AIMNet2Calculator") as MockCalc: + MockCalc.return_value = MagicMock() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + calc.setup( + model="aimnet2-rxn", + model_dir=str(tmp_path), + device="cpu", + ncores=1, + coulomb_method="dsf", + coulomb_cutoff=4.6, + ) + rxn_warnings = [ + w for w in caught if "aimnet2-rxn training cutoff" in str(w.message) + ] + assert rxn_warnings == [] + + def test_rxn_without_coulomb_method_no_warn(self, monkeypatch, tmp_path): + """aimnet2-rxn without --coulomb-method ignores --coulomb-cutoff + entirely; warning should NOT fire on the default 15.0.""" + monkeypatch.setattr("torch.cuda.is_available", lambda: False) + rxn_path = tmp_path / "aimnet2-rxn_0.pt" + rxn_path.write_bytes(b"") + calc = Aimnet2Calc() + with patch.object(Aimnet2Calc, "get_model_file", return_value=rxn_path): + with patch("oet.calculator.aimnet2.AIMNet2Calculator") as MockCalc: + MockCalc.return_value = MagicMock() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + calc.setup( + model="aimnet2-rxn", + model_dir=str(tmp_path), + device="cpu", + ncores=1, + ) + rxn_warnings = [ + w for w in caught if "aimnet2-rxn training cutoff" in str(w.message) + ] + assert rxn_warnings == [] + class TestReleaseHook: """Server-side eviction calls release(); verify it clears state.""" @@ -273,8 +372,6 @@ def test_release_then_setup_with_different_args(self, monkeypatch): (release() clears _setup_args so the args-mismatch check passes). Verifies the lifecycle: setup() -> release() -> setup() -> ... """ - from unittest.mock import patch - monkeypatch.setattr("torch.cuda.is_available", lambda: False) calc = Aimnet2Calc() # Pre-populate with first-setup state @@ -303,8 +400,6 @@ def test_setup_post_ctor_failure_clears_calc(self, monkeypatch): device, setup() must release device-resident state before re-raising (otherwise a repeat-bad-config server loop accumulates VRAM). """ - from unittest.mock import patch - monkeypatch.setattr("torch.cuda.is_available", lambda: False) calc = Aimnet2Calc()