From d9c0db1cd0d6c363c25f1fba2350423b70febc7b Mon Sep 17 00:00:00 2001 From: AsymmetryChou <181240085@smail.nju.edu.cn> Date: Sun, 26 Jul 2026 15:57:27 +0800 Subject: [PATCH 1/3] Feat: fix left-/right-most blocksize in BTD when using cache se --- dpnegf/negf/lead_property.py | 180 ++++++++++++++++++++++++- dpnegf/negf/negf_hamiltonian_init.py | 188 ++++++++++++++++++++++++++- dpnegf/runner/NEGF.py | 14 +- 3 files changed, 374 insertions(+), 8 deletions(-) diff --git a/dpnegf/negf/lead_property.py b/dpnegf/negf/lead_property.py index 22f509a..bf1513c 100644 --- a/dpnegf/negf/lead_property.py +++ b/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 \ No newline at end of file + 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_`` + groups, each holding ``k___`` 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__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 \ No newline at end of file diff --git a/dpnegf/negf/negf_hamiltonian_init.py b/dpnegf/negf/negf_hamiltonian_init.py index b064ffb..c81255b 100644 --- a/dpnegf/negf/negf_hamiltonian_init.py +++ b/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 compute_edge, compute_blocks 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 = self._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 @@ -776,6 +818,140 @@ def get_block_tridiagonal(self,HK,SK,structase:ase.Atoms,leftmost_size:int,right return hd, hu, hl, sd, su, sl, subblocks + def _constrained_subblocks(self, HK0, SK0, leftmost_size:int, rightmost_size:int): + """Partition the device into block-tridiagonal subblocks with fixed edges. + + Cache-driven BTD requires the first and last subblocks to keep exactly + ``leftmost_size`` / ``rightmost_size`` so the reused lead self-energy stays + dimension-compatible (see :meth:`LeadProperty.HDL_reduced`). Only the + interior is grown freely. + + The greedy splitter :func:`split_btd.compute_blocks` is run on the combined + ``|H| + |S|`` sparsity with the requested edge sizes as the leftmost / + rightmost blocks. Unlike the optimized splitter, it honours the requested + outer sizes *exactly* whenever they are compatible with the sparsity: it + pins ``blocks[0] == nL`` and, walking from the right, ``blocks[-1] == nR``. + When the requested size is too small to contain the block's coupling, the + greedy algorithm grows that edge block instead — which we detect and reject, + because the cached self-energy would then be dimension-incompatible. + + Parameters + ---------- + HK0, SK0 : torch.Tensor + Gamma-point Hamiltonian / overlap of the device, shape (D, D). + leftmost_size, rightmost_size : int + The exact required first / last block sizes. + + Returns + ------- + list[int] + Subblock sizes with ``subblocks[0] == leftmost_size`` and + ``subblocks[-1] == rightmost_size`` summing to ``D``. + + Raises + ------ + ValueError + If the requested edge sizes cannot form a valid block-tridiagonal + layout (too large for the device, no room for the interior, the greedy + splitter could not honour an edge size, or the couplings extend beyond + neighbouring blocks). + """ + D = int(HK0.shape[0]) + nL = int(leftmost_size) + nR = int(rightmost_size) + + if nL <= 0 or nR <= 0: + raise ValueError(f"Cache-driven BTD edge sizes must be positive, got leftmost={nL}, rightmost={nR}.") + if nL > D or nR > D: + raise ValueError(f"Cache-driven BTD edge sizes (leftmost={nL}, rightmost={nR}) exceed the device basis size {D}.") + + # Build the combined sparsity mask once; used for the greedy split and the + # final block-tridiagonality validation. + mask = (HK0.abs() + SK0.abs()) != 0 + + if nL + nR >= D: + # No interior room: the device is (at most) the two edge blocks. Valid + # only if they exactly tile the device (BTD is trivial for <=2 blocks). + if nL + nR != D: + raise ValueError( + f"Cache-driven BTD edge sizes (leftmost={nL}, rightmost={nR}) leave " + f"no room for the interior and do not tile the device basis {D}.") + subblocks = [nL, nR] + else: + edge, edge1 = compute_edge(mask.detach().cpu().numpy()) + subblocks = [int(b) for b in compute_blocks(nL, nR, edge, edge1)] + # The greedy splitter grows an edge block when the requested size cannot + # contain its coupling. A grown edge means the cached self-energy shape + # would not match the first/last Hamiltonian block, so reject it. + if subblocks[0] != nL: + raise ValueError( + f"Cache-driven leftmost block size {nL} is incompatible with the " + f"device sparsity: the block-tridiagonal splitter requires at least " + f"{subblocks[0]} orbitals in the first block. The cached self-energy " + f"cannot be reused for this device.") + if subblocks[-1] != nR: + raise ValueError( + f"Cache-driven rightmost block size {nR} is incompatible with the " + f"device sparsity: the block-tridiagonal splitter requires at least " + f"{subblocks[-1]} orbitals in the last block. The cached self-energy " + f"cannot be reused for this device.") + + subblocks = [int(b) for b in subblocks] + if sum(subblocks) != D: + raise ValueError( + f"Cache-driven BTD partition {subblocks} sums to {sum(subblocks)} but the " + f"device basis size is {D}.") + + if not self._validate_block_tridiagonal(mask, subblocks): + raise ValueError( + f"Cache-driven BTD partition with fixed edges (leftmost={nL}, " + f"rightmost={nR}) is not block-tridiagonal: the Hamiltonian/overlap " + f"couples non-neighbouring blocks. The cached self-energy edge sizes " + f"are incompatible with the current device sparsity.") + + return subblocks + + @staticmethod + def _validate_block_tridiagonal(mask, subblocks) -> bool: + """Check that ``mask`` is block-tridiagonal under the given partition. + + A partition is block-tridiagonal iff no nonzero entry couples blocks that + are more than one apart. Using the per-row sparsity profile (last nonzero + column per row, an O(D) quantity, cf. :func:`split_btd.compute_edge`), a + row belonging to block ``i`` must not reach past the end of block ``i+1``, + and by symmetry of the coupling structure this also bounds the lower part. + + Parameters + ---------- + mask : torch.Tensor + Boolean / 0-1 matrix marking nonzero entries of ``|H| + |S|``. + subblocks : list[int] + Candidate diagonal block sizes. + + Returns + ------- + bool + True if the layout is block-tridiagonal. + """ + mat = mask.detach().cpu().numpy() + edge, edge1 = compute_edge(mat) # edge[i]: 1 + last nonzero col in row i + bounds = np.cumsum([0] + list(subblocks)) # block boundaries + for bi in range(len(subblocks)): + r0, r1 = bounds[bi], bounds[bi + 1] + # Rightmost column any row in this block may reach: end of next block. + allowed = bounds[min(bi + 2, len(subblocks))] + if r1 > r0 and int(edge[r0:r1].max()) > allowed: + return False + # Lower part: leftmost column reachable is start of previous block. + # edge1 is the 180-degree-rotated profile; reuse the same bound from + # the other side to catch couplings below the sub-diagonal blocks. + allowed_lo = bounds[len(subblocks)] - bounds[max(bi - 1, 0)] + rr0 = bounds[len(subblocks)] - r1 + rr1 = bounds[len(subblocks)] - r0 + if rr1 > rr0 and int(edge1[rr0:rr1].max()) > allowed_lo: + return False + return True + def get_hs_device(self, kpoint=[0,0,0], V=None, block_tridiagonal=False, only_subblocks=False): """ get the device Hamiltonian and overlap matrix at a specific kpoint diff --git a/dpnegf/runner/NEGF.py b/dpnegf/runner/NEGF.py index 4faf305..c75cf9a 100644 --- a/dpnegf/runner/NEGF.py +++ b/dpnegf/runner/NEGF.py @@ -78,6 +78,9 @@ def __init__(self, se_cache = dict(se_opts.get("cache", {}) or {}) self.use_saved_se = se_cache.get("use_saved", False) self.self_energy_save_path = se_cache.get("save_path", None) + # Detected self-energy cache format + # defaults to HDF5 + self.se_cache_format = "h5" se_parallel = dict(se_opts.get("parallel", {}) or {}) self.cpu_budget = se_parallel.get("cpu_budget", None) self.n_workers = se_parallel.get("n_workers", -1) @@ -181,13 +184,19 @@ def __init__(self, stru_options=self.stru_options, unit = self.unit, results_path=self.results_path, - torch_device = torch.device("cpu")) + torch_device = torch.device("cpu"), + use_saved_se = self.use_saved_se, + self_energy_save_path = self.self_energy_save_path) # 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, block_tridiagnal=self.block_tridiagonal, plot_blocks=self.plot_blocks,\ useBloch=self.useBloch,bloch_factor=self.bloch_factor, use_saved_HS=self.use_saved_HS, saved_HS_path=self.saved_HS_path) + self.self_energy_save_path = \ + self.negf_hamiltonian.self_energy_save_path # update the self_energy_save_path in case it is None before + self.se_cache_format = \ + self.negf_hamiltonian.self_energy_cache_format # update the se_cache_format in case it is None before profiler.stop() output_path = os.path.join(self.results_path, "profile_report_ham_init.html") with open(output_path, 'w') as report_file: @@ -749,6 +758,7 @@ def negf_compute(self,scf_require=False,Vbias=None): eta_lead=self.eta_lead, method=self.sgf_solver, save_path=self.self_energy_save_path, + save_format=self.se_cache_format, se_info_display=self.se_info_display ) @@ -800,6 +810,7 @@ def negf_compute(self,scf_require=False,Vbias=None): eta_lead=self.eta_lead, method=self.sgf_solver, save_path=self.self_energy_save_path, + save_format=self.se_cache_format, se_info_display=self.se_info_display ) seL_list.append(self.deviceprop.lead_L.se) @@ -914,6 +925,7 @@ def negf_compute(self,scf_require=False,Vbias=None): eta_lead=self.eta_lead, method=self.sgf_solver, save_path=self.self_energy_save_path, + save_format=self.se_cache_format, se_info_display=self.se_info_display ) From 8ed4cbead1fc60ba84e4b2377b9c304eccd6fd48 Mon Sep 17 00:00:00 2001 From: AsymmetryChou <181240085@smail.nju.edu.cn> Date: Sun, 26 Jul 2026 15:57:46 +0800 Subject: [PATCH 2/3] Feat: add unit test for cache_driven BTD --- dpnegf/tests/test_cache_driven_btd.py | 354 ++++++++++++++++++++++++++ 1 file changed, 354 insertions(+) create mode 100644 dpnegf/tests/test_cache_driven_btd.py diff --git a/dpnegf/tests/test_cache_driven_btd.py b/dpnegf/tests/test_cache_driven_btd.py new file mode 100644 index 0000000..3ed2e05 --- /dev/null +++ b/dpnegf/tests/test_cache_driven_btd.py @@ -0,0 +1,354 @@ +"""Tests for cache-driven block-tridiagonal (BTD) edge matching. + +The optimized BTD splitter treats the leftmost/rightmost block sizes only as +lower bounds and re-optimizes cut points from the interior sparsity, so two +devices with an identical electrode region but a slightly different scattering +region could end up with different first/last block sizes. Because the lead +self-energy is pinned to those outer block sizes (see +``LeadProperty.HDL_reduced``), a cached self-energy from one run became +dimension-incompatible with the next. + +This module tests the two pieces of the fix: + +* ``inspect_self_energy_cache`` — infers the required first/last device block + sizes from a saved self-energy cache (HDF5 preferred, legacy PTH fallback). +* ``NEGFHamiltonianInit._constrained_subblocks`` / + ``_validate_block_tridiagonal`` — partition the device with the first/last + block sizes pinned exactly, growing only the interior, and reject layouts + incompatible with the device sparsity. +""" + +import os +import inspect + +import numpy as np +import pytest +import torch + +from dpnegf.negf.negf_hamiltonian_init import NEGFHamiltonianInit +from dpnegf.negf.lead_property import inspect_self_energy_cache, write_to_hdf5 +from dpnegf.negf.split_btd import ( + split_into_subblocks, + split_into_subblocks_optimized, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _make_btd_device(blocks, coupling=0.5, onsite=1.0): + """Build a Hamiltonian/overlap pair that is block-tridiagonal for ``blocks``. + + On-site terms populate the diagonal; nearest-block couplings populate the + super/sub block bands. The result is genuinely block-tridiagonal for the + given partition and its refinements. + """ + D = sum(blocks) + H = np.eye(D) * onsite + bnd = np.cumsum([0] + list(blocks)) + for bi in range(len(blocks) - 1): + r0, r1 = bnd[bi], bnd[bi + 1] + c0, c1 = bnd[bi + 1], bnd[bi + 2] + H[r0:r1, c0:c1] = coupling + H[c0:c1, r0:r1] = coupling + H = torch.tensor(H, dtype=torch.float64) + S = torch.eye(D, dtype=torch.float64) + return H, S + + +def _bare_ham(): + """A NEGFHamiltonianInit instance without running __init__. + + ``_constrained_subblocks`` only reaches ``self`` to call the static + ``_validate_block_tridiagonal``, so an uninitialized instance is enough to + exercise the partitioning logic without building a model. + """ + return object.__new__(NEGFHamiltonianInit) + + +def test_fixed_edge_sizes_are_not_public_constructor_arguments(): + parameters = inspect.signature(NEGFHamiltonianInit.__init__).parameters + assert "fixed_leftmost_size" not in parameters + assert "fixed_rightmost_size" not in parameters + + +def test_fresh_btd_uses_historical_splitter_path(): + ham = _bare_ham() + ham._self_energy_cache_edge_sizes = None + ham.results_path = None + H, S = _make_btd_device([3, 2, 2, 3]) + HK = H.unsqueeze(0) + SK = S.unsqueeze(0) + + expected = split_into_subblocks_optimized(HK[0], 3, 3) + if expected[0] < 3 or expected[-1] < 3: + expected = split_into_subblocks(HK[0], 3, 3) + + *_, actual = ham.get_block_tridiagonal( + HK, SK, structase=None, leftmost_size=3, rightmost_size=3) + assert actual == expected + + +def _write_h5_cache(path, tab, size, ks, es): + """Write a valid HDF5 self-energy cache with a single invariant size.""" + for k in ks: + for e in es: + se = torch.eye(size, dtype=torch.complex128) * (e + 1j) + write_to_hdf5(path, k, e, se) + + +# --------------------------------------------------------------------------- +# inspect_self_energy_cache +# --------------------------------------------------------------------------- +def test_inspect_cache_valid_h5(tmp_path): + d = str(tmp_path) + ks = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.5]] + es = [-1.0, 0.0, 1.0] + _write_h5_cache(os.path.join(d, "self_energy_leadL.h5"), "lead_L", 4, ks, es) + _write_h5_cache(os.path.join(d, "self_energy_leadR.h5"), "lead_R", 6, ks, es) + + nL, nR, fmt = inspect_self_energy_cache(d) + assert (nL, nR, fmt) == (4, 6, "h5") + + +def test_inspect_cache_valid_pth(tmp_path): + d = str(tmp_path) + torch.save(torch.eye(3, dtype=torch.complex128), + os.path.join(d, "se_lead_L_k0.0000_0.0000_0.0000_E0.000000.pth")) + torch.save(torch.eye(3, dtype=torch.complex128), + os.path.join(d, "se_lead_L_k0.0000_0.0000_0.0000_E1.000000.pth")) + torch.save(torch.eye(5, dtype=torch.complex128), + os.path.join(d, "se_lead_R_k0.0000_0.0000_0.0000_E0.000000.pth")) + + nL, nR, fmt = inspect_self_energy_cache(d) + assert (nL, nR, fmt) == (3, 5, "pth") + + +def test_inspect_cache_h5_precedence(tmp_path): + d = str(tmp_path) + ks = [[0.0, 0.0, 0.0]] + es = [0.0] + # HDF5 says (4, 6); PTH says (3, 5). HDF5 must win. + _write_h5_cache(os.path.join(d, "self_energy_leadL.h5"), "lead_L", 4, ks, es) + _write_h5_cache(os.path.join(d, "self_energy_leadR.h5"), "lead_R", 6, ks, es) + torch.save(torch.eye(3, dtype=torch.complex128), + os.path.join(d, "se_lead_L_k0.0000_0.0000_0.0000_E0.000000.pth")) + torch.save(torch.eye(5, dtype=torch.complex128), + os.path.join(d, "se_lead_R_k0.0000_0.0000_0.0000_E0.000000.pth")) + + nL, nR, fmt = inspect_self_energy_cache(d) + assert (nL, nR, fmt) == (4, 6, "h5") + + +def test_inspect_cache_missing_lead(tmp_path): + d = str(tmp_path) + _write_h5_cache(os.path.join(d, "self_energy_leadL.h5"), "lead_L", 4, + [[0.0, 0.0, 0.0]], [0.0]) + # Only left lead present. + with pytest.raises(ValueError, match="[Ii]ncomplete"): + inspect_self_energy_cache(d) + + +def test_inspect_cache_absent(tmp_path): + with pytest.raises(ValueError, match="No self-energy cache"): + inspect_self_energy_cache(str(tmp_path)) + + +def test_inspect_cache_missing_dir(): + with pytest.raises(ValueError, match="does not exist"): + inspect_self_energy_cache(os.path.join(os.sep, "no", "such", "dir_xyz")) + + +def test_inspect_cache_empty_h5(tmp_path): + d = str(tmp_path) + import h5py + # Two present but empty HDF5 files. + h5py.File(os.path.join(d, "self_energy_leadL.h5"), "w").close() + h5py.File(os.path.join(d, "self_energy_leadR.h5"), "w").close() + with pytest.raises(ValueError, match="empty"): + inspect_self_energy_cache(d) + + +def test_inspect_cache_non_square(tmp_path): + d = str(tmp_path) + import h5py + with h5py.File(os.path.join(d, "self_energy_leadL.h5"), "w") as f: + g = f.require_group("E_0.00000000") + g.create_dataset("k_0.0_0.0_0.0", data=np.zeros((3, 4), dtype=np.complex128)) + _write_h5_cache(os.path.join(d, "self_energy_leadR.h5"), "lead_R", 5, + [[0.0, 0.0, 0.0]], [0.0]) + with pytest.raises(ValueError, match="non-square"): + inspect_self_energy_cache(d) + + +def test_inspect_cache_mixed_shape(tmp_path): + d = str(tmp_path) + import h5py + with h5py.File(os.path.join(d, "self_energy_leadL.h5"), "w") as f: + f.require_group("E_0.00000000").create_dataset( + "k_0.0_0.0_0.0", data=np.eye(4, dtype=np.complex128)) + f.require_group("E_1.00000000").create_dataset( + "k_0.0_0.0_0.0", data=np.eye(5, dtype=np.complex128)) # size drift + _write_h5_cache(os.path.join(d, "self_energy_leadR.h5"), "lead_R", 6, + [[0.0, 0.0, 0.0]], [0.0]) + with pytest.raises(ValueError, match="mixed matrix sizes"): + inspect_self_energy_cache(d) + + +# --------------------------------------------------------------------------- +# _constrained_subblocks / _validate_block_tridiagonal +# --------------------------------------------------------------------------- +def test_constrained_endpoints_preserved(): + ham = _bare_ham() + H, S = _make_btd_device([3, 2, 2, 3]) + sub = ham._constrained_subblocks(H, S, 3, 3) + assert sub[0] == 3 and sub[-1] == 3 + assert sum(sub) == H.shape[0] + + +def test_constrained_endpoints_stable_under_interior_change(): + """Same electrodes (nL=nR=3), different scattering interiors, same endpoints. + + This is the whole point of the fix: cached self-energy edge sizes stay + reusable across scattering-region variants. + """ + ham = _bare_ham() + for interior in ([2, 2], [5], [2, 4], [3, 3, 2]): + H, S = _make_btd_device([3] + interior + [3]) + sub = ham._constrained_subblocks(H, S, 3, 3) + assert sub[0] == 3, (interior, sub) + assert sub[-1] == 3, (interior, sub) + assert sum(sub) == H.shape[0] + assert NEGFHamiltonianInit._validate_block_tridiagonal( + (H.abs() + S.abs()) != 0, sub) + + +def test_constrained_full_coverage_and_tridiagonal(): + ham = _bare_ham() + H, S = _make_btd_device([4, 3, 3, 4]) + sub = ham._constrained_subblocks(H, S, 4, 4) + assert sum(sub) == H.shape[0] + mask = (H.abs() + S.abs()) != 0 + assert NEGFHamiltonianInit._validate_block_tridiagonal(mask, sub) + + +def test_constrained_two_block_exact_tiling(): + ham = _bare_ham() + H, S = _make_btd_device([5, 5]) + sub = ham._constrained_subblocks(H, S, 5, 5) + assert sub == [5, 5] + + +def test_validate_rejects_non_tridiagonal_partition(): + # [3,2,2,3] device: partition [2,2,2,2,2] splits the first natural block and + # makes block 0 couple block 2 -> not block-tridiagonal. + H, S = _make_btd_device([3, 2, 2, 3]) + mask = (H.abs() + S.abs()) != 0 + assert not NEGFHamiltonianInit._validate_block_tridiagonal(mask, [2, 2, 2, 2, 2]) + assert NEGFHamiltonianInit._validate_block_tridiagonal(mask, [3, 2, 2, 3]) + + +# --------------------------------------------------------------------------- +# Expected failures +# --------------------------------------------------------------------------- +def test_constrained_edge_exceeds_device(): + ham = _bare_ham() + H, S = _make_btd_device([3, 2, 3]) + with pytest.raises(ValueError, match="exceed"): + ham._constrained_subblocks(H, S, H.shape[0] + 1, 3) + + +def test_constrained_non_positive_edge(): + ham = _bare_ham() + H, S = _make_btd_device([3, 2, 3]) + with pytest.raises(ValueError, match="positive"): + ham._constrained_subblocks(H, S, 0, 3) + + +def test_constrained_no_interior_no_tiling(): + ham = _bare_ham() + H, S = _make_btd_device([3, 2, 3]) # D = 8 + # 6 + 6 > 8 but 6 + 6 != 8 -> cannot tile. + with pytest.raises(ValueError, match="no room|tile"): + ham._constrained_subblocks(H, S, 6, 6) + + +def test_constrained_edge_too_small_for_coupling(): + """A [5,5] device cannot honour nL=3: the coupling forces a larger first block. + + The greedy splitter would grow the edge to 7, which the constrained path + must reject because the cached (3x3) self-energy would not fit. + """ + ham = _bare_ham() + H, S = _make_btd_device([5, 5]) + with pytest.raises(ValueError, match="incompatible"): + ham._constrained_subblocks(H, S, 3, 3) + + +# --------------------------------------------------------------------------- +# End-to-end integration: run once to populate the cache, then reuse it with +# cache-driven BTD and confirm the endpoint blocks match the cached self-energy +# and no self-energy is recomputed. +# --------------------------------------------------------------------------- +@pytest.fixture(scope='session') +def root_directory(request): + return str(request.config.rootdir) + + +def _read_subblocks(results_path): + import h5py + with h5py.File(os.path.join(results_path, "HS_device.h5"), "r") as f: + return [int(x) for x in np.array(f["subblocks"])] + + +def test_cache_driven_btd_end_to_end(root_directory, tmp_path): + import json + from dpnegf.entrypoints.run import run + + data = os.path.join(root_directory, "dpnegf/tests/data/test_negf/test_negf_run") + ckpt = os.path.join(data, "nnsk_C_new.json") + stru = os.path.join(data, "chain.vasp") + base = json.load(open(os.path.join(data, "negf_chain_new.json"))) + + # Block-tridiagonal, tiny grid, single SE worker for determinism. + base["task_options"]["block_tridiagonal"] = True + base["task_options"]["emin"] = -0.05 + base["task_options"]["emax"] = 0.05 + base["task_options"]["espacing"] = 0.05 + base["task_options"]["n_cpus"] = 1 + + # --- run 1: fresh, populate the self-energy cache --- + in1 = str(tmp_path / "run1.json"); json.dump(base, open(in1, "w")) + out1 = str(tmp_path / "out1") + run(INPUT=in1, init_model=ckpt, structure=stru, output=out1, + log_level=20, log_path=os.path.join(out1, "log.txt")) + + res1 = os.path.join(out1, "results") + se_dir = os.path.join(res1, "self_energy") + assert os.path.isfile(os.path.join(se_dir, "self_energy_leadL.h5")) + assert os.path.isfile(os.path.join(se_dir, "self_energy_leadR.h5")) + + nL, nR, fmt = inspect_self_energy_cache(se_dir) + assert fmt == "h5" + sub1 = _read_subblocks(res1) + assert sub1[0] == nL and sub1[-1] == nR + + # --- run 2: reuse the cache; endpoints must match the cached SE dims --- + base2 = json.loads(json.dumps(base)) + base2["task_options"]["use_saved_se"] = True + base2["task_options"]["self_energy_save_path"] = se_dir + in2 = str(tmp_path / "run2.json"); json.dump(base2, open(in2, "w")) + out2 = str(tmp_path / "out2") + run(INPUT=in2, init_model=ckpt, structure=stru, output=out2, + log_level=20, log_path=os.path.join(out2, "log.txt")) + + res2 = os.path.join(out2, "results") + sub2 = _read_subblocks(res2) + # The core guarantee: cache-driven endpoints equal the cached self-energy size. + assert sub2[0] == nL and sub2[-1] == nR + assert os.path.exists(os.path.join(res2, "negf.out.pth")) + + log2 = open(os.path.join(out2, "log.txt")).read() + assert "Cache-driven BTD active" in log2 + assert "Using saved self-energy" in log2 + From cdd25a687e38488759c947f73514578be9b2aaef Mon Sep 17 00:00:00 2001 From: AsymmetryChou <181240085@smail.nju.edu.cn> Date: Sun, 26 Jul 2026 16:20:05 +0800 Subject: [PATCH 3/3] Refac: move constrained_subblocks to split_btd.py --- dpnegf/negf/negf_hamiltonian_init.py | 138 +------------------------- dpnegf/negf/split_btd.py | 135 +++++++++++++++++++++++++ dpnegf/tests/test_cache_driven_btd.py | 45 ++++----- 3 files changed, 156 insertions(+), 162 deletions(-) diff --git a/dpnegf/negf/negf_hamiltonian_init.py b/dpnegf/negf/negf_hamiltonian_init.py index c81255b..58d308d 100644 --- a/dpnegf/negf/negf_hamiltonian_init.py +++ b/dpnegf/negf/negf_hamiltonian_init.py @@ -17,7 +17,7 @@ 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 compute_edge, compute_blocks +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 @@ -772,7 +772,7 @@ def get_block_tridiagonal(self,HK,SK,structase:ase.Atoms,leftmost_size:int,right f"contain the full device-lead coupling.") log.error(msg=msg) raise ValueError(msg) - subblocks = self._constrained_subblocks( + 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) @@ -818,140 +818,6 @@ def get_block_tridiagonal(self,HK,SK,structase:ase.Atoms,leftmost_size:int,right return hd, hu, hl, sd, su, sl, subblocks - def _constrained_subblocks(self, HK0, SK0, leftmost_size:int, rightmost_size:int): - """Partition the device into block-tridiagonal subblocks with fixed edges. - - Cache-driven BTD requires the first and last subblocks to keep exactly - ``leftmost_size`` / ``rightmost_size`` so the reused lead self-energy stays - dimension-compatible (see :meth:`LeadProperty.HDL_reduced`). Only the - interior is grown freely. - - The greedy splitter :func:`split_btd.compute_blocks` is run on the combined - ``|H| + |S|`` sparsity with the requested edge sizes as the leftmost / - rightmost blocks. Unlike the optimized splitter, it honours the requested - outer sizes *exactly* whenever they are compatible with the sparsity: it - pins ``blocks[0] == nL`` and, walking from the right, ``blocks[-1] == nR``. - When the requested size is too small to contain the block's coupling, the - greedy algorithm grows that edge block instead — which we detect and reject, - because the cached self-energy would then be dimension-incompatible. - - Parameters - ---------- - HK0, SK0 : torch.Tensor - Gamma-point Hamiltonian / overlap of the device, shape (D, D). - leftmost_size, rightmost_size : int - The exact required first / last block sizes. - - Returns - ------- - list[int] - Subblock sizes with ``subblocks[0] == leftmost_size`` and - ``subblocks[-1] == rightmost_size`` summing to ``D``. - - Raises - ------ - ValueError - If the requested edge sizes cannot form a valid block-tridiagonal - layout (too large for the device, no room for the interior, the greedy - splitter could not honour an edge size, or the couplings extend beyond - neighbouring blocks). - """ - D = int(HK0.shape[0]) - nL = int(leftmost_size) - nR = int(rightmost_size) - - if nL <= 0 or nR <= 0: - raise ValueError(f"Cache-driven BTD edge sizes must be positive, got leftmost={nL}, rightmost={nR}.") - if nL > D or nR > D: - raise ValueError(f"Cache-driven BTD edge sizes (leftmost={nL}, rightmost={nR}) exceed the device basis size {D}.") - - # Build the combined sparsity mask once; used for the greedy split and the - # final block-tridiagonality validation. - mask = (HK0.abs() + SK0.abs()) != 0 - - if nL + nR >= D: - # No interior room: the device is (at most) the two edge blocks. Valid - # only if they exactly tile the device (BTD is trivial for <=2 blocks). - if nL + nR != D: - raise ValueError( - f"Cache-driven BTD edge sizes (leftmost={nL}, rightmost={nR}) leave " - f"no room for the interior and do not tile the device basis {D}.") - subblocks = [nL, nR] - else: - edge, edge1 = compute_edge(mask.detach().cpu().numpy()) - subblocks = [int(b) for b in compute_blocks(nL, nR, edge, edge1)] - # The greedy splitter grows an edge block when the requested size cannot - # contain its coupling. A grown edge means the cached self-energy shape - # would not match the first/last Hamiltonian block, so reject it. - if subblocks[0] != nL: - raise ValueError( - f"Cache-driven leftmost block size {nL} is incompatible with the " - f"device sparsity: the block-tridiagonal splitter requires at least " - f"{subblocks[0]} orbitals in the first block. The cached self-energy " - f"cannot be reused for this device.") - if subblocks[-1] != nR: - raise ValueError( - f"Cache-driven rightmost block size {nR} is incompatible with the " - f"device sparsity: the block-tridiagonal splitter requires at least " - f"{subblocks[-1]} orbitals in the last block. The cached self-energy " - f"cannot be reused for this device.") - - subblocks = [int(b) for b in subblocks] - if sum(subblocks) != D: - raise ValueError( - f"Cache-driven BTD partition {subblocks} sums to {sum(subblocks)} but the " - f"device basis size is {D}.") - - if not self._validate_block_tridiagonal(mask, subblocks): - raise ValueError( - f"Cache-driven BTD partition with fixed edges (leftmost={nL}, " - f"rightmost={nR}) is not block-tridiagonal: the Hamiltonian/overlap " - f"couples non-neighbouring blocks. The cached self-energy edge sizes " - f"are incompatible with the current device sparsity.") - - return subblocks - - @staticmethod - def _validate_block_tridiagonal(mask, subblocks) -> bool: - """Check that ``mask`` is block-tridiagonal under the given partition. - - A partition is block-tridiagonal iff no nonzero entry couples blocks that - are more than one apart. Using the per-row sparsity profile (last nonzero - column per row, an O(D) quantity, cf. :func:`split_btd.compute_edge`), a - row belonging to block ``i`` must not reach past the end of block ``i+1``, - and by symmetry of the coupling structure this also bounds the lower part. - - Parameters - ---------- - mask : torch.Tensor - Boolean / 0-1 matrix marking nonzero entries of ``|H| + |S|``. - subblocks : list[int] - Candidate diagonal block sizes. - - Returns - ------- - bool - True if the layout is block-tridiagonal. - """ - mat = mask.detach().cpu().numpy() - edge, edge1 = compute_edge(mat) # edge[i]: 1 + last nonzero col in row i - bounds = np.cumsum([0] + list(subblocks)) # block boundaries - for bi in range(len(subblocks)): - r0, r1 = bounds[bi], bounds[bi + 1] - # Rightmost column any row in this block may reach: end of next block. - allowed = bounds[min(bi + 2, len(subblocks))] - if r1 > r0 and int(edge[r0:r1].max()) > allowed: - return False - # Lower part: leftmost column reachable is start of previous block. - # edge1 is the 180-degree-rotated profile; reuse the same bound from - # the other side to catch couplings below the sub-diagonal blocks. - allowed_lo = bounds[len(subblocks)] - bounds[max(bi - 1, 0)] - rr0 = bounds[len(subblocks)] - r1 - rr1 = bounds[len(subblocks)] - r0 - if rr1 > rr0 and int(edge1[rr0:rr1].max()) > allowed_lo: - return False - return True - def get_hs_device(self, kpoint=[0,0,0], V=None, block_tridiagonal=False, only_subblocks=False): """ get the device Hamiltonian and overlap matrix at a specific kpoint diff --git a/dpnegf/negf/split_btd.py b/dpnegf/negf/split_btd.py index fc4d909..179f245 100644 --- a/dpnegf/negf/split_btd.py +++ b/dpnegf/negf/split_btd.py @@ -726,6 +726,141 @@ def compute_blocks(left_block, right_block, edge, edge1): return out[:n].tolist() +def constrained_subblocks(HK0, SK0, leftmost_size:int, rightmost_size:int): + """Partition the device into block-tridiagonal subblocks with fixed edges. + + Cache-driven BTD requires the first and last subblocks to keep exactly + ``leftmost_size`` / ``rightmost_size`` so the reused lead self-energy stays + dimension-compatible (see :meth:`LeadProperty.HDL_reduced`). Only the + interior is grown freely. + + The greedy splitter :func:`compute_blocks` is run on the combined + ``|H| + |S|`` sparsity with the requested edge sizes as the leftmost / + rightmost blocks. Unlike the optimized splitter, it honours the requested + outer sizes *exactly* whenever they are compatible with the sparsity: it + pins ``blocks[0] == nL`` and, walking from the right, ``blocks[-1] == nR``. + When the requested size is too small to contain the block's coupling, the + greedy algorithm grows that edge block instead — which we detect and reject, + because the cached self-energy would then be dimension-incompatible. + + Parameters + ---------- + HK0, SK0 : torch.Tensor + Gamma-point Hamiltonian / overlap of the device, shape (D, D). + leftmost_size, rightmost_size : int + The exact required first / last block sizes. + + Returns + ------- + list[int] + Subblock sizes with ``subblocks[0] == leftmost_size`` and + ``subblocks[-1] == rightmost_size`` summing to ``D``. + + Raises + ------ + ValueError + If the requested edge sizes cannot form a valid block-tridiagonal + layout (too large for the device, no room for the interior, the greedy + splitter could not honour an edge size, or the couplings extend beyond + neighbouring blocks). + """ + D = int(HK0.shape[0]) + nL = int(leftmost_size) + nR = int(rightmost_size) + + if nL <= 0 or nR <= 0: + raise ValueError(f"Cache-driven BTD edge sizes must be positive, got leftmost={nL}, rightmost={nR}.") + if nL > D or nR > D: + raise ValueError(f"Cache-driven BTD edge sizes (leftmost={nL}, rightmost={nR}) exceed the device basis size {D}.") + + # Build the combined sparsity mask once; used for the greedy split and the + # final block-tridiagonality validation. + mask = (HK0.abs() + SK0.abs()) != 0 + + if nL + nR >= D: + # No interior room: the device is (at most) the two edge blocks. Valid + # only if they exactly tile the device (BTD is trivial for <=2 blocks). + if nL + nR != D: + raise ValueError( + f"Cache-driven BTD edge sizes (leftmost={nL}, rightmost={nR}) leave " + f"no room for the interior and do not tile the device basis {D}.") + subblocks = [nL, nR] + else: + edge, edge1 = compute_edge(mask.detach().cpu().numpy()) + subblocks = [int(b) for b in compute_blocks(nL, nR, edge, edge1)] + # The greedy splitter grows an edge block when the requested size cannot + # contain its coupling. A grown edge means the cached self-energy shape + # would not match the first/last Hamiltonian block, so reject it. + if subblocks[0] != nL: + raise ValueError( + f"Cache-driven leftmost block size {nL} is incompatible with the " + f"device sparsity: the block-tridiagonal splitter requires at least " + f"{subblocks[0]} orbitals in the first block. The cached self-energy " + f"cannot be reused for this device.") + if subblocks[-1] != nR: + raise ValueError( + f"Cache-driven rightmost block size {nR} is incompatible with the " + f"device sparsity: the block-tridiagonal splitter requires at least " + f"{subblocks[-1]} orbitals in the last block. The cached self-energy " + f"cannot be reused for this device.") + + subblocks = [int(b) for b in subblocks] + if sum(subblocks) != D: + raise ValueError( + f"Cache-driven BTD partition {subblocks} sums to {sum(subblocks)} but the " + f"device basis size is {D}.") + + if not validate_block_tridiagonal(mask, subblocks): + raise ValueError( + f"Cache-driven BTD partition with fixed edges (leftmost={nL}, " + f"rightmost={nR}) is not block-tridiagonal: the Hamiltonian/overlap " + f"couples non-neighbouring blocks. The cached self-energy edge sizes " + f"are incompatible with the current device sparsity.") + + return subblocks + + +def validate_block_tridiagonal(mask, subblocks) -> bool: + """Check that ``mask`` is block-tridiagonal under the given partition. + + A partition is block-tridiagonal iff no nonzero entry couples blocks that + are more than one apart. Using the per-row sparsity profile (last nonzero + column per row, an O(D) quantity, cf. :func:`compute_edge`), a + row belonging to block ``i`` must not reach past the end of block ``i+1``, + and by symmetry of the coupling structure this also bounds the lower part. + + Parameters + ---------- + mask : torch.Tensor + Boolean / 0-1 matrix marking nonzero entries of ``|H| + |S|``. + subblocks : list[int] + Candidate diagonal block sizes. + + Returns + ------- + bool + True if the layout is block-tridiagonal. + """ + mat = mask.detach().cpu().numpy() + edge, edge1 = compute_edge(mat) # edge[i]: 1 + last nonzero col in row i + bounds = np.cumsum([0] + list(subblocks)) # block boundaries + for bi in range(len(subblocks)): + r0, r1 = bounds[bi], bounds[bi + 1] + # Rightmost column any row in this block may reach: end of next block. + allowed = bounds[min(bi + 2, len(subblocks))] + if r1 > r0 and int(edge[r0:r1].max()) > allowed: + return False + # Lower part: leftmost column reachable is start of previous block. + # edge1 is the 180-degree-rotated profile; reuse the same bound from + # the other side to catch couplings below the sub-diagonal blocks. + allowed_lo = bounds[len(subblocks)] - bounds[max(bi - 1, 0)] + rr0 = bounds[len(subblocks)] - r1 + rr1 = bounds[len(subblocks)] - r0 + if rr1 > rr0 and int(edge1[rr0:rr1].max()) > allowed_lo: + return False + return True + + def show_blocks(subblocks, input_mat, results_path): """This is a script for visualizing the sparsity pattern and a block-tridiagonal structure of a matrix. diff --git a/dpnegf/tests/test_cache_driven_btd.py b/dpnegf/tests/test_cache_driven_btd.py index 3ed2e05..f4fbaaf 100644 --- a/dpnegf/tests/test_cache_driven_btd.py +++ b/dpnegf/tests/test_cache_driven_btd.py @@ -12,8 +12,8 @@ * ``inspect_self_energy_cache`` — infers the required first/last device block sizes from a saved self-energy cache (HDF5 preferred, legacy PTH fallback). -* ``NEGFHamiltonianInit._constrained_subblocks`` / - ``_validate_block_tridiagonal`` — partition the device with the first/last +* ``dpnegf.negf.split_btd.constrained_subblocks`` / + ``validate_block_tridiagonal`` — partition the device with the first/last block sizes pinned exactly, growing only the interior, and reject layouts incompatible with the device sparsity. """ @@ -30,6 +30,8 @@ from dpnegf.negf.split_btd import ( split_into_subblocks, split_into_subblocks_optimized, + constrained_subblocks, + validate_block_tridiagonal, ) @@ -59,9 +61,8 @@ def _make_btd_device(blocks, coupling=0.5, onsite=1.0): def _bare_ham(): """A NEGFHamiltonianInit instance without running __init__. - ``_constrained_subblocks`` only reaches ``self`` to call the static - ``_validate_block_tridiagonal``, so an uninitialized instance is enough to - exercise the partitioning logic without building a model. + Used to exercise ``get_block_tridiagonal`` on synthetic Hamiltonians + without building a model. """ return object.__new__(NEGFHamiltonianInit) @@ -196,12 +197,11 @@ def test_inspect_cache_mixed_shape(tmp_path): # --------------------------------------------------------------------------- -# _constrained_subblocks / _validate_block_tridiagonal +# constrained_subblocks / validate_block_tridiagonal # --------------------------------------------------------------------------- def test_constrained_endpoints_preserved(): - ham = _bare_ham() H, S = _make_btd_device([3, 2, 2, 3]) - sub = ham._constrained_subblocks(H, S, 3, 3) + sub = constrained_subblocks(H, S, 3, 3) assert sub[0] == 3 and sub[-1] == 3 assert sum(sub) == H.shape[0] @@ -212,30 +212,27 @@ def test_constrained_endpoints_stable_under_interior_change(): This is the whole point of the fix: cached self-energy edge sizes stay reusable across scattering-region variants. """ - ham = _bare_ham() for interior in ([2, 2], [5], [2, 4], [3, 3, 2]): H, S = _make_btd_device([3] + interior + [3]) - sub = ham._constrained_subblocks(H, S, 3, 3) + sub = constrained_subblocks(H, S, 3, 3) assert sub[0] == 3, (interior, sub) assert sub[-1] == 3, (interior, sub) assert sum(sub) == H.shape[0] - assert NEGFHamiltonianInit._validate_block_tridiagonal( + assert validate_block_tridiagonal( (H.abs() + S.abs()) != 0, sub) def test_constrained_full_coverage_and_tridiagonal(): - ham = _bare_ham() H, S = _make_btd_device([4, 3, 3, 4]) - sub = ham._constrained_subblocks(H, S, 4, 4) + sub = constrained_subblocks(H, S, 4, 4) assert sum(sub) == H.shape[0] mask = (H.abs() + S.abs()) != 0 - assert NEGFHamiltonianInit._validate_block_tridiagonal(mask, sub) + assert validate_block_tridiagonal(mask, sub) def test_constrained_two_block_exact_tiling(): - ham = _bare_ham() H, S = _make_btd_device([5, 5]) - sub = ham._constrained_subblocks(H, S, 5, 5) + sub = constrained_subblocks(H, S, 5, 5) assert sub == [5, 5] @@ -244,33 +241,30 @@ def test_validate_rejects_non_tridiagonal_partition(): # makes block 0 couple block 2 -> not block-tridiagonal. H, S = _make_btd_device([3, 2, 2, 3]) mask = (H.abs() + S.abs()) != 0 - assert not NEGFHamiltonianInit._validate_block_tridiagonal(mask, [2, 2, 2, 2, 2]) - assert NEGFHamiltonianInit._validate_block_tridiagonal(mask, [3, 2, 2, 3]) + assert not validate_block_tridiagonal(mask, [2, 2, 2, 2, 2]) + assert validate_block_tridiagonal(mask, [3, 2, 2, 3]) # --------------------------------------------------------------------------- # Expected failures # --------------------------------------------------------------------------- def test_constrained_edge_exceeds_device(): - ham = _bare_ham() H, S = _make_btd_device([3, 2, 3]) with pytest.raises(ValueError, match="exceed"): - ham._constrained_subblocks(H, S, H.shape[0] + 1, 3) + constrained_subblocks(H, S, H.shape[0] + 1, 3) def test_constrained_non_positive_edge(): - ham = _bare_ham() H, S = _make_btd_device([3, 2, 3]) with pytest.raises(ValueError, match="positive"): - ham._constrained_subblocks(H, S, 0, 3) + constrained_subblocks(H, S, 0, 3) def test_constrained_no_interior_no_tiling(): - ham = _bare_ham() H, S = _make_btd_device([3, 2, 3]) # D = 8 # 6 + 6 > 8 but 6 + 6 != 8 -> cannot tile. with pytest.raises(ValueError, match="no room|tile"): - ham._constrained_subblocks(H, S, 6, 6) + constrained_subblocks(H, S, 6, 6) def test_constrained_edge_too_small_for_coupling(): @@ -279,10 +273,9 @@ def test_constrained_edge_too_small_for_coupling(): The greedy splitter would grow the edge to 7, which the constrained path must reject because the cached (3x3) self-energy would not fit. """ - ham = _bare_ham() H, S = _make_btd_device([5, 5]) with pytest.raises(ValueError, match="incompatible"): - ham._constrained_subblocks(H, S, 3, 3) + constrained_subblocks(H, S, 3, 3) # ---------------------------------------------------------------------------