Merge branch 'main' of https://github.com/QG-phy/DeePTB into hotfix_tojson

Author: QG-phy

dptb/nn/nnsk.py

@@ -1021,6 +1021,8 @@ def to_json(self, version=2, basisref=None):
             rev_line = self.idp_sk.transform_bond(jan, ian)
             for orbpair, slices in self.idp_sk.orbpair_maps.items():
                 fiorb, fjorb = orbpair.split("-")
+                if fiorb not in self.idp_sk.full_basis_to_basis[iasym] or fjorb not in self.idp_sk.full_basis_to_basis[jasym]:
+                    continue
                 iorb = self.idp_sk.full_basis_to_basis[iasym].get(fiorb)
                 jorb = self.idp_sk.full_basis_to_basis[jasym].get(fjorb)
 
@@ -1064,6 +1066,8 @@ def to_json(self, version=2, basisref=None):
                 rev_line = self.idp_sk.transform_bond(jan, ian)
                 for orbpair, slices in self.idp_sk.orbpair_maps.items():
                     fiorb, fjorb = orbpair.split("-")
+                    if fiorb not in self.idp_sk.full_basis_to_basis[iasym] or fjorb not in self.idp_sk.full_basis_to_basis[jasym]:
+                        continue
                     iorb = self.idp_sk.full_basis_to_basis[iasym].get(fiorb)
                     jorb = self.idp_sk.full_basis_to_basis[jasym].get(fjorb)
 
@@ -1167,6 +1171,8 @@ def to_json(self, version=2, basisref=None):
             soc_param = {}
             for asym in self.idp_sk.type_names:
                 for fiorb, slices in self.idp_sk.sksoc_maps.items():
+                    if fiorb not in self.idp_sk.full_basis_to_basis[asym]:
+                        continue
                     iorb = self.idp_sk.full_basis_to_basis[asym][fiorb]
                     if fiorb not in self.idp_sk.full_basis_to_basis[asym]:
                         continue

dptb/postprocess/bandstructure/band.py

