Construct 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:

import os
import torch
import torchani
import ase

Now let’s read constants from constant file and construct AEV computer.

try:
    path = os.path.dirname(os.path.realpath(__file__))
except NameError:
    path = os.getcwd()
const_file = os.path.join(path, '../torchani/resources/ani-1x_8x/rHCNO-5.2R_16-3.5A_a4-8.params')  # noqa: E501
consts = torchani.neurochem.Constants(const_file)
aev_computer = torchani.AEVComputer(**consts)

Now let’s read self energies and construct energy shifter.

sae_file = os.path.join(path, '../torchani/resources/ani-1x_8x/sae_linfit.dat')  # noqa: E501
energy_shifter = torchani.neurochem.load_sae(sae_file)

Now let’s read a whole ensemble of models.

model_prefix = os.path.join(path, '../torchani/resources/ani-1x_8x/train')  # noqa: E501
ensemble = torchani.neurochem.load_model_ensemble(consts.species, model_prefix, 8)  # noqa: E501

Or alternatively a single model.

model_dir = os.path.join(path, '../torchani/resources/ani-1x_8x/train0/networks')  # noqa: E501
model = torchani.neurochem.load_model(consts.species, model_dir)

You can create the pipeline of computing energies: (Coordinates) -[AEVComputer]-> (AEV) -[Neural Network]-> (Raw energies) -[EnergyShifter]-> (Final energies) From using either the ensemble or a single model:

nnp1 = torchani.nn.Sequential(aev_computer, ensemble, energy_shifter)
nnp2 = torchani.nn.Sequential(aev_computer, model, energy_shifter)
print(nnp1)
print(nnp2)

Sequential(
  (0): AEVComputer()
  (1): Ensemble(
    (0-7): 8 x ANIModel(
      (H): Sequential(
        (0): Linear(in_features=384, out_features=160, bias=True)
        (1): CELU(alpha=0.1)
        (2): Linear(in_features=160, out_features=128, bias=True)
        (3): CELU(alpha=0.1)
        (4): Linear(in_features=128, out_features=96, bias=True)
        (5): CELU(alpha=0.1)
        (6): Linear(in_features=96, out_features=1, bias=True)
      )
      (C): Sequential(
        (0): Linear(in_features=384, out_features=144, bias=True)
        (1): CELU(alpha=0.1)
        (2): Linear(in_features=144, out_features=112, bias=True)
        (3): CELU(alpha=0.1)
        (4): Linear(in_features=112, out_features=96, bias=True)
        (5): CELU(alpha=0.1)
        (6): Linear(in_features=96, out_features=1, bias=True)
      )
      (N): Sequential(
        (0): Linear(in_features=384, out_features=128, bias=True)
        (1): CELU(alpha=0.1)
        (2): Linear(in_features=128, out_features=112, bias=True)
        (3): CELU(alpha=0.1)
        (4): Linear(in_features=112, out_features=96, bias=True)
        (5): CELU(alpha=0.1)
        (6): Linear(in_features=96, out_features=1, bias=True)
      )
      (O): Sequential(
        (0): Linear(in_features=384, out_features=128, bias=True)
        (1): CELU(alpha=0.1)
        (2): Linear(in_features=128, out_features=112, bias=True)
        (3): CELU(alpha=0.1)
        (4): Linear(in_features=112, out_features=96, bias=True)
        (5): CELU(alpha=0.1)
        (6): Linear(in_features=96, out_features=1, bias=True)
      )
    )
  )
  (2): EnergyShifter()
)
Sequential(
  (0): AEVComputer()
  (1): ANIModel(
    (H): Sequential(
      (0): Linear(in_features=384, out_features=160, bias=True)
      (1): CELU(alpha=0.1)
      (2): Linear(in_features=160, out_features=128, bias=True)
      (3): CELU(alpha=0.1)
      (4): Linear(in_features=128, out_features=96, bias=True)
      (5): CELU(alpha=0.1)
      (6): Linear(in_features=96, out_features=1, bias=True)
    )
    (C): Sequential(
      (0): Linear(in_features=384, out_features=144, bias=True)
      (1): CELU(alpha=0.1)
      (2): Linear(in_features=144, out_features=112, bias=True)
      (3): CELU(alpha=0.1)
      (4): Linear(in_features=112, out_features=96, bias=True)
      (5): CELU(alpha=0.1)
      (6): Linear(in_features=96, out_features=1, bias=True)
    )
    (N): Sequential(
      (0): Linear(in_features=384, out_features=128, bias=True)
      (1): CELU(alpha=0.1)
      (2): Linear(in_features=128, out_features=112, bias=True)
      (3): CELU(alpha=0.1)
      (4): Linear(in_features=112, out_features=96, bias=True)
      (5): CELU(alpha=0.1)
      (6): Linear(in_features=96, out_features=1, bias=True)
    )
    (O): Sequential(
      (0): Linear(in_features=384, out_features=128, bias=True)
      (1): CELU(alpha=0.1)
      (2): Linear(in_features=128, out_features=112, bias=True)
      (3): CELU(alpha=0.1)
      (4): Linear(in_features=112, out_features=96, bias=True)
      (5): CELU(alpha=0.1)
      (6): Linear(in_features=96, out_features=1, bias=True)
    )
  )
  (2): EnergyShifter()
)

You can also create an ASE calculator using the ensemble or single model:

calculator1 = torchani.ase.Calculator(consts.species, nnp1)
calculator2 = torchani.ase.Calculator(consts.species, nnp2)
print(calculator1)
print(calculator1)

<torchani.ase.Calculator object at 0x7f691d6df640>
<torchani.ase.Calculator object at 0x7f691d6df640>

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]]],
                           requires_grad=True)
species = consts.species_to_tensor(['C', 'H', 'H', 'H', 'H']).unsqueeze(0)
methane = ase.Atoms(['C', 'H', 'H', 'H', 'H'], positions=coordinates.squeeze().detach().numpy())

Now let’s compute energies using the ensemble directly:

energy = nnp1((species, coordinates)).energies
derivative = torch.autograd.grad(energy.sum(), coordinates)[0]
force = -derivative
print('Energy:', energy.item())
print('Force:', force.squeeze())

Energy: -40.45902210758762
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]])

And using the ASE interface of the ensemble:

methane.set_calculator(calculator1)
print('Energy:', methane.get_potential_energy() / ase.units.Hartree)
print('Force:', methane.get_forces() / ase.units.Hartree)

Energy: -40.45902210758762
Force: [[ 0.03062528 -0.13160461 -0.05265211]
 [-0.12927507  0.16388635 -0.07736803]
 [ 0.08563145 -0.04288925  0.04082095]
 [ 0.02681219  0.00601404  0.03809873]
 [-0.01379385  0.00459347  0.05110047]]

We can do the same thing with the single model:

energy = nnp2((species, coordinates)).energies
derivative = torch.autograd.grad(energy.sum(), coordinates)[0]
force = -derivative
print('Energy:', energy.item())
print('Force:', force.squeeze())

methane.set_calculator(calculator2)
print('Energy:', methane.get_potential_energy() / ase.units.Hartree)
print('Force:', methane.get_forces() / ase.units.Hartree)

Energy: -40.46280037151413
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]])
Energy: -40.46280037151413
Force: [[ 0.05614968 -0.12697159 -0.05413406]
 [-0.1400753   0.15523373 -0.07533834]
 [ 0.0753209  -0.0374411   0.03949669]
 [ 0.02423098  0.00235087  0.03344228]
 [-0.01562626  0.00682809  0.05653343]]