drop E3statistics when model does is valid and does not have a scale shift

Author: floatingCatty

dptb/data/dataset/_default_dataset.py

@@ -401,6 +401,10 @@ def E3statistics(self, model: torch.nn.Module=None, decay=False):
 
         if self.data[AtomicDataDict.EDGE_FEATURES_KEY].abs().sum() < 1e-7:
             return None
+    
+        if model is not None:
+            if not isinstance(model.node_prediction_h, torch.nn.Module):
+                return None
 
         typed_dataset = idp(self.data.clone().to_dict())
         e3h = E3Hamiltonian(basis=idp.basis, decompose=True)

dptb/data/dataset/lmdb_dataset.py

@@ -158,6 +158,10 @@ def E3statistics(self, model: torch.nn.Module=None):
 
         if not self.get_Hamiltonian and not self.get_DM:
             return None
+
+        if model is not None:
+            if not isinstance(model.node_prediction_h, torch.nn.Module):
+                return None
 
         assert self.transform is not None
         idp = self.transform

dptb/nnops/loss.py

@@ -364,7 +364,7 @@ def forward(self, data: AtomicDataDict, ref_data: AtomicDataDict):
             elif batch.max() >= 1:
                 slices = [data["__slices__"]["pos"][i]-data["__slices__"]["pos"][i-1] for i in range(1,len(data["__slices__"]["pos"]))]
                 slices = [0] + slices
-                ndiag_batch = torch.stack([i.shape[0] for i in self.idp.mask_to_ndiag[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()].split(slices)])
+                ndiag_batch = torch.IntTensor([i.shape[0] for i in self.idp.mask_to_ndiag[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()].split(slices)], device=self.device)
                 ndiag_batch = torch.cumsum(ndiag_batch, dim=0)
                 mu = torch.stack([mu[ndiag_batch[i]:ndiag_batch[i+1]].mean() for i in range(len(ndiag_batch)-1)])
                 ss = (ref_data[AtomicDataDict.NODE_OVERLAP_KEY].sum(dim=-1) * ref_data[AtomicDataDict.NODE_OVERLAP_KEY]).sum(dim=-1)
@@ -456,7 +456,7 @@ def forward(self, data: AtomicDataDict, ref_data: AtomicDataDict):
             elif batch.max() >= 1:
                 slices = [data["__slices__"]["pos"][i]-data["__slices__"]["pos"][i-1] for i in range(1,len(data["__slices__"]["pos"]))]
                 slices = [0] + slices
-                ndiag_batch = torch.stack([i.shape[0] for i in self.idp.mask_to_ndiag[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()].split(slices)])
+                ndiag_batch = torch.IntTensor([i.shape[0] for i in self.idp.mask_to_ndiag[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()].split(slices)], device=self.device)
                 ndiag_batch = torch.cumsum(ndiag_batch, dim=0)
                 mu = torch.stack([mu[ndiag_batch[i]:ndiag_batch[i+1]].mean() for i in range(len(ndiag_batch)-1)])
                 ss = (ref_data[AtomicDataDict.NODE_OVERLAP_KEY].sum(dim=-1) * ref_data[AtomicDataDict.NODE_OVERLAP_KEY]).sum(dim=-1)
@@ -536,7 +536,7 @@ def forward(self, data: AtomicDataDict, ref_data: AtomicDataDict):
             elif batch.max() >= 1:
                 slices = [data["__slices__"]["pos"][i]-data["__slices__"]["pos"][i-1] for i in range(1,len(data["__slices__"]["pos"]))]
                 slices = [0] + slices
-                ndiag_batch = torch.stack([i.shape[0] for i in self.idp.mask_to_ndiag[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()].split(slices)])
+                ndiag_batch = torch.IntTensor([i.shape[0] for i in self.idp.mask_to_ndiag[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()].split(slices)], device=self.device)
                 ndiag_batch = torch.cumsum(ndiag_batch, dim=0)
                 mu = torch.stack([mu[ndiag_batch[i]:ndiag_batch[i+1]].mean() for i in range(len(ndiag_batch)-1)])
                 ss = (ref_data[AtomicDataDict.NODE_OVERLAP_KEY].sum(dim=-1) * ref_data[AtomicDataDict.NODE_OVERLAP_KEY]).sum(dim=-1)
@@ -682,7 +682,7 @@ def forward(self, data: AtomicDataDict, ref_data: AtomicDataDict):
             elif batch.max() >= 1:
                 slices = [data["__slices__"]["pos"][i]-data["__slices__"]["pos"][i-1] for i in range(1,len(data["__slices__"]["pos"]))]
                 slices = [0] + slices
-                ndiag_batch = torch.stack([i.shape[0] for i in self.idp.mask_to_ndiag[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()].split(slices)])
+                ndiag_batch = torch.IntTensor([i.shape[0] for i in self.idp.mask_to_ndiag[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()].split(slices)], device=self.device)
                 ndiag_batch = torch.cumsum(ndiag_batch, dim=0)
                 mu = torch.stack([mu[ndiag_batch[i]:ndiag_batch[i+1]].mean() for i in range(len(ndiag_batch)-1)])
                 ss = (ref_data[AtomicDataDict.NODE_OVERLAP_KEY].sum(dim=-1) * ref_data[AtomicDataDict.NODE_OVERLAP_KEY]).sum(dim=-1)
@@ -811,7 +811,7 @@ def __call__(self, data: AtomicDataDict, ref_data: AtomicDataDict, running_avg:
             elif batch.max() >= 1:
                 slices = [data["__slices__"]["pos"][i]-data["__slices__"]["pos"][i-1] for i in range(1,len(data["__slices__"]["pos"]))]
                 slices = [0] + slices
-                ndiag_batch = torch.stack([i.shape[0] for i in self.idp.mask_to_ndiag[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()].split(slices)])
+                ndiag_batch = torch.IntTensor([i.shape[0] for i in self.idp.mask_to_ndiag[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()].split(slices)], device=self.device)
                 ndiag_batch = torch.cumsum(ndiag_batch, dim=0)
                 mu = torch.stack([mu[ndiag_batch[i]:ndiag_batch[i+1]].mean() for i in range(len(ndiag_batch)-1)])
                 ss = (ref_data[AtomicDataDict.NODE_OVERLAP_KEY].sum(dim=-1) * ref_data[AtomicDataDict.NODE_OVERLAP_KEY]).sum(dim=-1)