Merge pull request #9 from So-Takamoto/batch_abc

Bugfix: batch calculation with abc=True

Author: corochann

.flexci/config.pbtxt

@@ -8,6 +8,9 @@ configs {
       disk: 10
       gpu: 1
     }
+    time_limit {
+      seconds: 1800
+    }
     command:
         "bash -x .flexci/pytest_script.sh"
   }

tests/test_torch_dftd3_calculator.py

@@ -72,7 +72,15 @@ def _assert_energy_force_stress_equal(calc1, calc2, atoms: Atoms):
 
 
 def _test_calc_energy_force_stress(
-    damping, xc, old, atoms, device="cpu", dtype=torch.float64, abc=False, cnthr=15.0
+    damping,
+    xc,
+    old,
+    atoms,
+    device="cpu",
+    dtype=torch.float64,
+    bidirectional=True,
+    abc=False,
+    cnthr=15.0,
 ):
     cutoff = 22.0  # Make test faster
     with tempfile.TemporaryDirectory() as tmpdirname:
@@ -95,6 +103,7 @@ def _test_calc_energy_force_stress(
             cutoff=cutoff,
             cnthr=cnthr,
             abc=abc,
+            bidirectional=bidirectional,
         )
         _assert_energy_force_stress_equal(dftd3_calc, torch_dftd3_calc, atoms)
 
@@ -201,13 +210,39 @@ def test_calc_energy_force_stress_with_dft():
 @pytest.mark.parametrize("atoms", _create_atoms())
 @pytest.mark.parametrize("device", ["cpu", "cuda:0"])
 @pytest.mark.parametrize("dtype", [torch.float64])
+@pytest.mark.parametrize("bidirectional", [True, False])
 @pytest.mark.parametrize("abc", [True])
-def test_calc_energy_force_stress_device_abc(damping, old, atoms, device, dtype, abc):
+def test_calc_energy_force_stress_device_abc(
+    damping, old, atoms, device, dtype, bidirectional, abc
+):
     """Test: check tri-partite calc with device, dtype dependency."""
     xc = "pbe"
-    _test_calc_energy_force_stress(
-        damping, xc, old, atoms, device=device, dtype=dtype, abc=abc, cnthr=7.0
-    )
+    if np.all(atoms.pbc) and bidirectional == False:
+        # TODO: bidirectional=False is not implemented for pbc now.
+        with pytest.raises(NotImplementedError):
+            _test_calc_energy_force_stress(
+                damping,
+                xc,
+                old,
+                atoms,
+                device=device,
+                dtype=dtype,
+                bidirectional=bidirectional,
+                abc=abc,
+                cnthr=7.0,
+            )
+    else:
+        _test_calc_energy_force_stress(
+            damping,
+            xc,
+            old,
+            atoms,
+            device=device,
+            dtype=dtype,
+            bidirectional=bidirectional,
+            abc=abc,
+            cnthr=7.0,
+        )
 
 
 if __name__ == "__main__":

tests/test_torch_dftd3_calculator_batch.py

@@ -13,13 +13,18 @@
 from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
 
 
-def _create_atoms() -> List[Atoms]:
+def _create_atoms() -> List[List[Atoms]]:
     """Initialization"""
     atoms = molecule("CH3CH2OCH3")
 
     slab = fcc111("Au", size=(2, 1, 3), vacuum=80.0)
     slab.pbc = np.array([True, True, True])
-    return [atoms, slab]
+
+    slab_wo_pbc = slab.copy()
+    slab_wo_pbc.pbc = np.array([False, False, False])
+
+    null = Atoms()
+    return [[atoms, slab], [atoms, slab_wo_pbc], [null]]
 
 
 def _assert_energy_equal_batch(calc1, atoms_list: List[Atoms]):
