AIMNet2

AIMNet2: machine-learned potentials from the Isayev group, through aimnet.

The codec calls aimnet.calculators.AIMNet2Calculator in the calling process; install it next to qccompute with pip install aimnet. Models download on first use and are loaded once per process, then reused for every later calculation with the same model and device.

Inputs map to AIMNet2 as follows:

• model.method is a model name from aimnet's registry: aimnet2 (ωB97M-D3, 14 elements, any charge), aimnet2-rxn (reactions, H C N O), aimnet2-nse (open-shell species), aimnet2-2025 (B97-3c), aimnet2-b973c, aimnet2-pd, or a member such as aimnet2-wb97m-d3_1; or a path to a model file. Models have no basis set.
• structure.charge goes to the model. structure.multiplicity goes to the model only for spin-aware models (aimnet2-nse); the others are closed-shell, so the codec raises for a multiplicity other than 1.
• keywords: device ("cpu", "cuda", "cuda:1"; CUDA when available by default) and validate_species (True by default: aimnet refuses elements or charges a model was not trained on). Nothing else is accepted. A model on a GPU declares one for qccompute, which gives each calculation of a list its own; an explicit index such as "cuda:1" opts out of that.
• cmdline_args and files are not accepted.

Calculation types: energy, gradient, and hessian (by automatic differentiation). Every result records the model's partial charges in atomic_charges. AIMNet2 predicts energies in eV; the codec converts them, and the forces and Hessian, to Hartree and bohr.

AIMNet2 does not move the molecule, so results are in the input's frame. No input field carries atom indices.

Energies and gradients of several structures can be computed as one batch: qccompute.compute does this for a list of inputs that share a model, keywords, and calculation type, evaluating them in one pass instead of one call each. Hessians are computed one structure at a time.

Examples

Each script runs as written. Copy it, delete the options you don't need, and run it with python script.py. They live in the repository's examples/ folder.

Single point: energy, gradient, Hessian

examples/aimnet2/single_point.py

"""AIMNet2 single point: energy, gradient, or Hessian from a machine-learned potential.

Requires: ``pip install qccompute aimnet``. AIMNet2 is a Python library, so nothing has
to be on PATH: the model runs inside this Python process, on the GPU when PyTorch sees
one. Weights download on first use; later calculations in the same process reuse the
loaded model, so an optimization or path pays the loading cost once.

How the codec talks to AIMNet2 (https://github.com/isayevlab/aimnetcentral):

- ``model.method`` is a model from aimnet's registry: "aimnet2" (ωB97M-D3, 14
  elements, any charge), "aimnet2-rxn" (reactions, H C N O), "aimnet2-nse" (radicals
  and other open shells), "aimnet2-2025" (B97-3c), ...; or a path to a model file.
  There is no basis set.
- Charge comes from ``structure``. Only "aimnet2-nse" takes a multiplicity; the other
  models are closed-shell and the codec refuses open shells for them.
- aimnet refuses elements and charges a model was not trained on; set
  ``validate_species`` to False to override (results are then undefined).
- Hessians come from automatic differentiation, not finite differences.
- Energies come back in Hartree and gradients in Hartree/bohr, converted from
  AIMNet2's eV and angstrom.
"""

from qccompute import compute
from qcdata import CalcType, Model, ProgramInput, Structure

calctype = CalcType.gradient  # or CalcType.energy, CalcType.hessian

# Coordinates are in bohr (Structure.from_xyz / Structure.open read angstrom files).
structure = Structure(
    symbols=["O", "H", "H"],
    geometry=[[0.0, 0.0, 0.0], [0.524, 1.687, 0.48], [1.146, -0.45, -1.354]],
    charge=0,
    multiplicity=1,  # other than 1 needs "aimnet2-nse"
)

prog_input = ProgramInput(
    program="aimnet2",
    calctype=calctype,
    structure=structure,
    model=Model(method="aimnet2"),  # or "aimnet2-rxn", "aimnet2-nse", "aimnet2-2025"
    keywords={
        "device": "cpu",  # or "cuda", "cuda:1"; leave out to use a GPU when present
        # "validate_species": False,  # allow elements/charges outside the training set
    },
)

output = compute(prog_input)

results = output.results
print("Energy (Hartree):", results.energy)
if results.gradient is not None:
    print("Gradient (Hartree/bohr):\n", results.gradient)
if results.hessian is not None:
    print("Hessian (Hartree/bohr^2):", results.hessian.shape)
print("Partial charges:", results.atomic_charges)
print("aimnet version:", results.provenance.program_version)

Many structures at once: batches

examples/aimnet2/many_calculations.py

"""AIMNet2 for many structures at once: a bond scan evaluated in batches.

Requires: what examples/aimnet2/single_point.py requires, and qccompute 0.21.0 or newer.

How batching works:

- Pass ``compute`` a list of inputs. Energy and gradient inputs with the same program,
  model, and keywords are evaluated in batches, one model call for up to
  ``get_codec("aimnet2").max_batch`` (32) structures, in this Python process. On a GPU
  that is many times faster than one call per structure.
- Each structure's results are what computing it alone gives, to the model's precision.
  Structures in one batch may differ in size, composition, and charge.
- Inputs with different models or keywords go in separate batches. Hessians are
  computed one structure at a time.
- A batch that runs out of GPU memory is split in half until it fits. A batch that
  fails otherwise is computed again one structure at a time, so each error belongs to
  its own input.
- The model stays loaded in this process, so the next ``compute`` call (the next step
  of a path optimization, say) does not load it again.
- ``output.execution.extras`` records ``batch_size`` and the structure's
  ``batch_index``; ``logs`` and ``wall_time`` are the whole batch's.
"""

import math
import time

from qccompute import compute
from qcdata import CalcType, Model, ProgramInput, Structure


def water(r: float) -> Structure:
    """Water with one O-H bond stretched to r bohr, the other at 1.81 bohr."""
    angle = math.radians(104.5)
    return Structure(
        symbols=["O", "H", "H"],
        geometry=[
            [0.0, 0.0, 0.0],
            [r, 0.0, 0.0],
            [1.81 * math.cos(angle), 1.81 * math.sin(angle), 0.0],
        ],
    )


distances = [1.4 + 0.1 * i for i in range(20)]  # bohr
inputs = [
    ProgramInput(
        program="aimnet2",
        calctype=CalcType.gradient,  # or CalcType.energy
        structure=water(r),
        model=Model(method="aimnet2"),
        keywords={
            # "device": "cpu",  # or "cuda"; leave out to use a GPU when present
        },
    )
    for r in distances
]

compute(inputs)  # load the model and warm up, so the timings compare evaluation only

start = time.perf_counter()
outputs = compute(inputs)  # one batch of 20
batched = time.perf_counter() - start

start = time.perf_counter()
alone = [compute(inp) for inp in inputs]  # 20 separate calls, for comparison
separate = time.perf_counter() - start

print("r (bohr)  energy (Hartree)  batch_index/batch_size")
for r, output in zip(distances, outputs, strict=True):
    extras = output.execution.extras
    print(
        f"{r:8.2f}  {output.results.energy:16.8f}  "
        f"{extras['batch_index']}/{extras['batch_size']}"
    )
largest = max(
    abs(a.results.energy - b.results.energy)
    for a, b in zip(alone, outputs, strict=True)
)
print(f"Batched: {batched:.3f} s; one at a time: {separate:.3f} s")
print(f"Largest energy difference from computing alone: {largest:.1e} Hartree")