Add support for specifying number of CPU cores

dpnegf/negf/lead_property.py

@@ -517,7 +517,7 @@ def _estimate_worker_memory(lead_L, lead_R, kpoint=None, temp_allocation_factor=
     return total_estimate
 
 
-def _get_safe_n_jobs(lead_L, lead_R, requested_n_jobs=-1, max_memory_fraction=0.7, min_workers=1, kpoint=None):
+def _get_safe_n_jobs(lead_L, lead_R, requested_n_jobs=-1, max_memory_fraction=0.7, min_workers=1, kpoint=None, n_cpus=None):
     """
     Calculate safe number of parallel workers based on available system memory.
 
@@ -539,7 +539,7 @@ def _get_safe_n_jobs(lead_L, lead_R, requested_n_jobs=-1, max_memory_fraction=0.
     int
         Safe number of parallel workers.
     """
-    cpu_count = os.cpu_count()
+    cpu_count = n_cpus if n_cpus is not None else os.cpu_count()
     if cpu_count is None or cpu_count < 1:
         cpu_count = 1
         log.warning("os.cpu_count() returned None or invalid value. Defaulting to 1 CPU core.")
@@ -592,7 +592,7 @@ def _get_safe_n_jobs(lead_L, lead_R, requested_n_jobs=-1, max_memory_fraction=0.
 
 
 def compute_all_self_energy(eta, lead_L, lead_R, kpoints_grid, energy_grid,
-                            self_energy_save_path=None, n_jobs=-1, batch_size=200):
+                            self_energy_save_path=None, n_jobs=-1, batch_size=200, n_cpus=None):
     """
     Computes and saves self-energy matrices for all combinations of k-points and energy values
     for left and right leads.
@@ -633,7 +633,7 @@ def compute_all_self_energy(eta, lead_L, lead_R, kpoints_grid, energy_grid,
     # Calculate safe number of workers based on available memory
     # Use first k-point for memory estimation
     sample_kpoint = kpoints_grid[0] if len(kpoints_grid) > 0 else None
-    safe_n_jobs = _get_safe_n_jobs(lead_L, lead_R, requested_n_jobs=n_jobs, kpoint=sample_kpoint)
+    safe_n_jobs = _get_safe_n_jobs(lead_L, lead_R, requested_n_jobs=n_jobs, kpoint=sample_kpoint, n_cpus=n_cpus)
     if n_jobs == -1:
         log.info(f"Auto-detected safe n_jobs={safe_n_jobs} based on available memory")
     elif safe_n_jobs < n_jobs:

dpnegf/runner/NEGF.py

@@ -53,14 +53,16 @@ def __init__(self,
                 out_current: bool=False,out_current_nscf: bool=False,out_ldos: bool=False,out_lcurrent: bool=False,
                 results_path: Optional[str]=None,
                 torch_device: Union[str, torch.device]=torch.device('cpu'),
-                AtomicData_options: Optional[dict]=None, 
+                AtomicData_options: Optional[dict]=None,
+                n_cpus: Optional[int]=None,
                 **kwargs):
 
 
-        # self.model = model # No need to set model as property for memory saving      
+        # self.model = model # No need to set model as property for memory saving
         self.results_path = results_path
         self.cdtype = torch.complex128
         self.torch_device = torch_device
+        self.n_cpus = n_cpus
 
         # get the parameters
         self.ele_T = ele_T
@@ -556,11 +558,11 @@ def prepare_self_energy(self, scf_require: bool) -> None:
                 #         self.deviceprop.lead_L.self_energy(kpoint=k, energy=e, eta_lead=self.eta_lead, save=True)
                 #         self.deviceprop.lead_R.self_energy(kpoint=k, energy=e, eta_lead=self.eta_lead, save=True)
                 compute_all_self_energy(self.eta_lead, self.deviceprop.lead_L, self.deviceprop.lead_R,
-                                        self.kpoints, self.density.integrate_range, self.self_energy_save_path)
+                                        self.kpoints, self.density.integrate_range, self.self_energy_save_path, n_cpus=self.n_cpus)
             elif not self.scf:
                 # In non-scf case, the self-energy of the leads is calculated for each energy point in the energy grid.
                 compute_all_self_energy(self.eta_lead, self.deviceprop.lead_L, self.deviceprop.lead_R,
-                                        self.kpoints, self.uni_grid, self.self_energy_save_path)
+                                        self.kpoints, self.uni_grid, self.self_energy_save_path, n_cpus=self.n_cpus)
         log.info(msg="-----------------------------------\n")
 
 

dpnegf/utils/argcheck.py

@@ -1064,7 +1064,8 @@ def negf():
         Argument("out_current", bool, optional=True, default=False, doc=doc_out_current),
         Argument("out_current_nscf", bool, optional=True, default=False, doc=doc_out_current_nscf),
         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("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().")
     ]
 
 def stru_options():

examples/atomic_chain_cli/input_files/negf_chain_new.json

@@ -2,6 +2,7 @@
     "task_options": {
         "task": "negf",
         "scf": false,
+        "n_cpus": 2,
         "block_tridiagonal": false,
         "ele_T": 300,
         "unit": "eV",