Fix block size handling in cache-driven BTD

dpnegf/negf/lead_property.py

@@ -1194,4 +1194,182 @@ def _has_saved_self_energy(root: str) -> bool:
         for pat in patterns:
             if any(p.rglob(pat)):
                 return True
-        return False
+        return False
+
+
+def _collect_h5_self_energy_size(h5_path: str, tab: str) -> int:
+    """Return the single invariant matrix size of an HDF5 self-energy cache.
+
+    The HDF5 layout written by :func:`write_to_hdf5` is a set of ``E_<energy>``
+    groups, each holding ``k_<kx>_<ky>_<kz>`` datasets of a square self-energy
+    matrix. This walks every (k, E) entry, requires every matrix to be square,
+    and requires a single invariant size across the whole file.
+
+    Parameters
+    ----------
+    h5_path : str
+        Path to ``self_energy_leadL.h5`` / ``self_energy_leadR.h5``.
+    tab : str
+        Lead tag used only for error messages.
+
+    Returns
+    -------
+    int
+        The invariant matrix size ``n`` (matrices are ``n x n``).
+    """
+    size = None
+    n_entries = 0
+    with h5py.File(h5_path, "r") as f:
+        for group_name in f.keys():
+            group = f[group_name]
+            if not isinstance(group, h5py.Group):
+                continue
+            for dset_name in group.keys():
+                shape = group[dset_name].shape
+                if len(shape) != 2 or shape[0] != shape[1]:
+                    raise ValueError(
+                        f"Self-energy cache {h5_path} contains a non-square matrix "
+                        f"of shape {tuple(shape)} at {group_name}/{dset_name} for {tab}."
+                    )
+                if size is None:
+                    size = int(shape[0])
+                elif int(shape[0]) != size:
+                    raise ValueError(
+                        f"Self-energy cache {h5_path} for {tab} has mixed matrix sizes "
+                        f"({size} vs {shape[0]}). A reusable cache must have a single "
+                        f"invariant size across all k-points and energies."
+                    )
+                n_entries += 1
+    if n_entries == 0 or size is None:
+        raise ValueError(
+            f"Self-energy cache {h5_path} for {tab} is empty (no saved matrices)."
+        )
+    return size
+
+
+def _collect_pth_self_energy_size(pth_paths: List[str], tab: str) -> int:
+    """Return the single invariant matrix size of a legacy PTH self-energy cache.
+
+    Each ``se_<tab>_k*_E*.pth`` file stores one square self-energy tensor. This
+    requires every file to hold a square matrix and a single invariant size
+    across all files for the lead.
+
+    Parameters
+    ----------
+    pth_paths : list[str]
+        Sorted list of PTH cache files for one lead.
+    tab : str
+        Lead tag used only for error messages.
+
+    Returns
+    -------
+    int
+        The invariant matrix size ``n`` (matrices are ``n x n``).
+    """
+    size = None
+    for path in pth_paths:
+        try:
+            se = torch.load(path, weights_only=False)
+        except Exception as e:  # pragma: no cover - IO error path
+            raise IOError(f"Failed to load self-energy cache {path} for {tab}.") from e
+        shape = tuple(se.shape)
+        if len(shape) != 2 or shape[0] != shape[1]:
+            raise ValueError(
+                f"Self-energy cache {path} for {tab} is not a square matrix "
+                f"(shape {shape})."
+            )
+        if size is None:
+            size = int(shape[0])
+        elif int(shape[0]) != size:
+            raise ValueError(
+                f"Legacy PTH self-energy cache for {tab} has mixed matrix sizes "
+                f"({size} vs {shape[0]}). A reusable cache must have a single "
+                f"invariant size across all k-points and energies."
+            )
+    if size is None:
+        raise ValueError(f"No PTH self-energy files found for {tab}.")
+    return size
+
+
+def inspect_self_energy_cache(save_path: str):
+    """Infer the required left/right device edge block sizes from a saved cache.
+
+    Reusing a cached lead self-energy across nearby scattering-region variants
+    only works if the block-tridiagonal (BTD) partition keeps the same first and
+    last device block sizes, because the lead self-energy shape is pinned to
+    those outer blocks (see :meth:`LeadProperty.HDL_reduced`). This inspects the
+    saved self-energy cache and returns the exact matrix sizes so the BTD
+    construction can preserve them.
+
+    Discovery order:
+
+    * a valid pair of standard HDF5 caches ``self_energy_leadL.h5`` /
+      ``self_energy_leadR.h5`` (preferred), else
+    * legacy per-(k, E) PTH files ``se_lead_L_k*_E*.pth`` /
+      ``se_lead_R_k*_E*.pth``.
+
+    If both cache styles are present, HDF5 is used and the choice is logged.
+
+    Parameters
+    ----------
+    save_path : str
+        Directory that holds the self-energy cache files.
+
+    Returns
+    -------
+    (nL, nR, fmt) : tuple[int, int, str]
+        Left and right invariant self-energy matrix sizes, and the detected
+        cache format (``"h5"`` or ``"pth"``).
+
+    Raises
+    ------
+    ValueError
+        If the cache is missing, empty, incomplete (only one lead), contains a
+        non-square matrix, or has mixed matrix shapes for a lead.
+    """
+    if save_path is None or not os.path.isdir(save_path):
+        raise ValueError(
+            f"Self-energy cache directory {save_path} does not exist; cannot infer "
+            f"cache-driven BTD edge block sizes."
+        )
+
+    h5_L = os.path.join(save_path, "self_energy_leadL.h5")
+    h5_R = os.path.join(save_path, "self_energy_leadR.h5")
+    has_h5 = os.path.isfile(h5_L) and os.path.isfile(h5_R)
+
+    pth_L = sorted(glob.glob(os.path.join(save_path, "se_lead_L_k*_E*.pth")))
+    pth_R = sorted(glob.glob(os.path.join(save_path, "se_lead_R_k*_E*.pth")))
+    has_pth = bool(pth_L) and bool(pth_R)
+
+    if has_h5:
+        if has_pth:
+            log.info(
+                "Both HDF5 and legacy PTH self-energy caches found in "
+                f"{save_path}; using HDF5 to infer BTD edge block sizes."
+            )
+        nL = _collect_h5_self_energy_size(h5_L, "lead_L")
+        nR = _collect_h5_self_energy_size(h5_R, "lead_R")
+        fmt = "h5"
+    elif has_pth:
+        nL = _collect_pth_self_energy_size(pth_L, "lead_L")
+        nR = _collect_pth_self_energy_size(pth_R, "lead_R")
+        fmt = "pth"
+    else:
+        # Distinguish partial caches (one lead only) from a fully absent cache.
+        only_h5 = os.path.isfile(h5_L) or os.path.isfile(h5_R)
+        only_pth = bool(pth_L) or bool(pth_R)
+        if only_h5 or only_pth:
+            raise ValueError(
+                f"Incomplete self-energy cache in {save_path}: both left and right "
+                f"lead caches are required to infer BTD edge block sizes."
+            )
+        raise ValueError(
+            f"No self-energy cache (self_energy_leadL.h5/self_energy_leadR.h5 or "
+            f"se_lead_L_k*_E*.pth/se_lead_R_k*_E*.pth) found in {save_path}."
+        )
+
+    log.info(
+        f"Inferred cache-driven BTD edge block sizes from {fmt} self-energy cache: "
+        f"leftmost={nL}, rightmost={nR}."
+    )
+    return nL, nR, fmt

