rename shift to shift_pos

Author: So-Takamoto

tests/functions_tests/test_triplets.py

@@ -13,12 +13,14 @@ def test_calc_triplets():
         dtype=torch.long,
         device=device,
     )
-    shift = torch.zeros((edge_index.shape[1], 3), dtype=torch.float32, device=device)
-    shift[:, 0] = torch.tensor(
+    shift_pos = torch.zeros((edge_index.shape[1], 3), dtype=torch.float32, device=device)
+    shift_pos[:, 0] = torch.tensor(
         [1, 2, 3, 4, 5, 6, -1, -2, -3, -4, -5, -6], dtype=torch.float32, device=device
     )
     # print("shift", shift.shape)
-    triplet_node_index, multiplicity, edge_jk, batch_triplets = calc_triplets(edge_index, shift)
+    triplet_node_index, multiplicity, edge_jk, batch_triplets = calc_triplets(
+        edge_index, shift_pos
+    )
     # print("triplet_node_index", triplet_node_index.shape, triplet_node_index)
     # print("multiplicity", multiplicity.shape, multiplicity)
     # print("triplet_shift", triplet_shift.shape, triplet_shift)
@@ -44,9 +46,9 @@ def test_calc_triplets():
     # shift for edge `i->j`, `i->k`, `j->k`.
     triplet_shift = torch.stack(
         [
-            -shift[edge_jk[:, 0]],
-            -shift[edge_jk[:, 1]],
-            shift[edge_jk[:, 0]] - shift[edge_jk[:, 1]],
+            -shift_pos[edge_jk[:, 0]],
+            -shift_pos[edge_jk[:, 1]],
+            shift_pos[edge_jk[:, 0]] - shift_pos[edge_jk[:, 1]],
         ],
         dim=1,
     )

torch_dftd/functions/dftd3.py

@@ -117,7 +117,7 @@ def edisp(
     cnthr: Optional[float] = None,
     batch: Optional[Tensor] = None,
     batch_edge: Optional[Tensor] = None,
-    shift: Optional[Tensor] = None,
+    shift_pos: Optional[Tensor] = None,
     pos: Optional[Tensor] = None,
     cell: Optional[Tensor] = None,
     r2=None,
@@ -146,7 +146,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 (Tensor or None): (n_atoms,) used to calculate 3-body term when abc=True
+        shift_pos (Tensor or None): (n_atoms,) 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):
@@ -267,7 +267,7 @@ def edisp(
         edge_index_abc = edge_index[:, within_cutoff]
         batch_edge_abc = None if batch_edge is None else batch_edge[within_cutoff]
         # c6_abc = c6[within_cutoff]
-        shift_abc = None if shift is None else shift[within_cutoff]
+        shift_abc = None if shift_pos is None else shift_pos[within_cutoff]
 
         n_atoms = Z.shape[0]
         if not bidirectional:
@@ -290,7 +290,7 @@ def edisp(
             batch_triplets: Optional[Tensor]
             triplet_node_index, multiplicity, edge_jk, batch_triplets = calc_triplets(
                 edge_index_abc,
-                shift=shift_abc,
+                shift_pos=shift_abc,
                 dtype=pos.dtype,
                 batch_edge=batch_edge_abc,
             )

torch_dftd/functions/distance.py

@@ -8,7 +8,7 @@ def calc_distances(
     pos: Tensor,
     edge_index: Tensor,
     cell: Optional[Tensor] = None,
-    shift: Optional[Tensor] = None,
+    shift_pos: Optional[Tensor] = None,
     eps=1e-20,
 ) -> Tensor:
     """Distance calculation function.
@@ -17,7 +17,7 @@ def calc_distances(
         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.
-        shift (Tensor): (n_edges, 3) position shift vectors of edges owing to the periodic boundary. It should be length unit.
+        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.
 
     Returns:
@@ -30,7 +30,7 @@ def calc_distances(
     Ri = pos[idx_i]
     Rj = pos[idx_j]
     if cell is not None:
-        Rj += shift
+        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)
     return Dij

torch_dftd/functions/triplets.py

@@ -7,15 +7,15 @@
 
 def calc_triplets(
     edge_index: Tensor,
-    shift: Optional[Tensor] = None,
+    shift_pos: Optional[Tensor] = None,
     dtype=torch.float32,
     batch_edge: Optional[Tensor] = None,
 ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
     """Calculate triplet edge index.
 
     Args:
         edge_index (Tensor): (2, n_edges) edge_index for graph. It must be bidirectional edge.
-        shift (Tensor or None): (n_edges, 3) used to calculate unique atoms when pbc=True.
+        shift_pos (Tensor or None): (n_edges, 3) used to calculate unique atoms when pbc=True.
         dtype: dtype for `multiplicity`
         batch_edge (Tensor or None): Specify batch indices for `edge_index`.
 
@@ -37,10 +37,10 @@ def calc_triplets(
     src = src[sort_inds]
     dst = dst[sort_inds]
 
-    if shift is None:
+    if shift_pos is None:
         edge_indices = torch.arange(src.shape[0], dtype=torch.long, device=edge_index.device)
     else:
-        edge_indices = torch.arange(shift.shape[0], dtype=torch.long, device=edge_index.device)
+        edge_indices = torch.arange(shift_pos.shape[0], dtype=torch.long, device=edge_index.device)
         edge_indices = edge_indices[is_larger][sort_inds]
 
     if batch_edge is None:

torch_dftd/nn/base_dftd_module.py

@@ -17,7 +17,7 @@ def calc_energy_batch(
         edge_index: Tensor,
         cell: Optional[Tensor] = None,
         pbc: Optional[Tensor] = None,
-        shift: Optional[Tensor] = None,
+        shift_pos: Optional[Tensor] = None,
         batch: Optional[Tensor] = None,
         batch_edge: Optional[Tensor] = None,
         damping: str = "zero",
@@ -32,7 +32,7 @@ def calc_energy_batch(
             edge_index (Tensor): (2, n_edges) edge index within cutoff
             cell (Tensor): (n_atoms, 3) cell size in angstrom, None for non periodic system.
             pbc (Tensor): (bs, 3) pbc condition, None for non periodic system.
-            shift (Tensor): (n_atoms, 3) shift vector
+            shift_pos (Tensor): (n_atoms, 3) shift vector (length unit).
             batch (Tensor): (n_atoms,) Specify which graph this atom belongs to
             batch_edge (Tensor): (n_edges, 3) Specify which graph this edge belongs to
             damping (str):
@@ -49,7 +49,7 @@ def calc_energy(
         edge_index: Tensor,
         cell: Optional[Tensor] = None,
         pbc: Optional[Tensor] = None,
-        shift: Optional[Tensor] = None,
+        shift_pos: Optional[Tensor] = None,
         batch: Optional[Tensor] = None,
         batch_edge: Optional[Tensor] = None,
         damping: str = "zero",
@@ -64,6 +64,7 @@ def calc_energy(
             edge_index (Tensor):
             cell (Tensor): cell size in angstrom, None for non periodic system.
             pbc (Tensor): pbc condition, None for non periodic system.
+            shift_pos (Tensor):  (n_atoms, 3) shift vector (length unit).
             batch (Tensor):
             batch_edge (Tensor):
             damping (str): damping method. "zero", "bj", "zerom", "bjm"
@@ -73,7 +74,7 @@ def calc_energy(
         """
         with torch.no_grad():
             E_disp = self.calc_energy_batch(
-                Z, pos, edge_index, cell, pbc, shift, batch, batch_edge, damping=damping
+                Z, pos, edge_index, cell, pbc, shift_pos, batch, batch_edge, damping=damping
             )
         if batch is None:
             return [{"energy": E_disp.item()}]
@@ -91,7 +92,7 @@ def calc_energy_and_forces(
         edge_index: Tensor,
         cell: Optional[Tensor] = None,
         pbc: Optional[Tensor] = None,
-        shift: Optional[Tensor] = None,
+        shift_pos: Optional[Tensor] = None,
         batch: Optional[Tensor] = None,
         batch_edge: Optional[Tensor] = None,
         damping: str = "zero",
@@ -103,6 +104,7 @@ def calc_energy_and_forces(
             pos (Tensor): atom positions in angstrom
             cell (Tensor): cell size in angstrom, None for non periodic system.
             pbc (Tensor): pbc condition, None for non periodic system.
+            shift_pos (Tensor):  (n_atoms, 3) shift vector (length unit).
             damping (str): damping method. "zero", "bj", "zerom", "bjm"
 
         Returns:
@@ -117,11 +119,11 @@ def calc_energy_and_forces(
             # We need to explicitly include this dependency to calculate cell gradient
             # for stress computation.
             # pos is assumed to be inside "cell", so relative position `rel_pos` lies between 0~1.
-            assert isinstance(shift, Tensor)
-            shift.requires_grad_(True)
+            assert isinstance(shift_pos, Tensor)
+            shift_pos.requires_grad_(True)
 
         E_disp = self.calc_energy_batch(
-            Z, pos, edge_index, cell, pbc, shift, batch, batch_edge, damping=damping
+            Z, pos, edge_index, cell, pbc, shift_pos, batch, batch_edge, damping=damping
         )
 
         E_disp.sum().backward()
@@ -140,7 +142,7 @@ def calc_energy_and_forces(
         if cell is not None:
             # stress = torch.mm(cell_grad, cell.T) / cell_volume
             # Get stress in Voigt notation (xx, yy, zz, yz, xz, xy)
-            assert isinstance(shift, Tensor)
+            assert isinstance(shift_pos, Tensor)
             voigt_left = [0, 1, 2, 1, 2, 0]
             voigt_right = [0, 1, 2, 2, 0, 1]
             if batch is None:
@@ -149,7 +151,8 @@ def calc_energy_and_forces(
                     (pos[:, voigt_left] * pos.grad[:, voigt_right]).to(torch.float64), dim=0
                 )
                 cell_grad += torch.sum(
-                    (shift[:, voigt_left] * shift.grad[:, voigt_right]).to(torch.float64), dim=0
+                    (shift_pos[:, voigt_left] * shift_pos.grad[:, voigt_right]).to(torch.float64),
+                    dim=0,
                 )
                 stress = cell_grad.to(cell.dtype) / cell_volume
                 results_list[0]["stress"] = stress.detach().cpu().numpy()
@@ -166,7 +169,7 @@ def calc_energy_and_forces(
                 cell_grad.scatter_add_(
                     0,
                     batch_edge.view(batch_edge.size()[0], 1).expand(batch_edge.size()[0], 6),
-                    (shift[:, voigt_left] * shift.grad[:, voigt_right]).to(torch.float64),
+                    (shift_pos[:, voigt_left] * shift_pos.grad[:, voigt_right]).to(torch.float64),
                 )
                 stress = cell_grad.to(cell.dtype) / cell_volume[:, None]
                 stress = stress.detach().cpu().numpy()

torch_dftd/nn/dftd2_module.py

@@ -48,19 +48,19 @@ def calc_energy_batch(
         edge_index: Tensor,
         cell: Optional[Tensor] = None,
         pbc: Optional[Tensor] = None,
-        shift: Optional[Tensor] = None,
+        shift_pos: Optional[Tensor] = None,
         batch: Optional[Tensor] = None,
         batch_edge: Optional[Tensor] = None,
         damping: str = "zero",
     ) -> Tensor:
         """Forward computation to calculate atomic wise dispersion energy"""
-        shift = pos.new_zeros((edge_index.size()[1], 3, 3)) if shift is None else shift
+        shift_pos = pos.new_zeros((edge_index.size()[1], 3, 3)) if shift_pos is None else shift_pos
         pos_bohr = pos / d3_autoang  # angstrom -> bohr
         if cell is None:
             cell_bohr: Optional[Tensor] = None
         else:
             cell_bohr = cell / d3_autoang  # angstrom -> bohr
-        shift_bohr = shift / d3_autoang  # angstrom -> bohr
+        shift_bohr = shift_pos / d3_autoang  # angstrom -> bohr
         r = calc_distances(pos_bohr, edge_index, cell_bohr, shift_bohr)
 
         # E_disp (n_graphs,): Energy in eV unit

torch_dftd/nn/dftd3_module.py

@@ -70,19 +70,19 @@ def calc_energy_batch(
         edge_index: Tensor,
         cell: Optional[Tensor] = None,
         pbc: Optional[Tensor] = None,
-        shift: Optional[Tensor] = None,
+        shift_pos: Optional[Tensor] = None,
         batch: Optional[Tensor] = None,
         batch_edge: Optional[Tensor] = None,
         damping: str = "zero",
     ) -> Tensor:
         """Forward computation to calculate atomic wise dispersion energy"""
-        shift = pos.new_zeros((edge_index.size()[1], 3, 3)) if shift is None else shift
+        shift_pos = pos.new_zeros((edge_index.size()[1], 3, 3)) if shift_pos is None else shift_pos
         pos_bohr = pos / d3_autoang  # angstrom -> bohr
         if cell is None:
             cell_bohr: Optional[Tensor] = None
         else:
             cell_bohr = cell / d3_autoang  # angstrom -> bohr
-        shift_bohr = shift / d3_autoang  # angstrom -> bohr
+        shift_bohr = shift_pos / d3_autoang  # angstrom -> bohr
         r = calc_distances(pos_bohr, edge_index, cell_bohr, shift_bohr)
         # E_disp (n_graphs,): Energy in eV unit
         E_disp = d3_autoev * edisp(
@@ -98,7 +98,7 @@ def calc_energy_batch(
             cnthr=self.cnthr / Bohr,
             batch=batch,
             batch_edge=batch_edge,
-            shift=shift_bohr,
+            shift_pos=shift_bohr,
             damping=damping,
             cutoff_smoothing=self.cutoff_smoothing,
             bidirectional=self.bidirectional,

torch_dftd/torch_dftd3_calculator.py

@@ -103,10 +103,12 @@ def _preprocess_atoms(self, atoms: Atoms) -> Dict[str, Optional[Tensor]]:
         pbc = torch.tensor(atoms.pbc, device=self.device)
         edge_index, S = self._calc_edge_index(pos, cell, pbc)
         if cell is None:
-            shift = S
+            shift_pos = S
         else:
-            shift = torch.mm(S, cell.detach())
-        input_dicts = dict(pos=pos, Z=Z, cell=cell, pbc=pbc, edge_index=edge_index, shift=shift)
+            shift_pos = torch.mm(S, cell.detach())
+        input_dicts = dict(
+            pos=pos, Z=Z, cell=cell, pbc=pbc, edge_index=edge_index, shift_pos=shift_pos
+        )
         return input_dicts
 
     def calculate(self, atoms=None, properties=["energy"], system_changes=all_changes):
@@ -155,7 +157,6 @@ def batch_calculate(self, atoms_list=None, properties=["energy"], system_changes
         # Calculator.calculate(self, atoms, properties, system_changes)
         input_dicts_list = [self._preprocess_atoms(atoms) for atoms in atoms_list]
         # --- Make batch ---
-        # pos=pos, Z=Z, cell=cell, pbc=pbc, edge_index=edge_index, shift=S
         n_nodes_list = [d["Z"].shape[0] for d in input_dicts_list]
         shift_index_array = torch.cumsum(torch.tensor([0] + n_nodes_list), dim=0)
         cell_batch = torch.stack(
@@ -171,7 +172,7 @@ def batch_calculate(self, atoms_list=None, properties=["energy"], system_changes
             pos=torch.cat([d["pos"] for d in input_dicts_list], dim=0),  # (n_nodes,)
             cell=cell_batch,  # (bs, 3, 3)
             pbc=torch.stack([d["pbc"] for d in input_dicts_list]),  # (bs, 3)
-            shift=torch.cat([d["shift"] for d in input_dicts_list], dim=0),  # (n_nodes,)
+            shift_pos=torch.cat([d["shift_pos"] for d in input_dicts_list], dim=0),  # (n_nodes,)
         )
 
         batch_dicts["edge_index"] = torch.cat(