more energy_grid_options

dpnegf/negf/lead_property.py

@@ -660,18 +660,33 @@ def compute_all_self_energy(eta, lead_L, lead_R, kpoints_grid, energy_grid,
     # Capture the parent's log level so loky workers (which start with a clean
     # logging state and the WARNING default) can match it when they reinit.
     parent_log_level = logging.getLogger().getEffectiveLevel()
+
+    thread_per_task = 1
+    if n_cpus:
+        thread_per_task_from_tasks = int(n_cpus / min(len(total_tasks), batch_size))
+        if thread_per_task_from_tasks > 1:
+            if len(total_tasks) <= batch_size:
+                log.info(f"n_cpus:{n_cpus} is twice more than total_tasks:{len(total_tasks)},")
+            else:
+                log.info(f"n_cpus:{n_cpus} is twice more than batch_size:{batch_size},")
+        thread_per_task_from_jobs = int(n_cpus / safe_n_jobs)
+        if thread_per_task_from_jobs > 1:
+            log.info(f"n_cpus:{n_cpus} is twice more than safe_n_jobs:{safe_n_jobs},")
+        thread_per_task = max(thread_per_task_from_jobs, thread_per_task_from_tasks)
+        log.info(f"allocating {thread_per_task} thread for each task.")
+
     if len(total_tasks) <= batch_size:
-        Parallel(n_jobs=safe_n_jobs, backend="loky")(
+        Parallel(n_jobs=min(safe_n_jobs, len(total_tasks)), backend="loky")(
             delayed(_self_energy_worker_blas1)(k, e, eta, leadL_pack, leadR_pack, 
-                                               self_energy_save_path, se_numba_jit, parent_log_level)
+                                               self_energy_save_path, se_numba_jit, parent_log_level, thread_per_task)
             for k, e in total_tasks
         )
     else:
         for i in range(0, len(total_tasks), batch_size):
             batch = total_tasks[i:i+batch_size]
             Parallel(n_jobs=safe_n_jobs, backend="loky")(
                 delayed(_self_energy_worker_blas1)(k, e, eta, leadL_pack, leadR_pack, 
-                                                   self_energy_save_path, se_numba_jit, parent_log_level)
+                                                   self_energy_save_path, se_numba_jit, parent_log_level, thread_per_task)
                 for k, e in batch
             )
 
@@ -875,7 +890,7 @@ def _self_energy_worker_pure(k, e, eta, leadL_pack, leadR_pack, self_energy_save
     write_to_hdf5(save_tmp_R, k, e, seR)
 
 
-def _self_energy_worker_blas1(k, e, eta, leadL_pack, leadR_pack, self_energy_save_path, se_numba_jit, log_level):
+def _self_energy_worker_blas1(k, e, eta, leadL_pack, leadR_pack, self_energy_save_path, se_numba_jit, log_level, thread_per_task):
     """Loky entry point that pins this worker's BLAS/LAPACK runtime to a
     single thread, then delegates to `_self_energy_worker_pure`.
 
@@ -886,7 +901,7 @@ def _self_energy_worker_blas1(k, e, eta, leadL_pack, leadR_pack, self_energy_sav
     `solve` / `inv` / matmul calls where single-threaded BLAS already wins
     per-call; outer joblib-level parallelism handles scaling.
     """
-    with threadpool_limits(limits=1, user_api='blas'):
+    with threadpool_limits(limits=thread_per_task, user_api='blas'):
         return _self_energy_worker_pure(
             k, e, eta, leadL_pack, leadR_pack,
             self_energy_save_path, se_numba_jit, log_level,
@@ -985,13 +1000,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 "