Constructing a TorchANI model from NeuroChem files#

This tutorial illustrates how to manually load model from NeuroChem files.

# To begin with, let's first import the modules we will use:
from pathlib import Path
import torch

from torchani.paths import neurochem_dir
from torchani.grad import energies_and_forces
from torchani.neurochem import (
    download_model_parameters,
    load_aev_computer_and_symbols,
    load_sae,
    load_member,
    load_ensemble,
    load_model_from_info_file,
)

First lets download all model parameters, by default they will be loaded into torchani.paths.NEUROCHEM, which is ~/.local/torchani/Neurochem.

root = neurochem_dir()
download_model_parameters()
Found existing files in directory, assuming params already present

Now let’s read constants from constant file and construct AEV computer, The sae’s and construct a SelfEnergy potential, and the networks to construct an ensemble

const_file = Path(root, "ani-1x_8x", "rHCNO-5.2R_16-3.5A_a4-8.params")
aev_computer, symbols = load_aev_computer_and_symbols(const_file)
model_prefix = Path(root, "ani-1x_8x", "train")
ensemble = load_ensemble(symbols, model_prefix, 8)
sae_file = Path(root, "ani-1x_8x", "sae_linfit.dat")
energy_shifter = load_sae(sae_file)

We can also load a single model from the ensemble

model_dir = Path(root, "ani-1x_8x", "train0", "networks")
member = load_member(symbols, model_dir)

We can also load an ANI model using neurochem

