test_ase.py

from ase.lattice.cubic import Diamond
from ase.md.langevin import Langevin
from ase import units, Atoms
from ase.calculators.test import numeric_force
import torch
import torchani
import unittest
import numpy
import itertools
import math
import os
import pickle

path = os.path.dirname(os.path.realpath(__file__))
N = 97
tol = 5e-5


def get_numeric_force(atoms, eps):
    fn = torch.zeros((len(atoms), 3))
    for i in range(len(atoms)):
        for j in range(3):
            fn[i, j] = numeric_force(atoms, i, j, eps)
    return fn


class TestASE(unittest.TestCase):

    def _testForce(self, pbc):
        atoms = Diamond(symbol="C", pbc=pbc)
        builtin = torchani.neurochem.Builtins()
        calculator = torchani.ase.Calculator(
            builtin.species, builtin.aev_computer,
            builtin.models, builtin.energy_shifter)
        atoms.set_calculator(calculator)
        dyn = Langevin(atoms, 5 * units.fs, 30000000 * units.kB, 0.002)
        dyn.run(100)
        f = torch.from_numpy(atoms.get_forces())
        fn = get_numeric_force(atoms, 0.001)
        df = (f - fn).abs().max()
        avgf = f.abs().mean()
        if avgf > 0:
            self.assertLess(df / avgf, 0.1)

    def testForceWithPBCEnabled(self):
        self._testForce(True)

    def testForceWithPBCDisabled(self):
        self._testForce(False)

    def testANIDataset(self):
        builtin = torchani.neurochem.Builtins()
        calculator = torchani.ase.Calculator(
            builtin.species, builtin.aev_computer,
            builtin.models, builtin.energy_shifter)
        default_neighborlist_calculator = torchani.ase.Calculator(
            builtin.species, builtin.aev_computer,
            builtin.models, builtin.energy_shifter, True)
        nnp = torch.nn.Sequential(
            builtin.aev_computer,
            builtin.models,
            builtin.energy_shifter
        )
        for i in range(N):
            datafile = os.path.join(path, 'test_data/ANI1_subset/{}'.format(i))
            with open(datafile, 'rb') as f:
                coordinates, species, _, _, _, _ = pickle.load(f)
                coordinates = coordinates[0]
                species = species[0]
                species_str = [builtin.consts.species[i] for i in species]

                atoms = Atoms(species_str, positions=coordinates)
                atoms.set_calculator(calculator)
                energy1 = atoms.get_potential_energy() / units.Hartree
                forces1 = atoms.get_forces() / units.Hartree

                atoms2 = Atoms(species_str, positions=coordinates)
                atoms2.set_calculator(default_neighborlist_calculator)
                energy2 = atoms2.get_potential_energy() / units.Hartree
                forces2 = atoms2.get_forces() / units.Hartree

                coordinates = torch.tensor(coordinates,
                                           requires_grad=True).unsqueeze(0)
                _, energy3 = nnp((torch.from_numpy(species).unsqueeze(0),
                                  coordinates))
                forces3 = -torch.autograd.grad(energy3.squeeze(),
                                               coordinates)[0].numpy()
                energy3 = energy3.item()

                self.assertLess(abs(energy1 - energy2), tol)
                self.assertLess(abs(energy1 - energy3), tol)

                diff_f12 = torch.tensor(forces1 - forces2).abs().max().item()
                self.assertLess(diff_f12, tol)
                diff_f13 = torch.tensor(forces1 - forces3).abs().max().item()
                self.assertLess(diff_f13, tol)

    def testForceAgainstDefaultNeighborList(self):
        atoms = Diamond(symbol="C", pbc=False)
        builtin = torchani.neurochem.Builtins()

        calculator = torchani.ase.Calculator(
            builtin.species, builtin.aev_computer,
            builtin.models, builtin.energy_shifter)
        default_neighborlist_calculator = torchani.ase.Calculator(
            builtin.species, builtin.aev_computer,
            builtin.models, builtin.energy_shifter, True)

        atoms.set_calculator(calculator)
        dyn = Langevin(atoms, 5 * units.fs, 50 * units.kB, 0.002)

        def test_energy(a=atoms):
            a = a.copy()
            a.set_calculator(calculator)
            e1 = a.get_potential_energy()
            a.set_calculator(default_neighborlist_calculator)
            e2 = a.get_potential_energy()
            self.assertLess(abs(e1 - e2), tol)

        dyn.attach(test_energy, interval=1)
        dyn.run(500)

    def testTranslationalInvariancePBC(self):
        atoms = Atoms('CH4', [[0, 0, 0],
                              [1, 0, 0],
                              [0, 1, 0],
                              [0, 0, 1],
                              [0, 1, 1]],
                      cell=[2, 2, 2], pbc=True)

        builtin = torchani.neurochem.Builtins()
        calculator = torchani.ase.Calculator(
            builtin.species, builtin.aev_computer,
            builtin.models, builtin.energy_shifter)
        atoms.set_calculator(calculator)
        e = atoms.get_potential_energy()

        for _ in range(100):
            positions = atoms.get_positions()
            translation = (numpy.random.rand(3) - 0.5) * 2
            atoms.set_positions(positions + translation)
            self.assertEqual(e, atoms.get_potential_energy())

    def assertTensorEqual(self, a, b):
        self.assertLess((a - b).abs().max().item(), 1e-6)

    def testPBCConnersSeeEachOther(self):
        species = torch.tensor([[0, 0]])
        neighborlist = torchani.ase.NeighborList(cell=[10, 10, 10], pbc=True)

        xyz1 = torch.tensor([0.1, 0.1, 0.1])
        xyz2s = [
            torch.tensor([9.9, 0.0, 0.0]),
            torch.tensor([0.0, 9.9, 0.0]),
            torch.tensor([0.0, 0.0, 9.9]),
            torch.tensor([9.9, 9.9, 0.0]),
            torch.tensor([0.0, 9.9, 9.9]),
            torch.tensor([9.9, 0.0, 9.9]),
            torch.tensor([9.9, 9.9, 9.9]),
        ]

        for xyz2 in xyz2s:
            coordinates = torch.stack([xyz1, xyz2]).unsqueeze(0)
            s, _, D = neighborlist(species, coordinates, 1)
            self.assertListEqual(list(s.shape), [1, 2, 1])
            neighbor_coordinate = D[0][0].squeeze() + xyz1
            mirror = xyz2
            for i in range(3):
                if mirror[i] > 5:
                    mirror[i] -= 10
            self.assertTensorEqual(neighbor_coordinate, mirror)

    def testPBCSurfaceSeeEachOther(self):
        species = torch.tensor([[0, 0]])
        neighborlist = torchani.ase.NeighborList(cell=[10, 10, 10], pbc=True)

        for i in range(3):
            xyz1 = torch.tensor([5.0, 5.0, 5.0])
            xyz1[i] = 0.1
            xyz2 = xyz1.clone()
            xyz2[i] = 9.9

            coordinates = torch.stack([xyz1, xyz2]).unsqueeze(0)
            s, _, D = neighborlist(species, coordinates, 1)
            self.assertListEqual(list(s.shape), [1, 2, 1])
            neighbor_coordinate = D[0][0].squeeze() + xyz1
            xyz2[i] = -0.1
            self.assertTensorEqual(neighbor_coordinate, xyz2)

    def testPBCEdgesSeeEachOther(self):
        species = torch.tensor([[0, 0]])
        neighborlist = torchani.ase.NeighborList(cell=[10, 10, 10], pbc=True)

        for i, j in itertools.combinations(range(3), 2):
            xyz1 = torch.tensor([5.0, 5.0, 5.0])
            xyz1[i] = 0.1
            xyz1[j] = 0.1
            for new_i, new_j in [[0.1, 9.9], [9.9, 0.1], [9.9, 9.9]]:
                xyz2 = xyz1.clone()
                xyz2[i] = new_i
                xyz2[j] = new_i

            coordinates = torch.stack([xyz1, xyz2]).unsqueeze(0)
            s, _, D = neighborlist(species, coordinates, 1)
            self.assertListEqual(list(s.shape), [1, 2, 1])
            neighbor_coordinate = D[0][0].squeeze() + xyz1

            if xyz2[i] > 5:
                xyz2[i] = -0.1
            if xyz2[j] > 5:
                xyz2[j] = -0.1

            self.assertTensorEqual(neighbor_coordinate, xyz2)

    def testNonRectangularPBCConnersSeeEachOther(self):
        species = torch.tensor([[0, 0]])
        neighborlist = torchani.ase.NeighborList(
            cell=[10, 10, 10 * math.sqrt(2), 90, 45, 90], pbc=True)

        xyz1 = torch.tensor([0.1, 0.1, 0.05])
        xyz2 = torch.tensor([10.0, 0.1, 0.1])
        mirror = torch.tensor([0.0, 0.1, 0.1])

        coordinates = torch.stack([xyz1, xyz2]).unsqueeze(0)
        s, _, D = neighborlist(species, coordinates, 1)
        self.assertListEqual(list(s.shape), [1, 2, 1])
        neighbor_coordinate = D[0][0].squeeze() + xyz1
        for i in range(3):
            if mirror[i] > 5:
                mirror[i] -= 10
        self.assertTensorEqual(neighbor_coordinate, mirror)


if __name__ == '__main__':
    unittest.main()