Exploring PyTorch + ANI + MD
PyTorch provides nice utilities for differentiation. ANI provides some interatomic potentials trained on some neural networks. Molecular Dynamics might be an interesting combination
import torch
import matplotlib.pyplot as plt
x = torch.ones((2,2), requires_grad=True)
A simple quadratic function
def sq_function(x):
return x**2
Since we have an array of 1s, the square won't look very interesting...
foo = sq_function(x)
foo
More interstingly, we can compute the gradient of this function.
To compute the gradient, the value/function needs to be a scalar, but this scalar could be computed from a bunch of other functions stemming from some independent variables (our tensor x). In this case, our final scalar looks like this, $ Y = x_0^2 + x_1^2 + x_2^2 + x_3^2 $. Taking the gradient means taking 4 partial derivatives for each input. Fortunately, the equation is simple to compute each partial derivative, $ \frac{\partial Y}{\partial x_i} = 2*x_i $, where $i = [0,4)$. Since this is an array of 1s, each partial derivative evaluates to 2
torch.autograd.grad(foo.sum(), x)
We've evaluated the function and its gradient at just one point, but we can use some numpy-esque functions to evaluate the square-function and its gradient at a range of points.
Yup, looks right to me
some_xvals = torch.arange(-12., 12., step=0.5, requires_grad=True)
some_yvals = sq_function(some_xvals)
fig, ax = plt.subplots(1,1)
ax.plot(some_xvals.detach().numpy(), some_yvals.detach().numpy())
ax.plot(some_xvals.detach().numpy(),
torch.autograd.grad(some_yvals.sum(), some_xvals)[0])
Slightly more book-keeping, 3x 1-D harmonic springs
Define an energy function as the sum of 3 harmonic springs
$ V(x, y, z) = V_x + V_y + V_z = (x-x_0)^2 + (y-y_0)^2 + (z-z_0)^2 $
The gradient, the 3 partial derivatives, are computed as such (being verbose with the chain rule)
$ \frac{\partial V}{\partial X} = 2 *(x-x_0) * 1 $
$\frac{\partial V}{\partial Y} = 2 *(y-y_0) * 1$
$\frac{\partial V}{\partial Z} = 2 *(z-z_0) * 1$
def harmonic_spring_3d(coord, origin=torch.tensor([0,0,0])):
V_x = (coord[0]-origin[0])**2
V_y = (coord[1]-origin[1])**2
V_z = (coord[2]-origin[2])**2
return V_x + V_y + V_z
We can evaluate the potential energy at 1 point, which involves computing the energy in 3 dimensions.
Our "anchor" will be the origin, and our endpoint will be (1,2,3)
$ 1^2 + 2^2 + 3^2 = 14 $
my_coords = torch.tensor([1.,2.,3.], requires_grad=True)
total_energy = harmonic_spring_3d(my_coords)
total_energy
Computing the gradient, partial derivatives in each direction, which is simply 2 times the distance in each dimension
$ \nabla \hat V = < 2*1, 2*2, 2*3 > = <2,4,6> $
torch.autograd.grad(total_energy, my_coords)
More involved: Lennard Jones
The Lennard-Jones potential describes the potential energy between two particles. Not the most accurate potential, but has been decent for a long time now. Some background information on the Lenanrd-Jones potential. For simplicity, assume $\epsilon =1$ and $\sigma=1$ in unitless quantities:
$ V_{LJ} = 4 * ( \frac{1}{r}^{12} - \frac{1}{r}^6) $
$ -\frac{\partial V}{\partial r} = -4 * (-12 * r^{-13} + 6 * r^{-7}) $
def lj(val):
return 4 * ((1/val)**12 - (1/val)**6)
r_values = torch.arange(0.1, 12., step=0.001, requires_grad=True)
energy = lj(r_values)
forces = -torch.autograd.grad(energy.sum(), r_values)[0]
For sanity check, we can confirm that energy reaches a critical point (local minimum) when the force is 0.
Also, this definitely looks like a LJ potential to me
import matplotlib.pyplot as plt
fig, ax = plt.subplots(1,1, dpi=100)
ax.plot(r_values.detach().numpy(), energy.detach().numpy(), label='energy')
ax.plot(r_values.detach().numpy(), forces.detach().numpy(), label='force')
ax.set_ylim([-2,1])
ax.legend()
ax.set_xlim([0,2])
ax.axhline(y=0, color='r', linestyle='--')
To begin, we have to define our elements (a tensor of atomic numbers). For the molecular mechanics people, each atom is identifiable by its element, and not one of many atom-types.
We have to define the positions (units of Angstrom), which is also a multi-dimensional tensor.
Load the model, specifying to convert the atomic numbers to indices suitable for ANI.
We can compute the energies and forces from the model. The energy comes from the model, but the force is obtained via an autograd call, observing that we are differentiating the sum of the forces, evaluating at the positions
import torchani
elements = torch.tensor([[6, 6]])
positions = torch.tensor([[[3.0, 3.0, 3.0],
[3.5, 3.5, 3.5]]], requires_grad=True)
model = torchani.models.ANI2x(periodic_table_index=True)
energy = model((elements, positions)).energies
forces = -1.0 * torch.autograd.grad(energy.sum(), positions)[0]
energy
forces
Going a step further, we can try to visualize the interaction potential by evaluating the energy at a variety of distances. We can also do some autodifferentiation to compute the forces.
In this example, we have 2 atoms that share X and Y coordinates, but pull them apart in the Z direction
all_z = torch.arange(3.0, 12.0, step=0.1)
all_energy = []
all_forces = []
for z in all_z:
# Generate a new set of positions
positions = torch.tensor([[[3.0, 3.0, 3.0],
[3.0, 3.0, z]]], requires_grad=True
)
# Compute energy
energy = model((elements, positions)).energies
# Compute force
forces = -1.0 * torch.autograd.grad(energy.sum(), positions)[0]
# Get the force vector on the first atom
one_atom_forces = forces[0,0]
# Compute the magnitude of this force vector
force_magnitude = torch.sqrt(torch.dot(one_atom_forces, one_atom_forces))
# Calculate the unit vector for this force vector,
# although it's a little unnecessary because the only distance is in the
# z direction
unit_vector_force = one_atom_forces/force_magnitude
# Get z-component of force vector
force_vector_z = unit_vector_force[2]*force_magnitude
# Some nans will form if the force magnitude is zero, but this
# is really just a 0 force vector
if torch.isnan(force_vector_z).any():
force_vector_z = 0.0
else:
force_vector_z = float(force_vector_z.detach().numpy())
# Accumulate
all_energy.append(float(energy.detach().numpy()))
all_forces.append(force_vector_z)
Hmmm... this does not resemble the Lennard-Jones potential (or basic chemistry for that matter)
fig, ax = plt.subplots(1,1, dpi=100)
ax.plot(all_z-3, all_energy)
ax.set_xlabel(r"Distance ($\AA$)")
ax.set_ylabel("Energy (Hartree)")
fig, ax = plt.subplots(1,1, dpi=100)
ax.plot(all_z-3, all_forces)
ax.set_xlabel(r"Distance ($\AA$)")
ax.set_ylabel("Force (Hartree / $\AA$)")
from mbuild.lib.recipes import Alkane
# The mBuild alkane recipe is mainly used to generate
# some particles and positions
cmpd = Alkane(n=5)
# Convert to mdtraj trajectory out of convenience for atomic numbers
traj = cmpd.to_trajectory()
# Periodic cell, from nm to angstrom
cell = torch.tensor(traj.unitcell_vectors[0]*10)
# We just need atomic numbers
species = torch.tensor([[
a.element.atomic_number for a in traj.top.atoms
]])
# Make tensor for coordinates
# Since we are differentiating WRT coordinates, we need the
# requires_grad=True
coordinates = torch.tensor(traj.xyz*10, requires_grad=True)
# PBC flag necessary for computing energies with periodic boundaries
pbc = torch.tensor([True, True, True], dtype=torch.bool)
energies = model((species, coordinates), cell=cell, pbc=pbc).energies
forces = -1.0 * (
torch.autograd.grad(energies.sum(), coordinates)[0]
)
energies
forces
To be continued ...
One might imagine trying to incorporate ANI potentials into MD simulations (which has been done in ASE). However, the torchani-API is general enough that you could use any number of computational chemistry packages to feed into torchani. The output is also general enough you could imagine trying to apply your own integrators and make your own simulation. But... from the weird 2-atom interatomic potentials, some of these methods might require some debugging.
Files and environment can be found here