Note
Go to the end to download the full example code.
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
We can also load an ANI model using neurochem
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:
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:
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]])