Note
Go to the end to download the full example code
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 0x7f28a943c5b0>
<torchani.ase.Calculator object at 0x7f28a943c5b0>
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]]
Total running time of the script: (0 minutes 2.026 seconds)