more

more

more

more

more

more

more

more

more

more

test

fix

more

more

docs/api.rst

@@ -1,6 +1,8 @@
-.. automodule:: nnp.pbc
+.. automodule:: nnp.so3
     :members:
 .. automodule:: nnp.md
     :members:
+.. automodule:: nnp.pbc
+    :members:
 .. automodule:: nnp.vib
     :members:

docs/conf.py

@@ -32,7 +32,8 @@
 sphinx_gallery_conf = {
     'examples_dirs': ['../tests/', '../nnp/'],
     'gallery_dirs': ['examples', 'code'],
-    'filename_pattern': r'.*\.py'
+    'filename_pattern': r'.*\.py',
+    'ignore_pattern': r'__init__\.py|test_.*',
 }
 
 intersphinx_mapping = {

docs/index.rst

@@ -4,7 +4,7 @@ Welcome to NNP's documentation!
 
 .. toctree::
     :maxdepth: 2
-    :caption: Getting Started
+    :caption: Introduction
 
     start
 

docs/start.rst

@@ -1,8 +1,25 @@
-Installation
+Introduction
 ============
 
+NNP is a python package that provide common tools used wo work with
+neural network potentials with PyTorch. Currently it includes tools
+to deal with periodic boundary condition, rotation in 3D space, analytical
+hessian, vibrational analysis, and molecular dynamics.
+
 To install NNP, just do
 
 .. code-block:: bash
 
     pip install nnp
+
+
+TODO:
+- [x] so3
+- [x] test so3
+- [ ] pbc
+- [ ] test pbc
+- [ ] md
+- [x] test md: autograd
+- [ ] test md: stress
+- [ ] vib
+- [ ] test vib

nnp/md.py

@@ -2,15 +2,15 @@
 Molecular Dynamics
 ==================
 
-The module `nnp.md` provide tools to run molecular dynamics with a potential
+The module ``nnp.md`` provide tools to run molecular dynamics with a potential
 defined by PyTorch.
 """
 
 import torch
 from torch import Tensor
-import ase.calculators.calculator
 from nnp import pbc
 from typing import Callable, Sequence
+import ase.calculators.calculator
 
 
 class Calculator(ase.calculators.calculator.Calculator):

nnp/pbc.py

@@ -2,32 +2,62 @@
 Periodic Boundary Conditions
 ============================
 
-This file contains
+The module ``nnp.pbc`` contains tools to deal with periodic boundary conditions.
 """
-
+###############################################################################
+# Let's first import all the packages we will use:
 import torch
 from torch import Tensor
 
 
+###############################################################################
+# The following function computes the number of repeats required to capture
+# all the neighbor atoms within the cutoff radius.
+def num_repeats(cell: Tensor, pbc: Tensor, cutoff: float) -> Tensor:
+    """Compute the number of repeats required along each cell vector to make
+    the original cell and repeated cells together form a large enough box
+    so that for each atom in the original cell, all its neighbor atoms within
+    the given cutoff distance are contained in the big box.
+
+    Arguments:
+        cell: tensor of shape ``(3, 3)`` of the three vectors defining unit cell:
+
+            .. code-block:: python
+
+                tensor([[x1, y1, z1],
+                        [x2, y2, z2],
+                        [x3, y3, z3]])
+
+        pbc: boolean vector of size 3 storing if pbc is enabled for that direction.
+        cutoff: the cutoff inside which atoms are considered as pairs
+
+    Returns:
+        numbers of repeats required to make the repeated box contains all the neighbors
+        of atoms in the ceter cell
+    """
+    reciprocal_cell = cell.inverse().t()
+    inv_distances = reciprocal_cell.norm(2, -1)
+    result = torch.ceil(cutoff * inv_distances).to(torch.long)
+    return torch.where(pbc, result, torch.tensor(0))
+
+
 def map2central(cell: Tensor, coordinates: Tensor, pbc: Tensor) -> Tensor:
     """Map atoms outside the unit cell into the cell using PBC.
 
     Arguments:
-        cell (:class:`torch.Tensor`): tensor of shape (3, 3) of the three
-            vectors defining unit cell:
+        cell: tensor of shape ``(3, 3)`` of the three vectors defining unit cell:
 
             .. code-block:: python
 
                 tensor([[x1, y1, z1],
                         [x2, y2, z2],
                         [x3, y3, z3]])
 
-        coordinates (:class:`torch.Tensor`): Tensor of shape ``(atoms, 3)``
-            or ``(molecules, atoms, 3)``.
-        pbc (:class:`torch.Tensor`): boolean vector of size 3 storing
-            if pbc is enabled for that direction.
+        coordinates: Tensor of shape ``(atoms, 3)`` or ``(molecules, atoms, 3)``.
+        pbc: boolean vector of size 3 storing if pbc is enabled for that direction.
+
     Returns:
-        :class:`torch.Tensor`: coordinates of atoms mapped back to unit cell.
+        coordinates of atoms mapped back to unit cell.
     """
     # Step 1: convert coordinates from standard cartesian coordinate to unit
     # cell coordinates

nnp/so3.py

@@ -0,0 +1,86 @@
+"""
+SO(3) Lie Group Operations
+==========================
+
+The module ``nnp.so3`` contains tools to rotate point clouds in 3D space.
+"""
+###############################################################################
+# Let's first import all the packages we will use:
+import torch
+from scipy.linalg import expm as scipy_expm
+from torch import Tensor
+
+
+###############################################################################
+# The following function implements rotation along an axis passing the origin.
+# Rotation in 3D is not as trivial as in 2D. Here we start from the equation of
+# motion. Let :math:`\vec{n}` be the unit vector pointing to
+# direction of the axis, then for an infinitesimal rotation, we have
+#
+# .. math::
+#   \frac{\mathrm{d}\vec{r}}{\mathrm{d}\theta}
+#       = \vec{n} \times \vec{r}
+#
+# That is
+#
+# .. math::
+#   \frac{\mathrm{d}r_i}{\mathrm{d}\theta} = \epsilon_{ijk} n_j r_k
+#
+# where :math:`\epsilon_{ijk}` is the Levi-Civita symbol, let
+# :math:`W_{ik}=\epsilon_{ijk} n_j`, then the above equation becomes
+# a matrix equation:
+#
+# .. math::
+#   \frac{\mathrm{d}\vec{r}}{\mathrm{d}\theta} = W \cdot \vec{r}
+#
+# It is not hard to see that :math:`W` is a skew-symmetric matrix.
+# From the above equation and the knowledge of linear algebra,
+# matrix Lie algebra/group, it is not hard to see that the set of
+# all rotation operations along the axis :math:`\vec{n}` is a one
+# parameter Lie group. And the skew-symmetric matrices together
+# with standard matrix commutator is a Lie algebra.This Lie group
+# and Lie algebra is connected by the exponential map. See Wikipedia
+# `Exponential map (Lie theory)`_ for more detail.
+#
+# .. _Exponential map (Lie theory):
+#   https://en.wikipedia.org/wiki/Exponential_map_(Lie_theory)
+#
+# So it is easy to tell that:
+#
+# .. math::
+#   \vec{r}\left(\theta\right) = \exp \left(\theta W\right) \cdot
+#       \vec{r}\left(0\right)
+#
+# where :math:`\vec{r}\left(0\right)` is the initial coordinates,
+# and :math:`\vec{r}\left(\theta\right)` is the final coordinates
+# after rotating :math:`\theta`.
+#
+# To implement, let's first define the Levi-Civita symbol:
+levi_civita = torch.zeros(3, 3, 3)
+levi_civita[0, 1, 2] = levi_civita[1, 2, 0] = levi_civita[2, 0, 1] = 1
+levi_civita[0, 2, 1] = levi_civita[2, 1, 0] = levi_civita[1, 0, 2] = -1
+
+
+###############################################################################
+# PyTorch does not have matrix exp, let's implement it here using scipy
+def expm(matrix: Tensor) -> Tensor:
+    # TODO: remove this part when pytorch support matrix_exp
+    ndarray = matrix.detach().cpu().numpy()
+    return torch.from_numpy(scipy_expm(ndarray)).to(matrix)
+
+
+###############################################################################
+# Now we are ready to implement the :math:`\exp \left(\theta W\right)`
+def rotate_along(axis: Tensor) -> Tensor:
+    r"""Compute group elements of rotating along an axis passing origin.
+
+    Arguments:
+        axis: a vector (x, y, z) whose direction specifies the axis of the rotation,
+            length specifies the radius to rotate, and sign specifies clockwise
+            or anti-clockwise.
+
+    Return:
+        the rotational matrix :math:`\exp{\left(\theta W\right)}`.
+    """
+    W = torch.einsum('ijk,j->ik', levi_civita.to(axis), axis)
+    return expm(W)

nnp/vib.py

@@ -14,29 +14,30 @@ def hessian(coordinates: Tensor, energies: Optional[Tensor] = None,
     """Compute analytical hessian from the energy graph or force graph.
 
     Arguments:
-        coordinates (:class:`torch.Tensor`): Tensor of shape `(molecules, atoms, 3)` or
-            ``(atoms, 3)``.
-        energies (:class:`torch.Tensor`): Tensor of shape `(molecules,)`, if specified,
-            then `forces` must be `None`. This energies must be computed from
-            `coordinates` in a graph.
-        forces (:class:`torch.Tensor`): Tensor of shape `(molecules, atoms, 3)`, if specified,
+        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.
     Returns:
-        :class:`torch.Tensor`: Tensor of shape `(molecules, 3A, 3A)` where A is the number of
-        atoms in each molecule
+        Tensor of shape `(molecules, 3A, 3A)` where A is the number of atoms in
+        each molecule
     """
     if energies is None and forces is None:
         raise ValueError()
     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]
+        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)
     return -torch.stack([
-        torch.autograd.grad(f.sum(), coordinates, retain_graph=True)[0].flatten(start_dim=1)
+        torch.autograd.grad(
+            [f.sum()], [coordinates], retain_graph=True)[0].flatten(start_dim=1)
         for f in force_components
     ], dim=1)
 
