add lennardjones to autograd module

mlp/ag/lennardjones.py

@@ -0,0 +1,44 @@
+"""A periodic Lennard Jones potential."""
+
+
+import autograd.numpy as np
+from autograd import elementwise_grad
+from mlp.ag.neighborlist import get_distances
+
+def energy(params, positions, cell, strain=np.zeros((3, 3))):
+    "Compute the energy of a Lennard-Jones system."
+
+    sigma = params.get('sigma', 1.0)
+    epsilon = params.get('epsilon', 1.0)
+
+    rc = 3 * sigma
+
+    e0 = 4 * epsilon * ((sigma / rc)**12 - (sigma / rc)**6)
+
+    r2 = get_distances(positions, cell, rc, 0.01, strain)**2
+
+    zeros = np.equal(r2, 0.0)
+    adjusted = np.where(zeros, np.ones_like(r2), r2)
+
+    c6 = np.where((r2 <= rc**2) & (r2 > 0.0), (sigma**2 / adjusted)**3, np.zeros_like(r2))
+    c6 = np.where(zeros, np.zeros_like(r2), c6)
+    energy = -e0 * (c6 != 0.0).sum()
+    c12 = c6**2
+    energy += np.sum(4 * epsilon * (c12 - c6))
+
+    return energy / 2
+
+
+def forces(params, positions, cell):
+    dEdR = elementwise_grad(energy, 1)
+    return -dEdR(params, positions, cell)
+
+
+def stress(params, positions, cell, strain=np.zeros((3, 3))):
+    dEdst = elementwise_grad(energy, 3)
+
+    volume = np.abs(np.linalg.det(cell))
+
+    der = dEdst(params, positions, cell, strain)
+    result = (der + der.T) / 2 / volume
+    return np.take(result, [0, 4, 8, 5, 2, 1])

mlp/tests/test_ag.py

@@ -2,7 +2,12 @@
 import autograd.numpy as np
 from ase.build import bulk
 from ase.neighborlist import NeighborList
+from ase.calculators.lj import LennardJones
+
 from mlp.ag.neighborlist import get_distances
+from mlp.ag.lennardjones import energy, forces, stress
+
+
 
 class TestNeighborList(unittest.TestCase):
 
@@ -30,3 +35,30 @@ def test0(self):
                 nns = inds.sum((1, 2))
 
                 self.assertTrue(np.all(nns_ase == nns))
+
+
+class TestLennardJones(unittest.TestCase):
+    def test_energy(self):
+
+        atoms = bulk('Cu', 'fcc', a=3.7).repeat((2, 2, 1))
+        atoms.rattle(0.02)
+        atoms.set_calculator(LennardJones())
+
+        ase_energy = atoms.get_potential_energy()
+        lj_energy = energy({}, atoms.positions, atoms.cell)
+
+        self.assertAlmostEqual(ase_energy, lj_energy)
+
+        lj_forces = forces({}, atoms.positions, atoms.cell)
+        self.assertTrue(np.allclose(atoms.get_forces(),
+                                     lj_forces))
+
+        lj_stress = stress({}, atoms.positions, atoms.cell)
+
+        self.assertTrue(np.allclose(atoms.get_stress(),
+                                     lj_stress))
+
+
+
+
+