Add direct-BTD and Refactor Hamiltonian HDF5 dataflow

dpnegf/negf/lead_property.py

@@ -309,8 +309,17 @@ def self_energy_cal(self,
                     = self.hamiltonian.get_hs_lead(kpoint, tab=self.tab, v=self.voltage)
                 self.voltage_old = self.voltage
                 self.kpoint = torch.tensor(kpoint)
-
-            HDL_reduced, SDL_reduced = self.HDL_reduced(self.HDLk, self.SDLk,subblocks)
+
+            metadata = getattr(
+                self.hamiltonian, "contact_metadata", {}
+            ).get(self.tab, {})
+            HDL_reduced, SDL_reduced = self.HDL_reduced(
+                self.HDLk, 
+                self.SDLk,
+                subblocks,
+                tab = self.tab,
+                reduced=metadata.get("contact_reduced", False),
+            )
 
             self.se, _ = selfEnergy(
                 ee=energy,
@@ -385,7 +394,8 @@ def self_energy_cal(self,
         return self.se
 
     @staticmethod
-    def HDL_reduced(HDL: torch.Tensor, SDL: torch.Tensor, subblocks: np.ndarray) -> torch.Tensor:
+    def HDL_reduced(HDL: torch.Tensor, SDL: torch.Tensor, subblocks: np.ndarray,
+                    tab=None, reduced=False,) -> torch.Tensor:
         '''This function takes in Hamiltonian/Overlap matrix between lead and device and reduces 
         it based on the subblocks results or non-zero range of the Hamiltonian matrix.
 
@@ -408,7 +418,37 @@ def HDL_reduced(HDL: torch.Tensor, SDL: torch.Tensor, subblocks: np.ndarray) ->
         assert len(SDL.shape) == 2, "The shape of SDL should be 2."
         assert HDL.shape == SDL.shape, "The shape of HDL and SDL should be the same."
 
-        HDL_nonzero_range = (HDL.nonzero().min(dim=0).values, HDL.nonzero().max(dim=0).values)
+        if reduced:
+            if tab not in ("lead_L", "lead_R") or subblocks is None:
+                raise ValueError(
+                    "Reduced contacts require explicit lead identity and "
+                    "device blocks."
+                )
+            expected = (
+                subblocks[0] if tab == "lead_L" else subblocks[-1]
+            )
+            if HDL.shape[0] != expected:
+                raise ValueError(
+                    "Reduced contact height does not match its device boundary."
+                )
+            return HDL, SDL
+
+        if tab is not None and subblocks is not None:
+            if tab == "lead_L":
+                return HDL[:subblocks[0]], SDL[:subblocks[0]]
+            if tab == "lead_R":
+                return HDL[-subblocks[-1]:], SDL[-subblocks[-1]:]
+            raise ValueError("Unknown lead identity.")
+        nonzero = ((HDL.abs() + SDL.abs()) != 0).nonzero()
+        if not len(nonzero):
+            raise ValueError(
+                "Cannot infer lead identity from an empty legacy contact."
+            )
+        HDL_nonzero_range = (
+            nonzero.min(dim=0).values,
+            nonzero.max(dim=0).values,
+        )
+
         if subblocks is None:
             cut_range = HDL_nonzero_range
         else:
@@ -892,6 +932,13 @@ def _precompute_lead_kdata(lead, kpoints_grid):
                 k, tab=lead.tab, v=lead.voltage
             )
             pack["kdata"][key] = _pack_lead_matrices(HLk, HLLk, HDLk, SLk, SLLk, SDLk, subblocks)
+            metadata = getattr(
+                lead.hamiltonian, "contact_metadata", {}
+            ).get(lead.tab, {})
+            pack["kdata"][key].update(
+                contact_reduced=metadata.get("contact_reduced", False),
+                lead_tab=lead.tab
+            )
         else:
             kpoints_bloch = bloch_unfolder.unfold_points(list(np.asarray(k, dtype=float).reshape(3)))
             bloch_entries = []
@@ -964,7 +1011,11 @@ def _compute_self_energy_from_pack(pack, k, e, eta_lead, method="Lopez-Sancho",
 
     if not pack["useBloch"]:
         HDL_reduced, SDL_reduced = LeadProperty.HDL_reduced(
-            entry["HDLk"], entry["SDLk"], subblocks
+            entry["HDLk"], 
+            entry["SDLk"], 
+            subblocks,
+            tab=entry.get("lead_tab"),
+            reduced=entry.get("contact_reduced", False)
         )
         se, _ = selfEnergy(
             ee=energy,

dpnegf/negf/split_btd.py

@@ -662,7 +662,7 @@ def compute_edge(mat):
     return edge, edge1
 
 
-def compute_blocks(left_block, right_block, edge, edge1):
+def compute_blocks(left_block, right_block, edge, edge1, *, use_jit=True):
     """This is an implementation of the greedy algorithm for
      computing block-tridiagonal representation of a matrix.
      The information regarding the input matrix is represented
@@ -722,7 +722,12 @@ def compute_blocks(left_block, right_block, edge, edge1):
     # let the JIT core return the actual length.
     size = edge_arr.shape[0]
     out = np.empty(size + 2, dtype=np.int64)
-    n = _compute_blocks_into(left_block, right_block, edge_arr, edge1_arr, out, 0)
+    kernel = (
+        _compute_blocks_into
+        if use_jit
+        else getattr(_compute_blocks_into, "py_func", _compute_blocks_into)
+    )
+    n = kernel(left_block, right_block, edge_arr, edge1_arr, out, 0)
     return out[:n].tolist()
 
 

dpnegf/runner/NEGF.py

@@ -72,6 +72,8 @@ def __init__(self,
                 AtomicData_options: Optional[dict]=None,
                 energy_grid: Optional[dict]=None,
                 conductance_options: Optional[dict]=None,
+                btd_initialization: str="dense",
+                direct_btd_max_mib: Optional[float]=None,
                 **kwargs):
 
         # Use an explicitly supplied energy grid; otherwise preserve
@@ -239,7 +241,9 @@ def __init__(self,
                                                         results_path=self.results_path,
                                                         torch_device = torch.device("cpu"),
                                                         use_saved_se = self.use_saved_se,
-                                                        self_energy_save_path = self.self_energy_save_path)
+                                                        self_energy_save_path = self.self_energy_save_path,
+                                                        btd_initialization = btd_initialization,
+                                                        direct_btd_max_mib = direct_btd_max_mib,)
             # if useBloch is None, structure_leads_fold,bloch_sorted_indices,bloch_R_lists = None,None,None
             struct_device, struct_leads,structure_leads_fold,bloch_sorted_indices,bloch_R_lists = \
                 self.negf_hamiltonian.initialize(kpoints=self.kpoints,