ensemble_model = load_model_from_info_file(root / "ani-1x_8x.info")
single_model = load_model_from_info_file(root / "ani-1x_8x.info", model_index=0)
print(ensemble_model)
print(single_model)
ANI(
  (neighborlist): AllPairs()
  (energy_shifter): SelfEnergy()
  (species_converter): SpeciesConverter()
  (potentials): ModuleDict(
    (nnp): NNPotential(
      (aev_computer): AEVComputer(
        #  out_dim=384
        #  radial_len=64 (16.67% of feats)
        #  angular_len=320 (83.33% of feats)
        num_species=4,
        strategy=pyaev,
        (radial): ANIRadial(
          #  num_feats=16
          #  num_shifts=16
          eta=16.0000,
          shifts=[0.9000, 1.1688, 1.4375, 1.7062, 1.9750, 2.2438, 2.5125, 2.7812, 3.0500, 3.3187, 3.5875, 3.8563, 4.1250, 4.3938, 4.6625, 4.9313],
          cutoff=5.2000,
          (cutoff_fn): CutoffCosine()
        )
        (angular): ANIAngular(
          #  num_feats=32
          #  num_shifts=4
          #  num_sections=8
          eta=8.0000,
          zeta=32.0000,
          shifts=[0.9000, 1.5500, 2.2000, 2.8500],
          sections=[0.1963, 0.5890, 0.9817, 1.3744, 1.7671, 2.1598, 2.5525, 2.9452],
          cutoff=3.5000,
          (cutoff_fn): CutoffCosine()
        )
        (neighborlist): AllPairs()
      )
      (neural_networks): Ensemble(
        (members): ModuleList(
          (0-7): 8 x ANINetworks(
            (atomics): ModuleDict(
              (H): AtomicNetwork(
                layer_dims=(384, 160, 128, 96, 1),
                activation=TightCELU(),
                bias=True,
                (layers): ModuleList(
                  (0): Linear(in_features=384, out_features=160, bias=True)
                  (1): Linear(in_features=160, out_features=128, bias=True)
                  (2): Linear(in_features=128, out_features=96, bias=True)
                )
                (final_layer): Linear(in_features=96, out_features=1, bias=True)
                (activation): TightCELU()
              )
              (C): AtomicNetwork(
                layer_dims=(384, 144, 112, 96, 1),
                activation=TightCELU(),
                bias=True,
                (layers): ModuleList(
                  (0): Linear(in_features=384, out_features=144, bias=True)
                  (1): Linear(in_features=144, out_features=112, bias=True)
                  (2): Linear(in_features=112, out_features=96, bias=True)
                )
                (final_layer): Linear(in_features=96, out_features=1, bias=True)
                (activation): TightCELU()
              )
              (N): AtomicNetwork(
                layer_dims=(384, 128, 112, 96, 1),
                activation=TightCELU(),
                bias=True,
                (layers): ModuleList(
                  (0): Linear(in_features=384, out_features=128, bias=True)
                  (1): Linear(in_features=128, out_features=112, bias=True)
                  (2): Linear(in_features=112, out_features=96, bias=True)
                )
                (final_layer): Linear(in_features=96, out_features=1, bias=True)
                (activation): TightCELU()
              )
              (O): AtomicNetwork(
                layer_dims=(384, 128, 112, 96, 1),
                activation=TightCELU(),
                bias=True,
                (layers): ModuleList(
                  (0): Linear(in_features=384, out_features=128, bias=True)
                  (1): Linear(in_features=128, out_features=112, bias=True)
                  (2): Linear(in_features=112, out_features=96, bias=True)
                )
                (final_layer): Linear(in_features=96, out_features=1, bias=True)
                (activation): TightCELU()
              )
            )
          )
        )
      )
    )
  )
)
ANI(
  (neighborlist): AllPairs()
  (energy_shifter): SelfEnergy()
  (species_converter): SpeciesConverter()
  (potentials): ModuleDict(
    (nnp): NNPotential(
      (aev_computer): AEVComputer(
        #  out_dim=384
        #  radial_len=64 (16.67% of feats)
        #  angular_len=320 (83.33% of feats)
        num_species=4,
        strategy=pyaev,
        (radial): ANIRadial(
          #  num_feats=16
          #  num_shifts=16
          eta=16.0000,
          shifts=[0.9000, 1.1688, 1.4375, 1.7062, 1.9750, 2.2438, 2.5125, 2.7812, 3.0500, 3.3187, 3.5875, 3.8563, 4.1250, 4.3938, 4.6625, 4.9313],
          cutoff=5.2000,
          (cutoff_fn): CutoffCosine()
        )
        (angular): ANIAngular(
          #  num_feats=32
          #  num_shifts=4
          #  num_sections=8
          eta=8.0000,
          zeta=32.0000,
          shifts=[0.9000, 1.5500, 2.2000, 2.8500],
          sections=[0.1963, 0.5890, 0.9817, 1.3744, 1.7671, 2.1598, 2.5525, 2.9452],
          cutoff=3.5000,
          (cutoff_fn): CutoffCosine()
        )
        (neighborlist): AllPairs()
      )
      (neural_networks): ANINetworks(
        (atomics): ModuleDict(
          (H): AtomicNetwork(
            layer_dims=(384, 160, 128, 96, 1),
            activation=TightCELU(),
            bias=True,
            (layers): ModuleList(
              (0): Linear(in_features=384, out_features=160, bias=True)
              (1): Linear(in_features=160, out_features=128, bias=True)
              (2): Linear(in_features=128, out_features=96, bias=True)
            )
            (final_layer): Linear(in_features=96, out_features=1, bias=True)
            (activation): TightCELU()
          )
          (C): AtomicNetwork(
            layer_dims=(384, 144, 112, 96, 1),
            activation=TightCELU(),
            bias=True,
            (layers): ModuleList(
              (0): Linear(in_features=384, out_features=144, bias=True)
              (1): Linear(in_features=144, out_features=112, bias=True)
              (2): Linear(in_features=112, out_features=96, bias=True)
            )
            (final_layer): Linear(in_features=96, out_features=1, bias=True)
            (activation): TightCELU()
          )
          (N): AtomicNetwork(
            layer_dims=(384, 128, 112, 96, 1),
            activation=TightCELU(),
            bias=True,
            (layers): ModuleList(
              (0): Linear(in_features=384, out_features=128, bias=True)
              (1): Linear(in_features=128, out_features=112, bias=True)
              (2): Linear(in_features=112, out_features=96, bias=True)
            )
            (final_layer): Linear(in_features=96, out_features=1, bias=True)
            (activation): TightCELU()
          )
          (O): AtomicNetwork(
            layer_dims=(384, 128, 112, 96, 1),
            activation=TightCELU(),
            bias=True,
            (layers): ModuleList(
              (0): Linear(in_features=384, out_features=128, bias=True)
              (1): Linear(in_features=128, out_features=112, bias=True)
              (2): Linear(in_features=112, out_features=96, bias=True)
            )
            (final_layer): Linear(in_features=96, out_features=1, bias=True)
            (activation): TightCELU()
          )
        )
      )
    )
  )
)

Now let’s define a methane molecule

coordinates = torch.tensor(
    [
        [
            [0.03192167, 0.00638559, 0.01301679],
            [-0.83140486, 0.39370209, -0.26395324],
            [-0.66518241, -0.84461308, 0.20759389],
            [0.45554739, 0.54289633, 0.81170881],
            [0.66091919, -0.16799635, -0.91037834],
        ]
    ],
)
species = torch.tensor([[6, 1, 1, 1, 1]], dtype=torch.long)

Now let’s compute energies using the ensemble directly:

energies, forces = energies_and_forces(ensemble_model, species, coordinates)
print("Energy:", energies.item())
print("Force:", forces.squeeze())
Energy: -40.45901870727539
Force: tensor([[ 0.031, -0.132, -0.053],
        [-0.129,  0.164, -0.077],
        [ 0.086, -0.043,  0.041],
        [ 0.027,  0.006,  0.038],
        [-0.014,  0.005,  0.051]])

We can do the same thing with the single model:

energies, forces = energies_and_forces(single_model, species, coordinates)
print("Energy:", energies.item())
print("Force:", forces.squeeze())
Energy: -40.46279525756836
Force: tensor([[ 0.056, -0.127, -0.054],
        [-0.140,  0.155, -0.075],
        [ 0.075, -0.037,  0.039],
        [ 0.024,  0.002,  0.033],
        [-0.016,  0.007,  0.057]])