vib

more

fix

more

.github/workflows/test-nightly-torch.yml

@@ -0,0 +1,31 @@
+name: test-nightly-torch
+
+on:
+  pull_request:
+  push:
+    branches:
+      - master
+  schedule:
+    - cron: '0 0 * * *'
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      max-parallel: 4
+      matrix:
+        python-version: [3.6, 3.7]
+
+    steps:
+    - uses: actions/checkout@v1
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v1
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Run tests
+      run: |
+        pip install --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+        pip install -r optional_requirements.txt
+        pip install .
+        pytest

README.md

@@ -6,6 +6,7 @@ Status:
 [![Actions Status](https://github.com/aiqm/nnp/workflows/mypy/badge.svg)](https://github.com/aiqm/nnp/actions)
 [![Actions Status](https://github.com/aiqm/nnp/workflows/docs/badge.svg)](https://github.com/aiqm/nnp/actions)
 [![Actions Status](https://github.com/aiqm/nnp/workflows/test/badge.svg)](https://github.com/aiqm/nnp/actions)
+[![Actions Status](https://github.com/aiqm/nnp/workflows/test-nightly-torch/badge.svg)](https://github.com/aiqm/nnp/actions)
 [![Actions Status](https://github.com/aiqm/nnp/workflows/deploy-test-pypi/badge.svg)](https://github.com/aiqm/nnp/actions)
 [![Actions Status](https://github.com/aiqm/nnp/workflows/deploy-pypi/badge.svg)](https://github.com/aiqm/nnp/actions)
 [![Actions Status](https://github.com/aiqm/nnp/workflows/deploy-docs/badge.svg)](https://github.com/aiqm/nnp/actions)

docs/start.rst

@@ -14,12 +14,13 @@ To install NNP, just do
 
 
 TODO:
+
 - [x] so3
 - [x] test so3
+- [x] vib
+- [ ] test vib
 - [ ] pbc
 - [ ] test pbc
 - [ ] md
 - [x] test md: autograd
 - [ ] test md: stress
-- [ ] vib
-- [ ] test vib

nnp/vib.py

@@ -1,66 +1,131 @@
 """
 Vibrational Analysis
 ====================
-"""
 
+The module ``nnp.vib`` contains tools to compute analytical hessian
+and do vibrational analysis.
+"""
+###############################################################################
+# Let's first import all the packages we will use:
 import torch
 from torch import Tensor
 from typing import NamedTuple, Optional
-import math
+
+
+###############################################################################
+# With ``torch.autograd`` automatic differentiation engine, computing analytical
+# hessian is very simple: Just compute the gradient of energies with respect
+# to forces, and then compute the gradient of each element of forces with respect
+# to coordinates again.
+#
+# The ``torch.autograd.grad`` returns a list of ``Optional[Tensor]``. To be
+# compatible with ``torch.jit``, we need to use assert statement to allow
+# ``torch.jit`` to do `type refinement`_.
+#
+# .. _type refinement:
+#   https://pytorch.org/docs/stable/jit.html#optional-type-refinement
+def _get_derivatives_not_none(x: Tensor, y: Tensor, retain_graph: Optional[bool] = None,
+                              create_graph: bool = False) -> Tensor:
+    ret = torch.autograd.grad(
+        [y.sum()], [x], retain_graph=retain_graph, create_graph=create_graph)[0]
+    assert ret is not None
+    return ret
 
 
 def hessian(coordinates: Tensor, energies: Optional[Tensor] = None,
             forces: Optional[Tensor] = None) -> Tensor:
     """Compute analytical hessian from the energy graph or force graph.
 
     Arguments:
-        coordinates: Tensor of shape `(molecules, atoms, 3)` or ``(atoms, 3)``.
-        energies: Tensor of shape `(molecules,)`, if specified, then `forces`
-            must be `None`. This energies must be computed from `coordinates`
-            in a graph.
-        forces: Tensor of shape `(molecules, atoms, 3)`, if specified,
-            then `energies` must be `None`. This forces must be computed from
-            `coordinates` in a graph.
+        coordinates: Tensor of shape `(molecules, atoms, 3)` or `(atoms, 3)`
+        energies: Tensor of shape `(molecules,)`, or scalar, if specified,
+            then `forces` must be `None`. This energies must be computed
+            from `coordinates` in a graph.
+        forces: Tensor of shape `(molecules, atoms, 3)` or `(atoms, 3)`,
+            if specified, then `energies` must be `None`. This forces must
+            be computed from `coordinates` in a graph.
+
     Returns:
-        Tensor of shape `(molecules, 3A, 3A)` where A is the number of atoms in
-        each molecule
+        Tensor of shape `(molecules, 3 * atoms, 3 * atoms)` or `(3 * atoms, 3 * atoms)`
     """
     if energies is None and forces is None:
-        raise ValueError()
+        raise ValueError('Energies or forces must be specified')
     if energies is not None and forces is not None:
         raise ValueError('Energies or forces can not be specified at the same time')
     if forces is None:
         assert energies is not None
-        forces = -torch.autograd.grad(
-            [energies.sum()], [coordinates], create_graph=True)[0]
-    flattened_force = forces.flatten(start_dim=1)
-    force_components = flattened_force.unbind(dim=1)
+        forces = -_get_derivatives_not_none(coordinates, energies, create_graph=True)
+    flattened_force = forces.flatten(start_dim=-2)
+    force_components = flattened_force.unbind(dim=-1)
     return -torch.stack([
-        torch.autograd.grad(
-            [f.sum()], [coordinates], retain_graph=True)[0].flatten(start_dim=1)
+        _get_derivatives_not_none(coordinates, f, retain_graph=True).flatten(start_dim=-2)
         for f in force_components
-    ], dim=1)
+    ], dim=-1)
 
 
+###############################################################################
+# Below are helper functions to compute vibrational frequencies and normal modes.
+# The normal modes and vibrational frquencies satisfies the following equation.
+#
+# .. math::
+#   H q = \omega^2 T q
+#
+# where :math:`H` is the Hessian matrix, :math:`q` is the normal coordinates, and
+#
+# .. math::
+#   T=\left[
+#    \begin{array}{ccccccc}
+#    m_{1}\\
+#    & m_{1}\\
+#    &  & m_{1}\\
+#    &  &  & m_{2}\\
+#    &  &  &  & m_{2}\\
+#    &  &  &  &  & m_{2}\\
+#    &  &  &  &  &  & \ddots
+#    \end{array}
+#    \right]
+#
+# is the mass for each coordinate.
+#
+# This is a generalized eigen problem, which is not immediately supported by PyTorch
+# So we solve this problem through Lowdin diagnolization:
+#
+# .. math::
+#   H q = \omega^2 T q \Longrightarrow H q = \omega^2 T^{\frac{1}{2}} T^{\frac{1}{2}} q
+#
+# Let :math:`q' = T^{\frac{1}{2}} q`, we then have
+#
+# .. math::
+#   T^{\frac{1}{2}} H T^{\frac{1}{2}} q' = \omega^2 q'
+#
+# this is a regular eigen problem
 class FreqsModes(NamedTuple):
-    freqs: Tensor
+    angular_frequencies: Tensor
     modes: Tensor
 
 
 def vibrational_analysis(masses: Tensor, hessian: Tensor) -> FreqsModes:
-    """Computing the vibrational wavenumbers from hessian."""
-    # Solving the eigenvalue problem: Hq = w^2 * T q
-    # where H is the Hessian matrix, q is the normal coordinates,
-    # T = diag(m1, m1, m1, m2, m2, m2, ....) is the mass
-    # 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'
-    inv_sqrt_mass = masses.rsqrt().repeat_interleave(3, dim=1)  # shape (molecule, 3 * atoms)
-    mass_scaled_hessian = hessian * inv_sqrt_mass.unsqueeze(1) * inv_sqrt_mass.unsqueeze(2)
-    mass_scaled_hessian = mass_scaled_hessian.squeeze(0)
+    """Computing the vibrational wavenumbers from hessian.
+
+    Arguments:
+        masses: Tensor of shape `(molecules, atoms)` or `(atoms,)`.
+        hessian: Tensor of shape `(molecules, 3 * atoms, 3 * atoms)` or
+            `(3 * atoms, 3 * atoms)`.
+
+    Returns:
+        A namedtuple `(angular_frequencies, modes)` where
+
+        angular_frequencies:
+            Tensor of shape `(molecules, 3 * atoms)` or `(3 * atoms,)`
+        modes:
+            Tensor of shape `(molecules, modes, atoms, 3)` or `(modes, atoms, 3)`
+            where `modes = 3 * atoms` is the number of normal modes.
+    """
+    inv_sqrt_mass = masses.rsqrt().repeat_interleave(3, dim=-1)
+    mass_scaled_hessian = hessian * inv_sqrt_mass.unsqueeze(-2) * inv_sqrt_mass.unsqueeze(-1)
     eigenvalues, eigenvectors = torch.symeig(mass_scaled_hessian, eigenvectors=True)
     angular_frequencies = eigenvalues.sqrt()
-    frequencies = angular_frequencies / (2 * math.pi)
-    modes = (eigenvectors.t() * inv_sqrt_mass).reshape(frequencies.numel(), -1, 3)
-    return FreqsModes(frequencies, modes)
+    modes = (eigenvectors.transpose(-1, -2) * inv_sqrt_mass.unsqueeze(-2))
+    new_shape = modes.shape[:-1] + (-1, 3)
+    modes = modes.reshape(new_shape)
+    return FreqsModes(angular_frequencies, modes)

tests/planetary_motion.py

@@ -117,4 +117,4 @@ def test_is_circle():
 
 
 if __name__ == '__main__':
-    pytest.main([sys.argv[0]])
+    pytest.main([sys.argv[0], '-v'])

tests/rotation3d.py

@@ -54,4 +54,4 @@ def test_rotated_unit_vectors():
 ###############################################################################
 # Now let's run all the tests
 if __name__ == '__main__':
-    pytest.main([sys.argv[0]])
+    pytest.main([sys.argv[0], '-v'])

tests/test_jit.py

@@ -3,16 +3,19 @@
 import sys
 # import nnp.so3 as so3
 import nnp.pbc as pbc
-# import nnp.vib as vib
+import nnp.vib as vib
 
+TORCH_TOO_OLD = torch.__version__ < "1.4.0.dev20191123"
 
+
+@pytest.mark.skipif(TORCH_TOO_OLD, reason="JIT compatibility requires new PyTorch")
 def test_script():
     # torch.jit.script(so3.rotate_along)
     torch.jit.script(pbc.num_repeats)
     torch.jit.script(pbc.map2central)
-    # torch.jit.script(vib.hessian)
-    # torch.jit.script(vib.vibrational_analysis)
+    torch.jit.script(vib.hessian)
+    torch.jit.script(vib.vibrational_analysis)
 
 
 if __name__ == '__main__':
-    pytest.main([sys.argv[0]])
+    pytest.main([sys.argv[0], '-v'])