Don't require to specify masses for vibrational analysis (#427)

* Add tuple of atomic masses

* Add get_atomic_masses function, to get masses from periodic table indices (znumbers)

* Fix bug that prevented using ASE together with periodic_table_index

* ASE has to know if the model has PTI or not

* Allow for models that don't have a periodic_table_index attribute

* now vibrational analysis can be computed without explicit masses if periodic_table_index=True

* Add atomic masses calculator to docs

* flake8

* flake8

* flake8
Co-authored-by: Gao, Xiang <qasdfgtyuiop@gmail.com>

docs/api.rst

@@ -47,6 +47,7 @@ Utilities
     :members:
 .. autofunction:: torchani.utils.hessian
 .. autofunction:: torchani.utils.vibrational_analysis
+.. autofunction:: torchani.utils.get_atomic_masses
 
 
 NeuroChem

examples/vibration_analysis.py

@@ -19,7 +19,8 @@ import math
 
 ###############################################################################
 # Let's now manually specify the device we want TorchANI to run:
-model = torchani.models.ANI1x().double()
+device = torch.device('cpu')
+model = torchani.models.ANI1x(periodic_table_index=True).to(device).double()
 
 ###############################################################################
 # Let's first construct a water molecule and do structure optimization:
@@ -36,19 +37,14 @@ opt.run(fmax=1e-6)
 ###############################################################################
 # Now let's extract coordinates and species from ASE to use it directly with
 # TorchANI:
-species = model.species_to_tensor(molecule.get_chemical_symbols()).unsqueeze(0)
+species = torch.tensor(molecule.get_atomic_numbers(), device=device, dtype=torch.long).unsqueeze(0)
 coordinates = torch.from_numpy(molecule.get_positions()).unsqueeze(0).requires_grad_(True)
 
 ###############################################################################
-# TorchANI needs to know the mass of each atom in amu in order to do vibration
-# analysis:
-element_masses = torch.tensor([
-    1.008,  # H
-    12.011,  # C
-    14.007,  # N
-    15.999,  # O
-], dtype=torch.double)
-masses = element_masses[species]
+# TorchANI needs the masses of elements in AMU to compute vibrations. The
+# masses in AMU can be obtained from a tensor with atomic numbers by using
+# this utility:
+masses = torchani.utils.get_atomic_masses(species)
 
 ###############################################################################
 # To do vibration analysis, we first need to generate a graph that computes

torchani/utils.py

@@ -357,6 +357,63 @@ def vibrational_analysis(masses, hessian, mode_type='MDU', unit='cm^-1'):
     return VibAnalysis(wavenumbers, modes, fconstants, rmasses)
 
 
+def get_atomic_masses(species):
+    r"""Convert a tensor of znumbers into a tensor of atomic masses
+
+    Atomic masses supported are the first 119 elements, and are taken from:
+
+    Atomic weights of the elements 2013 (IUPAC Technical Report). Meija, J.,
+    Coplen, T., Berglund, M., et al. (2016). Pure and Applied Chemistry, 88(3), pp.
+    265-291. Retrieved 30 Nov. 2016, from doi:10.1515/pac-2015-0305
+
+    They are all consistent with those used in ASE
+
+    Arguments:
+        species (:class:`torch.Tensor`): tensor with atomic numbers
+
+    Returns:
+        :class:`torch.Tensor`: Tensor of dtype :class:`torch.double`, with
+        atomic masses, with the same shape as the input.
+    """
+    # Note that there should not be any atoms with index zero, because that is
+    # not an element
+    assert len((species == 0).nonzero()) == 0
+    default_atomic_masses = torch.tensor(
+            [0.        ,   1.008     ,   4.002602  ,   6.94      , # noqa
+             9.0121831 ,  10.81      ,  12.011     ,  14.007     , # noqa
+            15.999     ,  18.99840316,  20.1797    ,  22.98976928, # noqa
+            24.305     ,  26.9815385 ,  28.085     ,  30.973762  , # noqa
+            32.06      ,  35.45      ,  39.948     ,  39.0983    , # noqa
+            40.078     ,  44.955908  ,  47.867     ,  50.9415    , # noqa
+            51.9961    ,  54.938044  ,  55.845     ,  58.933194  , # noqa
+            58.6934    ,  63.546     ,  65.38      ,  69.723     , # noqa
+            72.63      ,  74.921595  ,  78.971     ,  79.904     , # noqa
+            83.798     ,  85.4678    ,  87.62      ,  88.90584   , # noqa
+            91.224     ,  92.90637   ,  95.95      ,  97.90721   , # noqa
+           101.07      , 102.9055    , 106.42      , 107.8682    , # noqa
+           112.414     , 114.818     , 118.71      , 121.76      , # noqa
+           127.6       , 126.90447   , 131.293     , 132.90545196, # noqa
+           137.327     , 138.90547   , 140.116     , 140.90766   , # noqa
+           144.242     , 144.91276   , 150.36      , 151.964     , # noqa
+           157.25      , 158.92535   , 162.5       , 164.93033   , # noqa
+           167.259     , 168.93422   , 173.054     , 174.9668    , # noqa
+           178.49      , 180.94788   , 183.84      , 186.207     , # noqa
+           190.23      , 192.217     , 195.084     , 196.966569  , # noqa
+           200.592     , 204.38      , 207.2       , 208.9804    , # noqa
+           208.98243   , 209.98715   , 222.01758   , 223.01974   , # noqa
+           226.02541   , 227.02775   , 232.0377    , 231.03588   , # noqa
+           238.02891   , 237.04817   , 244.06421   , 243.06138   , # noqa
+           247.07035   , 247.07031   , 251.07959   , 252.083     , # noqa
+           257.09511   , 258.09843   , 259.101     , 262.11      , # noqa
+           267.122     , 268.126     , 271.134     , 270.133     , # noqa
+           269.1338    , 278.156     , 281.165     , 281.166     , # noqa
+           285.177     , 286.182     , 289.19      , 289.194     , # noqa
+           293.204     , 293.208     , 294.214], # noqa
+           dtype=torch.double, device=species.device)
+    masses = default_atomic_masses[species]
+    return masses
+
+
 __all__ = ['pad_atomic_properties', 'present_species', 'hessian',
            'vibrational_analysis', 'strip_redundant_padding',
-           'ChemicalSymbolsToInts']
+           'ChemicalSymbolsToInts', 'get_atomic_masses']