* Forces by automatic differentiation in molecular simulation
:PROPERTIES:
:categories: simulation, autograd
:date: 2017/11/14 21:06:22
:updated: 2017/11/14 21:06:22
:org-url: http://kitchingroup.cheme.cmu.edu/org/2017/11/14/Forces-by-automatic-differentiation-in-molecular-simulation.org
:permalink: http://kitchingroup.cheme.cmu.edu/blog/2017/11/14/Forces-by-automatic-differentiation-in-molecular-simulation/index.html
:END:
In molecular simulation we often use a potential to compute the total energy of a system. For example, we might use something simple like a [[https://en.wikipedia.org/wiki/Lennard-Jones_potential][Lennard-Jones]] potential. If we have a potential function, e.g. $E = V(R)$ where $R$ are the positions of the atoms, then we know the forces on the atoms are defined by $f = -\frac{dV}{dR}$. For simple functions, you can derive the derivative pretty easily, but these functions quickly get complicated. In this post, we consider [[https://en.wikipedia.org/wiki/Automatic_differentiation][automatic differentiation]] as implemented by [[https://github.com/HIPS/autograd][autograd]]. This is neither symbolic nor numerical differentiation. The gist is that the program uses the chain rule to evaluate derivatives. Here we do not delve into how it is done, we just see how it might help us in molecular simulation.
For reference, here is a result from the LennardJones calculator in ASE.
#+NAME: hotel-red-skylark-august
#+BEGIN_SRC ipython
from ase.calculators.lj import LennardJones
from ase.cluster.icosahedron import Icosahedron
atoms = Icosahedron('Ar', noshells=2, latticeconstant=3)
atoms.set_calculator(LennardJones())
atoms.rattle(0.5)
print('LJ: ', atoms.get_potential_energy())
#+END_SRC
#+RESULTS: hotel-red-skylark-august
:RESULTS:
#+BEGIN_EXAMPLE
('LJ: ', -3.3553466825679812)
#+END_EXAMPLE
:END:
First, we define a function for the Lennard-Jones potential. I adapted the code [[https://wiki.fysik.dtu.dk/ase/_modules/ase/calculators/lj.html#LennardJones][here]] to implement a function that calculates the Lennard-Jones energy for a cluster of atoms with no periodic boundary conditions. /Instead/ of using numpy directly, we import it from the autograd package which puts thin wrappers around the functions to enable the derivative calculations.
#+NAME: violet-fourteen-emma-utah
#+BEGIN_SRC ipython
import autograd.numpy as np
def energy(positions):
"Compute the energy of a Lennard-Jones system."
natoms = len(positions)
sigma = 1.0
epsilon = 1.0
rc = 3 * sigma
e0 = 4 * epsilon * ((sigma / rc)**12 - (sigma / rc)**6)
energy = 0.0
for a1 in range(natoms):
for j in range(a1 + 1, natoms):
r2 = np.sum((positions[a1] - positions[j])**2)
if r2 <= rc**2:
c6 = (sigma**2 / r2)**3
energy -= e0
c12 = c6**2
energy += 4 * epsilon * (c12 - c6)
return energy
#+END_SRC
Here is our function in action, and it produces the same result as the ASE calculator. So far there is nothing new.
#+NAME: uniform-virginia-yellow-xray
#+BEGIN_SRC ipython
print('our func: ', energy(atoms.positions))
#+END_SRC
#+RESULTS: uniform-virginia-yellow-xray
:RESULTS:
#+BEGIN_EXAMPLE
('our func: ', -3.3553466825679803)
#+END_EXAMPLE
:END:
Now, we look at the forces from the ASE calculator. If you look at the ASE [[https://wiki.fysik.dtu.dk/ase/_modules/ase/calculators/lj.html#LennardJones][code]] you will see that the formula for forces was analytically derived and accumulated in the loop.
#+NAME: alabama-spring-georgia-solar
#+BEGIN_SRC ipython
np.set_printoptions(precision=3, suppress=True)
print(atoms.get_forces())
#+END_SRC
#+RESULTS: alabama-spring-georgia-solar
:RESULTS:
#+BEGIN_EXAMPLE
[[ 0.545 1.667 0.721]
[-0.068 0.002 0.121]
[-0.18 0.018 -0.121]
[ 0.902 -0.874 -0.083]
[ 0.901 -0.937 -1.815]
[ 0.243 -0.19 0.063]
[-0.952 -1.776 -0.404]
[-0.562 1.822 1.178]
[-0.235 0.231 0.081]
[-0.023 0.204 -0.294]
[ 0.221 -0.342 -0.425]
[-5.385 -6.017 1.236]
[ 4.593 6.193 -0.258]]
#+END_EXAMPLE
:END:
Now we look at how to use autograd for this purpose. We want an element-wise gradient of the total energy with respect to the positions. autograd returns functions, so we wrap it in another function so we can take the negative of that function.
#+NAME: beryllium-pasta-juliet-table
#+BEGIN_SRC ipython
from autograd import elementwise_grad
def forces(pos):
dEdR = elementwise_grad(energy)
return -dEdR(pos)
print(forces(atoms.positions))
#+END_SRC
#+RESULTS: beryllium-pasta-juliet-table
:RESULTS:
#+BEGIN_EXAMPLE
[[ 0.545 1.667 0.721]
[-0.068 0.002 0.121]
[-0.18 0.018 -0.121]
[ 0.902 -0.874 -0.083]
[ 0.901 -0.937 -1.815]
[ 0.243 -0.19 0.063]
[-0.952 -1.776 -0.404]
[-0.562 1.822 1.178]
[-0.235 0.231 0.081]
[-0.023 0.204 -0.294]
[ 0.221 -0.342 -0.425]
[-5.385 -6.017 1.236]
[ 4.593 6.193 -0.258]]
#+END_EXAMPLE
:END:
Here we show the results are the same from both approaches.
#+NAME: illinois-violet-helium-island
#+BEGIN_SRC ipython
print(np.allclose(atoms.get_forces(), forces(atoms.positions)))
#+END_SRC
#+RESULTS: illinois-violet-helium-island
:RESULTS:
#+BEGIN_EXAMPLE
True
#+END_EXAMPLE
:END:
Wow. We got forces without deriving a derivative, or using numerical finite differences, across loops, and conditionals. That is pretty awesome. You can easily modify the potential function now, without the need to rederive the force derivatives! This is an idea worth exploring further. In principle, it should be possible to include periodic boundary conditions and use autograd to compute stresses too. Maybe that will be a future post.