TorchANI

TorchANI: the ANI family of machine-learned potentials.

The codec calls TorchANI (https://github.com/aiqm/torchani) in the calling process; install it next to qccompute with pip install torchani. Parameters download on first use, and each model is loaded once per process and reused.

Inputs map to TorchANI as follows:

• model.method names a model in torchani.models: ANI2x (ωB97X/6-31G(d), H C N O F S Cl), ANI1x, ANI1ccx (CCSD(T)/CBS, H C N O), ANI2xr, ANI2dr, ANImbis, or SnnANI2xr. Case and dashes are ignored, so ani-2x works too. Models have no basis set.
• ANI models describe neutral closed-shell molecules only, so the codec raises for a charge other than 0 or a multiplicity other than 1.
• keywords: device ("cpu", "cuda"; CUDA when available by default), dtype ("float64" by default, for smooth gradients and Hessians; or "float32", faster), and model_index (one member of the ensemble instead of its average). 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). TorchANI predicts energies in Hartree with coordinates in angstrom; the codec converts gradients and Hessians to bohr. ANImbis also fills atomic_charges.

TorchANI 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/torchani/single_point.py

"""TorchANI single point: energy, gradient, or Hessian from an ANI potential.

Requires: ``pip install qccompute torchani``. TorchANI 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. Parameters download on first use; later calculations in the same process
reuse the loaded model.

How the codec talks to TorchANI (https://github.com/aiqm/torchani):

- ``model.method`` names a model in ``torchani.models``: "ANI2x" (ωB97X/6-31G(d);
  H C N O F S Cl), "ANI1x", "ANI1ccx" (CCSD(T)/CBS; H C N O), "ANI2xr", "ANI2dr",
  "ANImbis", "SnnANI2xr". Case and dashes are ignored. There is no basis set.
- ANI describes neutral closed-shell molecules only; the codec refuses anything else.
- ``dtype`` is "float64" by default, for smooth gradients and Hessians; "float32" is
  faster on GPUs.
- ANI models are ensembles; ``model_index`` picks one member instead of the average.
- Hessians come from automatic differentiation. Energies are in Hartree.
"""

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,  # ANI: neutral closed-shell only
    multiplicity=1,
)

prog_input = ProgramInput(
    program="torchani",
    calctype=calctype,
    structure=structure,
    model=Model(method="ANI2x"),  # or "ANI1ccx", "ANI1x", ...
    keywords={
        "device": "cpu",  # or "cuda"; leave out to use a GPU when present
        "dtype": "float64",  # or "float32"
        # "model_index": 0,  # one member of the ensemble instead of its average
    },
)

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("torchani version:", results.provenance.program_version)

Many structures at once: batches

examples/torchani/many_calculations.py

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

Requires: what examples/torchani/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("torchani").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="torchani",
        calctype=CalcType.gradient,  # or CalcType.energy
        structure=water(r),
        model=Model(method="ANI2x"),
        keywords={
            "dtype": "float64",
            # "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")