"docs/how to/use-transferbench.rst" did not exist on "c02da24bbe4dc2f22443a31fcf1adad3b452a515"
test_ase.py 3.02 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
import os

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


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

26
    def setUp(self):
27
        self.model = torchani.models.ANI1x(model_index=0).double()
28

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

42
43
44
    def testWithNumericalStressWithPBCEnabled(self):
        filename = os.path.join(path, '../tools/generate-unit-test-expect/others/Benzene.cif')
        benzene = read(filename)
45
        calculator = self.model.ase()
46
47
48
49
50
51
52
53
54
55
56
57
58
59
        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
60

61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
class TestASEWithPTI(unittest.TestCase):

    def setUp(self):
        self.model_pti = torchani.models.ANI1x(periodic_table_index=True).double()
        self.model = torchani.models.ANI1x().double()

    def testEqualEnsemblePTI(self):
        calculator_pti = self.model_pti.ase()
        calculator = self.model.ase()
        atoms = Diamond(symbol="C", pbc=True)
        atoms_pti = Diamond(symbol="C", pbc=True)
        atoms.set_calculator(calculator)
        atoms_pti.set_calculator(calculator_pti)
        self.assertEqual(atoms.get_potential_energy(), atoms_pti.get_potential_energy())

    def testEqualOneModelPTI(self):
        calculator_pti = self.model_pti[0].ase()
        calculator = self.model[0].ase()
        atoms = Diamond(symbol="C", pbc=True)
        atoms_pti = Diamond(symbol="C", pbc=True)
        atoms.set_calculator(calculator)
        atoms_pti.set_calculator(calculator_pti)
        self.assertEqual(atoms.get_potential_energy(), atoms_pti.get_potential_energy())


Gao, Xiang's avatar
Gao, Xiang committed
86
87
if __name__ == '__main__':
    unittest.main()