Merge pull request #25 from AsymmetryChou/btd_acc

Refactor and test vectorized BTD

Author: AsymmetryChou

dpnegf/negf/lead_property.py

@@ -487,7 +487,7 @@ def _estimate_worker_memory(lead_L, lead_R, kpoint=None, temp_allocation_factor=
         Estimated memory in bytes per worker.
     """
     # Base overhead for Python process + libraries (interpreter, numpy, scipy, torch, etc.)
-    base_overhead = 300 * 1024 * 1024  # 300 MB
+    base_overhead = 100 * 1024 * 1024  # 100 MB
 
     matrix_bytes = 0
 
@@ -522,7 +522,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, n_cpus=None):
+def _get_safe_n_jobs(lead_L, lead_R, requested_n_jobs=-1, max_memory_fraction=0.9, min_workers=1, kpoint=None, n_cpus=None):
     """
     Calculate safe number of parallel workers based on available system memory.
 
@@ -533,7 +533,7 @@ def _get_safe_n_jobs(lead_L, lead_R, requested_n_jobs=-1, max_memory_fraction=0.
     requested_n_jobs : int
         User-requested n_jobs. -1 means auto-detect.
     max_memory_fraction : float
-        Maximum fraction of available memory to use. Default 0.7.
+        Maximum fraction of available memory to use. Default 0.9.
     min_workers : int
         Minimum number of workers to use. Default 1.
     kpoint : array-like, optional

dpnegf/negf/negf_hamiltonian_init.py

@@ -141,7 +141,8 @@ def __init__(self,
             atom_norbs.append(int(self.model.idp.atom_norb[model.idp.chemical_symbol_to_type[atom.symbol]]))
         self.atom_norbs = atom_norbs
 
-    def initialize(self, kpoints, block_tridiagnal=False,useBloch=False,bloch_factor=None,\
+    def initialize(self, kpoints, block_tridiagnal=False, plot_blocks=False, \
+                   useBloch=False,bloch_factor=None,\
                    use_saved_HS=False, saved_HS_path=None):
         '''This function initializes the structure and Hamiltonian for a system with optional block tridiagonal
         and Bloch factor parameters.
@@ -152,6 +153,8 @@ def initialize(self, kpoints, block_tridiagnal=False,useBloch=False,bloch_factor
             Kpoints in the Brillouin zone
         block_tridiagnal, optional
              a boolean flag that determines whether the Hamiltonian matrix should be stored in a block-tridiagonal form. 
+        plot_blocks, optional
+            a boolean flag that determines whether to visualize the block structure using `show_blocks`.
         useBloch, optional
             a boolean flag that determines whether Bloch boundary conditions should be used in the lead self energy calculations. 
         bloch_factor
@@ -199,7 +202,7 @@ def initialize(self, kpoints, block_tridiagnal=False,useBloch=False,bloch_factor
             log.info(msg="=="*40)
         else:
             self.saved_HS_path = self.results_path
-            self.Hamiltonian_initialized(kpoints,useBloch,bloch_factor,block_tridiagnal,\
+            self.Hamiltonian_initialized(kpoints,useBloch,bloch_factor,block_tridiagnal,plot_blocks,\
                                                  lead_atom_range,structure_leads,structure_leads_fold)
             log.info(msg="--"*40)
             log.info(msg=f"The Hamiltonian has been initialized by model.")
@@ -211,7 +214,8 @@ def initialize(self, kpoints, block_tridiagnal=False,useBloch=False,bloch_factor
 
 
     def Hamiltonian_initialized(self,kpoints:List[List[float]],useBloch:bool,bloch_factor:List[int],\
-                                block_tridiagnal:bool,lead_atom_range:dict,structure_leads:Atoms,structure_leads_fold:Atoms):
+                                block_tridiagnal:bool, plot_blocks:bool,\
+                                lead_atom_range:dict,structure_leads:Atoms,structure_leads_fold:Atoms):
         '''This function initializes the Hamiltonian for a device with leads, handling various calculations
         and checks along the way.
 
@@ -236,6 +240,8 @@ def Hamiltonian_initialized(self,kpoints:List[List[float]],useBloch:bool,bloch_f
         should be stored in a block-tridiagonal form. If `block_tridiagnal` is set to `True`, the Hamiltonian
         matrix will be stored in a block-tridiagonal form. Otherwise, the Hamiltonian matrix will be stored in
         a full matrix format.
+        plot_blocks : bool
+            The `plot_blocks` parameter is a boolean flag that determines whether to visualize the block structure
         lead_atom_range : dict
             The `lead_atom_range` parameter indicates the range of leads. The key of the dictionary is the
         lead name, and the value is a list containing the start and end indices of the lead atoms.
@@ -432,7 +438,8 @@ def Hamiltonian_initialized(self,kpoints:List[List[float]],useBloch:bool,bloch_f
             leftmost_size = coupling_width['lead_L']
             rightmost_size = coupling_width['lead_R']
             hd, hu, hl, sd, su, sl, subblocks = self.get_block_tridiagonal(HD*self.h_factor,SD.cdouble(),self.structase,\
-                                                                leftmost_size,rightmost_size)
+                                                                leftmost_size,rightmost_size,
+                                                                plot_blocks=plot_blocks)
             HS_device.update({"hd":hd, "hu":hu, "hl":hl, "sd":sd, "su":su, "sl":sl, \
                               "subblocks":subblocks, "block_tridiagonal":True})
 
@@ -637,7 +644,7 @@ def get_lead_structure(self,kk,natom,useBloch=False,bloch_factor=None):
 
         return stru_lead, stru_lead_fold, bloch_sorted_indice, bloch_R_list
 
-    def get_block_tridiagonal(self,HK,SK,structase:ase.Atoms,leftmost_size:int,rightmost_size:int):
+    def get_block_tridiagonal(self,HK,SK,structase:ase.Atoms,leftmost_size:int,rightmost_size:int,plot_blocks=False):
         """
         Block-tridiagonalizes the Hamiltonian (HK) and overlap (SK) matrices for a given atomic structure.
         This method splits the input matrices into block tridiagonal form based on the atomic structure along the z-axis.
@@ -671,6 +678,8 @@ def get_block_tridiagonal(self,HK,SK,structase:ase.Atoms,leftmost_size:int,right
             List of lower-diagonal blocks of SK for each k-point.
         subblocks : list
             List of subblock sizes (number of orbitals in each block).
+        plot_blocks : bool
+            Whether to visualize the block structure using `show_blocks`.
         Notes
         -----
         - Uses optimized or fallback block splitting depending on the structure.
@@ -726,7 +735,9 @@ def get_block_tridiagonal(self,HK,SK,structase:ase.Atoms,leftmost_size:int,right
         log.info(msg="               occupation of subblocks: {} %".format(num_total/(HK[0].shape[0]**2)*100))     
 
         subblocks = subblocks[1:]
-        show_blocks(subblocks,HK[0],self.results_path)
+
+        if plot_blocks:
+            show_blocks(subblocks,HK[0],self.results_path)
 
         return hd, hu, hl, sd, su, sl, subblocks