@@ -48,7 +49,6 @@ class FreqsModes(NamedTuple):
 
 def vibrational_analysis(masses: Tensor, hessian: Tensor) -> FreqsModes:
     """Computing the vibrational wavenumbers from hessian."""
-    assert hessian.shape[0] == 1, 'Currently only supporting computing one molecule a time'
     # 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
@@ -58,8 +58,6 @@ def vibrational_analysis(masses: Tensor, hessian: Tensor) -> FreqsModes:
     # 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)
-    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, eigenvectors = torch.symeig(mass_scaled_hessian, eigenvectors=True)
     angular_frequencies = eigenvalues.sqrt()

optional_requirements.txt

@@ -1,2 +1,3 @@
 ase
 pytest
+scipy

tests/planetary_motion.py

@@ -22,6 +22,7 @@
 from ase.md.verlet import VelocityVerlet
 import nnp.md as md
 import pytest
+import sys
 import math
 import matplotlib.pyplot as plt
 
@@ -116,4 +117,4 @@ def test_is_circle():
 
 
 if __name__ == '__main__':
-    pytest.main()
+    pytest.main([sys.argv[0]])

tests/rotation3d.py

@@ -0,0 +1,57 @@
+"""
+Rotation in 3D Space
+====================
+
+This tutorial demonstrates how rotate a group of points in 3D space.
+"""
+###############################################################################
+# To begin with, we must understand that rotation is a linear transformation.
+# The set of all rotations forms the group SO(3). It is a Lie group that each
+# group element is described by an orthogonal 3x3 matrix.
+#
+# That is, if you have two points :math:`\vec{r}_1` and :math:`\vec{r}_2`, and
+# you want to rotate the two points along the same axis for the same number of
+# degrees, then there is a single orthogonal matrix :math:`R` that no matter
+# the value of :math:`\vec{r}_1` and :math:`\vec{r}_2`, their rotation is always
+# :math:`R\cdot\vec{r}_1` and :math:`R\cdot\vec{r}_2`.
+#
+# Let's import libraries first
+import math
+import torch
+import pytest
+import sys
+import nnp.so3 as so3
+
+###############################################################################
+# Let's first take a look at a special case: rotating the three unit vectors
+# ``(1, 0, 0)``, ``(0, 1, 0)`` and ``(0, 0, 1)`` along the diagonal for 120
+# degree, this should permute these points:
+rx = torch.tensor([1, 0, 0])
+ry = torch.tensor([0, 1, 0])
+rz = torch.tensor([0, 0, 1])
+points = torch.stack([rx, ry, rz])
+
+###############################################################################
+# Now let's compute the rotation matrix
+axis = torch.ones(3) / math.sqrt(3) * (math.pi * 2 / 3)
+R = so3.rotate_along(axis)
+
+###############################################################################
+# Note that we need to do matrix-vector product for all these unit vectors
+# so we need to do a transpose in order to use ``@`` operator.
+rotated = (R @ points.float().t()).t()
+print(rotated)
+
+
+###############################################################################
+# After this rotation, the three vector permutes: rx->ry, ry->rz, rz->rx. Let's
+# programmically check that. This check will be run by pytest later.
+def test_rotated_unit_vectors():
+    expected = torch.stack([ry, rz, rx]).float()
+    assert torch.allclose(rotated, expected, atol=1e-5)
+
+
+###############################################################################
+# Now let's run all the tests
+if __name__ == '__main__':
+    pytest.main([sys.argv[0]])

tests/test_jit.py

@@ -0,0 +1,18 @@
+import torch
+import pytest
+import sys
+# import nnp.so3 as so3
+import nnp.pbc as pbc
+# import nnp.vib as vib
+
+
+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)
+
+
+if __name__ == '__main__':
+    pytest.main([sys.argv[0]])