@@ -60,7 +65,15 @@ def _assert_energy_force_stress_equal_batch(calc1, atoms_list: List[Atoms]):
 
 
 def _test_calc_energy_force_stress(
-    damping, xc, old, atoms_list, device="cpu", dtype=torch.float64, abc=False, cnthr=15.0
+    damping,
+    xc,
+    old,
+    atoms_list,
+    device="cpu",
+    dtype=torch.float64,
+    bidirectional=True,
+    abc=False,
+    cnthr=15.0,
 ):
     cutoff = 22.0  # Make test faster
     torch_dftd3_calc = TorchDFTD3Calculator(
@@ -72,41 +85,68 @@ def _test_calc_energy_force_stress(
         cutoff=cutoff,
         cnthr=cnthr,
         abc=abc,
+        bidirectional=bidirectional,
     )
     _assert_energy_force_stress_equal_batch(torch_dftd3_calc, atoms_list)
 
 
 @pytest.mark.parametrize("damping,old", damping_method_list)
+@pytest.mark.parametrize("atoms_list", _create_atoms())
 @pytest.mark.parametrize("device", ["cpu", "cuda:0"])
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
-def test_calc_energy_device_batch(damping, old, device, dtype):
+def test_calc_energy_device_batch(damping, old, atoms_list, device, dtype):
     """Test2-1: check device, dtype dependency. with only various damping method."""
     xc = "pbe"
-    atoms_list = _create_atoms()
     _test_calc_energy(damping, xc, old, atoms_list, device=device, dtype=dtype)
 
 
 @pytest.mark.parametrize("damping,old", damping_method_list)
+@pytest.mark.parametrize("atoms_list", _create_atoms())
 @pytest.mark.parametrize("device", ["cpu", "cuda:0"])
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
-def test_calc_energy_force_stress_device_batch(damping, old, device, dtype):
+def test_calc_energy_force_stress_device_batch(damping, old, atoms_list, device, dtype):
     """Test2-2: check device, dtype dependency. with only various damping method."""
     xc = "pbe"
-    atoms_list = _create_atoms()
     _test_calc_energy_force_stress(damping, xc, old, atoms_list, device=device, dtype=dtype)
 
 
 @pytest.mark.parametrize("damping,old", damping_method_list)
+@pytest.mark.parametrize("atoms_list", _create_atoms())
 @pytest.mark.parametrize("device", ["cpu", "cuda:0"])
+@pytest.mark.parametrize("bidirectional", [True, False])
 @pytest.mark.parametrize("dtype", [torch.float64])
-def test_calc_energy_force_stress_device_batch_abc(damping, old, device, dtype):
-    """Test2-2: check device, dtype dependency. with only various damping method."""
+def test_calc_energy_force_stress_device_batch_abc(
+    damping, old, atoms_list, device, bidirectional, dtype
+):
+    """Test2-3: check device, dtype dependency. with only various damping method."""
     xc = "pbe"
     abc = True
-    atoms_list = _create_atoms()
-    _test_calc_energy_force_stress(
-        damping, xc, old, atoms_list, device=device, dtype=dtype, cnthr=7.0
-    )
+    if any([np.all(atoms.pbc) for atoms in atoms_list]) and bidirectional == False:
+        # TODO: bidirectional=False is not implemented for pbc now.
+        with pytest.raises(NotImplementedError):
+            _test_calc_energy_force_stress(
+                damping,
+                xc,
+                old,
+                atoms_list,
+                device=device,
+                dtype=dtype,
+                bidirectional=bidirectional,
+                abc=abc,
+                cnthr=7.0,
+            )
+    else:
+        _test_calc_energy_force_stress(
+            damping,
+            xc,
+            old,
+            atoms_list,
+            device=device,
+            dtype=dtype,
+            bidirectional=bidirectional,
+            abc=abc,
+            cnthr=7.0,
+        )
 
 
 if __name__ == "__main__":

torch_dftd/functions/dftd2.py

@@ -72,7 +72,10 @@ def edisp_d2(
         g = e6.sum()[None]
     else:
         # (n_graphs,)
-        n_graphs = int(batch[-1]) + 1
+        if batch.size()[0] == 0:
+            n_graphs = 1
+        else:
+            n_graphs = int(batch[-1]) + 1
         g = e6.new_zeros((n_graphs,))
         g.scatter_add_(0, batch_edge, e6)
 

torch_dftd/functions/dftd3.py

@@ -250,7 +250,10 @@ def edisp(
         g = e68.sum()[None]
     else:
         # (n_graphs,)
-        n_graphs = int(batch[-1]) + 1
+        if batch.size()[0] == 0:
+            n_graphs = 1
+        else:
+            n_graphs = int(batch[-1]) + 1
         g = e68.new_zeros((n_graphs,))
         g.scatter_add_(0, batch_edge, e68)
 
@@ -262,29 +265,40 @@ def edisp(
         # r_abc = r[within_cutoff]
         # r2_abc = r2[within_cutoff]
         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]
 
         n_atoms = Z.shape[0]
         if not bidirectional:
             # (2, n_edges) -> (2, n_edges * 2)
-            edge_index_abc = torch.cat([edge_index_abc, edge_index_abc.flip(dims=[1])], dim=1)
+            edge_index_abc = torch.cat([edge_index_abc, edge_index_abc.flip(dims=[0])], dim=1)
+            # (n_edges, ) -> (n_edges * 2, )
+            batch_edge_abc = (
+                None
+                if batch_edge_abc is None
+                else torch.cat([batch_edge_abc, batch_edge_abc], dim=0)
+            )
+            # (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)
             triplet_node_index, multiplicity, triplet_shift, batch_triplets = calc_triplets(
-                edge_index_abc, shift=shift_abc, dtype=pos.dtype, batch_edge=batch_edge
+                edge_index_abc, shift=shift_abc, dtype=pos.dtype, batch_edge=batch_edge_abc
             )
+            batch_triplets = None if batch_edge is None else batch_triplets
 
         # Apply `cnthr` cutoff threshold for r_kj
         idx_j, idx_k = triplet_node_index[:, 1], triplet_node_index[:, 2]
         ts2 = triplet_shift[:, 2]
-        r_jk = calc_distances(pos, torch.stack([idx_j, idx_k], dim=0), cell, ts2)
+        r_jk = calc_distances(pos, torch.stack([idx_j, idx_k], dim=0), cell, ts2, batch_triplets)
         kj_within_cutoff = r_jk <= cnthr
 
         triplet_node_index = triplet_node_index[kj_within_cutoff]
-        multiplicity, triplet_shift = (
+        multiplicity, triplet_shift, batch_triplets = (
             multiplicity[kj_within_cutoff],
             triplet_shift[kj_within_cutoff],
+            None if batch_triplets is None else batch_triplets[kj_within_cutoff],
         )
 
         idx_i, idx_j, idx_k = (
@@ -294,8 +308,8 @@ def edisp(
         )
         ts0, ts1, ts2 = triplet_shift[:, 0], triplet_shift[:, 1], triplet_shift[:, 2]
 
-        r_ij = calc_distances(pos, torch.stack([idx_i, idx_j], dim=0), cell, ts0)
-        r_ik = calc_distances(pos, torch.stack([idx_i, idx_k], dim=0), cell, ts1)
+        r_ij = calc_distances(pos, torch.stack([idx_i, idx_j], dim=0), cell, ts0, batch_triplets)
+        r_ik = calc_distances(pos, torch.stack([idx_i, idx_k], dim=0), cell, ts1, batch_triplets)
         r_jk = r_jk[kj_within_cutoff]
 
         Zti, Ztj, Ztk = Z[idx_i], Z[idx_j], Z[idx_k]
@@ -328,9 +342,9 @@ def edisp(
         if batch_edge is None:
             e6abc = e3.sum()
             g += e6abc
-            print("g", g)
-            print("e6abc", e6abc)
+            # print("g", g)
+            # print("e6abc", e6abc)
         else:
             g.scatter_add_(0, batch_triplets, e3)
-            print("g", g)
+            # print("g", g)
     return g  # (n_graphs,)

torch_dftd/nn/base_dftd_module.py

@@ -78,7 +78,10 @@ def calc_energy(
         if batch is None:
             return [{"energy": E_disp.item()}]
         else:
-            n_graphs = int(batch[-1]) + 1
+            if batch.size()[0] == 0:
+                n_graphs = 1
+            else:
+                n_graphs = int(batch[-1]) + 1
             return [{"energy": E_disp[i].item()} for i in range(n_graphs)]
 
     def calc_energy_and_forces(
@@ -133,7 +136,10 @@ def calc_energy_and_forces(
         if batch is None:
             results_list = [{"energy": E_disp.item(), "forces": forces.cpu().numpy()}]
         else:
-            n_graphs = int(batch[-1]) + 1
+            if batch.size()[0] == 0:
+                n_graphs = 1
+            else:
+                n_graphs = int(batch[-1]) + 1
             results_list = [{"energy": E_disp[i].item()} for i in range(n_graphs)]
             for i in range(n_graphs):
                 results_list[i]["forces"] = forces[batch == i].cpu().numpy()

torch_dftd/torch_dftd3_calculator.py

@@ -166,8 +166,10 @@ def batch_calculate(self, atoms_list=None, properties=["energy"], system_changes
         )
         batch_dicts = dict(
             Z=torch.cat([d["Z"] for d in input_dicts_list], dim=0),  # (n_nodes,)
-            pos=torch.cat([d["pos"] for d in input_dicts_list], dim=0),  # (n_nodes,)
-            cell=cell_batch,  # (bs, 3, 3)
+            pos=torch.cat([d["pos"] for d in input_dicts_list], dim=0).requires_grad_(
+                True
+            ),  # (n_nodes,)
+            cell=cell_batch.requires_grad_(True),  # (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,)
         )