niklashoelter · Oct 30, 2026 · Oct 29, 2026 · Oct 30, 2026 · Oct 30, 2026
diff --git a/AGENTS.md b/AGENTS.md
@@ -12,7 +12,7 @@ Guidance for AI coding agents working on the GPUMA repository. (Agents that
 | `src/gpuma/cli_report.py` | Exit codes and the `--json` summary |
 | `src/gpuma/config.py`, `config_schema.py` | Config defaults, validation, JSON Schema |
 | `src/gpuma/optimizer.py`, `autobatching.py` | ASE and torch-sim optimization, batching |
-| `src/gpuma/models/` | Model backends (Fairchem UMA, ORB-v3, SevenNet) and device handling |
+| `src/gpuma/models/` | Model backends (Fairchem UMA, ORB-v3, SevenNet, AIMNet2) and device handling |
 | `src/gpuma/conformer_generation/` | SMILES embedding (nvMolKit on GPU, morfeus on CPU), RMSD pruning |
 | `src/gpuma/utils/reporting.py` | Opt-in diagnostic events used by `--json` |
 | `src/gpuma/doctor.py`, `agent_skills.py` | `gpuma doctor` and `gpuma skills install` |
@@ -43,4 +43,8 @@ Guidance for AI coding agents working on the GPUMA repository. (Agents that
 - The `--json` summary is a contract for agents. Only add keys; bump
   `SUMMARY_SCHEMA_VERSION` in `cli_report.py` for anything else.
 - No backward-compatibility shims unless asked for.
+- The AIMNet2 L-BFGS warning (`_warn_batch_optimizer_memory` in `optimizer.py`)
+  is temporary: drop it, and its docs and skill lines, once an `aimnet` release
+  contains isayevlab/aimnetcentral#150. Bump the pin, then rerun a large mixed
+  L-BFGS job to confirm.
 - Don't add AI co-author trailers or "generated with" footers to commits or PRs.
diff --git a/README.md b/README.md
@@ -17,14 +17,16 @@
 GPUMA is a minimalist Python toolkit for facile and rapid high-throughput molecular geometry optimization
 using machine-learning interatomic potentials (MLIPs).
 
-Three model backends are supported out of the box:
+Four model backends are supported out of the box:
 
 - **Fairchem UMA** ([UMA/OMol25](https://arxiv.org/abs/2505.08762)) — the default backend.
 - **ORB-v3** ([orbital-materials/orb-models](https://github.com/orbital-materials/orb-models)).
 - **SevenNet** ([SevenNet](https://github.com/MDIL-SNU/SevenNet)), including multi-modal checkpoints such as `7net-omni`.
+- **AIMNet2** ([isayevlab/aimnetcentral](https://github.com/isayevlab/aimnetcentral)): `aimnet2`, the open-shell `aimnet2-nse` and `aimnet2-2025`.
 
-All three backends support optional DFT-D3(BJ) dispersion correction. It is off by
-default and rejected for the default `uma-s-1p2p1-ft-d4rev`, which already includes dispersion.
+Fairchem UMA, ORB-v3 and SevenNet support optional DFT-D3(BJ) dispersion correction. It is
+off by default and rejected for the default `uma-s-1p2p1-ft-d4rev` and the AIMNet2 models,
+which already include dispersion.
 
 GPUMA is especially designed for batch optimizations of many structures (conformer ensembles, datasets) on GPU,
 ensuring efficient parallelization and maximum GPU utilization by leveraging the [torch-sim library](https://arxiv.org/abs/2508.06628).
@@ -45,7 +47,7 @@ Feedback and improvements are always welcome!
 ### Option 1: Install from PyPI (recommended)
 
 This installs `gpuma` together with all dependencies (including the
-Fairchem UMA, ORB-v3, and SevenNet backends, as well as the nvMolKit GPU
+Fairchem UMA, ORB-v3, SevenNet and AIMNet2 backends, as well as the nvMolKit GPU
 conformer-embedding backend). Every dependency is pinned to the exact tested
 version, so install gpuma into a fresh environment.
 At the moment, installation and tests have only been
@@ -134,6 +136,7 @@ Also check the [examples/](examples) folder in the repository for sample config
 - [examples/config.json](examples/config.json) – Fairchem UMA configuration
 - [examples/config_orb.json](examples/config_orb.json) – ORB-v3 configuration
 - [examples/config_sevennet.json](examples/config_sevennet.json) – SevenNet configuration (multi-modal `7net-omni`)
+- [examples/config_aimnet.json](examples/config_aimnet.json) – AIMNet2 configuration (`aimnet2`)
 - [examples/example_single_optimization.py](examples/example_single_optimization.py) – single-structure optimization from Python
 - [examples/example_ensemble_optimization.py](examples/example_ensemble_optimization.py) – ensemble/multi-structure optimization from Python
 - [examples/example_dispersion.py](examples/example_dispersion.py) – DFT-D3(BJ) correction with Fairchem and ORB (single + batch)
@@ -190,9 +193,16 @@ When a run is started from SMILES, an RDKit/MMFF force field (via nvMolKit on GP
 - SevenNet: models are downloaded automatically on first use; multi-modal
   checkpoints (e.g. `7net-omni`, `7net-mf-ompa`) also require a
   `model_modal` fidelity (see [docs/config.md](docs/config.md)).
+- AIMNet2: models are downloaded automatically on first use. They cover 14
+  elements, and only `aimnet2-nse` handles open-shell structures; other inputs
+  are rejected before the model loads (see [docs/config.md](docs/config.md)).
+  Use the `fire` batch optimizer: `lbfgs` can run out of GPU memory on large
+  AIMNet2 jobs.
+- BFGS: `"batch_optimizer": "bfgs"` keeps a dense Hessian per structure and can
+  run out of GPU memory on large batch jobs with any backend; prefer `lbfgs` or `fire`.
 - D3 dispersion correction: off by default. Set `"d3_correction": true` to
   enable DFT-D3(BJ) for models trained without dispersion. It is rejected for
-  `uma-s-1p2p1-ft-d4rev` and logs a warning for the other UMA models, ORB
+  `uma-s-1p2p1-ft-d4rev` and AIMNet2 and logs a warning for the other UMA models, ORB
   `*_omol` models and SevenNet `omol25_*` modals (see [docs/config.md](docs/config.md)).
 
 ## License

diff --git a/docs/cli.md b/docs/cli.md
@@ -99,6 +99,22 @@ gpuma batch --xyz-dir examples/example_input_xyzs/multi_xyz_dir/ --output batch_
 You can also select SevenNet with the `--model-type` flag
 (`--model-type sevennet` or `--model-type 7net`), just like the other backends.
 
+### Using AIMNet2 models
+
+Set `"model_type": "aimnet"` (alias `"aimnet2"`) and `"model_name"` to
+`"aimnet2"`, `"aimnet2-nse"` (open-shell) or `"aimnet2-2025"` (see
+`examples/config_aimnet.json`), or pass `--model-type aimnet`. Inputs outside
+the model's range (elements, open-shell on a closed-shell model, H–H contacts
+below 0.8 Å) are rejected with exit code 3 before the model loads:
+
+```bash
+# Batch optimization from a multi-XYZ file with AIMNet2
+gpuma batch --multi-xyz examples/example_input_xyzs/multi_xyz_file.xyz --output batch_aimnet.xyz --config examples/config_aimnet.json
+
+# Open-shell structure with AIMNet2-NSE (set "model_name": "aimnet2-nse" in the config)
+gpuma optimize --xyz examples/example_input_xyzs/butene_triplet.xyz --multiplicity 3 --output triplet.xyz --config my_nse_config.json
+```
+
 **Note:**
 - If `--config` is not specified, `config.json` in the current directory is loaded by default.
   If that file does not exist either, GPUMA runs on its built-in defaults and logs a warning

diff --git a/docs/config.md b/docs/config.md
@@ -97,11 +97,28 @@ For multi-modal checkpoints such as `7net-omni` or `7net-mf-ompa`, a
 }
 ```
 
+## Minimal Example (JSON, AIMNet2)
+
+Use `aimnet2-nse` for open-shell structures; `aimnet2` and `aimnet2-2025` are
+closed-shell only.
+
+```json
+{
+  "model": {
+    "model_type": "aimnet",
+    "model_name": "aimnet2"
+  },
+  "technical": {
+    "device": "cuda"
+  }
+}
+```
+
 ## D3 Dispersion Correction
 
 D3 is off by default and adds GPU memory per batch; enable it only for models
-trained without dispersion. `uma-s-1p2p1-ft-d4rev` already includes dispersion
-and rejects D3. The other UMA models, the ORB `*_omol` models and SevenNet's
+trained without dispersion. `uma-s-1p2p1-ft-d4rev` and the AIMNet2 models
+already include dispersion and reject D3. The other UMA models, the ORB `*_omol` models and SevenNet's
 `omol25_*` modals are trained on ωB97M-V, which also includes it, so gpuma logs
 a warning when D3 is enabled for them.
 
@@ -183,7 +200,7 @@ technical:
 | Parameter | Default | Description |
 |---|---|---|
 | `batch_optimization_mode` | `"batch"` | `"sequential"` (ASE per structure) or `"batch"` (torch-sim GPU-accelerated). If a batch fails, its structures are retried one at a time, so one bad input fails only itself |
-| `batch_optimizer` | `"fire"` | Optimizer for both single and batch modes: `"fire"`, `"gradient_descent"`, `"lbfgs"`, or `"bfgs"`. In single-structure (ASE) mode, maps to ASE FIRE/BFGS/LBFGS (`"gradient_descent"` falls back to FIRE). Defaults to `"fire"` if unset; an unknown value is an error |
+| `batch_optimizer` | `"fire"` | Optimizer for both single and batch modes: `"fire"`, `"gradient_descent"`, `"lbfgs"`, or `"bfgs"`. In single-structure (ASE) mode, maps to ASE FIRE/BFGS/LBFGS (`"gradient_descent"` falls back to FIRE). Defaults to `"fire"` if unset; an unknown value is an error. In batch mode, `"bfgs"` keeps a dense Hessian per structure that the memory estimate leaves out and can run out of GPU memory on large jobs, and `"lbfgs"` can do the same with AIMNet2; gpuma warns for both |
 | `charge` | `0` | Total charge of the system (inferred from SMILES, not overridden). Must be an integer |
 | `multiplicity` | `1` | Spin multiplicity of the system. A multiplicity the electron count can't have (e.g. a singlet radical) is an error |
 | `force_convergence_criterion` | `5e-2` | Force convergence threshold (eV/A). Used for both single and batch modes |
@@ -194,14 +211,14 @@ technical:
 
 | Parameter | Default | Description |
 |---|---|---|
-| `model_type` | `"fairchem"` | Backend: `"fairchem"` / `"uma"`, `"orb"` / `"orb-v3"`, or `"sevennet"` / `"7net"` |
+| `model_type` | `"fairchem"` | Backend: `"fairchem"` / `"uma"`, `"orb"` / `"orb-v3"`, `"sevennet"` / `"7net"`, or `"aimnet"` / `"aimnet2"` |
 | `model_name` | `"uma-s-1p2p1-ft-d4rev"` | Model identifier (see [Available Models](#available-models)) |
 | `model_modal` | `null` | SevenNet multi-modal fidelity selector. Valid values depend on the checkpoint (`7net-omni`: e.g. `"omol25_high"`, `"omol25_low"`, `"spice"`, `"qcml"`; `7net-mf-ompa`: `"mpa"`, `"omat24"`). Required for multi-modal checkpoints; ignored by single-modal checkpoints and by other backends |
-| `model_path` | `null` | Local checkpoint path (Fairchem and SevenNet; overrides `model_name`). A path that doesn't exist, or one set for ORB, is an error |
-| `model_cache_dir` | `null` | Directory for cached model downloads (Fairchem only; ORB and SevenNet use their own caches, so setting it for them is an error). A directory that can't be created is an error |
+| `model_path` | `null` | Local checkpoint path (Fairchem and SevenNet; overrides `model_name`). A path that doesn't exist, or one set for ORB or AIMNet2, is an error |
+| `model_cache_dir` | `null` | Directory for cached model downloads (Fairchem and AIMNet2, which gpuma points there via `AIMNET_CACHE_DIR`; ORB and SevenNet use their own caches, so setting it for them is an error). A directory that can't be created is an error |
 | `huggingface_token` | `null` | HuggingFace token for model access |
 | `huggingface_token_file` | `null` | File path to read the HF token from |
-| `d3_correction` | `false` | Enable D3 dispersion correction. Supported for the ORB, Fairchem/UMA, and SevenNet backends; rejected for `uma-s-1p2p1-ft-d4rev` |
+| `d3_correction` | `false` | Enable D3 dispersion correction. Supported for the ORB, Fairchem/UMA, and SevenNet backends; rejected for `uma-s-1p2p1-ft-d4rev` and AIMNet2 |
 | `d3_functional` | `"PBE"` | DFT functional for D3 correction (ORB/Fairchem use the orb-models BJ-damping table; SevenNet uses its own native D3) |
 | `d3_damping` | `"BJ"` | Damping scheme for D3 correction (`"BJ"` or `"BJM"`). SevenNet supports only `"BJ"` |
 
@@ -220,7 +237,7 @@ technical:
 | `max_memory_padding` | `0.95` | Fraction of GPU memory the autobatcher is allowed to fill during calibration. Lower = more headroom, smaller batches. Must be between 0.05 and 0.95 |
 | `memory_scaling_factor` | `1.75` | Factor by which the autobatcher grows the probe size during calibration. Larger = faster calibration but coarser final batch size; smaller = slower but tighter. A job runs as one batch without probing only if its forward-pass peak times this factor fits in `max_memory_padding` of free memory. Must be > 1 |
 | `max_atoms_to_try` | `100000` | Upper bound on the autobatcher's calibration probe size (atoms) |
-| `memory_scales_with` | `"auto"` | Per-structure memory metric for batch sizing: `"n_atoms"`, `"n_edges"` or `"auto"`. `"auto"` uses `n_atoms` for UMA, whose large fixed memory per structure makes edge counts overcharge small molecules, and `n_edges` for ORB and SevenNet |
+| `memory_scales_with` | `"auto"` | Per-structure memory metric for batch sizing: `"n_atoms"`, `"n_edges"` or `"auto"`. `"auto"` uses `n_atoms` for UMA, whose large fixed memory per structure makes edge counts overcharge small molecules, and for AIMNet2, whose Coulomb term grows with the square of the batch's atom count; `n_edges` for ORB and SevenNet |
 | `steps_between_swaps` | `1` | Optimization steps between batch swaps in the in-flight autobatcher. `1` is fastest on this codebase's screenings (uma-s/uma-m/orb); higher values are monotonically slower. Must be at least 1 |
 | `logging_level` | `"INFO"` | Logging verbosity: `"DEBUG"`, `"INFO"`, `"WARNING"`, `"ERROR"`, `"CRITICAL"` |
 
@@ -281,11 +298,27 @@ instead be supplied via `model_path`.
 | `7net-omni-i8` | SevenNet Omni i8 (multi-modal) |
 | `7net-omni-i12` | SevenNet Omni i12 (multi-modal) |
 
-> D3 dispersion correction can be enabled for any ORB, Fairchem/UMA, or SevenNet model with `d3_correction: true`, except `uma-s-1p2p1-ft-d4rev`. gpuma warns when it is enabled for UMA, ORB `*_omol` models or SevenNet `omol25_*` modals; SevenNet supports only BJ damping.
+### AIMNet2 (`model_type: "aimnet"`)
+
+All three models cover H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br and I, handle
+charged molecules and include D3 dispersion. gpuma loads ensemble member 0 of
+each. Before loading the model, gpuma rejects structures with other elements, an
+open-shell structure (multiplicity above 1) on `aimnet2` or `aimnet2-2025`, and,
+for those two models, an H–H distance below 0.8 Å with at least one heavy atom
+present: upstream reports a spurious minimum where two hydrogens on one atom
+collapse ([aimnetcentral#139](https://github.com/isayevlab/aimnetcentral/issues/139)).
+
+| Model Name | Description |
+|---|---|
+| `aimnet2` | ωB97M-D3/def2-TZVPP, closed-shell; the general-purpose default |
+| `aimnet2-nse` | ωB97M-D3/def2-TZVPP (unrestricted), open-shell: uses the multiplicity |
+| `aimnet2-2025` | B97-3c, closed-shell; improved non-covalent interactions |
+
+> D3 dispersion correction can be enabled for any ORB, Fairchem/UMA, or SevenNet model with `d3_correction: true`, except `uma-s-1p2p1-ft-d4rev`; AIMNet2 models already include it. gpuma warns when it is enabled for UMA, ORB `*_omol` models or SevenNet `omol25_*` modals; SevenNet supports only BJ damping.
 
 These model names are also available programmatically as
-`gpuma.AVAILABLE_FAIRCHEM_MODELS`, `gpuma.AVAILABLE_ORB_MODELS`, and
-`gpuma.AVAILABLE_SEVENNET_MODELS`.
+`gpuma.AVAILABLE_FAIRCHEM_MODELS`, `gpuma.AVAILABLE_ORB_MODELS`,
+`gpuma.AVAILABLE_SEVENNET_MODELS`, and `gpuma.AVAILABLE_AIMNET_MODELS`.
 
 ---
 

diff --git a/docs/index.md b/docs/index.md
@@ -9,14 +9,16 @@
 GPUMA is a minimalist Python toolkit for facile and rapid high-throughput molecular geometry optimization
 based on machine-learning interatomic potentials (MLIPs).
 
-Three model backends are supported out of the box:
+Four model backends are supported out of the box:
 
 - **Fairchem UMA** ([UMA/OMol25](https://arxiv.org/abs/2505.08762)) — the default backend.
 - **ORB-v3** ([orbital-materials/orb-models](https://github.com/orbital-materials/orb-models)) — included in core dependencies.
 - **SevenNet** ([SevenNet](https://github.com/MDIL-SNU/SevenNet)) — included in core dependencies, with support for multi-modal checkpoints such as `7net-omni`.
+- **AIMNet2** ([isayevlab/aimnetcentral](https://github.com/isayevlab/aimnetcentral)) — included in core dependencies: `aimnet2`, the open-shell `aimnet2-nse` and `aimnet2-2025`.
 
-All three backends support optional DFT-D3(BJ) dispersion correction. It is off by
-default and rejected for the default `uma-s-1p2p1-ft-d4rev`, which already includes dispersion.
+Fairchem UMA, ORB-v3 and SevenNet support optional DFT-D3(BJ) dispersion correction. It is
+off by default and rejected for the default `uma-s-1p2p1-ft-d4rev` and the AIMNet2 models,
+which already include dispersion.
 
 GPUMA is especially designed for batch optimizations of many structures (conformer ensembles, datasets) on GPU,
 ensuring efficient parallelization and maximum GPU utilization by leveraging the [torch-sim library](https://arxiv.org/abs/2508.06628).
@@ -40,7 +42,7 @@ Feedback and improvements are always welcome!
 ### Option 1: Install from PyPI (recommended)
 
 This installs `gpuma` together with all dependencies (including the
-Fairchem UMA, ORB-v3, and SevenNet backends, plus the nvMolKit GPU
+Fairchem UMA, ORB-v3, SevenNet and AIMNet2 backends, plus the nvMolKit GPU
 conformer-embedding backend).
 Every dependency is pinned to the exact tested version, so install gpuma into
 a fresh environment.
@@ -102,6 +104,14 @@ checkpoints (e.g. `7net-omni`, `7net-mf-ompa`) additionally require a
 `"model_modal"` fidelity such as `"omol25_high"` (see [Configuration](config.md)
 and `examples/config_sevennet.json`).
 
+### AIMNet2 support
+
+AIMNet2 models are included in the standard installation. Set
+`"model_type": "aimnet"` (alias `"aimnet2"`) and a `"model_name"` of
+`"aimnet2"`, `"aimnet2-nse"` or `"aimnet2-2025"` (see [Configuration](config.md)
+and `examples/config_aimnet.json`). They cover 14 elements (H, B, C, N, O, F,
+Si, P, S, Cl, As, Se, Br, I); only `aimnet2-nse` handles open-shell structures.
+
 
 ### Option 2: Install from source
 
@@ -150,9 +160,16 @@ When a run is started from SMILES, an RDKit/MMFF force field (via nvMolKit on GP
 - ORB-v3: set `"model_type": "orb"` in the `model` section of your config.
 - SevenNet: set `"model_type": "sevennet"` in the `model` section of your
   config; multi-modal checkpoints also need a `"model_modal"` fidelity.
+- AIMNet2: set `"model_type": "aimnet"`. Structures with other elements, an
+  open-shell structure on a closed-shell model, or an H–H distance below 0.8 Å
+  (for `aimnet2` and `aimnet2-2025`) are rejected before the model loads.
+  Use the `fire` batch optimizer: `lbfgs` can run out of GPU memory on large
+  AIMNet2 jobs.
+- BFGS: `"batch_optimizer": "bfgs"` keeps a dense Hessian per structure and can
+  run out of GPU memory on large batch jobs with any backend; prefer `lbfgs` or `fire`.
 - D3 dispersion correction: off by default. Enable it with
   `"d3_correction": true` for models trained without dispersion. It is
-  rejected for `uma-s-1p2p1-ft-d4rev` and logs a warning for the other UMA
+  rejected for `uma-s-1p2p1-ft-d4rev` and the AIMNet2 models and logs a warning for the other UMA
   models, ORB `*_omol` models and SevenNet `omol25_*` modals (see
   [Configuration](config.md)).
 

diff --git a/docs/install.md b/docs/install.md
@@ -34,7 +34,7 @@ version is replaced by the default PyPI build:
 pip install gpuma
 ```
 
-This pulls in all model backends (Fairchem UMA, ORB-v3, and SevenNet) as
+This pulls in all model backends (Fairchem UMA, ORB-v3, SevenNet and AIMNet2) as
 core dependencies — there is no separate optional extra to install.
 
 ### Pinned dependencies

diff --git a/docs/reference.md b/docs/reference.md
@@ -10,7 +10,7 @@ Documentation for the core functions in `gpuma`.
 
     - `gpuma.models` — model dispatchers (`load_calculator`, `load_torchsim_model`)
       and the `AVAILABLE_*_MODELS` registries, with backend-specific loaders under
-      `models/fairchem.py`, `models/orb.py`, `models/sevennet.py`.
+      `models/fairchem.py`, `models/orb.py`, `models/sevennet.py`, `models/aimnet.py`.
     - `gpuma.conformer_generation` — SMILES → 3D structure/ensemble embedding
       (`embed.py`, `mol_utils.py`).
     - `gpuma.utils` — I/O (`io_handler.py`), logging (`logging_utils.py`) and