TorchANI

Several core classes have been renamed or have updated method signatures in TorchANI 2:

1. torchani.nn.ANIModel has been renamed to torchani.nn.ANINetworks. While the old name is temporarily supported, it is slated for removal.
2. torchani.AEVComputer initialization has changed. If you need the old signature, use the torchani.aev.AEVComputer.from_constants constructor.
3. Component Workflow: The interaction between SpeciesConverter, AEVComputer, and ANINetworks has been streamlined. SpeciesConverter now typically takes atomic_nums directly, and AEVComputer no longer requires the same tuple-based input as before.

Separating Neighborlists: You can now separate the AEVComputer and torchani.neighbors.Neighborlist parts of a calculation. This is useful if you want to reuse a computed neighborlist in other modules (like pair potentials).

import torchani

# New workflow for components
converter = torchani.nn.SpeciesConverter("...")
aevc = torchani.AEVComputer("...")
ani_nets = torchani.ANINetworks("...")
ensemble = torchani.Ensemble("...")

idxs = converter(atomic_nums)
aevs = aevc(idxs, coords, cell, pbc)
energies = ani_nets(idxs, aevs)

# Separating Neighborlist
neighborlist = torchani.neighbors.AllPairs()
# ... setup aevc and converter ...
cutoff = aevc.radial.cutoff
neighbors = neighborlist(cutoff, idxs, coords, cell, pbc)
aevc = aevc.compute_from_neighbors(idxs, neighbors)

Instead of calling models directly with a tuple of (species_indices, coords), it is recommended to use the torchani.single_point utility. This approach is more robust as it allows models to return multiple quantities (like charges, forces, or hessians) in a dictionary format and handles gradient requirements automatically.

To obtain forces and hessians, pass forces=True and hessians=True to the function.

import torchani
from torchani import single_point

# Setup
atomic_nums = torch.tensor([[1, 6, 6, 1]])
coords = torch.tensor(...) # your coordinates
model = torchani.models.ANI1x()

# Standard usage
result = single_point(model, atomic_nums, coords)
energies = result["energies"]

# Advanced usage (forces and hessians)
result = single_point(model, atomic_nums, coords, forces=True, hessians=True)
atomic_charges = result["atomic_charges"]  # Only for models that support this
energies = result["energies"]
forces = result["forces"]
hessians = result["hessians"]

The use of torchani.nn.Sequential is highly discouraged due to being error-prone and verbose. Instead, use factory functions in torchani.arch or the torchani.arch.Assembler class to build models ready for training.

Option 1: Simple Factory Functions

Use torchani.arch.simple_ani for a quick setup with random weights.

Option 2: Using Assembler for Customization

Use torchani.arch.Assembler to fine-tune components like radial/angular terms, strategies (e.g., cuAEV), and ground state atomic energies (GSAEs).

Option 3: Custom torch.nn.Module

For maximum flexibility, inherit from torch.nn.Module and manually compose the components.

# Option 1: Simple factory
from torchani.arch import simple_ani
model = simple_ani(("H", "C", "N", "O", "S"), lot="wb97x-631gd")

# Option 2: Assembler
import torchani
asm = torchani.arch.Assembler()
asm.set_symbols(("H", "C", "N", "O"))
asm.set_aev_computer(radial="ani2x", angular="ani2x", strategy="cuaev")
asm.set_atomic_networks(ctor="ani2x")
asm.set_gsaes_as_self_energies("wb97x-631gd")
model = asm.assemble()

# Option 3: Custom Module
import torchani
from torch.nn import Module

class Model(Module):
    def __init__(self):
        super().__init__()
        self.converter = torchani.nn.SpeciesConverter("...")
        self.neighborlist = torchani.neighbors.AllPairs("...")
        self.aevc = torchani.aev.AEVComputer("...")
        self.nn = torchani.nn.ANINetworks("...")
        self.shifter = torchani.sae.SelfEnergy("...")

    def forward(self, atomic_nums, coords, cell, pbc):
        idxs = self.converter(atomic_nums)
        cutoff = self.aevc.radial.cutoff
        neighbors = self.neighborlist(cutoff, idxs, coords, cell, pbc)
        aevs = self.aevc.compute_from_neighbors(idxs, neighbors)
        return self.nn(idxs, aevs) + self.shifter(idxs)

model = Model()

When working on an HPC cluster, you must first start an interactive session with GPU access.

Example session commands:

• Hipergator: srun -p gpu --ntasks=1 --cpus-per-task=2 --gpus=geforce:1 --time=02:00:00 --mem=10gb --pty -u bash -i
• Bridges2: srun -p GPU-small --ntasks=1 --cpus-per-task=5 --gpus=1 --time=02:00:00 --mem=20gb --pty -u bash -i
• Expanse: srun -p gpu-shared --ntasks=1 --account=cwr109 --cpus-per-task=1 --gpus=1 --time=01:00:00 --mem=10gb --pty -u bash -i
• Moria: srun --ntasks=1 --cpus-per-task=2 --gres=gpu:1 --time=02:00:00 --mem=10gb --pty -u bash -i

Installation steps in the session:

1. Create environment: conda create -f ./environment.yml
2. Clone repo: git clone https://github.com/roitberg-group/torchani_sandbox.git
3. Install base: pip install -v --no-deps --no-build-isolation --editable .
4. Install with all-GPU support: pip install -v --no-deps --no-build-isolation --editable . --global-option="--ext-all-sms"