Test on structural optimization (#153)

.gitattributes

 *.h5 filter=lfs diff=lfs merge=lfs -text
 tests/test_data/ANI1_subset/* filter=lfs diff=lfs merge=lfs -text
 tests/test_data/NIST/* filter=lfs diff=lfs merge=lfs -text
-tools/diverse_test_set/result.json filter=lfs diff=lfs merge=lfs -text
+tests/test_data/NeuroChemOptimized/* filter=lfs diff=lfs merge=lfs -text
+tools/generate-unit-test-expect/nist-dataset/result.json filter=lfs diff=lfs merge=lfs -text

tests/test_structure_optim.py

+import os
+import unittest
+import torch
+import torchani
+import copy
+import pickle
+from ase.optimize import BFGS
+
+
+path = os.path.dirname(os.path.realpath(__file__))
+
+
+class TestStructureOptimization(unittest.TestCase):
+
+    def setUp(self):
+        self.tolerance = 1e-6
+        self.builtin = torchani.neurochem.Builtins()
+        self.calculator = torchani.ase.Calculator(
+            self.builtin.species, self.builtin.aev_computer,
+            self.builtin.models, self.builtin.energy_shifter)
+
+    def testRMSE(self):
+        datafile = os.path.join(path, 'test_data/NeuroChemOptimized/all')
+        with open(datafile, 'rb') as f:
+            all_atoms = pickle.load(f)
+            for atoms in all_atoms:
+                old_coordinates = copy.deepcopy(atoms.get_positions())
+                old_coordinates = torch.from_numpy(old_coordinates)
+                atoms.set_calculator(self.calculator)
+                opt = BFGS(atoms)
+                opt.run()
+                coordinates = atoms.get_positions()
+                coordinates = torch.from_numpy(coordinates)
+                distances = (old_coordinates - coordinates).norm(dim=1)
+                rmse = distances.mean()
+                self.assertLess(rmse, self.tolerance)
+
+
+if __name__ == '__main__':
+    unittest.main()

tools/generate-unit-test-expect/ANI1.py

+import os
+import pickle
+import pyanitools
+from neurochem_calculator import NeuroChem, path
+import tqdm
+
+
+neurochem = NeuroChem()
+
+# generate expect for ANI1 subset
+mol_count = 0
+for i in [1, 2, 3, 4]:
+    data_file = os.path.join(
+        path, '../../dataset/ani_gdb_s0{}.h5'.format(i))
+    adl = pyanitools.anidataloader(data_file)
+    for data in tqdm.tqdm(adl, desc='ANI1: {} heavy atoms'.format(i)):
+        coordinates = data['coordinates'][:10, :]
+        pickleobj = neurochem(coordinates, data['species'])
+        dumpfile = os.path.join(
+            path, '../../tests/test_data/ANI1_subset/{}'.format(mol_count))
+        with open(dumpfile, 'wb') as f:
+            pickle.dump(pickleobj, f)
+        mol_count += 1

tools/neurochem_calculator.py → tools/generate-unit-test-expect/neurochem_calculator.py

@@ -5,7 +5,7 @@ import ase_interface
 import numpy
 
 path = os.path.dirname(os.path.realpath(__file__))
-builtin_path = os.path.join(path, '../torchani/resources/ani-1x_dft_x8ens/')
+builtin_path = os.path.join(path, '../../torchani/resources/ani-1x_dft_x8ens/')
 const_file = os.path.join(builtin_path, 'rHCNO-5.2R_16-3.5A_a4-8.params')
 sae_file = os.path.join(builtin_path, 'sae_linfit.dat')
 network_dir = os.path.join(builtin_path, 'train0/networks/')
@@ -13,12 +13,16 @@ radial_length = 64
 conv_au_ev = 27.21138505
 all_species = ['H', 'C', 'N', 'O']
 species_indices = {all_species[i]: i for i in range(len(all_species))}
+nc = pyNeuroChem.molecule(const_file, sae_file, network_dir, 0)
 
 
-class NeuroChem:
+def calc():
+    calc = ase_interface.ANI(False)
+    calc.setnc(nc)
+    return calc
+
 
-    def __init__(self):
-        self.nc = pyNeuroChem.molecule(const_file, sae_file, network_dir, 0)
+class NeuroChem:
 
     def _get_radial_part(self, fullaev):
         return fullaev[:, :, :radial_length]
@@ -29,8 +33,7 @@ class NeuroChem:
     def _per_conformation(self, coordinates, species):
         atoms = coordinates.shape[0]
         mol = ase.Atoms(''.join(species), positions=coordinates)
-        mol.set_calculator(ase_interface.ANI(False))
-        mol.calc.setnc(self.nc)
+        mol.set_calculator(calc())
         energy = mol.get_potential_energy() / conv_au_ev
         aevs = [mol.calc.nc.atomicenvironments(j) for j in range(atoms)]
         force = mol.get_forces() / conv_au_ev

tools/generate-unit-test-expect.py → tools/generate-unit-test-expect/nist.py

 import os
 import pickle
-import pyanitools
 import numpy
 from neurochem_calculator import NeuroChem, path
 import json
@@ -9,27 +8,9 @@ import random
 
 
 neurochem = NeuroChem()
-
-# generate expect for ANI1 subset
-mol_count = 0
-for i in [1, 2, 3, 4]:
-    data_file = os.path.join(
-        path, '../dataset/ani_gdb_s0{}.h5'.format(i))
-    adl = pyanitools.anidataloader(data_file)
-    for data in tqdm.tqdm(adl, desc='ANI1: {} heavy atoms'.format(i)):
-        coordinates = data['coordinates'][:10, :]
-        pickleobj = neurochem(coordinates, data['species'])
-        dumpfile = os.path.join(
-            path, '../tests/test_data/ANI1_subset/{}'.format(mol_count))
-        with open(dumpfile, 'wb') as f:
-            pickle.dump(pickleobj, f)
-        mol_count += 1
-
-
-# generate expect for NIST
 keep_ratio = 0.1  # reduce the size of generated file by discarding
 mol_count = 0
-with open(os.path.join(path, 'diverse_test_set/result.json')) as f:
+with open(os.path.join(path, 'nist-dataset/result.json')) as f:
     pickle_objects = []
     for i in tqdm.tqdm(json.load(f), desc='NIST'):
         if random.random() > keep_ratio:
@@ -46,6 +27,6 @@ with open(os.path.join(path, 'diverse_test_set/result.json')) as f:
         mol_count += 1
 
     dumpfile = os.path.join(
-        path, '../tests/test_data/NIST/all')
+        path, '../../tests/test_data/NIST/all')
     with open(dumpfile, 'wb') as f:
         pickle.dump(pickle_objects, f)

tools/generate-unit-test-expect/structure-optim.py

+import os
+import pickle
+import json
+import tqdm
+import random
+import ase
+from ase.optimize import BFGS
+from neurochem_calculator import calc, path
+
+keep_ratio = 0.01  # reduce the size of generated file by discarding
+mol_count = 0
+with open(os.path.join(path, 'nist-dataset/result.json')) as f:
+    pickle_objects = []
+    for i in tqdm.tqdm(json.load(f), desc='NIST'):
+        if random.random() > keep_ratio:
+            continue
+        atoms = i['atoms']
+        natoms = len(atoms)
+        species = []
+        coordinates = []
+        for atype, x, y, z in atoms:
+            species.append(atype)
+            coordinates.append([x, y, z])
+        mol = ase.Atoms(species, positions=coordinates)
+        mol.set_calculator(calc())
+        opt = BFGS(mol, logfile='/dev/null')
+        opt.run()
+        mol.set_calculator(None)
+        pickle_objects.append(mol)
+        mol_count += 1
+
+    dumpfile = os.path.join(
+        path, '../../tests/test_data/NeuroChemOptimized/all')
+    with open(dumpfile, 'wb') as f:
+        pickle.dump(pickle_objects, f)