Support torch.jit.script

.flexci/build_and_push.sh

@@ -32,10 +32,10 @@ docker_build_and_push() {
 WAIT_PIDS=""
 
 # PyTorch 1.5 + Python 3.6
-docker_build_and_push torch15 \
+docker_build_and_push torch19 \
     --build-arg base_image="nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04" \
-    --build-arg python_version="3.6.12" \
-    --build-arg pip_packages="torch==1.5.* torchvision==0.6.* ${TEST_PIP_PACKAGES}" &
+    --build-arg python_version="3.7.9" \
+    --build-arg pip_packages="torch==1.9.1 ${TEST_PIP_PACKAGES}" &
 WAIT_PIDS="$! ${WAIT_PIDS}"
 
 # Wait until the build complete.

.flexci/pytest_script.sh

@@ -2,7 +2,7 @@
 set -eu
 
 #IMAGE=pytorch/pytorch:1.5.1-cuda10.1-cudnn7-devel
-IMAGE=asia.gcr.io/pfn-public-ci/torch-dftd-ci:torch15
+IMAGE=asia.gcr.io/pfn-public-ci/torch-dftd-ci:torch19
 
 
 main() {
@@ -20,7 +20,7 @@ main() {
   docker run --runtime=nvidia --rm --volume="$(pwd)":/workspace -w /workspace \
     ${IMAGE} \
     bash -x -c "pip install flake8 pytest pytest-cov pytest-xdist pytest-benchmark && \
-      pip install cupy-cuda102 pytorch-pfn-extras!=0.5.0 && \
+      pip install cupy-cuda102 pytorch-pfn-extras==0.4.2 && \
       pip install -e .[develop] && \
       pysen run lint && \
       pytest --cov=torch_dftd -n $(nproc) -m 'not slow' tests &&

README.md

@@ -34,14 +34,14 @@ print(f"forces {forces}")
 ## Dependency
 
 The library is tested under following environment.
- - python: 3.6
+ - python: 3.7
  - CUDA: 10.2
 ```bash
-torch==1.5.1
+torch==1.9.1
 ase==3.21.1
 # Below is only for 3-body term
-cupy-cuda102==8.6.0
-pytorch-pfn-extras==0.3.2
+cupy-cuda102==9.5.0
+pytorch-pfn-extras==0.4.2
 ```
 
 ## Development tips

torch_dftd/dftd3_xc_params.py

@@ -535,7 +535,7 @@ def get_dftd3_default_params(
         rs6 = 1.1
         s18 = 0.0
         alp = 20.0
-        rs18 = None  # Not used.
+        rs18 = 0.0  # It is DUMMY value. Not used.
         if xc == "b-lyp":
             s6 = 1.2
         elif xc == "b-p":

torch_dftd/functions/dftd2.py

@@ -6,6 +6,7 @@
 from torch_dftd.functions.smoothing import poly_smoothing
 
 
+@torch.jit.script
 def edisp_d2(
     Z: Tensor,
     r: Tensor,
@@ -44,7 +45,7 @@ def edisp_d2(
     r2 = r ** 2
     r6 = r2 ** 3
 
-    idx_i, idx_j = edge_index
+    idx_i, idx_j = edge_index[0], edge_index[1]
     # compute all necessary quantities
     Zi = Z[idx_i]  # (n_edges,)
     Zj = Z[idx_j]
@@ -71,6 +72,7 @@ def edisp_d2(
         # (1,)
         g = e6.sum()[None]
     else:
+        assert batch is not None
         # (n_graphs,)
         if batch.size()[0] == 0:
             n_graphs = 1

torch_dftd/functions/dftd3.py

@@ -9,15 +9,16 @@
 
 # conversion factors used in grimme d3 code
 
-d3_autoang = 0.52917726  # for converting distance from bohr to angstrom
-d3_autoev = 27.21138505  # for converting a.u. to eV
+d3_autoang: float = 0.52917726  # for converting distance from bohr to angstrom
+d3_autoev: float = 27.21138505  # for converting a.u. to eV
 
 d3_k1 = 16.000
 d3_k2 = 4 / 3
 d3_k3 = -4.000
 d3_maxc = 5  # maximum number of coordination complexes
 
 
+@torch.jit.script
 def _ncoord(
     Z: Tensor,
     r: Tensor,
@@ -53,7 +54,7 @@ def _ncoord(
     Zi = Z[idx_i]
     Zj = Z[idx_j]
     rco = rcov[Zi] + rcov[Zj]  # (n_edges,)
-    rr = rco.type(r.dtype) / r
+    rr = rco.to(r.dtype) / r
     damp = 1.0 / (1.0 + torch.exp(-k1 * (rr - 1.0)))
     if cutoff is not None and cutoff_smoothing == "poly":
         damp *= poly_smoothing(r, cutoff)
@@ -66,6 +67,7 @@ def _ncoord(
     return g  # (n_atoms,)
 
 
+@torch.jit.script
 def _getc6(
     Zi: Tensor, Zj: Tensor, nci: Tensor, ncj: Tensor, c6ab: Tensor, k3: float = d3_k3
 ) -> Tensor:
@@ -84,7 +86,7 @@ def _getc6(
     """
     # gather the relevant entries from the table
     # c6ab (95, 95, 5, 5, 3) --> c6ab_ (n_edges, 5, 5, 3)
-    c6ab_ = c6ab[Zi, Zj].type(nci.dtype)
+    c6ab_ = c6ab[Zi, Zj].to(nci.dtype)
     # calculate c6 coefficients
 
     # cn0, cn1, cn2 (n_edges, 5, 5)
@@ -104,6 +106,7 @@ def _getc6(
     return c6
 
 
+@torch.jit.script
 def edisp(
     Z: Tensor,
     r: Tensor,
@@ -120,12 +123,12 @@ def edisp(
     shift_pos: Optional[Tensor] = None,
     pos: Optional[Tensor] = None,
     cell: Optional[Tensor] = None,
-    r2=None,
-    r6=None,
-    r8=None,
-    k1=d3_k1,
-    k2=d3_k2,
-    k3=d3_k3,
+    r2: Optional[Tensor] = None,
+    r6: Optional[Tensor] = None,
+    r8: Optional[Tensor] = None,
+    k1: float = d3_k1,
+    k2: float = d3_k2,
+    k3: float = d3_k3,
     cutoff_smoothing: str = "none",
     damping: str = "zero",
     bidirectional: bool = False,
@@ -146,7 +149,7 @@ def edisp(
         cnthr (float or None): cutoff distance for coordination number calculation in **bohr**
         batch (Tensor or None): (n_atoms,)
         batch_edge (Tensor or None): (n_edges,)
-        shift_pos (Tensor or None): (n_atoms,) used to calculate 3-body term when abc=True
+        shift_pos (Tensor or None): (n_edges, 3) used to calculate 3-body term when abc=True
         pos (Tensor): (n_atoms, 3) position in **bohr**
         cell (Tensor): (3, 3) cell size in **bohr**
         r2 (Tensor or None):
@@ -171,7 +174,7 @@ def edisp(
     if r8 is None:
         r8 = r6 * r2
 
-    idx_i, idx_j = edge_index
+    idx_i, idx_j = edge_index[0], edge_index[1]
     # compute all necessary quantities
     Zi = Z[idx_i]  # (n_edges,)
     Zj = Z[idx_j]
@@ -192,7 +195,7 @@ def edisp(
     ncj = nc[idx_j]
     c6 = _getc6(Zi, Zj, nci, ncj, c6ab=c6ab, k3=k3)  # c6 coefficients
 
-    c8 = 3 * c6 * r2r4[Zi].type(c6.dtype) * r2r4[Zj].type(c6.dtype)  # c8 coefficient
+    c8 = 3 * c6 * r2r4[Zi].to(c6.dtype) * r2r4[Zj].to(c6.dtype)  # c8 coefficient
 
     s6 = params["s6"]
     s8 = params["s18"]
@@ -250,6 +253,7 @@ def edisp(
         g = e68.to(torch.float64).sum()[None]
     else:
         # (n_graphs,)
+        assert batch is not None
         if batch.size()[0] == 0:
             n_graphs = 1
         else:
@@ -261,6 +265,7 @@ def edisp(
         g *= 2.0
 
     if abc:
+        assert cnthr is not None
         within_cutoff = r <= cnthr
         # r_abc = r[within_cutoff]
         # r2_abc = r2[within_cutoff]
@@ -282,12 +287,7 @@ def edisp(
             # (n_edges, ) -> (n_edges * 2, )
             shift_abc = None if shift_abc is None else torch.cat([shift_abc, -shift_abc], dim=0)
         with torch.no_grad():
-            # triplet_node_index, triplet_edge_index = calc_triplets_cycle(edge_index_abc, n_atoms, shift=shift_abc)
-            # Type hinting
-            triplet_node_index: Tensor
-            multiplicity: Tensor
-            edge_jk: Tensor
-            batch_triplets: Optional[Tensor]
+            assert pos is not None
             triplet_node_index, multiplicity, edge_jk, batch_triplets = calc_triplets(
                 edge_index_abc,
                 shift_pos=shift_abc,
@@ -303,7 +303,6 @@ def edisp(
         )
         r_jk = calc_distances(pos, torch.stack([idx_j, idx_k], dim=0), cell, shift_jk)
         kj_within_cutoff = r_jk <= cnthr
-        del shift_jk
 
         triplet_node_index = triplet_node_index[kj_within_cutoff]
         multiplicity, edge_jk, batch_triplets = (
@@ -355,5 +354,6 @@ def edisp(
             e6abc = e3.to(torch.float64).sum()
             g += e6abc
         else:
+            assert batch_triplets is not None
             g.scatter_add_(0, batch_triplets, e3.to(torch.float64))
     return g  # (n_graphs,)

torch_dftd/functions/distance.py

@@ -4,19 +4,21 @@
 from torch import Tensor
 
 
+@torch.jit.script
 def calc_distances(
     pos: Tensor,
     edge_index: Tensor,
     cell: Optional[Tensor] = None,
     shift_pos: Optional[Tensor] = None,
-    eps=1e-20,
+    eps: float = 1e-20,
 ) -> Tensor:
     """Distance calculation function.
 
     Args:
         pos (Tensor): (n_atoms, 3) atom positions.
         edge_index (Tensor): (2, n_edges) edge_index for graph.
         cell (Tensor): cell size, None for non periodic system.
+            This it NOT USED now, it is left for backward compatibility.
         shift_pos (Tensor): (n_edges, 3) position shift vectors of edges owing to the periodic boundary. It should be length unit.
         eps (float): Small float value to avoid NaN in backward when the distance is 0.
 
@@ -25,11 +27,11 @@ def calc_distances(
 
     """
 
-    idx_i, idx_j = edge_index
+    idx_i, idx_j = edge_index[0], edge_index[1]
     # calculate interatomic distances
     Ri = pos[idx_i]
     Rj = pos[idx_j]
-    if cell is not None:
+    if shift_pos is not None:
         Rj += shift_pos
     # eps is to avoid Nan in backward when Dij = 0 with sqrt.
     Dij = torch.sqrt(torch.sum((Ri - Rj) ** 2, dim=-1) + eps)

torch_dftd/functions/smoothing.py

@@ -2,12 +2,13 @@
 from torch import Tensor
 
 
+@torch.jit.script
 def poly_smoothing(r: Tensor, cutoff: float) -> Tensor:
     """Computes a smooth step from 1 to 0 starting at 1 bohr before the cutoff
 
     Args:
         r (Tensor): (n_edges,)
-        cutoff (float): ()
+        cutoff (float): cutoff length
 
     Returns:
         r (Tensor): Smoothed `r`