more energy_grid_options

dpnegf/negf/lead_property.py

@@ -985,13 +985,21 @@ def merge_hdf5_files(tmp_dir, output_path, pattern, remove=True):
 
 
 def _has_saved_self_energy(root: str) -> bool:
-        from pathlib import Path
-        p = Path(root) if root is not None else None
-        if p is None or not p.exists():
-            return False
-
-        patterns = ("*.h5", "*.pth")
-        for pat in patterns:
-            if any(p.rglob(pat)):
-                return True
-        return False
+    from pathlib import Path
+
+    p = Path(root) if root is not None else None
+    if p is None or not p.exists():
+        return False
+
+    for file in p.rglob("*"):
+        if not file.is_file():
+            continue
+
+        if file.suffix.lower() not in {".pth", ".h5"}:
+            continue
+
+        name = file.name.lower()
+        if "self_energy" in name or "se" in name:
+            return True
+
+    return False

dpnegf/runner/NEGF.py

@@ -40,23 +40,25 @@ def __init__(self,
                 model: torch.nn.Module,
                 structure: Union[AtomicData, ase.Atoms, str],
                 ele_T: float,
-                emin: float, emax: float, espacing: float,
                 density_options: dict,
                 unit: str,
                 scf: bool, poisson_options: dict,
                 stru_options: dict,eta_lead: float,eta_device: float,
                 block_tridiagonal: bool,
                 sgf_solver: str,
+                emin: float=None, emax: float=None, espacing: float=None,
                 e_fermi: float=None,
+                plot_blocks: bool=False,
                 use_saved_HS: bool=False, saved_HS_path: str=None,
                 use_saved_se: bool=False, self_energy_save_path: str=None, 
                 se_info_display: bool=False, se_numba_jit: Optional[bool]=None,
                 out_tc: bool=False,out_dos: bool=False,out_density: bool=False,out_potential: bool=False,
                 out_current: bool=False,out_current_nscf: bool=False,out_ldos: bool=False,out_lcurrent: bool=False,
-                results_path: Optional[str]=None, plot_blocks: Optional[bool]=False,
+                results_path: Optional[str]=None,
                 rgf_device: Union[str, torch.device]='cpu',
                 AtomicData_options: Optional[dict]=None,
                 n_cpus: Optional[int]=None,
+                energy_grid_options: Optional[dict]=None,
                 e_batch_size: Optional[int]=None,
                 **kwargs):
 