dpnegf/negf/negf_hamiltonian_init.py

@@ -17,7 +17,9 @@
 from dpnegf.negf.bloch import Bloch
 from dpnegf.negf.sort_btd import sort_lexico, sort_projection, sort_capacitance
 from dpnegf.negf.split_btd import show_blocks,split_into_subblocks,split_into_subblocks_optimized
+from dpnegf.negf.split_btd import constrained_subblocks
 from dpnegf.negf.negf_utils import natsorted
+from dpnegf.negf.lead_property import inspect_self_energy_cache
 
 '''
 a Hamiltonian object  that initializes and manipulates device and  lead Hamiltonians for NEGF
@@ -59,8 +61,10 @@ def __init__(self,
                  pbc_negf: List[bool],
                  stru_options:dict, 
                  unit: str,
+                 use_saved_se: bool=False,
+                 self_energy_save_path: Optional[str]=None,
                  results_path:Optional[str]=None,
-                 torch_device: Union[str, torch.device]=torch.device('cpu')
+                 torch_device: Union[str, torch.device]=torch.device('cpu'),
                  ) -> None:
 
         # TODO: add dtype and device setting to the model
@@ -104,6 +108,10 @@ def __init__(self,
         self.results_path = results_path
         self.saved_HS_path = None
         self.subblocks = None
+        self.use_saved_se = use_saved_se
+        self.self_energy_save_path = self_energy_save_path
+        self.self_energy_cache_format = "h5"  # default to h5 format for self-energy cache
+        self._self_energy_cache_edge_sizes = None
 
         self.h2k = HR2HK(
             idp=model.idp, 
@@ -225,6 +233,17 @@ def initialize(self, kpoints, block_tridiagnal=False, plot_blocks=False, \
                 bloch_sorted_indices[kk],bloch_R_lists[kk] = self.get_lead_structure(kk,n_proj_atom_lead,\
                                 useBloch=useBloch,bloch_factor=bloch_factor) 
 
+        if self.use_saved_se:
+            if self.self_energy_save_path is None:
+                self.self_energy_save_path = os.path.join(self.results_path, "self_energy")
+            left_size, right_size, self.self_energy_cache_format = \
+                inspect_self_energy_cache(self.self_energy_save_path)
+            if block_tridiagnal and not use_saved_HS:
+                self._self_energy_cache_edge_sizes = (left_size, right_size)
+                log.info(msg="Cache-driven BTD active: using left/right block sizes "
+                        "{} and {} read directly from the {} self-energy cache."
+                        .format(left_size, right_size, self.self_energy_cache_format))
+
         # Hamiltonian initialization
         if use_saved_HS:
             if saved_HS_path is None:
@@ -477,6 +496,9 @@ def Hamiltonian_initialized(self,kpoints:List[List[float]],useBloch:bool,bloch_f
                                                                 plot_blocks=plot_blocks)
             HS_device.update({"hd":hd, "hu":hu, "hl":hl, "sd":sd, "su":su, "sl":sl, \
                               "subblocks":subblocks, "block_tridiagonal":True})
+            if self._self_energy_cache_edge_sizes is not None:
+                log.info(msg="Cache-driven BTD active: using left/right-most block sizes {} and {} from self-energy cache {}."
+                        .format(self._self_energy_cache_edge_sizes[0], self._self_energy_cache_edge_sizes[1], self.self_energy_cache_format))
 
         self.subblocks = subblocks
         # torch.save(HS_device, os.path.join(self.results_path, "HS_device.pth"))
@@ -697,6 +719,8 @@ def get_block_tridiagonal(self,HK,SK,structase:ase.Atoms,leftmost_size:int,right
             Number of orbitals in the leftmost block. If None, it is determined from the structure.
         rightmost_size : int or None
             Number of orbitals in the rightmost block. If None, it is determined from the structure.
+        plot_blocks : bool, optional
+            Whether to visualize the block structure using `show_blocks`. Default is False.
         Returns
         -------
         hd : list
@@ -733,11 +757,29 @@ def get_block_tridiagonal(self,HK,SK,structase:ase.Atoms,leftmost_size:int,right
             rightmost_atoms_index = np.where(structase.positions[:,2]==max(structase.positions[:,2]))[0]
             rightmost_size = sum([self.atom_norbs[rightmost_atoms_index[i]] for i in range(len(rightmost_atoms_index))])
 
-        subblocks = split_into_subblocks_optimized(HK[0],leftmost_size,rightmost_size)
-        if subblocks[0] < leftmost_size or subblocks[-1] < rightmost_size:
-            subblocks = split_into_subblocks(HK[0],leftmost_size,rightmost_size)
-            log.info(msg="The optimized block tridiagonalization is not successful, \
-                     the original block tridiagonalization is used.")
+        if self._self_energy_cache_edge_sizes is not None:
+            cache_leftmost_size, cache_rightmost_size = \
+                self._self_energy_cache_edge_sizes
+            if cache_leftmost_size < leftmost_size:
+                msg = (f"The cached left self-energy size {cache_leftmost_size} is smaller "
+                       f"than the lead_L coupling width {leftmost_size}; it cannot "
+                       f"contain the full device-lead coupling.")
+                log.error(msg=msg)
+                raise ValueError(msg)
+            if cache_rightmost_size < rightmost_size:
+                msg = (f"The cached right self-energy size {cache_rightmost_size} is smaller "
+                       f"than the lead_R coupling width {rightmost_size}; it cannot "
+                       f"contain the full device-lead coupling.")
+                log.error(msg=msg)
+                raise ValueError(msg)
+            subblocks = constrained_subblocks(
+                HK[0], SK[0], cache_leftmost_size, cache_rightmost_size)
+        else:
+            subblocks = split_into_subblocks_optimized(HK[0],leftmost_size,rightmost_size)
+            if subblocks[0] < leftmost_size or subblocks[-1] < rightmost_size:
+                subblocks = split_into_subblocks(HK[0],leftmost_size,rightmost_size)
+                log.info(msg="The optimized block tridiagonalization is not successful, \
+                         the original block tridiagonalization is used.")
         subblocks = [0]+subblocks