update deephe3 module

SharpLonde · Nov 30, 2023 · 2db61f2 · 2db61f2
1 parent ebf326f
commit 2db61f2
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 18 deletions.
diff --git a/dptb/data/interfaces/abacus.py b/dptb/data/interfaces/abacus.py
@@ -32,14 +32,14 @@ def __init__(self):
         #     3: [0, 2, 4, 6],
         # }
 
-        minus_dict = {
-            1: [0, 2],
-            2: [1, 3],
-            3: [0, 2, 4, 6],
-        }
+        # minus_dict = {
+        #     1: [0, 2],
+        #     2: [1, 3],
+        #     3: [0, 2, 4, 6],
+        # }
 
-        for k, v in minus_dict.items():
-            self.Us_abacus2deeptb[k][v] *= -1  # add phase (-1)^m
+        # for k, v in minus_dict.items():
+        #     self.Us_abacus2deeptb[k][v] *= -1  # add phase (-1)^m
 
     def get_U(self, l):
         if l > 3:

diff --git a/dptb/nn/deeptb.py b/dptb/nn/deeptb.py
@@ -215,7 +215,7 @@ def forward(self, data: AtomicDataDict.Type):
             data = self.node_prediction_h(data)
             data = self.edge_prediction_h(data)
 
-        data = self.hamiltonian(data)
+        # data = self.hamiltonian(data)
 
         if self.overlap:
             data = self.edge_prediction_s(data)

diff --git a/dptb/nn/hamiltonian.py b/dptb/nn/hamiltonian.py
@@ -160,11 +160,9 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
                     HR = data[self.node_field][:, self.idp.nodetype_maps[opairtype]]
                     HR = HR.reshape(n_node, -1, nL, nR).permute(0,2,3,1)# shape (N, nL, nR, n_pair)
 
-                    # rme = torch.sum(self.cgbasis[opairtype][None,:,:,:,None] * \
-                    #     HR[:,:,:,None,:], dim=(1,2)) # shape (N, n_rme, n_pair)
-                    # rme = rme.transpose(1,2).reshape(n_node, -1)
-
-                    rme = HR.permute(0,3,1,2).reshape(n_node, -1)
+                    rme = torch.sum(self.cgbasis[opairtype][None,:,:,:,None] * \
+                        HR[:,:,:,None,:], dim=(1,2)) # shape (N, n_rme, n_pair)
+                    rme = rme.transpose(1,2).reshape(n_node, -1)
 
                     # the onsite block doesnot have rotation
                     data[self.node_field][:, self.idp.nodetype_maps[opairtype]] = rme

diff --git a/dptb/nnops/_loss.py b/dptb/nnops/_loss.py
@@ -169,11 +169,15 @@ def forward(self, data: AtomicDataDict, ref_data: AtomicDataDict):
         # data[AtomicDataDict.NODE_FEATURES_KEY].masked_fill(~self.idp.mask_to_nrme[data[AtomicDataDict.ATOM_TYPE_KEY]], 0.)
         # data[AtomicDataDict.EDGE_FEATURES_KEY].masked_fill(~self.idp.mask_to_erme[data[AtomicDataDict.EDGE_TYPE_KEY]], 0.)
 
-        node_mean = ref_data[AtomicDataDict.NODE_FEATURES_KEY].mean(dim=-1, keepdim=True)
-        edge_mean = ref_data[AtomicDataDict.EDGE_FEATURES_KEY].mean(dim=-1, keepdim=True)
-        node_weight = 1/((ref_data[AtomicDataDict.NODE_FEATURES_KEY]-node_mean).norm(dim=-1, keepdim=True)+1e-5)
-        edge_weight = 1/((ref_data[AtomicDataDict.EDGE_FEATURES_KEY]-edge_mean).norm(dim=-1, keepdim=True)+1e-5)
-
+        # node_mean = ref_data[AtomicDataDict.NODE_FEATURES_KEY].mean(dim=-1, keepdim=True)
+        # edge_mean = ref_data[AtomicDataDict.EDGE_FEATURES_KEY].mean(dim=-1, keepdim=True)
+        # node_weight = 1/((ref_data[AtomicDataDict.NODE_FEATURES_KEY]-node_mean).norm(dim=-1, keepdim=True)+1e-5)
+        # edge_weight = 1/((ref_data[AtomicDataDict.EDGE_FEATURES_KEY]-edge_mean).norm(dim=-1, keepdim=True)+1e-5)
+
+        node_mean = 0.
+        edge_mean = 0.
+        node_weight = 1.
+        edge_weight = 1.
 
         pre = (node_weight*(data[AtomicDataDict.NODE_FEATURES_KEY]-node_mean))[self.idp.mask_to_nrme[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()]]
         tgt = (node_weight*(ref_data[AtomicDataDict.NODE_FEATURES_KEY]-node_mean))[self.idp.mask_to_nrme[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()]]