Add Benzene and Tripeptide MD trajectory to unit test (#213)

tests/test_energies.py

@@ -4,7 +4,6 @@ import unittest
 import os
 import pickle
 import math
-import random
 
 
 path = os.path.dirname(os.path.realpath(__file__))
@@ -42,6 +41,42 @@ class TestEnergies(unittest.TestCase):
                 max_diff = (energies - energies_).abs().max().item()
                 self.assertLess(max_diff, self.tolerance)
 
+    @unittest.skipIf(True, "WIP")
+    def testBenzeneMD(self):
+        tolerance = 1e-5
+        for i in range(100):
+            datafile = os.path.join(path, 'test_data/benzene-md/{}.dat'.format(i))
+            with open(datafile, 'rb') as f:
+                coordinates, species, _, _, energies, _, cell, pbc \
+                    = pickle.load(f)
+                coordinates = torch.from_numpy(coordinates).float().unsqueeze(0)
+                species = torch.from_numpy(species).unsqueeze(0)
+                cell = torch.from_numpy(cell).float()
+                pbc = torch.from_numpy(pbc)
+                coordinates = torchani.utils.map2central(cell, coordinates, pbc)
+                coordinates = self.transform(coordinates)
+                species = self.transform(species)
+                energies = self.transform(energies)
+                _, energies_ = self.model((species, coordinates, cell, pbc))
+                max_diff = (energies - energies_).abs().max().item()
+                self.assertLess(max_diff, tolerance)
+
+    def testTripeptideMD(self):
+        tolerance = 2e-4
+        for i in range(100):
+            datafile = os.path.join(path, 'test_data/tripeptide-md/{}.dat'.format(i))
+            with open(datafile, 'rb') as f:
+                coordinates, species, _, _, energies, _, _, _ \
+                    = pickle.load(f)
+                coordinates = torch.from_numpy(coordinates).float().unsqueeze(0)
+                species = torch.from_numpy(species).unsqueeze(0)
+                coordinates = self.transform(coordinates)
+                species = self.transform(species)
+                energies = self.transform(energies)
+                _, energies_ = self.model((species, coordinates))
+                max_diff = (energies - energies_).abs().max().item()
+                self.assertLess(max_diff, tolerance)
+
     def testPadding(self):
         species_coordinates = []
         energies = []
@@ -80,19 +115,5 @@ class TestEnergies(unittest.TestCase):
                 self.assertLess(max_diff / math.sqrt(natoms), self.tolerance)
 
 
-class TestEnergiesASEComputer(TestEnergies):
-
-    def setUp(self):
-        super(TestEnergiesASEComputer, self).setUp()
-
-    def transform(self, x):
-        """To reduce the size of test cases for faster test speed"""
-        return x[:2, ...]
-
-    def random_skip(self):
-        """To reduce the size of test cases for faster test speed"""
-        return random.random() < 0.95
-
-
 if __name__ == '__main__':
     unittest.main()

tests/test_forces.py

@@ -3,7 +3,6 @@ import torchani
 import unittest
 import os
 import pickle
-import random
 
 path = os.path.dirname(os.path.realpath(__file__))
 N = 97
@@ -68,6 +67,48 @@ class TestForce(unittest.TestCase):
             max_diff = (forces + derivative).abs().max().item()
             self.assertLess(max_diff, self.tolerance)
 
+    @unittest.skipIf(True, "WIP")
+    def testBenzeneMD(self):
+        tolerance = 1e-6
+        for i in range(100):
+            datafile = os.path.join(path, 'test_data/benzene-md/{}.dat'.format(i))
+            with open(datafile, 'rb') as f:
+                coordinates, species, _, _, _, forces, cell, pbc \
+                    = pickle.load(f)
+                coordinates = torch.from_numpy(coordinates).float().unsqueeze(0).requires_grad_(True)
+                species = torch.from_numpy(species).unsqueeze(0)
+                cell = torch.from_numpy(cell).float()
+                pbc = torch.from_numpy(pbc)
+                forces = torch.from_numpy(forces)
+                coordinates = torchani.utils.map2central(cell, coordinates, pbc)
+                coordinates = self.transform(coordinates)
+                species = self.transform(species)
+                forces = self.transform(forces)
+                _, energies_ = self.model((species, coordinates, cell, pbc))
+                derivative = torch.autograd.grad(energies_.sum(),
+                                                 coordinates)[0]
+                max_diff = (forces + derivative).abs().max().item()
+                self.assertLess(max_diff, tolerance)
+
+    def testTripeptideMD(self):
+        tolerance = 2e-6
+        for i in range(100):
+            datafile = os.path.join(path, 'test_data/tripeptide-md/{}.dat'.format(i))
+            with open(datafile, 'rb') as f:
+                coordinates, species, _, _, _, forces, _, _ \
+                    = pickle.load(f)
+                coordinates = torch.from_numpy(coordinates).float().unsqueeze(0).requires_grad_(True)
+                species = torch.from_numpy(species).unsqueeze(0)
+                forces = torch.from_numpy(forces)
+                coordinates = self.transform(coordinates)
+                species = self.transform(species)
+                forces = self.transform(forces)
+                _, energies_ = self.model((species, coordinates))
+                derivative = torch.autograd.grad(energies_.sum(),
+                                                 coordinates)[0]
+                max_diff = (forces + derivative).abs().max().item()
+                self.assertLess(max_diff, tolerance)
+
     def testNIST(self):
         datafile = os.path.join(path, 'test_data/NIST/all')
         with open(datafile, 'rb') as f:
@@ -86,19 +127,5 @@ class TestForce(unittest.TestCase):
                 self.assertLess(max_diff, self.tolerance)
 
 
-class TestForceASEComputer(TestForce):
-
-    def setUp(self):
-        super(TestForceASEComputer, self).setUp()
-
-    def transform(self, x):
-        """To reduce the size of test cases for faster test speed"""
-        return x[:2, ...]
-
-    def random_skip(self):
-        """To reduce the size of test cases for faster test speed"""
-        return random.random() < 0.95
-
-
 if __name__ == '__main__':
     unittest.main()

tools/generate-unit-test-expect/benzene-md.py

+import ase
+import ase.io
+import ase.optimize
+import ase.md.velocitydistribution
+import ase.md.verlet
+import os
+from neurochem_calculator import NeuroChem, path
+import torchani
+import pickle
+
+
+neurochem = NeuroChem()
+
+molecule = ase.io.read('others/Benzene.cif')
+
+temp = 300 * ase.units.kB
+stepsize = ase.units.fs
+steps = int(10000 * ase.units.fs / stepsize)
+
+ase.md.velocitydistribution.MaxwellBoltzmannDistribution(molecule, temp, force_temp=True)
+ase.md.velocitydistribution.Stationary(molecule)
+ase.md.velocitydistribution.ZeroRotation(molecule)
+
+print("Initial temperature from velocities %.2f" % molecule.get_temperature())
+
+molecule.set_calculator(torchani.models.ANI1ccx().to('cuda').ase())
+
+dyn = ase.md.verlet.VelocityVerlet(
+    molecule,
+    stepsize,
+    logfile='-',
+)
+
+counter = 0
+data_dir = os.path.join(path, '../../tests/test_data/benzene-md/')
+
+
+def dump_neurochem_data(molecule=molecule):
+    global counter
+    filename = os.path.join(data_dir, '{}.dat'.format(counter))
+    ret = neurochem(molecule)
+    with open(filename, 'wb') as f:
+        pickle.dump(ret, f)
+    counter += 1
+
+
+dyn.attach(dump_neurochem_data, interval=100)
+dyn.run(steps)

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_8x/')
 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/')
@@ -25,10 +25,16 @@ def calc():
 class NeuroChem:
 
     def _get_radial_part(self, fullaev):
-        return fullaev[:, :, :radial_length]
+        if len(fullaev.shape) == 3:
+            return fullaev[:, :, :radial_length]
+        assert len(fullaev.shape) == 2
+        return fullaev[:, :radial_length]
 
     def _get_angular_part(self, fullaev):
-        return fullaev[:, :, radial_length:]
+        if len(fullaev.shape) == 3:
+            return fullaev[:, :, radial_length:]
+        assert len(fullaev.shape) == 2
+        return fullaev[:, radial_length:]
 
     def _per_conformation(self, coordinates, species):
         atoms = coordinates.shape[0]
@@ -40,7 +46,27 @@ class NeuroChem:
         aevs = numpy.stack(aevs)
         return aevs, energy, force
 
-    def __call__(self, coordinates, species):
+    def __call__(self, args):
+        if len(args) == 2:
+            return self.from_coordinates_and_species(*args)
+        return self.from_atoms_obj(args)
+
+    def from_atoms_obj(self, mol):
+        natoms = len(mol)
+        mol = mol.copy()
+        mol.set_calculator(calc())
+        energy = mol.get_potential_energy() / conv_au_ev
+        aevs = [mol.calc.nc.atomicenvironments(j) for j in range(natoms)]
+        aevs = numpy.stack(aevs)
+        force = mol.get_forces() / conv_au_ev
+        cell = mol.get_cell(complete=True)
+        pbc = mol.get_pbc().astype(numpy.uint8)
+        coordinates = mol.get_positions()
+        species = numpy.array([species_indices[i] for i in mol.get_chemical_symbols()])
+        return coordinates, species, self._get_radial_part(aevs), \
+            self._get_angular_part(aevs), energy, force, cell, pbc
+
+    def from_coordinates_and_species(self, coordinates, species):
         if len(coordinates.shape) == 2:
             coordinates = coordinates.reshape(1, -1, 3)
         conformations = coordinates.shape[0]

tools/generate-unit-test-expect/others/Benzene.cif

+
+#######################################################################
+#
+#                 Cambridge Crystallographic Data Centre
+#                                CCDC 
+#
+#######################################################################
+#
+# If this CIF has been generated directly or indirectly from an entry in the 
+# Cambridge Structural Database, then it will include bibliographic, chemical, 
+# crystal, experimental, refinement or atomic coordinate data resulting from 
+# the CCDC's data processing and validation procedures. Files generated from 
+# CSD entries are Copyright 2012 Cambridge Crystallographic Data Centre. They 
+# may be used in bona fide research applications only, and may not be copied or 
+# further disseminated in any form, whether machine-readable or not, except for 
+# the purposes of generating routine backup copies on your local computer 
+# system.
+# 
+# Files arising from any other source may also contain material that is the 
+# copyright of third parties, including the originator, and you should check 
+# with the originator concerning the permitted uses of the information 
+# contained in this CIF.
+#
+# For further information on the CCDC and the free tools enCIFer and Mercury
+# for validating and visualising CIF files, please visit https://urldefense.proofpoint.com/v2/url?u=http-3A__www.ccdc.cam.ac.uk&d=DwIGAg&c=sJ6xIWYx-zLMB3EPkvcnVg&r=dl7Zd5Rzbdvo14I2ndQf4w&m=orkB7wHPEppuovTbB3oqDN_oQdvAKyxgBSb_iCN15FM&s=ILVnBGtKq4Zt42mg_pLsCGetPJ0_KzFzTbICq_zg04Y&e=
+#
+#######################################################################
+
+data_BENZEN11
+_audit_creation_date             2006-03-08
+_database_code_depnum_ccdc_archive 'CCDC 298305'
+_journal_coeditor_code           'IUCr AV5045'
+_chemical_formula_moiety         'C6 H6'
+_chemical_name_systematic        Benzene
+_journal_coden_Cambridge         622
+_journal_volume                  62
+_journal_year                    2006
+_journal_page_first              94
+_journal_name_full               'Acta Crystallogr.,Sect.B:Struct.Sci.'
+loop_
+_publ_author_name
+A.Budzianowski
+A.Katrusiak
+_chemical_absolute_configuration unk
+_diffrn_ambient_temperature      296
+_diffrn_ambient_pressure         'at 0.70 GPa'
+_exptl_crystal_density_meas      1.157
+_exptl_crystal_density_diffrn    1.157
+#These two values have been output from a single CSD field.
+_refine_ls_R_factor_gt           0.053
+_refine_ls_wR_factor_gt          0.053
+_diffrn_radiation_probe          x-ray
+_symmetry_cell_setting           orthorhombic
+_symmetry_space_group_name_H-M   'P b c a'
+_symmetry_Int_Tables_number      61
+loop_
+_symmetry_equiv_pos_site_id
+_symmetry_equiv_pos_as_xyz
+1 x,y,z
+2 1/2-x,-y,1/2+z
+3 -x,1/2+y,1/2-z
+4 1/2+x,1/2-y,-z
+5 -x,-y,-z
+6 1/2+x,y,1/2-z
+7 x,1/2-y,1/2+z
+8 1/2-x,1/2+y,z
+_cell_length_a                   7.287(6)
+_cell_length_b                   9.20(2)
+_cell_length_c                   6.688(9)
+_cell_angle_alpha                90
+_cell_angle_beta                 90
+_cell_angle_gamma                90
+_cell_volume                     448.366
+_exptl_crystal_colour            colorless
+_cell_formula_units_Z            4
+loop_
+_atom_site_label
+_atom_site_type_symbol
+_atom_site_fract_x
+_atom_site_fract_y
+_atom_site_fract_z
+C1 C -0.0537(8) 0.1425(9) 0.0097(12)
+H1 H -0.085(6) 0.246(7) 0.034(8)
+C2 C 0.0840(7) 0.0924(10) 0.1373(10)
+H2 H 0.140(6) 0.156(6) 0.219(8)
+C3 C -0.1343(7) 0.0521(9) -0.1235(12)
+H3 H -0.220(6) 0.080(6) -0.204(9)
+C1D C 0.0537 -0.1425 -0.0097
+H1D H 0.085 -0.246 -0.034
+C2D C -0.084 -0.0924 -0.1373
+H2D H -0.14 -0.156 -0.219
+C3D C 0.1343 -0.0521 0.1235
+H3D H 0.22 -0.08 0.204
+
+#END
+
+

tools/generate-unit-test-expect/tripeptide-md.py

+import ase
+import ase.io
+import ase.optimize
+import ase.md.velocitydistribution
+import ase.md.verlet
+import os
+from neurochem_calculator import NeuroChem, path
+import torchani
+import pickle
+
+
+neurochem = NeuroChem()
+
+molecule = ase.io.read(os.path.join(path, 'tripeptide/tripeptide-000.ipt_optimization.xyz'))
+
+temp = 300 * ase.units.kB
+stepsize = 0.25 * ase.units.fs
+steps = int(10000 * ase.units.fs / stepsize)
+
+ase.md.velocitydistribution.MaxwellBoltzmannDistribution(molecule, temp, force_temp=True)
+ase.md.velocitydistribution.Stationary(molecule)
+ase.md.velocitydistribution.ZeroRotation(molecule)
+
+print("Initial temperature from velocities %.2f" % molecule.get_temperature())
+
+molecule.set_calculator(torchani.models.ANI1ccx().to('cuda').ase())
+
+dyn = ase.md.verlet.VelocityVerlet(
+    molecule,
+    stepsize,
+    logfile='-',
+)
+
+counter = 0
+data_dir = os.path.join(path, '../../tests/test_data/tripeptide-md/')
+
+
+def dump_neurochem_data(molecule=molecule):
+    global counter
+    filename = os.path.join(data_dir, '{}.dat'.format(counter))
+    ret = neurochem(molecule)
+    with open(filename, 'wb') as f:
+        pickle.dump(ret, f)
+    counter += 1
+
+
+dyn.attach(dump_neurochem_data, interval=400)
+dyn.run(steps)

torchani/resources/ani-1ccx_8x.info

+ani-1ccx_8x/rHCNO-5.2R_16-3.5A_a4-8.params
+ani-1ccx_8x/sae_linfit.dat
+ani-1ccx_8x/train
+8
+

torchani/resources/ani-1x_8x.info

+ani-1x_8x/rHCNO-5.2R_16-3.5A_a4-8.params
+ani-1x_8x/sae_linfit.dat
+ani-1x_8x/train
+8