test_ase.py 2.04 KB
Newer Older
Gao, Xiang's avatar
Gao, Xiang committed
1
2
from ase.lattice.cubic import Diamond
from ase.md.langevin import Langevin
3
from ase.md.nptberendsen import NPTBerendsen
4
from ase import units
5
from ase.io import read
Gao, Xiang's avatar
Gao, Xiang committed
6
7
8
9
from ase.calculators.test import numeric_force
import torch
import torchani
import unittest
10
11
12
13
14
import os

path = os.path.dirname(os.path.realpath(__file__))
N = 97
tol = 5e-5
Gao, Xiang's avatar
Gao, Xiang committed
15
16
17


def get_numeric_force(atoms, eps):
18
    fn = torch.zeros((len(atoms), 3), dtype=torch.double)
Gao, Xiang's avatar
Gao, Xiang committed
19
20
21
22
23
24
25
26
    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):

27
28
29
    def setUp(self):
        self.model = torchani.models.ANI1x().double()[0]

30
31
    def testWithNumericalForceWithPBCEnabled(self):
        atoms = Diamond(symbol="C", pbc=True)
32
        calculator = self.model.ase()
Gao, Xiang's avatar
Gao, Xiang committed
33
34
35
36
37
38
39
        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()
Gao, Xiang's avatar
Gao, Xiang committed
40
41
42
        if avgf > 0:
            self.assertLess(df / avgf, 0.1)

43
44
45
    def testWithNumericalStressWithPBCEnabled(self):
        filename = os.path.join(path, '../tools/generate-unit-test-expect/others/Benzene.cif')
        benzene = read(filename)
46
        calculator = self.model.ase()
47
48
49
50
51
52
53
54
55
56
57
58
59
60
        benzene.set_calculator(calculator)
        dyn = NPTBerendsen(benzene, timestep=0.1 * units.fs,
                           temperature=300 * units.kB,
                           taut=0.1 * 1000 * units.fs, pressure=1.01325,
                           taup=1.0 * 1000 * units.fs, compressibility=4.57e-5)

        def test_stress():
            stress = benzene.get_stress()
            numerical_stress = calculator.calculate_numerical_stress(benzene)
            diff = torch.from_numpy(stress - numerical_stress).abs().max().item()
            self.assertLess(diff, tol)
        dyn.attach(test_stress, interval=30)
        dyn.run(120)

Gao, Xiang's avatar
Gao, Xiang committed
61
62
63

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