@@ -207,8 +207,12 @@ def get_bands(self, data: Union[AtomicData, ase.Atoms, str], kpath_kwargs: dict,
         else:
             log.error('Error, now, kline_type only support ase_kpath, abacus, or vasp.')
             raise ValueError
+
+        override_overlap = None
+        if kpath_kwargs.get("override_overlap", None):
+            override_overlap = kpath_kwargs["override_overlap"]
 
-        data, eigenvalues = self.get_eigs(data=data, klist=klist, pbc=pbc, AtomicData_options=AtomicData_options)
+        data, eigenvalues = self.get_eigs(data=data, klist=klist, pbc=pbc, AtomicData_options=AtomicData_options, override_overlap=override_overlap)
 
 
         # get the E_fermi from data

dptb/postprocess/elec_struc_cal.py

@@ -1,3 +1,5 @@
+import os
+import h5py
 import numpy as np
 from ase.io import read
 import ase
@@ -7,7 +9,7 @@
 from typing import Optional
 import logging
 log = logging.getLogger(__name__)
-from dptb.data import AtomicData, AtomicDataDict
+from dptb.data import AtomicData, AtomicDataDict, block_to_feature
 from dptb.nn.energy import Eigenvalues
 from dptb.utils.argcheck import get_cutoffs_from_model_options
 from copy import deepcopy
@@ -66,7 +68,12 @@ def __init__ (
             )
         r_max, er_max, oer_max  = get_cutoffs_from_model_options(model.model_options)
         self.cutoffs = {'r_max': r_max, 'er_max': er_max, 'oer_max': oer_max}
-    def get_data(self,data: Union[AtomicData, ase.Atoms, str],pbc:Union[bool,list]=None, device: Union[str, torch.device]=None, AtomicData_options:dict=None):
+    def get_data(self,
+                 data: Union[AtomicData, ase.Atoms, str],
+                 pbc:Union[bool,list]=None,
+                 device: Union[str, torch.device]=None,
+                 AtomicData_options:dict=None,
+                 override_overlap:Optional[str]=None):
         '''The function `get_data` takes input data in the form of a string, ase.Atoms object, or AtomicData
         object, processes it accordingly, and returns the AtomicData class.
         
@@ -81,6 +88,7 @@ def get_data(self,data: Union[AtomicData, ase.Atoms, str],pbc:Union[bool,list]=N
         device : Union[str, torch.device]
             The `device` parameter in the `get_data` function is used to specify the device on which the data
         should be processed. If no device is provided, it defaults to `self.device`.
+        override_overlap : the path for overlap.h5 to use and override overlap matrix from model.
         
         Returns
         -------
@@ -130,7 +138,30 @@ def get_data(self,data: Union[AtomicData, ase.Atoms, str],pbc:Union[bool,list]=N
             data = data
         else:
             raise ValueError('data should be either a string, ase.Atoms, or AtomicData')
-        
+
+        if isinstance(override_overlap, str):
+            assert os.path.exists(override_overlap), "Overlap file not found."
+            overlap_blocks = h5py.File(override_overlap, "r")
+            if len(overlap_blocks) != 1:
+                log.info('Overlap file contains more than one overlap matrix, only first will be used.')
+            if self.overlap:
+                log.warning('override_overlap is enabled while model contains overlap, override_overlap will be used.')
+            if "0" in overlap_blocks:
+                overlaps = overlap_blocks["0"]
+            else:
+                overlaps = overlap_blocks["1"]
+            block_to_feature(data, self.model.idp, blocks=False, overlap_blocks=overlaps)
+            if not self.overlap:
+                self.eigv = Eigenvalues(
+                    idp=self.model.idp,
+                    device=self.device,
+                    s_edge_field=AtomicDataDict.EDGE_OVERLAP_KEY,
+                    s_node_field=AtomicDataDict.NODE_OVERLAP_KEY,
+                    s_out_field=AtomicDataDict.OVERLAP_KEY,
+                    dtype=self.model.dtype,
+                )
+            overlap_blocks.close()
+
         if device is None:
             device = self.device
         data = AtomicData.to_AtomicDataDict(data.to(device))
@@ -139,7 +170,12 @@ def get_data(self,data: Union[AtomicData, ase.Atoms, str],pbc:Union[bool,list]=N
         return data
 
 
-    def get_eigs(self, data: Union[AtomicData, ase.Atoms, str], klist: np.ndarray, pbc:Union[bool,list]=None, AtomicData_options:dict=None):
+    def get_eigs(self,
+                 data: Union[AtomicData, ase.Atoms, str],
+                 klist: np.ndarray,
+                 pbc:Union[bool,list]=None,
+                 AtomicData_options:dict=None,
+                 override_overlap:Optional[str]=None):
         '''This function calculates eigenvalues for Hk at specified k-points.
         
         Parameters
@@ -152,20 +188,27 @@ def get_eigs(self, data: Union[AtomicData, ase.Atoms, str], klist: np.ndarray, p
         AtomicData_options : dict
             The `AtomicData_options` parameter is a dictionary that contains options for configuring the
         `AtomicData` object.
+        override_overlap : the path for overlap.h5 to use and override overlap matrix from model.
         
         Returns
         -------
             The function `get_eigs` returns the loaded data and the energy eigenvalues as a numpy array.
         
         '''
 
-        data  = self.get_data(data=data, pbc=pbc, device=self.device,AtomicData_options=AtomicData_options)
+        data  = self.get_data(data=data, pbc=pbc, device=self.device,AtomicData_options=AtomicData_options, override_overlap=override_overlap)
         # set the kpoint of the AtomicData
         data[AtomicDataDict.KPOINT_KEY] = \
             torch.nested.as_nested_tensor([torch.as_tensor(klist, dtype=self.model.dtype, device=self.device)])
+        if isinstance(override_overlap, str):
+            override_overlap_edge = data[AtomicDataDict.EDGE_OVERLAP_KEY]
+            override_overlap_node = data[AtomicDataDict.NODE_OVERLAP_KEY]
         # get the eigenvalues
         data = self.model(data)
-        if self.overlap == True:
+        if isinstance(override_overlap, str):
+            data[AtomicDataDict.EDGE_OVERLAP_KEY] = override_overlap_edge
+            data[AtomicDataDict.NODE_OVERLAP_KEY] = override_overlap_node
+        if self.overlap or isinstance(override_overlap, str):
             assert data.get(AtomicDataDict.EDGE_OVERLAP_KEY) is not None
         data = self.eigv(data)
 

dptb/tests/data/e3_band/data/Si64.0/atomic_numbers.dat

@@ -0,0 +1,64 @@
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14
+14

dptb/tests/data/e3_band/data/Si64.0/basis.dat

@@ -0,0 +1 @@
+{'Si': '1s1p'}

dptb/tests/data/e3_band/data/Si64.0/cell.dat

@@ -0,0 +1,3 @@
+1.088740444183349609e+01 0.000000000000000000e+00 0.000000000000000000e+00
+0.000000000000000000e+00 1.088740444183349609e+01 0.000000000000000000e+00
+0.000000000000000000e+00 0.000000000000000000e+00 1.088740444183349609e+01

dptb/tests/data/e3_band/data/Si64.0/hamiltonians.h5

@@ -0,0 +1,3 @@
+1.000000000000000000e+00
+1.000000000000000000e+00
+1.000000000000000000e+00