Using custom NNPs with Auto3D

Auto3D (>= 2.3.0) is compatible with any jitable NNPs. This notebook demonstrates how to wrapper and jit a custom NNP to a specific format that Auto3D can use.

import os, sys
root = os.path.dirname(os.path.dirname(os.path.abspath("__file__")))

import torch
import torchani
import Auto3D
from Auto3D import Auto3DOptions, main

print(Auto3D.__version__)

# Below is the template of the wrapper that you need to implement with your custom NNP.
class userNNP(torch.nn.Module):
    def __init__(self):
        super(userNNP, self).__init__()
        """This is an example NNP model that can be used with Auto3D.
        You can initialize an NNP model however you want,
        just make sure that:
            - It contains the coord_pad and species_pad attributes
              (These values will be used when processing the molecules in batch.)
            - The signature of the forward method is the same as below.
        """
        # Here I constructed an example NNP using ANI2x.
        # In your case, you can replace this with your own NNP model.
        self.model = torchani.models.ANI2x(periodic_table_index=True)

        self.coord_pad = 0  # int, the padding value for coordinates
        self.species_pad = -1  # int, the padding value for species.

    def forward(self,
                species: torch.Tensor,
                coords: torch.Tensor,
                charges: torch.Tensor) -> torch.Tensor:
        """
        Your NNP should take species, coords, and charges as input
        and return the energies of the molecules.

        species contains the atomic numbers of the atoms in the molecule: [B, N]
        where B is the batch size, N is the number of atoms in the largest molecule.

        coords contains the coordinates of the atoms in the molecule: [B, N, 3]
        where B is the batch size, N is the number of atoms in the largest molecule,
        and 3 represents the x, y, z coordinates.

        charges contains the molecular charges: [B]

        The forward function returns the energies of the molecules: [B],
        output energy unit: eV"""

        # an example for computing molecular energy, replace with your NNP model
        energies = self.model((species, coords)).energies * 27.211386245988
        return energies

# initialize and jit the wrapper with your NNP model
myNNP = userNNP()
myNNP_jit = torch.jit.script(myNNP)

# save the model to a file for later use
model_path = os.path.join(root, 'myNNP.pt')
torch.jit.save(myNNP_jit, model_path)

/home/jack/miniconda3/envs/py39/lib/python3.9/site-packages/torchani/resources/

# Now you can run Auto3D with your custom NNP model.
# Simply pass the model_path to the optimizing_engine argument

smi_path = os.path.join(root, "example/files/smiles.smi")  # You can specify the path to your file here
config = Auto3DOptions(path=smi_path, k=1, optimizing_engine=model_path, use_gpu=True, gpu_idx=0)
out = main(config)
print(out)