Update example

[moradza]

Example code crashes. Convertor should be the last one to be replaced, bc TorchANISymmetryFunctions, TorchANIBatchedNN, and TorchANIEnergyShifter assume a torchani converter, which is not a tuple

[raimis]

Could you elaborate what is the crash?

• The same order is used in OptimizedTorchANI: https://github.com/openmm/NNPOps/blob/ee338426d5a73938753fd8beb43e3dc866d4b4f9/src/pytorch/OptimizedTorchANI.py#L33-L56
• There are tests for that: https://github.com/openmm/NNPOps/blob/master/src/pytorch/TestOptimizedTorchANI.py
• But there are no any crashes: https://github.com/openmm/NNPOps/actions/runs/3257791064

[moradza]

OptimizedTorchANI is working fine (no problem with it), and tests only cover OptimizedTorchANI. It does not cover using TorchANISpeciesConverter, TorchANISymmetryFunctions, TorchANIBatchedNN, and TorchANIEnergyShifter.

For direct usage of TorchANI*, it is assumed that the species converter has output format of TorchANI for species.

import torch
import torchani

from NNPOps.SpeciesConverter import TorchANISpeciesConverter
from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
from NNPOps.BatchedNN import TorchANIBatchedNN
from NNPOps.EnergyShifter import TorchANIEnergyShifter

from NNPOps import OptimizedTorchANI

device = torch.device('cuda')

# Load a molecule
molecule = mdtraj.load('molecule.mol2')
species = torch.tensor([[atom.element.atomic_number for atom in molecule.top.atoms]], device=device)
positions = torch.tensor(molecule.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)

# Construct ANI-2x and replace its operations with the optimized ones
nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
nnp.species_converter = TorchANISpeciesConverter(nnp.species_converter, species).to(device)
nnp.aev_computer = TorchANISymmetryFunctions(nnp.species_converter, nnp.aev_computer, species).to(device)
nnp.neural_networks = TorchANIBatchedNN(nnp.species_converter, nnp.neural_networks, species).to(device)
nnp.energy_shifter = TorchANIEnergyShifter(nnp.species_converter, nnp.energy_shifter, species).to(device)

# Compute energy and forces
energy = nnp((species, positions)).energies
energy.backward()
forces = -positions.grad.clone()

    nnp.aev_computer = TorchANISymmetryFunctions(nnp.species_converter, nnp.aev_computer, species).to(device)
  File "/home/amoradzadeh/anaconda3/envs/tmp-nnpops/lib/python3.10/site-packages/NNPOps/SymmetryFunctions.py", line 90, in __init__
    species = converter((atomicNumbers, torch.empty(0))).species[0].tolist()
AttributeError: 'tuple' object has no attribute 'species'

[raimis]

Which TorchANI version are you using?

[moradza]

torchani 2.2 pyh9f0ad1d_0 conda-forge

[raimis]

You should use 2.2.2. 2.2 is known to be broken.

[moradza]

It is not related to torchANI version. I am closing this because user can infer it from the OptimizedTorchANI code.

Let's go through OptimizedTorchANI:

class OptimizedTorchANI(torch.nn.Module):

    from torchani.models import BuiltinModel

    def __init__(self, model: BuiltinModel, atomicNumbers: Tensor) -> None:

        super().__init__()

        # Optimize the components of an ANI model
        _self.species_converter_ = TorchANISpeciesConverter(**model.species_converter**, atomicNumbers) 
        # ---> Here you are using torchANI species convertor
        self.aev_computer = TorchANISymmetryFunctions(**model.species_converter**, model.aev_computer, atomicNumbers) 
        self.neural_networks = TorchANIBatchedNN(**model.species_converter**, model.neural_networks, atomicNumbers)
        self.energy_shifter = TorchANIEnergyShifter(**model.species_converter**, model.energy_shifter, atomicNumbers)
        # <--- 

    def forward(self, species_coordinates: Tuple[Tensor, Tensor],
                cell: Optional[Tensor] = None,
                pbc: Optional[Tensor] = None) -> SpeciesEnergies:

        species_coordinates = self.species_converter(species_coordinates)
        species_aevs = self.aev_computer(species_coordinates, cell=cell, pbc=pbc)
        species_energies = self.neural_networks(species_aevs)
        species_energies = self.energy_shifter(species_energies)

        return species_energies

Now take a look into above code, species converter is nnpops type and not torchani.

# Construct ANI-2x and replace its operations with the optimized ones
nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
**nnp.species_converter** = TorchANISpeciesConverter(**nnp.species_converter**, species).to(device)
nnp.aev_computer = TorchANISymmetryFunctions(**nnp.species_converter**, nnp.aev_computer, species).to(device)
nnp.neural_networks = TorchANIBatchedNN(**nnp.species_converter**, nnp.neural_networks, species).to(device)
nnp.energy_shifter = TorchANIEnergyShifter(**nnp.species_converter**, nnp.energy_shifter, species).to(device)