diff --git a/README.md b/README.md index 42aeae8..83b3135 100644 --- a/README.md +++ b/README.md @@ -65,6 +65,16 @@ 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 + +`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 All scripts must be executable as: diff --git a/pyproject.toml b/pyproject.toml index 6943fcc..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] @@ -115,9 +118,20 @@ 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' +# 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/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/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..b167e9f 100755 --- a/src/oet/calculator/aimnet2.py +++ b/src/oet/calculator/aimnet2.py @@ -11,8 +11,8 @@ Main function """ -import shutil import sys +import warnings from argparse import ArgumentParser from pathlib import Path from typing import Any @@ -21,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 @@ -43,32 +48,148 @@ 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[tuple[str, Any]] | None = None def get_calculator(self) -> AIMNet2Calculator: """ @@ -81,148 +202,331 @@ 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 _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 get_model_file(model: str, model_dir: str) -> Path: + 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_wb97m", or filename, e.g. "aimnet2_wb97m_0.jpt", 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 `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(): - 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. - 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 - # make sure the directory exists - model_path = Path(model_dir) - if model_path.exists() and not model_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) - else: - raise FileExistsError(f'Path "{model_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 "/" 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 - # 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) - # finally return the path - return model_path - - def setup(self, model: str, model_dir: str, device: str) -> None: - """ - Sets the calculator. Does nothing, if it is already set. + 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, + 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. Parameters ---------- - model: str - Model that the calculator should have - model_dir: str - Path to the model files - device: str - device to use + 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. """ - 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) + # 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") + + # 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 {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, + ) + + # 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, + needs_coulomb=self._TRISTATE_MAP[coulomb], + needs_dispersion=self._TRISTATE_MAP[dispersion], + 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 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. + + 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 + parser : ArgumentParser + Parser that should be extended. """ + # --- model selection --------------------------------------------- parser.add_argument( "-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. " + "For open-shell or charged systems: use aimnet2-nse " + "(other families are closed-shell-trained)." + ), ) parser.add_argument( "-p", @@ -231,7 +535,7 @@ def extend_parser(cls, parser: ArgumentParser) -> None: 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", @@ -239,33 +543,113 @@ def extend_parser(cls, parser: ArgumentParser) -> None: 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=( + "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." + ), + ) - 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,34 +661,67 @@ 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 ---------- 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] = { + "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. + 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": { - "coord": [coordinates], - "numbers": [numbers], - "charge": [charge], - "mult": [mult], - }, + "data": data, "forces": dograd, "stress": False, "hessian": False, @@ -336,9 +753,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 +769,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 +783,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") - 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): + # --- 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") + # 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( - 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/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_client.py b/tests/aimnet2/test_aiment2_client.py index c83c4cc..6fda0c0 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.647682538153e01 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.582629635076e01 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.568258700191e01 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..519d3c2 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.647682538153e01 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.582629635076e01 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.568258700191e01 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..293cfc6 --- /dev/null +++ b/tests/aimnet2/test_aiment2_v02_contract.py @@ -0,0 +1,634 @@ +"""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 +import warnings +from unittest.mock import MagicMock, patch + +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 + + 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).""" + + 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", + ) + + 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.""" + + 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() -> ... + """ + 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). + """ + 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 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]}\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}" + # 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}" + )