Merge pull request #287 from AsymmetryChou/fix_ovp_np

fix: Ensure proper tensor conversion for numpy solver in Eigenvalues …

Author: AsymmetryChou

dptb/nn/energy.py

@@ -83,23 +83,28 @@ def forward(self,
         for i in range(int(np.ceil(num_k / nk))):
             data[AtomicDataDict.KPOINT_KEY] = kpoints[i*nk:(i+1)*nk]
             data = self.h2k(data)
+            h_transformed_np = None
             if self.overlap:
                 data = self.s2k(data)
                 if eig_solver == 'torch':
                     chklowt = torch.linalg.cholesky(data[self.s_out_field])
                     chklowtinv = torch.linalg.inv(chklowt)
                     data[self.h_out_field] = (chklowtinv @ data[self.h_out_field] @ torch.transpose(chklowtinv,dim0=1,dim1=2).conj())
                 elif eig_solver == 'numpy':
-                    chklowt = np.linalg.cholesky(data[self.s_out_field].detach().numpy())
+                    s_np = data[self.s_out_field].detach().cpu().numpy()
+                    h_np = data[self.h_out_field].detach().cpu().numpy()
+                    chklowt = np.linalg.cholesky(s_np)
                     chklowtinv = np.linalg.inv(chklowt)
-                    data[self.h_out_field] = (chklowtinv @ data[self.h_out_field].detach().numpy() @ np.transpose(chklowtinv,(0,2,1)).conj())
-            else:
-                data[self.h_out_field] = data[self.h_out_field]
-            
+                    h_transformed_np = chklowtinv @ h_np @ np.transpose(chklowtinv,(0,2,1)).conj()
+
             if eig_solver == 'torch':
                 eigvals.append(torch.linalg.eigvalsh(data[self.h_out_field]))
             elif eig_solver == 'numpy':
-                eigvals.append(torch.from_numpy(np.linalg.eigvalsh(a=data[self.h_out_field])))
+                if h_transformed_np is None:
+                    h_transformed_np = data[self.h_out_field].detach().cpu().numpy()
+                eigvals_np = np.linalg.eigvalsh(a=h_transformed_np)
+                # Preserve dtype by converting to the Hamiltonian's original dtype
+                eigvals.append(torch.from_numpy(eigvals_np).to(dtype=self.h2k.dtype, device=self.h2k.device))
 
         data[self.out_field] = torch.nested.as_nested_tensor([torch.cat(eigvals, dim=0)])
         if nested: