Add vibrational mode support (#228)

examples/vibration_analysis.py

@@ -69,5 +69,9 @@ print(hessian.shape)
 # We are now ready to compute vibrational frequencies. The output has unit
 # cm^-1. Since there are in total 9 degree of freedom, there are in total 9
 # frequencies. Only the frequencies of the 3 vibrational modes are interesting.
-freq = torchani.utils.vibrational_analysis(masses, hessian)[-3:]
-print(freq)
+freq, modes = torchani.utils.vibrational_analysis(masses, hessian)
+freq = freq[-3:]
+modes = modes[-3:]
+
+print('Frequencies (cm^-1):', freq)
+print('Modes:', modes)

tests/test_vibrational.py

@@ -30,16 +30,27 @@ class TestVibrational(unittest.TestCase):
         # compute vibrational frequencies by ASE
         vib = ase.vibrations.Vibrations(molecule)
         vib.run()
-        freq = torch.tensor([numpy.real(x) for x in vib.get_frequencies()[-3:]])
+        freq = torch.tensor([numpy.real(x) for x in vib.get_frequencies()[6:]])
+        modes = []
+        for j in range(6, 6 + len(freq)):
+            modes.append(vib.get_mode(j))
+        modes = torch.tensor(modes)
         # compute vibrational by torchani
         species = model.species_to_tensor(molecule.get_chemical_symbols()).unsqueeze(0)
         coordinates = torch.from_numpy(molecule.get_positions()).unsqueeze(0).requires_grad_(True)
         _, energies = model((species, coordinates))
         hessian = torchani.utils.hessian(coordinates, energies=energies)
-        freq2 = torchani.utils.vibrational_analysis(masses[species], hessian)[-3:].float()
+        freq2, modes2 = torchani.utils.vibrational_analysis(masses[species], hessian)
+        freq2 = freq2[6:].float()
+        modes2 = modes2[6:]
         ratio = freq2 / freq
         self.assertLess((ratio - 1).abs().max(), 0.02)
 
+        diff1 = (modes - modes2).abs().max(dim=-1).values.max(dim=-1).values
+        diff2 = (modes + modes2).abs().max(dim=-1).values.max(dim=-1).values
+        diff = torch.where(diff1 < diff2, diff1, diff2)
+        self.assertLess(diff.max(), 0.02)
+
 
 if __name__ == '__main__':
     unittest.main()

torchani/utils.py

@@ -248,18 +248,19 @@ def vibrational_analysis(masses, hessian, unit='cm^-1'):
     # We solve this eigenvalue problem through Lowdin diagnolization:
     # Hq = w^2 * Tq ==> Hq = w^2 * T^(1/2) T^(1/2) q
     # Letting q' = T^(1/2) q, we then have
-    # T^(-1/2) H T^(1/2) q' = w^2 * q'
+    # T^(-1/2) H T^(-1/2) q' = w^2 * q'
     inv_sqrt_mass = (1 / masses.sqrt()).repeat_interleave(3, dim=1)  # shape (molecule, 3 * atoms)
     mass_scaled_hessian = hessian * inv_sqrt_mass.unsqueeze(1) * inv_sqrt_mass.unsqueeze(2)
     if mass_scaled_hessian.shape[0] != 1:
         raise ValueError('The input should contain only one molecule')
     mass_scaled_hessian = mass_scaled_hessian.squeeze(0)
-    eigenvalues = torch.symeig(mass_scaled_hessian).eigenvalues
+    eigenvalues, eigenvectors = torch.symeig(mass_scaled_hessian, eigenvectors=True)
     angular_frequencies = eigenvalues.sqrt()
     frequencies = angular_frequencies / (2 * math.pi)
     # converting from sqrt(hartree / (amu * angstrom^2)) to cm^-1
     wavenumbers = frequencies * 17092
-    return wavenumbers
+    modes = (eigenvectors.t() * inv_sqrt_mass).reshape(frequencies.numel(), -1, 3)
+    return wavenumbers, modes
 
 
 __all__ = ['pad', 'pad_coordinates', 'present_species', 'hessian',