@@ -68,6 +70,7 @@ def __init__(self,
             rgf_device = torch.device(rgf_device)
         self.rgf_device = rgf_device
         self.n_cpus = n_cpus
+
         self.e_batch_size = e_batch_size
 
         # The RGF q-loop allocates/frees many small slabs; with the default
@@ -89,7 +92,16 @@ def __init__(self,
         self.kBT = Boltzmann * self.ele_T / eV2J # change to eV
         self.e_fermi = e_fermi
         self.eta_lead = eta_lead; self.eta_device = eta_device
-        self.emin = emin; self.emax = emax; self.espacing = espacing
+
+        assert energy_grid_options or (espacing and emin and emax), "target energy grid not defined"
+        self.energy_grid_options = {} if energy_grid_options is None else energy_grid_options
+        self.emin = self.energy_grid_options.get("emin") if self.energy_grid_options.get("emin") is not None else emin
+        self.emax = self.energy_grid_options.get("emax") if self.energy_grid_options.get("emax") is not None else emax
+        self.espacing = self.energy_grid_options.get("espacing") if self.energy_grid_options.get("espacing") is not None else espacing
+        self.esteps = self.energy_grid_options.get("esteps", None)
+        self.force_zero_point = self.energy_grid_options.get("force_zero_point", False)
+        self.force_espacing = self.energy_grid_options.get("force_espacing", False)
+        self.energies = self.energy_grid_options.get("energies", None)
         self.stru_options = stru_options
         self.poisson_options = poisson_options
         if e_fermi is None:
@@ -132,6 +144,14 @@ def __init__(self,
             self.kpoints,self.wk = kmesh_sampling_negf(self.stru_options["kmesh"], 
                                                        self.stru_options["gamma_center"],
                                                      self.stru_options["time_reversal_symmetry"])
+
+        if 'override_overlap' in kwargs:
+            assert isinstance(kwargs['override_overlap'], str)
+            self.override_overlap = kwargs['override_overlap']
+            log.info(msg="Using external calculated overlap overriding!")
+        else:
+            self.override_overlap = None
+
         log.info(msg="------ k-point for NEGF -----")
         log.info(msg="Gamma Center: {0}".format(self.stru_options["gamma_center"]))
         log.info(msg="Time Reversal: {0}".format(self.stru_options["time_reversal_symmetry"]))
@@ -198,7 +218,7 @@ def __init__(self,
         e_fermi = {}; chemiPot = {}
         # calculate Fermi level
         if  self.e_fermi is None:        
-            elec_cal = ElecStruCal(model=model,device=torch.device("cpu"))
+            elec_cal = ElecStruCal(model=model, device=torch.device("cpu"), override_overlap=self.override_overlap)
             nel_atom_lead = self.get_nel_atom_lead(
                                 struct_leads, 
                                 charge={lead_tag: self.stru_options[lead_tag].get("charge", 0) for lead_tag in ["lead_L", "lead_R"]}
@@ -361,8 +381,68 @@ def generate_energy_grid(self):
             cal_int_grid = True
 
         if self.out_dos or self.out_tc or self.out_current_nscf or self.out_ldos:
-            # Energy gird is set relative to Fermi level
-            self.uni_grid = torch.linspace(start=self.emin, end=self.emax, steps=int((self.emax-self.emin)/self.espacing))
+            if self.energies is not None:
+                self.uni_grid = torch.as_tensor(self.energies, dtype=torch.get_default_dtype())
+            else:
+                emin = float(self.emin)
+                emax = float(self.emax)
+                espacing = float(self.espacing)
+                if espacing <= 0:
+                    raise ValueError("espacing should be positive")
+                if emax < emin:
+                    raise ValueError("emax should be no less than emin")
+
+                has_esteps = self.esteps is not None
+                if has_esteps:
+                    energy_steps = int(self.esteps)
+                    if energy_steps < 1:
+                        raise ValueError("esteps should be a positive integer")
+                else:
+                    energy_steps = int((emax - emin) / espacing) + 1
+
+                if self.force_zero_point and self.force_espacing:
+                    if energy_steps == 1:
+                        emin = 0.0
+                        emax = 0.0
+                    elif has_esteps:
+                        zero_index = int(np.clip(round((0.0 - emin) / espacing), 0, energy_steps - 1))
+                        emin = -zero_index * espacing
+                        emax = emin + espacing * (energy_steps - 1)
+                        if emin > float(self.emin) or emax < float(self.emax):
+                            raise ValueError("esteps is too small to keep espacing, include zero, and cover [emin, emax]")
+                    else:
+                        left_steps = int(np.ceil(max(0.0, -emin) / espacing - 1e-12))
+                        right_steps = int(np.ceil(max(0.0, emax) / espacing - 1e-12))
+                        energy_steps = max(energy_steps, left_steps + right_steps + 1)
+                        emin = -left_steps * espacing
+                        emax = emin + espacing * (energy_steps - 1)
+                elif self.force_zero_point:
+                    if energy_steps == 1:
+                        emin = 0.0
+                        emax = 0.0
+                    else:
+                        zero_index = int(np.clip(round((0.0 - emin) / (emax - emin) * (energy_steps - 1)), 0, energy_steps - 1))
+                        left_spacing = abs(float(self.emin)) / zero_index if zero_index > 0 else 0.0
+                        right_spacing = abs(float(self.emax)) / (energy_steps - 1 - zero_index) if zero_index < energy_steps - 1 else 0.0
+                        espacing = max(left_spacing, right_spacing, espacing if left_spacing == 0.0 and right_spacing == 0.0 else 0.0)
+                        emin = -zero_index * espacing
+                        emax = (energy_steps - 1 - zero_index) * espacing
+                elif self.force_espacing and energy_steps > 1:
+                    target_width = espacing * (energy_steps - 1)
+                    if target_width < emax - emin:
+                        if has_esteps:
+                            raise ValueError("esteps is too small to keep espacing and cover [emin, emax]")
+                        energy_steps = int(np.ceil((emax - emin) / espacing - 1e-12)) + 1
+                        target_width = espacing * (energy_steps - 1)
+                    center = 0.5 * (emin + emax)
+                    emin = center - 0.5 * target_width
+                    emax = center + 0.5 * target_width
+
+                # Energy gird is set relative to Fermi level
+                if self.force_espacing and energy_steps > 1:
+                    self.uni_grid = emin + espacing * torch.arange(energy_steps, dtype=torch.get_default_dtype())
+                else:
+                    self.uni_grid = torch.linspace(start=emin, end=emax, steps=energy_steps)
 
         if cal_pole and  self.density_options["method"] == "Ozaki":
             self.poles, self.residues = ozaki_residues(M_cut=self.density_options["M_cut"])
@@ -447,6 +527,8 @@ def compute(self,
         output_path = os.path.join(self.results_path, "profile_report_negf.html")
         with open(output_path, 'w') as report_file:
             report_file.write(profiler.output_html())
+
+        log.info(msg="dpnegf compute completed.")
 
         return None
 
@@ -571,7 +653,8 @@ def prepare_self_energy(self, scf_require: bool) -> None:
         # self energy calculation
         log.info(msg="------Self-energy calculation------")
         if  self.self_energy_save_path is None:
-            self.self_energy_save_path = os.path.join(self.results_path, "self_energy") 
+            self.self_energy_save_path = os.path.join(self.results_path, "self_energy")
+        self.self_energy_save_path = os.path.abspath(self.self_energy_save_path)
         os.makedirs(self.self_energy_save_path, exist_ok=True)
 
         if self.use_saved_se:

dpnegf/utils/argcheck.py

@@ -955,6 +955,25 @@ def normalize_test(data):
 
 
 
+def energy_grid_options():
+    doc_espacing = ""
+    doc_emin = ""
+    doc_emax = ""
+    doc_esteps = ""
+    doc_force_zero_point = ""
+    doc_force_espacing = ""
+    doc_energies = ""
+
+    return [
+        Argument("espacing", [int, float], optional=True, default=None, doc=doc_espacing),
+        Argument("emin", [int, float], optional=True, default=None, doc=doc_emin),
+        Argument("emax", [int, float], optional=True, default=None, doc=doc_emax),
+        Argument("esteps", [int, None], optional=True, default=None, doc=doc_esteps),
+        Argument("force_zero_point", bool, optional=True, default=False, doc=doc_force_zero_point),
+        Argument("force_espacing", bool, optional=True, default=False, doc=doc_force_espacing),
+        Argument("energies", [list, None], optional=True, default=None, doc=doc_energies),
+    ]
+
 
 def tbtrans_negf():
     doc_scf = ""
@@ -990,9 +1009,10 @@ def tbtrans_negf():
         Argument("stru_options", dict, optional=False, sub_fields=stru_options(), doc=doc_stru_options),
         Argument("poisson_options", dict, optional=True, default={}, sub_fields=[], sub_variants=[poisson_options()], doc=doc_poisson_options),
         Argument("sgf_solver", str, optional=True, default="Sancho-Rubio", doc=doc_sgf_solver),
-        Argument("espacing", [int, float], optional=False, doc=doc_espacing),
-        Argument("emin", [int, float], optional=False, doc=doc_emin),
-        Argument("emax", [int, float], optional=False, doc=doc_emax),
+        Argument("espacing", [int, float], optional=True, doc=doc_espacing),
+        Argument("emin", [int, float], optional=True, doc=doc_emin),
+        Argument("emax", [int, float], optional=True, doc=doc_emax),
+        Argument("energy_grid_options", dict, optional=True, default={}, sub_fields=energy_grid_options(), doc=""),
         Argument("e_fermi", [int, float], optional=False, doc=doc_e_fermi),
         Argument("density_options", dict, optional=True, default={}, sub_fields=[], sub_variants=[density_options()], doc=doc_density_options),
         Argument("eta_lead", [int, float], optional=True, default=1e-5, doc=doc_eta_lead),
@@ -1035,6 +1055,7 @@ def negf():
     doc_out_lcurrent = ""
     doc_density_options = ""
     doc_out_potential = ""
+    doc_override_overlap = ""
 
     return [
         Argument("scf", bool, optional=True, default=False, doc=doc_scf),
@@ -1052,9 +1073,10 @@ def negf():
         Argument("self_energy_save_path", str, optional=True, default=None, doc="the directory to save the self energy or load the self energy"),
         Argument("se_info_display", bool, optional=True, default=False, doc="whether to display the self energy information"),
         Argument("se_numba_jit", [bool, None], optional=True, default=None, doc="whether to use numba JIT for self energy calculation"),
-        Argument("espacing", [int, float], optional=False, doc=doc_espacing),
-        Argument("emin", [int, float], optional=False, doc=doc_emin),
-        Argument("emax", [int, float], optional=False, doc=doc_emax),
+        Argument("espacing", [int, float], optional=True, doc=doc_espacing),
+        Argument("emin", [int, float], optional=True, doc=doc_emin),
+        Argument("emax", [int, float], optional=True, doc=doc_emax),
+        Argument("energy_grid_options", dict, optional=True, default={}, sub_fields=energy_grid_options(), doc=""),
         Argument("e_batch_size", [int, float], optional=True, default=None, doc="the batch size for energy points in NEGF calculation"),
         Argument("e_fermi", [int, float], optional=True, default=None ,doc=doc_e_fermi),
         Argument("density_options", dict, optional=True, default={}, sub_fields=[], sub_variants=[density_options()], doc=doc_density_options),
@@ -1069,6 +1091,7 @@ def negf():
         Argument("out_ldos", bool, optional=True, default=False, doc=doc_out_ldos),
         Argument("out_lcurrent", bool, optional=True, default=False, doc=doc_out_lcurrent),
         Argument("n_cpus", [int, None], optional=True, default=None, doc="Number of CPU cores for parallel self-energy calculation. Default None uses os.cpu_count()."),
+        Argument("override_overlap", str, optional=True, doc=doc_override_overlap),
         Argument("rgf_device", str, optional=True, default="cpu",
                  doc="Device used only for the RGF (recursive Green's function) step. "
                      "'cpu' (default) or 'cuda'. Hamiltonian initialization always runs "

dpnegf/utils/elec_struc_cal.py

@@ -23,7 +23,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.
@@ -68,12 +69,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.
 
@@ -88,7 +95,8 @@ 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
         -------
@@ -139,6 +147,7 @@ def get_data(self,
         else:
             raise ValueError('data should be either a string, ase.Atoms, or AtomicData')
 
+        override_overlap = None if override_overlap == False else override_overlap if override_overlap else self.override_overlap
         if isinstance(override_overlap, str):
             assert os.path.exists(override_overlap), "Overlap file not found."
             overlap_blocks = h5py.File(override_overlap, "r")
@@ -189,14 +198,16 @@ def get_eigs(self,
         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.
+        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 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] = \