add argument for TBSystem to change eig_solver

dptb/postprocess/elec_struc_cal.py

@@ -24,7 +24,8 @@ class ElecStruCal(object):
     def __init__ (
             self, 
             model: torch.nn.Module,
-            device: Union[str, torch.device]=None
+            device: Union[str, torch.device]=None,
+            **kwargs
             ):
         '''It initializes ElecStruCal object with a neural network model, optional results path, GUI
         usage flag, and device information, and sets up eigenvalues  based on model properties.
@@ -69,12 +70,18 @@ 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}
+
+        if 'override_overlap' in kwargs and isinstance(kwargs['override_overlap'], str):
+            self.override_overlap = kwargs['override_overlap']
+        else:
+            self.override_overlap = 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):
+                 override_overlap:Union[str,bool,None]=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.
 
@@ -89,14 +96,15 @@ def get_data(self,
         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.
+        override_overlap : the path for overlap.h5 to use and override overlap matrix from model. If None, will try
+            to use self.override_overlap; If False, will not try anything.
 
         Returns
         -------
             the loaded AtomicData object.
 
         '''
-        if override_overlap is not None:
+        if override_overlap is not False and override_overlap or self.override_overlap:
             if not self.overlap:
                 self.eigv = Eigenvalues(
                     idp=self.model.idp,
@@ -112,7 +120,7 @@ def get_data(self,
             device=device if device else self.device,
             pbc=pbc,
             AtomicData_options=AtomicData_options,
-            override_overlap=override_overlap
+            override_overlap=None if override_overlap == False else override_overlap if override_overlap else self.override_overlap
         )
 
 
@@ -121,7 +129,7 @@ def get_eigs(self,
                  klist: np.ndarray,
                  pbc:Union[bool,list]=None,
                  AtomicData_options:dict=None,
-                 override_overlap:Optional[str]=None,
+                 override_overlap:Union[str,bool,None]=None,
                  eig_solver:Optional[str]=None):
         '''This function calculates eigenvalues for Hk at specified k-points.
 
@@ -142,7 +150,8 @@ def get_eigs(self,
             The function `get_eigs` returns the loaded data and the energy eigenvalues as a numpy array.
 
         '''
-
+
+        override_overlap = None if override_overlap == False else override_overlap if override_overlap else self.override_overlap
         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] = \

dptb/postprocess/unified/calculator.py

@@ -59,7 +59,7 @@ def get_hR(self, atomic_data: dict) -> Tuple[Any, Any]:
         pass
 
     @abstractmethod
-    def get_eigenvalues(self, atomic_data: dict) -> Tuple[dict, torch.Tensor]:
+    def get_eigenvalues(self, atomic_data: dict, **kwargs) -> Tuple[dict, torch.Tensor]:
         """
         Calculate eigenvalues for the given atomic data.
 
@@ -225,7 +225,12 @@ def get_hR(self, atomic_data):
     def get_eigenvalues(self, 
                         atomic_data: dict, 
                         nk: Optional[int]=None,
-                        solver: Optional[str]=None) -> Tuple[dict, torch.Tensor]:
+                        solver: Optional[str]=None,
+                        **kwargs) -> Tuple[dict, torch.Tensor]:
+        if not nk and kwargs.get("nk"):
+            nk = kwargs.get("nk")
+        if not solver and kwargs.get("eig_solver") or kwargs.get("solver"):
+            solver = kwargs.get("eig_solver") if kwargs.get("eig_solver") else kwargs.get("solver")
         # 1. Get Hamiltonian
         atomic_data = self.model_forward(atomic_data)
 

dptb/postprocess/unified/properties/band.py

@@ -221,7 +221,7 @@ def set_kpath(self, method: str, **kwargs):
         k_tensor = torch.as_tensor(self._k_points, dtype=self._system._calculator.dtype, device=self._system._calculator.device)
         self._system._atomic_data[AtomicDataDict.KPOINT_KEY] = torch.nested.as_nested_tensor([k_tensor])
 
-    def compute(self):
+    def compute(self, **kwargs):
         """
         Compute the band structure using the configured K-path and store result in system.
         """
@@ -232,7 +232,9 @@ def compute(self):
         data = self._system._atomic_data
 
         # Calculate
-        data, eigs = self._system.calculator.get_eigenvalues(data)
+        data, eigs = self._system.calculator.get_eigenvalues(data,
+                                                             nk=kwargs.get("nk", None),
+                                                             eig_solver=kwargs.get("eig_solver", None))
 
         # Extract results
         eigenvalues = eigs.detach().cpu().numpy() # [Nk, Nb]

dptb/postprocess/unified/system.py

@@ -240,7 +240,9 @@ def get_efermi(self, kmesh: List[int],
                                    device=self.calculator.device)
         data = self._atomic_data.copy()             
         data[AtomicDataDict.KPOINT_KEY] = torch.nested.as_nested_tensor([k_tensor])
-        data, eigs = self.calculator.get_eigenvalues(data)
+        data, eigs = self.calculator.get_eigenvalues(data,
+                                                     nk=kwargs.get("nk", None),
+                                                     eig_solver=kwargs.get("eig_solver", None))
 
         calculated_efermi = self.estimate_efermi_e(
                         eigenvalues=eigs.detach().numpy(),
@@ -294,7 +296,10 @@ def estimate_efermi_e(self, eigenvalues,
         )    
 
 
-    def get_bands(self, kpath_config: Optional[dict] = None, reuse: Optional[bool]=True, **kwargs):
+    def get_bands(self,
+                  kpath_config: Optional[dict] = None,
+                  reuse: Optional[bool]=True,
+                  **kwargs):
         # 计算能带，返回 bands
         # bands 应该是一个类，也有属性。bands.kpoints, bands.eigenvalues, bands.klabels, bands.kticks, 也有函数 bands.plot()
         if self.has_bands and reuse:
@@ -303,7 +308,8 @@ def get_bands(self, kpath_config: Optional[dict] = None, reuse: Optional[bool]=T
             assert kpath_config is not None, "kpath_config must be provided if bands not calculated."
             self._bands = BandAccessor(self)
             self._bands.set_kpath(**kpath_config)
-            self._bands.compute()
+            self._bands.compute(nk=kwargs.get("nk", None),
+                                eig_solver=kwargs.get("eig_solver", None))
             self.has_bands = True
             return self._bands