Add eigenvalue solver option to structure calculations (#46)

* Feat: update elec_stru_cal.py as same as dptb, with numpy eigsolver choice

* Feat: add eig_solver option to stru_options and update example JSON

Author: AsymmetryChou

dpnegf/runner/NEGF.py

@@ -178,11 +178,12 @@ def __init__(self,
             for lead_tag in ["lead_L", "lead_R"]:
                 log.info(msg="-----Calculating Fermi level for {0}-----".format(lead_tag))
                 _, e_fermi[lead_tag]  = elec_cal.get_fermi_level(data=struct_leads[lead_tag], 
-                                nel_atom = nel_atom_lead[lead_tag],
-                                meshgrid=self.stru_options[lead_tag]["kmesh_lead_Ef"],
-                                AtomicData_options=AtomicData_options,
-                                smearing_method=self.stru_options.get("e_fermi_smearing", "FD"),
-                                temp=100.0)
+                                                                nel_atom = nel_atom_lead[lead_tag],
+                                                                meshgrid=self.stru_options[lead_tag]["kmesh_lead_Ef"],
+                                                                AtomicData_options=AtomicData_options,
+                                                                smearing_method=self.stru_options.get("e_fermi_smearing", "FD"),
+                                                                temp=100.0,
+                                                                eig_solver=self.stru_options.get("eig_solver", "torch"),)
         else:
             e_fermi["lead_L"] = self.e_fermi
             e_fermi["lead_R"] = self.e_fermi

dpnegf/tests/test_get_fermi.py

@@ -1,29 +1,37 @@
-import pytest
-from dpnegf.utils.elec_struc_cal import ElecStruCal
+from dptb.postprocess.elec_struc_cal import ElecStruCal
+
+# from dptb.postprocess.bandstructure.band import Band
 from dptb.nn.build import build_model
+import os
+from pathlib import Path
 
+rootdir = os.path.join(Path(os.path.abspath(__file__)).parent, "data")
 
-@pytest.fixture(scope='session', autouse=True)
-def root_directory(request):
-    """
-    :return:
-    """
-    return str(request.config.rootdir)
 
-def test_get_fermi(root_directory):
-    ckpt = f"{root_directory}/dpnegf/tests/data/test_get_fermi/nnsk.best.pth"  #  'hopping': {'method': 'poly2exp', 'rs': 5.0, 'w': 0.6},
-    stru_data = f"{root_directory}/dpnegf/tests/data/test_get_fermi/PRIMCELL.vasp"
+def test_get_fermi():
+    ckpt = f"{rootdir}/test_get_fermi/nnsk.best.pth"  #  'hopping': {'method': 'poly2exp', 'rs': 5.0, 'w': 0.6},
+    stru_data = f"{rootdir}/test_get_fermi/PRIMCELL.vasp"
 
     model = build_model(checkpoint=ckpt)
     nel_atom = {"Au":11}
 
     elec_cal = ElecStruCal(model=model,device='cpu')
     _, efermi =elec_cal.get_fermi_level(data=stru_data, 
                     nel_atom = nel_atom,smearing_method='FD',
-                meshgrid=[30,30,30])
-    assert abs(efermi  + 3.2257686853408813) < 1e-6
+                    meshgrid=[30,30,30],eig_solver='torch')
+    assert abs(efermi  + 3.2262574434280395) < 1e-3
+
+    _, efermi =elec_cal.get_fermi_level(data=stru_data, 
+                    nel_atom = nel_atom,smearing_method='Gaussian',
+                    meshgrid=[30,30,30],eig_solver='torch')
+    assert abs(efermi  + 3.2262574434280395) < 1e-3
+
+    _, efermi =elec_cal.get_fermi_level(data=stru_data, 
+                    nel_atom = nel_atom,smearing_method='FD',
+                    meshgrid=[30,30,30],eig_solver='numpy')
+    assert abs(efermi  + 3.2262574434280395) < 1e-3
 
     _, efermi =elec_cal.get_fermi_level(data=stru_data, 
                     nel_atom = nel_atom,smearing_method='Gaussian',
-                meshgrid=[30,30,30])
-    assert abs(efermi  + 3.2267462015151978) < 1e-6
+                    meshgrid=[30,30,30],eig_solver='numpy')
+    assert abs(efermi  + 3.2262574434280395) < 1e-3

dpnegf/utils/argcheck.py

@@ -1076,6 +1076,7 @@ def stru_options():
     doc_gamma_center=""
     doc_time_reversal_symmetry=""
     doc_e_fermi_smearing="The smearing method for Fermi level."
+    doc_eig_solver="The eigenvalue solver to use."
     doc_nel_atom = "The number of electrons in each element."
     return [
         Argument("device", dict, optional=False, sub_fields=device(), doc=doc_device),
@@ -1086,6 +1087,7 @@ def stru_options():
         Argument("gamma_center", list, optional=True, default=True, doc=doc_gamma_center),
         Argument("time_reversal_symmetry", list, optional=True, default=True, doc=doc_time_reversal_symmetry),
         Argument("e_fermi_smearing", str, optional=True, default="FD", doc=doc_e_fermi_smearing),
+        Argument("eig_solver", str, optional=True, default="torch", doc=doc_eig_solver),
         Argument("nel_atom", [dict,None], optional=True, default=None, doc=doc_nel_atom)
     ]
 

dpnegf/utils/elec_struc_cal.py

@@ -1,12 +1,15 @@
+import os
+import h5py
 import numpy as np
 from ase.io import read
 import ase
 import numpy as np
 from typing import Union
 import torch
+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 dpnegf.utils.argcheck import get_cutoffs_from_model_options
 from copy import deepcopy
@@ -65,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.
         
@@ -80,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
         -------
@@ -129,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))
@@ -138,7 +170,13 @@ 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,
+                 eig_solver:Optional[str]=None):
         '''This function calculates eigenvalues for Hk at specified k-points.
         
         Parameters
@@ -151,28 +189,36 @@ 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)
+        data = self.eigv(data, eig_solver=eig_solver)
 
         return data, data[AtomicDataDict.ENERGY_EIGENVALUE_KEY][0].detach().cpu().numpy()
 
     def get_fermi_level(self, data: Union[AtomicData, ase.Atoms, str], nel_atom: dict, \
-                        meshgrid: list = None, klist: np.ndarray=None, pbc:Union[bool,list]=None,AtomicData_options:dict=None,
-                        q_tol:float=1e-10,smearing_method:str='FD',temp:float=300):
+                        meshgrid: list = None, klist: np.ndarray=None, pbc:Union[bool,list]=None,
+                        AtomicData_options:dict=None, q_tol:float=1e-10, smearing_method:str='FD', 
+                        temp:float=300,eig_solver:Optional[str]='torch'):
         '''This function calculates the Fermi level based on provided data with iteration method, electron counts per atom, and
         optional parameters like specific k-points and eigenvalues.
         
@@ -233,13 +279,14 @@ def get_fermi_level(self, data: Union[AtomicData, ase.Atoms, str], nel_atom: dic
 
         # eigenvalues would be used if provided, otherwise the eigenvalues would be calculated from the model on the specified k-points
         if not AtomicDataDict.ENERGY_EIGENVALUE_KEY in data:
-            data, eigs = self.get_eigs(data=data, klist=klist, pbc=pbc, AtomicData_options=AtomicData_options) 
+            data, eigs = self.get_eigs(data=data, klist=klist, pbc=pbc, 
+                                       AtomicData_options=AtomicData_options,
+                                       eig_solver=eig_solver) 
             log.info('Getting eigenvalues from the model.')
         else:
             log.info('The eigenvalues are already in data. will use them.')
             eigs = data[AtomicDataDict.ENERGY_EIGENVALUE_KEY][0].detach().cpu().numpy()
-
-            
+        
         if nel_atom is not None:
             atomtype_list = data[AtomicDataDict.ATOM_TYPE_KEY].flatten().tolist()
             atomtype_symbols = np.asarray(self.model.idp.type_names)[atomtype_list].tolist()

examples/atomic_chain_api/input_files/negf_chain_new.json

@@ -19,6 +19,7 @@
             "time_reversal_symmetry": true,
             "nel_atom": {"C": 1.0},
             "kmesh":[1,1,1],
+            "eig_solver":"numpy",
             "pbc":[false, false, false],
             "device":{
                 "id":"4-8",