From fb3f3f6fd324d2c9ce6baed9abe65828bc77d79f Mon Sep 17 00:00:00 2001 From: AsymmetryChou <181240085@smail.nju.edu.cn> Date: Thu, 1 Oct 2026 19:20:12 +0800 Subject: [PATCH 1/3] Feat: add btd_initialization=direct to avoid full H matrix --- dpnegf/runner/NEGF.py | 6 +- dpnegf/tests/test_direct_btd.py | 446 ++++++++++++++ dpnegf/tests/test_direct_btd_2d.py | 282 +++++++++ dpnegf/utils/argcheck.py | 19 +- dpnegf/utils/direct_btd.py | 929 +++++++++++++++++++++++++++++ 5 files changed, 1680 insertions(+), 2 deletions(-) create mode 100644 dpnegf/tests/test_direct_btd.py create mode 100644 dpnegf/tests/test_direct_btd_2d.py create mode 100644 dpnegf/utils/direct_btd.py diff --git a/dpnegf/runner/NEGF.py b/dpnegf/runner/NEGF.py index 5f9084e..9fc5e69 100644 --- a/dpnegf/runner/NEGF.py +++ b/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, diff --git a/dpnegf/tests/test_direct_btd.py b/dpnegf/tests/test_direct_btd.py new file mode 100644 index 0000000..ba2f5f3 --- /dev/null +++ b/dpnegf/tests/test_direct_btd.py @@ -0,0 +1,446 @@ +import copy +import json +import os +from pathlib import Path +import subprocess +import sys + +import h5py +import numpy as np +import torch +import pytest +from ase.io import read +from dptb.data import AtomicData, AtomicDataDict +from dptb.data.transforms import OrbitalMapper +from dptb.nn.build import build_model +from dptb.nn.hr2hk import HR2HK + +from dpnegf.negf.negf_hamiltonian_init import NEGFHamiltonianInit +from dpnegf.negf.lead_property import LeadProperty +from dpnegf.negf.split_btd import compute_edge +from dpnegf.utils.direct_btd import ( + LocalBlocks, + assemble, + partition_profiles, +) + +ROOT = Path(__file__).resolve().parents[2] + +@pytest.fixture(autouse=True) +def restore_torch_settings(): + original_dtype = torch.get_default_dtype() + original_threads = torch.get_num_threads() + yield + torch.set_default_dtype(original_dtype) + torch.set_num_threads(original_threads) + + +def dense_from_parts(parts): + sizes = [block.shape[0] for block in parts[0]] + bounds = np.cumsum([0] + sizes) + matrix = torch.zeros( + (sum(sizes), sum(sizes)), + dtype=parts[0][0].dtype, + ) + for index in range(len(sizes)): + section = slice(bounds[index], bounds[index + 1]) + matrix[section, section] = parts[0][index] + if index < len(sizes) - 1: + neighbour = slice(bounds[index + 1], bounds[index + 2]) + matrix[section, neighbour] = parts[1][index] + matrix[neighbour, section] = parts[2][index] + return matrix + + +@pytest.mark.parametrize("overlap", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_local_matches_hr2hk_mixed_masks_phases_duplicates(overlap, dtype): + mapper = OrbitalMapper( + {"C": ["2s", "2p"], "H": ["1s"]}, + method="e3tb", + ) + mapper.get_orbpair_maps() + generator = torch.Generator().manual_seed(19) + data = { + AtomicDataDict.ATOM_TYPE_KEY: torch.tensor( + [ + [mapper.chemical_symbol_to_type[symbol]] + for symbol in ("C", "H", "C") + ] + ), + AtomicDataDict.EDGE_INDEX_KEY: torch.tensor( + [[0, 1, 1, 2, 0, 0], [1, 0, 2, 1, 1, 0]] + ), + AtomicDataDict.EDGE_CELL_SHIFT_KEY: torch.tensor( + [ + [1, 0, 0], + [-1, 0, 0], + [0, 1, 0], + [0, -1, 0], + [1, 0, 0], + [1, 0, 0], + ], + dtype=dtype, + ), + AtomicDataDict.KPOINT_KEY: torch.tensor( + [[0, 0, 0], [0.17, -0.23, 0]], + dtype=dtype, + ), + } + edge_key = ( + AtomicDataDict.EDGE_OVERLAP_KEY + if overlap + else AtomicDataDict.EDGE_FEATURES_KEY + ) + node_key = ( + AtomicDataDict.NODE_OVERLAP_KEY + if overlap + else AtomicDataDict.NODE_FEATURES_KEY + ) + data[edge_key] = torch.randn( + (6, mapper.reduced_matrix_element), + generator=generator, + dtype=dtype, + ) + data[node_key] = torch.randn( + (3, mapper.reduced_matrix_element), + generator=generator, + dtype=dtype, + ) + + local = LocalBlocks(data, mapper, overlap=overlap) + reference = HR2HK( + idp=mapper, + edge_field=edge_key, + node_field=node_key, + overlap=overlap, + dtype=dtype, + )(dict(data))[AtomicDataDict.HAMILTONIAN_KEY] + + tolerance = 1e-6 if dtype == torch.float32 else 1e-14 + for k_index, kpoint in enumerate(data[AtomicDataDict.KPOINT_KEY]): + result = assemble(local, kpoint, (0, 9), [3, 3, 3], {}) + torch.testing.assert_close( + dense_from_parts(result), + reference[k_index], + atol=tolerance, + rtol=tolerance, + ) + + data[AtomicDataDict.NODE_SOC_SWITCH_KEY] = torch.tensor([True, False]) + with pytest.raises(ValueError, match="SOC"): + LocalBlocks(data, mapper) + + +def test_profiles_fixed_boundaries_and_overlap_only(): + matrix = np.eye(12) + matrix[0, 5] = matrix[5, 0] = 1 + matrix[5, 8] = matrix[8, 5] = 1 + + edge, reverse = compute_edge(matrix) + blocks = partition_profiles(edge, reverse, 3, 2, fixed=True) + + assert blocks[0] == 3 + assert blocks[-1] == 2 + assert sum(blocks) == 12 + + matrix[0, 11] = matrix[11, 0] = 1 + edge, reverse = compute_edge(matrix) + with pytest.raises(ValueError, match="incompatible"): + partition_profiles(edge, reverse, 3, 2, fixed=True) + + with pytest.raises(ValueError, match="overlap"): + partition_profiles(edge, reverse, 8, 8, fixed=True) + + +def test_reduced_contacts_explicit_side(): + h_contact = torch.zeros((3, 4)) + s_contact = torch.ones_like(h_contact) + actual = LeadProperty.HDL_reduced( + h_contact, + s_contact, + [2, 6, 3], + tab="lead_R", + reduced=True, + ) + + assert actual[0] is h_contact + assert actual[1] is s_contact + + with pytest.raises(ValueError, match="height"): + LeadProperty.HDL_reduced( + h_contact, + s_contact, + [2, 6, 4], + tab="lead_R", + reduced=True, + ) + + +def cnt_initializer(path, mode, dtype="float64"): + path.mkdir(parents=True, exist_ok=True) + torch.set_default_dtype(getattr(torch, dtype)) + options = json.loads( + (ROOT / "examples/CNT/input.json").read_text() + )["task_options"]["stru_options"] + + model = build_model( + checkpoint=str(ROOT / "examples/CNT/nnsk_dftb.json"), + common_options={"dtype": dtype, "device": "cpu"}, + ) + return NEGFHamiltonianInit( + model=model, + AtomicData_options={"r_max": 4.6}, + structure=read(ROOT / "examples/CNT/cnt7_0.xyz"), + block_tridiagonal=True, + pbc_negf=options["pbc"], + stru_options=copy.deepcopy(options), + unit="eV", + results_path=str(path), + btd_initialization=mode, + ) + +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +def test_cnt_cache_equivalence_and_no_global_hr2hk( + tmp_path, monkeypatch, dtype +): + torch.set_num_threads(1) + reference = cnt_initializer(tmp_path / "dense", "dense", dtype) + kpoints = [[0, 0, 0], [0.17, 0.23, 0]] + with torch.no_grad(): + reference.initialize(kpoints, block_tridiagnal=True) + direct = cnt_initializer(tmp_path / "direct", "direct", dtype) + direct._self_energy_cache_edge_sizes = ( + reference.subblocks[0], + reference.subblocks[-1], + ) + forward = HR2HK.forward + + def guarded_forward(transform, data): + assert len(data[AtomicDataDict.ATOM_TYPE_KEY]) <= 56 + return forward(transform, data) + + monkeypatch.setattr(HR2HK, "forward", guarded_forward) + original_zeros = torch.zeros + original_eye = torch.eye + forbidden_sizes = { + sum(direct.atom_norbs), + sum(direct.device_norbs), + } + + def guarded_zeros(*shape, **kwargs): + dimensions = ( + shape[0] + if len(shape) == 1 + and isinstance(shape[0], (tuple, list, torch.Size)) + else shape + ) + assert not ( + len(dimensions) >= 2 + and dimensions[-1] == dimensions[-2] + and dimensions[-1] in forbidden_sizes + ) + return original_zeros(*shape, **kwargs) + + def guarded_eye(size, *args, **kwargs): + assert size not in forbidden_sizes + return original_eye(size, *args, **kwargs) + + with monkeypatch.context() as allocation_guard: + def forbidden_resident_assembly(*args, **kwargs): + raise AssertionError( + "Production initialization must stream BTD blocks." + ) + + allocation_guard.setattr( + "dpnegf.utils.direct_btd.assemble", + forbidden_resident_assembly, + ) + allocation_guard.setattr(torch, "zeros", guarded_zeros) + allocation_guard.setattr(torch, "eye", guarded_eye) + direct.initialize(kpoints, block_tridiagnal=True) + assert direct.subblocks[0] == reference.subblocks[0] + assert direct.subblocks[-1] == reference.subblocks[-1] + tolerance = 2e-6 if dtype == "float32" else 1e-13 + for kpoint in kpoints: + reference_parts = reference.get_hs_device( + kpoint, V=0, block_tridiagonal=True + ) + direct_parts = direct.get_hs_device( + kpoint, V=0, block_tridiagonal=True + ) + for indices in ((0, 5, 2), (1, 3, 4)): + expected = dense_from_parts( + [reference_parts[index] for index in indices] + ) + actual = dense_from_parts( + [direct_parts[index] for index in indices] + ) + torch.testing.assert_close( + actual, expected, atol=tolerance, rtol=tolerance + ) + for tab in ("lead_L", "lead_R"): + expected = reference.get_hs_lead(kpoint, tab, 0) + actual = direct.get_hs_lead(kpoint, tab, 0) + reduced_h, reduced_s = LeadProperty.HDL_reduced( + expected[2], + expected[5], + reference.subblocks, + tab=tab, + ) + for index in (0, 1, 3, 4): + torch.testing.assert_close( + actual[index], + expected[index], + atol=tolerance, + rtol=tolerance, + ) + torch.testing.assert_close( + actual[2], reduced_h, atol=tolerance, rtol=tolerance + ) + torch.testing.assert_close( + actual[5], reduced_s, atol=tolerance, rtol=tolerance + ) + if dtype == "float64": + expected = direct.get_hs_lead(kpoints[0], "lead_L", 0) + torch.set_default_dtype(torch.float32) + actual = direct.get_hs_lead(kpoints[0], "lead_L", 0) + torch.set_default_dtype(torch.float64) + for expected_matrix, actual_matrix in zip(expected, actual): + torch.testing.assert_close( + actual_matrix, expected_matrix, atol=0, rtol=0 + ) + with h5py.File(tmp_path / "direct/HS_device.h5", "r") as cache: + assert cache.attrs["layout_version"] == 2 + assert cache.attrs["complete"] + assert "Hall" not in cache + with pytest.raises(ValueError, match="Bloch"): + direct.initialize( + kpoints, block_tridiagnal=True, useBloch=True + ) + with pytest.raises(ValueError, match="plot_blocks"): + direct.initialize( + kpoints, block_tridiagnal=True, plot_blocks=True + ) + with h5py.File(tmp_path / "direct/HS_lead_R.h5", "a") as cache: + cache.attrs["cache_id"] = "different-physical-cache" + with pytest.raises(ValueError, match="identities"): + direct.get_hs_lead(kpoints[0], "lead_R", 0) + + +def test_orthogonal_chain_and_budget(tmp_path): + fixture = ROOT / "dpnegf/tests/data/test_negf/test_negf_run" + options = json.loads( + (fixture / "negf_chain_new.json").read_text() + )["task_options"]["stru_options"] + matrices = [] + for mode in ("dense", "direct"): + path = tmp_path / mode + path.mkdir() + torch.set_default_dtype(torch.float64) + model = build_model( + checkpoint=str(fixture / "nnsk_C_new.json"), + common_options={"dtype": "float64", "device": "cpu"}, + ) + initializer = NEGFHamiltonianInit( + model, + {"r_max": 2.0}, + str(fixture / "chain.vasp"), + True, + options["pbc"], + options, + "Hartree", + results_path=str(path), + btd_initialization=mode, + ) + with torch.no_grad(): + initializer.initialize( + [[0, 0, 0]], block_tridiagnal=True + ) + parts = initializer.get_hs_device( + V=0, block_tridiagonal=True + ) + matrices.append( + dense_from_parts([parts[index] for index in (0, 5, 2)]) + ) + torch.testing.assert_close( + dense_from_parts([parts[index] for index in (1, 3, 4)]), + torch.eye( + sum(initializer.subblocks), dtype=torch.complex128 + ), + ) + torch.testing.assert_close(*matrices, atol=1e-12, rtol=1e-12) + initializer.direct_btd_max_mib = 1e-12 + with pytest.raises(MemoryError, match="No matrices allocated"): + initializer.initialize([[0, 0, 0]], block_tridiagnal=True) + + +def test_cnt_runner_transport_and_cache_reuse(tmp_path): + results = [] + for mode in ("dense", "direct", "reuse_se", "reuse_hs"): + config = json.loads((ROOT / "examples/CNT/input.json").read_text()) + config["dtype"] = "float64" + task = config["task_options"] + task["btd_initialization"] = ( + "dense" if mode == "dense" else "direct" + ) + for key in ("emin", "emax", "espacing"): + task.pop(key) + task["energy_grid"] = { + "method": "uniform", + "emin": -1.0, + "emax": 1.0, + "num_points": 31, + } + task["e_fermi"] = 0.0 + task["output_options"] = {"tc": True} + task["self_energy_options"]["numba_jit"] = False + task["self_energy_options"]["parallel"] = { + "cpu_budget": 1, + "n_workers": 1, + "blas_threads": 1, + } + if mode == "reuse_se": + task["self_energy_options"]["cache"] = { + "use_saved": True, + "save_path": str(tmp_path / "dense/results/self_energy"), + } + if mode == "reuse_hs": + task["hs_cache"] = { + "use_saved": True, + "save_path": str(tmp_path / "direct/results"), + } + task["rgf_options"] = {"device": "cpu", "e_batch_size": 3} + config_path = tmp_path / f"{mode}.json" + config_path.write_text(json.dumps(config)) + output = tmp_path / mode + command = [ + sys.executable, + "-c", + ( + "from dpnegf.entrypoints.run import run; " + f"run(INPUT={str(config_path)!r}, " + f"init_model={str(ROOT / 'examples/CNT/nnsk_dftb.json')!r}, " + f"structure={str(ROOT / 'examples/CNT/cnt7_0.xyz')!r}, " + f"output={str(output)!r}, log_level=20, " + f"log_path={str(output / 'run.log')!r})" + ), + ] + completed = subprocess.run( + command, + cwd=ROOT, + env=dict(os.environ, MPLCONFIGDIR=str(tmp_path / "matplotlib")), + capture_output=True, + text=True, + timeout=180, + ) + assert completed.returncode == 0, completed.stderr + results.append( + torch.load( + output / "results/negf.out.pth", weights_only=False + )["T_avg"] + ) + for actual in results[1:]: + torch.testing.assert_close( + actual, results[0], atol=1e-10, rtol=1e-10 + ) \ No newline at end of file diff --git a/dpnegf/tests/test_direct_btd_2d.py b/dpnegf/tests/test_direct_btd_2d.py new file mode 100644 index 0000000..73bb557 --- /dev/null +++ b/dpnegf/tests/test_direct_btd_2d.py @@ -0,0 +1,282 @@ +import json +from pathlib import Path + +import numpy as np +import pytest +import torch +from dptb.data import AtomicData, AtomicDataDict +from dptb.nn.hr2hk import HR2HK + +from dpnegf.tests.test_direct_btd import dense_from_parts + +ROOT = Path(__file__).parent.parent.parent + +def example_structure(material, repeat): + """Extend the device while preserving both original lead contacts.""" + import numpy as np + from ase.io import read + + structure = read(ROOT / f"examples/{material}/stru_negf.xyz") + options = json.loads( + (ROOT / f"examples/{material}/negf.json").read_text() + )["task_options"]["stru_options"] + original_device = structure[32:64] + displacement = ( + structure.positions[64:80] - structure.positions[32:48] + ).mean(axis=0) + assert np.allclose( + structure.positions[64:80] - structure.positions[32:48], + displacement, + atol=1e-6, + ) + extended = structure[:32] + for image_index in range(repeat): + image = original_device.copy() + image.positions += image_index * displacement + extended += image + right = structure[64:].copy() + right.positions += (repeat - 1) * displacement + extended += right + cell = structure.cell.copy() + cell[2] += (repeat - 1) * displacement + extended.set_cell(cell) + extended.set_pbc(structure.pbc) + options["device"]["id"] = f"32-{32 + 32 * repeat}" + options["lead_R"]["id"] = ( + f"{32 + 32 * repeat}-{64 + 32 * repeat}" + ) + return extended, options + + +def prepare_example(material, mode, repeat, output, dtype="float64"): + import torch + from ase.io import write + from dptb.nn.build import build_model + + from dpnegf.negf.negf_hamiltonian_init import NEGFHamiltonianInit + from dpnegf.utils.argcheck import get_cutoffs_from_model_options + + torch.set_default_dtype(getattr(torch, dtype)) + structure, options = example_structure(material, repeat) + output.mkdir(parents=True, exist_ok=True) + write(output / "structure.xyz", structure) + model = build_model( + checkpoint=str( + ROOT + / f"examples/{material}/train/train_out/checkpoint/nnsk.ep3000.pth" + ), + common_options={"dtype": dtype, "device": "cpu"}, + ) + cutoffs = dict( + zip( + ("r_max", "er_max", "oer_max"), + get_cutoffs_from_model_options(model.model_options), + ) + ) + return NEGFHamiltonianInit( + model, + cutoffs, + structure, + True, + options["pbc"], + options, + "eV", + results_path=str(output), + btd_initialization=mode, + ) + + +def initialize_example( + material, mode, repeat, output, kpoints, dtype="float64" +): + import torch + + initializer = prepare_example( + material, mode, repeat, output, dtype + ) + with torch.no_grad(): + initializer.initialize(kpoints, block_tridiagnal=True) + return initializer + +@pytest.mark.parametrize("material", ["graphene", "hBN"]) +def test_2d_example_dense_equivalence(material, tmp_path, monkeypatch): + original_dtype = torch.get_default_dtype() + original_threads = torch.get_num_threads() + torch.set_num_threads(1) + try: + kpoints = [[0.0, 0.0, 0.0], [0.0, 0.23, 0.0]] + dense = initialize_example( + material, "dense", 1, tmp_path / "dense", kpoints + ) + forward = HR2HK.forward + + def lead_only(transform, data): + assert len(data[AtomicDataDict.ATOM_TYPE_KEY]) <= 32 + return forward(transform, data) + + monkeypatch.setattr(HR2HK, "forward", lead_only) + direct = initialize_example( + material, "direct", 1, tmp_path / "direct", kpoints + ) + full_size = sum(direct.subblocks) + for kpoint in kpoints: + old_parts = dense.get_hs_device( + kpoint, V=0, block_tridiagonal=True + ) + new_parts = direct.get_hs_device( + kpoint, V=0, block_tridiagonal=True + ) + for indices in ((0, 5, 2), (1, 3, 4)): + expected = dense_from_parts( + [old_parts[index] for index in indices] + ) + actual = dense_from_parts( + [new_parts[index] for index in indices] + ) + torch.testing.assert_close( + actual, expected, atol=1e-12, rtol=1e-12 + ) + for tab in ("lead_L", "lead_R"): + expected = dense.get_hs_lead(kpoint, tab, 0) + actual = direct.get_hs_lead(kpoint, tab, 0) + for index in (0, 1, 3, 4): + torch.testing.assert_close( + actual[index], + expected[index], + atol=1e-12, + rtol=1e-12, + ) + for index in (2, 5): + padded = torch.zeros( + (full_size, actual[index].shape[1]), + dtype=torch.complex128, + ) + section = ( + slice(0, actual[index].shape[0]) + if tab == "lead_L" + else slice( + full_size - actual[index].shape[0], + full_size, + ) + ) + padded[section] = actual[index] + torch.testing.assert_close( + padded, + expected[index], + atol=1e-12, + rtol=1e-12, + ) + gamma = direct.get_hs_device( + kpoints[0], V=0, block_tridiagonal=True + )[0] + non_gamma = direct.get_hs_device( + kpoints[1], V=0, block_tridiagonal=True + )[0] + assert max( + float((left - right).abs().max()) + for left, right in zip(gamma, non_gamma) + ) > 1e-3 + finally: + torch.set_default_dtype(original_dtype) + torch.set_num_threads(original_threads) + + +@pytest.mark.parametrize("material", ["graphene", "hBN"]) +def test_extended_sheet_contacts_and_periodicity(material): + original, _ = example_structure(material, 1) + extended, options = example_structure(material, 10) + assert len(extended) == len(original) + 9 * 32 + assert options["pbc"] == [False, True, False] + np.testing.assert_allclose(extended.cell[:2], original.cell[:2]) + np.testing.assert_allclose( + extended.positions[:32], original.positions[:32] + ) + shift = extended.positions[-32:] - original.positions[-32:] + np.testing.assert_allclose( + shift, np.broadcast_to(shift[0], shift.shape), atol=1e-7 + ) + assert ( + extended.get_chemical_symbols()[-32:] + == original.get_chemical_symbols()[-32:] + ) + + +@pytest.mark.parametrize("material", ["graphene", "hBN"]) +def test_2d_runner_transmission(material, tmp_path): + from dpnegf.entrypoints.run import run + + original_dtype = torch.get_default_dtype() + original_threads = torch.get_num_threads() + torch.set_num_threads(1) + outputs = [] + try: + for mode in ("dense", "direct"): + config = json.loads( + (ROOT / f"examples/{material}/negf.json").read_text() + ) + config["dtype"] = "float64" + task = config["task_options"] + for key in ("emin", "emax", "espacing"): + task.pop(key, None) + task["energy_grid"] = { + "method": "uniform", + "emin": -8.0, + "emax": 2.0, + "num_points": 11, + } + task["e_fermi"] = 0.0 + task["stru_options"]["compute_band_edges"] = False + task["stru_options"]["kmesh"] = [1, 3, 1] + task["output_options"] = {"tc": True} + task["self_energy_options"] = { + "solver": "Sancho-Rubio", + "numba_jit": False, + "parallel": { + "n_workers": 1, + "cpu_budget": 1, + "blas_threads": 1, + }, + } + task["rgf_options"] = { + "device": "cpu", + "e_batch_size": 4, + } + task["btd_initialization"] = mode + config_path = tmp_path / f"{mode}.json" + config_path.write_text(json.dumps(config)) + output = tmp_path / mode + torch.set_default_dtype(torch.float64) + run( + INPUT=str(config_path), + init_model=str( + ROOT + / f"examples/{material}/train/train_out/checkpoint/nnsk.ep3000.pth" + ), + structure=str(ROOT / f"examples/{material}/stru_negf.xyz"), + output=str(output), + log_level=20, + log_path=str(output / "run.log"), + ) + outputs.append( + torch.load( + output / "results/negf.out.pth", + weights_only=False, + ) + ) + torch.testing.assert_close( + outputs[0]["T_avg"], + outputs[1]["T_avg"], + atol=1e-9, + rtol=1e-9, + ) + assert len(outputs[0]["T_k"]) > 1 + for key in outputs[0]["T_k"]: + torch.testing.assert_close( + outputs[0]["T_k"][key], + outputs[1]["T_k"][key], + atol=1e-9, + rtol=1e-9, + ) + finally: + torch.set_default_dtype(original_dtype) + torch.set_num_threads(original_threads) \ No newline at end of file diff --git a/dpnegf/utils/argcheck.py b/dpnegf/utils/argcheck.py index 609c327..78672d3 100644 --- a/dpnegf/utils/argcheck.py +++ b/dpnegf/utils/argcheck.py @@ -198,6 +198,8 @@ def negf(): Argument("ele_T", [float, int], optional=False, doc=doc_ele_T), Argument("unit", str, optional=True, default="Hartree", doc=doc_unit), Argument("hs_cache", dict, optional=True, default={}, sub_fields=hs_cache_options(), doc=doc_hs_cache), + Argument("btd_initialization", str, optional=True, default="dense", doc="Whether to initialize the block-tridiagonalization with a dense or direct method. Default: 'dense'."), + Argument("direct_btd_max_mib", [float, int, None], optional=True, default=None, doc="Maximum memory in MiB for direct BTD initialization. If None, no limit is applied."), Argument("scf_options", dict, optional=True, default={}, sub_fields=[], sub_variants=[scf_options()], doc=doc_scf_options), Argument("stru_options", dict, optional=False, sub_fields=stru_options(), doc=doc_stru_options), Argument("poisson_options", dict, optional=True, default={}, sub_fields=[], sub_variants=[poisson_options()], doc=doc_poisson_options), @@ -570,7 +572,22 @@ def normalize_run(data): if data.get("task_options", {}).get("task") == "negf": validate_energy_options(data["task_options"]) - + task = data["task_options"] + if task["btd_initialization"] not in ["dense", "direct"]: + raise ValueError(f"Invalid btd_initialization: {task['btd_initialization']}. Must be 'dense' or 'direct'.") + if task["btd_initialization"] == "direct": + if not task["block_tridiagonal"] or task["plot_blocks"]: + raise ValueError("btd_initialization='direct' requires block_tridiagonal=True and plot_blocks=False.") + if any( + value.get("useBloch", False) + for key, value in task["stru_options"].items() + if key.startswith("lead") + ): + raise ValueError("btd_initialization='direct' does not support Bloch expansion in leads.") + if ( + task["direct_btd_max_mib"] is not None + and task["direct_btd_max_mib"] <= 0): + raise ValueError("direct_btd_max_mib must be positive.") return data diff --git a/dpnegf/utils/direct_btd.py b/dpnegf/utils/direct_btd.py new file mode 100644 index 0000000..1f379ea --- /dev/null +++ b/dpnegf/utils/direct_btd.py @@ -0,0 +1,929 @@ +import logging +import os +import re +import uuid + +import h5py +import numpy as np +import torch +from dptb.data import AtomicData, AtomicDataDict +from dptb.utils.constants import anglrMId + +from dpnegf.negf.split_btd import compute_blocks + +log = logging.getLogger(__name__) + + + + + +class LocalBlocks: + """Expand local triangular features without constructing dense H/S.""" + + def __init__(self, data, idp, overlap=False): + soc = data.get(AtomicDataDict.NODE_SOC_SWITCH_KEY, False) + if torch.as_tensor(soc).any(): + raise ValueError("Direct BTD does not support SOC.") + idp.get_orbpair_maps() + edge_key = ( + AtomicDataDict.EDGE_OVERLAP_KEY + if overlap + else AtomicDataDict.EDGE_FEATURES_KEY + ) + node_key = ( + AtomicDataDict.NODE_OVERLAP_KEY + if overlap + else AtomicDataDict.NODE_FEATURES_KEY + ) + edge_features = data[edge_key] + node_features = data.get(node_key) + if node_features is None: + raise ValueError(f"Direct BTD requires DeePTB field {node_key}.") + atom_types = data[AtomicDataDict.ATOM_TYPE_KEY].flatten() + self.masks = [idp.mask_to_basis[int(atom_type)] for atom_type in atom_types] + self.offsets = np.cumsum( + [0] + [int(mask.sum()) for mask in self.masks] + ) + self.edges = data[AtomicDataDict.EDGE_INDEX_KEY].cpu().numpy() + self.shifts = data[AtomicDataDict.EDGE_CELL_SHIFT_KEY] + self.edge_orders = [ + np.argsort(endpoints, kind="stable") for endpoints in self.edges + ] + self.edge_starts = [ + np.concatenate( + ( + [0], + np.cumsum( + np.bincount(endpoints, minlength=len(atom_types)) + ), + ) + ) + for endpoints in self.edges + ] + full_norb = idp.full_basis_norb + self.nodes = node_features.new_zeros( + (len(atom_types), full_norb, full_norb) + ) + self.hoppings = edge_features.new_zeros( + (len(edge_features), full_norb, full_norb) + ) + row_start = 0 + for row_index, row_orbital in enumerate(idp.full_basis): + row_size = ( + 2 * anglrMId[re.findall(r"[a-zA-Z]+", row_orbital)[0]] + 1 + ) + column_start = 0 + for column_index, column_orbital in enumerate(idp.full_basis): + column_size = ( + 2 + * anglrMId[ + re.findall(r"[a-zA-Z]+", column_orbital)[0] + ] + + 1 + ) + if row_index <= column_index: + factor = 0.5 if row_orbital == column_orbital else 1.0 + pair_slice = idp.orbpair_maps[ + row_orbital + "-" + column_orbital + ] + matrix_slice = ( + slice(row_start, row_start + row_size), + slice(column_start, column_start + column_size), + ) + self.nodes[:, matrix_slice[0], matrix_slice[1]] = ( + factor + * node_features[:, pair_slice].reshape( + -1, row_size, column_size + ) + ) + self.hoppings[:, matrix_slice[0], matrix_slice[1]] = ( + factor + * edge_features[:, pair_slice].reshape( + -1, row_size, column_size + ) + ) + column_start += column_size + row_start += row_size + + def contributions(self): + """Yield unhermitianized onsite/edge blocks in HR2HK order.""" + for atom_index, block in enumerate(self.nodes): + mask = self.masks[atom_index] + yield atom_index, atom_index, block[mask][:, mask], None + for edge_index, (source, target) in enumerate(self.edges.T): + block = self.hoppings[edge_index][self.masks[source]][ + :, self.masks[target] + ] + yield int(source), int(target), block, edge_index + + def phases(self, kpoint): + kpoint = torch.as_tensor( + kpoint, dtype=self.nodes.dtype, device=self.nodes.device + ) + return torch.exp(-2j * torch.pi * (self.shifts @ kpoint)) + + def pairs(self, kpoint, atoms=None, phases=None): + """Yield selected atom rows with only one reduced row resident.""" + complex_dtype = ( + torch.complex64 + if self.nodes.dtype == torch.float32 + else torch.complex128 + ) + if phases is None: + phases = self.phases(kpoint) + if atoms is None: + atoms = range(len(self.nodes)) + if self.nodes.device.type == "cpu" and not torch.is_grad_enabled(): + yield from self._cpu_pairs(atoms, phases) + return + for source in atoms: + mask = self.masks[source] + onsite = self.nodes[source][mask][:, mask].to(complex_dtype) + halves = [{source: onsite}, {}] + for direction in (0, 1): + start, end = self.edge_starts[direction][source:source + 2] + for edge_index in self.edge_orders[direction][start:end]: + row, column = self.edges[:, edge_index] + if direction == 1 and row == column: + continue + target = int(column if direction == 0 else row) + value = self.hoppings[edge_index][self.masks[row]][ + :, self.masks[column] + ] + value = value.to(complex_dtype) * phases[edge_index] + if target in halves[direction]: + halves[direction][target] += value + else: + halves[direction][target] = value + for target in sorted(halves[0].keys() | halves[1].keys()): + value = halves[0].get(target) + reverse = value if target == source else halves[1].get(target) + if value is None: + value = reverse.mH + elif reverse is not None: + value = value + reverse.mH + yield source, target, value + + def _cpu_pairs(self, atoms, phases): + """Use zero-copy CPU views to avoid small Torch allocations.""" + nodes = self.nodes.numpy() + hoppings = self.hoppings.numpy() + phases = phases.numpy() + masks = [mask.numpy() for mask in self.masks] + for source in atoms: + mask = masks[source] + onsite = nodes[source][mask][:, mask].astype(phases.dtype) + halves = [{source: onsite}, {}] + for direction in (0, 1): + start, end = self.edge_starts[direction][source:source + 2] + for edge_index in self.edge_orders[direction][start:end]: + row, column = self.edges[:, edge_index] + if direction == 1 and row == column: + continue + target = int(column if direction == 0 else row) + value = hoppings[edge_index][masks[row]][:, masks[column]] + value = value.astype(phases.dtype) * phases[edge_index] + if target in halves[direction]: + halves[direction][target] += value + else: + halves[direction][target] = value + for target in sorted(halves[0].keys() | halves[1].keys()): + value = halves[0].get(target) + reverse = value if target == source else halves[1].get(target) + if value is None: + value = reverse.T.conj() + elif reverse is not None: + value = value + reverse.T.conj() + yield source, target, torch.from_numpy(value) + + +def partition_profiles(edge, reverse_edge, left, right, fixed=False): + """Partition exact H/S union profiles, preserving fixed boundaries.""" + size = len(edge) + if min(left, right) <= 0 or max(left, right) > size: + raise ValueError( + "Direct BTD boundary sizes must lie inside the device." + ) + if fixed and left + right > size: + raise ValueError("Fixed cache boundaries overlap.") + edge = np.maximum.accumulate(edge) + reverse_edge = np.maximum.accumulate(reverse_edge) + blocks = compute_blocks( + left, right, edge, reverse_edge, use_jit=False + ) + if fixed and (blocks[0] != left or blocks[-1] != right): + raise ValueError( + "Fixed self-energy cache boundaries are incompatible with H/S " + "connectivity." + ) + bounds = np.cumsum([0] + blocks) + for block_index in range(len(blocks)): + end = bounds[block_index + 1] + if edge[end - 1] > bounds[min(block_index + 2, len(blocks))]: + raise ValueError( + "Direct BTD partition has non-neighbouring H/S connections." + ) + reverse_end = size - bounds[block_index] + if ( + reverse_edge[reverse_end - 1] + > size - bounds[max(block_index - 1, 0)] + ): + raise ValueError( + "Direct BTD partition has non-neighbouring H/S connections." + ) + return blocks + + +def _window_add( + output, + value, + row_start, + column_start, + row_window, + column_window, +): + + row_lo = max(row_start, row_window[0]) + row_hi = min(row_start + value.shape[0], row_window[1]) + column_lo = max(column_start, column_window[0]) + column_hi = min(column_start + value.shape[1], column_window[1]) + + if row_lo < row_hi and column_lo < column_hi: + output[ + row_lo - row_window[0]:row_hi - row_window[0], + column_lo - column_window[0]:column_hi - column_window[0], + ] += value[ + row_lo - row_start:row_hi - row_start, + column_lo - column_start:column_hi - column_start, + ] + + +def assemble(local, kpoint, device_window, blocks, lead_windows): + """Assemble one k point, splitting atoms at orbital boundaries.""" + + + bounds = np.cumsum([device_window[0]] + list(blocks)) + dtype = ( + torch.complex64 + if local.nodes.dtype == torch.float32 + else torch.complex128 + ) + + kwargs = dict(dtype=dtype, device=local.nodes.device) + diagonal = [ + torch.zeros((size, size), **kwargs) for size in blocks + ] + upper = [ + torch.zeros((left, right), **kwargs) + for left, right in zip(blocks[:-1], blocks[1:]) + ] + lower = [ + torch.zeros((right, left), **kwargs) + for left, right in zip(blocks[:-1], blocks[1:]) + ] + + contacts = {} + principal = {} + + for tab, window in lead_windows.items(): + + boundary = 0 if tab == "lead_L" else len(blocks) - 1 + contacts[tab] = torch.zeros( + (blocks[boundary], window[1] - window[0]), **kwargs + ) + principal[tab] = torch.zeros( + (window[1] - window[0],) * 2, **kwargs + ) + + + for source, target, value in local.pairs(kpoint): + + row_start, row_end = local.offsets[source:source + 2] + column_start, column_end = local.offsets[target:target + 2] + + if row_start < device_window[1] and row_end > device_window[0]: + first_row = max( + 0, np.searchsorted(bounds, row_start, side="right") - 1 + ) + last_row = min( + len(blocks), + np.searchsorted(bounds, row_end, side="left"), + ) + first_column = max( + 0, np.searchsorted(bounds, column_start, side="right") - 1 + ) + last_column = min( + len(blocks), + np.searchsorted(bounds, column_end, side="left"), + ) + for row_block in range(first_row, last_row): + for column_block in range(first_column, last_column): + + if row_block == column_block: + output = diagonal[row_block] + + elif column_block == row_block + 1: + output = upper[row_block] + + elif row_block == column_block + 1: + output = lower[column_block] + + else: + row_slice = slice( + max(0, bounds[row_block] - row_start), + min( + value.shape[0], + bounds[row_block + 1] - row_start, + ), + ) + column_slice = slice( + max(0, bounds[column_block] - column_start), + min( + value.shape[1], + bounds[column_block + 1] - column_start, + ), + ) + + if torch.any(value[row_slice, column_slice] != 0): + raise ValueError( + "Non-neighbouring block contribution would " + "be lost." + ) + + continue + + _window_add( + output, + value, + row_start, + column_start, + bounds[row_block:row_block + 2], + bounds[column_block:column_block + 2], + ) + + for tab, window in lead_windows.items(): + boundary = 0 if tab == "lead_L" else len(blocks) - 1 + _window_add( + contacts[tab], + value, + row_start, + column_start, + bounds[boundary:boundary + 2], + window, + ) + + for tab, window in lead_windows.items(): + _window_add( + principal[tab], + value, + row_start, + column_start, + window, + window, + ) + + return diagonal, upper, lower, contacts, principal + +def _write_complex(group, name, value): + array = value.cpu().numpy() + group.create_dataset(f"{name}_real", data=array.real) + group.create_dataset(f"{name}_imag", data=array.imag) + + +def _build_profiles(local_blocks, kpoints, offsets, device_window, lead_windows): + size = int(device_window[1] - device_window[0]) + edge = np.arange(1, size + 1) + reverse_edge = edge.copy() + + contact_extents = {"lead_L": 1, "lead_R": 1} + + + for local in local_blocks: + if local is None: + continue + + for kpoint in kpoints: + + for source, target, value in local.pairs(kpoint): + + rows, columns = np.nonzero(value.cpu().numpy()) + rows = rows + offsets[source] + columns = columns + offsets[target] + + device_rows = ( + (rows >= device_window[0]) + & (rows < device_window[1]) + ) + + inside = ( + device_rows + & (columns >= device_window[0]) + & (columns < device_window[1]) + ) + + row_indices = rows[inside] - device_window[0] + column_indices = columns[inside] - device_window[0] + + np.maximum.at(edge, row_indices, column_indices + 1) + np.maximum.at( + reverse_edge, + size - 1 - row_indices, + size - column_indices, + ) + + for tab, window in lead_windows.items(): + + selected = ( + device_rows + & (columns >= window[0]) + & (columns < window[1]) + ) + + if selected.any(): + + if tab == "lead_L": + depth = rows[selected].max() - device_window[0] + 1 + else: + depth = device_window[1] - rows[selected].min() + + contact_extents[tab] = max( + contact_extents[tab], int(depth) + ) + + return edge, reverse_edge, contact_extents + +def write_block_rows( + cache, + prefix, + local, + kpoint, + k_index, + device_window, + blocks, + lead_windows, + factor=1.0, + orthogonal=False, +): + """Write one diagonal and at most two neighboring blocks at a time.""" + + bounds = np.cumsum([device_window[0]] + list(blocks)) + + dtype = ( + torch.complex64 + if local.nodes.dtype == torch.float32 + else torch.complex128 + ) + + + kwargs = dict(dtype=dtype, device=local.nodes.device) + phases = None if orthogonal else local.phases(kpoint) + maximum = max(blocks) + workspace = torch.empty((3, maximum, maximum), **kwargs) + + contacts = {} + principal = {} + + for tab, window in lead_windows.items(): + + boundary = 0 if tab == "lead_L" else len(blocks) - 1 + width = window[1] - window[0] + contacts[tab] = torch.zeros( + (blocks[boundary], width), **kwargs + ) + + principal[tab] = torch.zeros((width, width), **kwargs) + if orthogonal: + principal[tab].diagonal().fill_(1) + + + for block_index, size in enumerate(blocks): + + row_window = bounds[block_index:block_index + 2] + neighbours = range( + max(0, block_index - 1), + min(len(blocks), block_index + 2), + ) + + outputs = {} + + for slot, column_block in enumerate(neighbours): + output = workspace[slot, :size, :blocks[column_block]] + output.zero_() + outputs[column_block] = output + + if orthogonal: + outputs[block_index].diagonal().fill_(1) + + else: + first_atom = ( + np.searchsorted(local.offsets, row_window[0], side="right") - 1 + ) + last_atom = np.searchsorted( + local.offsets, row_window[1], side="left" + ) + for source, target, value in local.pairs( + kpoint, range(first_atom, last_atom), phases=phases + ): + + row_start = local.offsets[source] + column_start = local.offsets[target] + for column_block, output in outputs.items(): + _window_add( + output, + value, + row_start, + column_start, + row_window, + bounds[column_block:column_block + 2], + ) + device_column_lo = max(column_start, device_window[0]) + device_column_hi = min( + column_start + value.shape[1], device_window[1] + ) + + allowed_lo = bounds[max(0, block_index - 1)] + allowed_hi = bounds[min(len(blocks), block_index + 2)] + row_slice = slice( + max(0, row_window[0] - row_start), + min(value.shape[0], row_window[1] - row_start), + ) + + for column_lo, column_hi in ( + ( + device_column_lo, + min(device_column_hi, allowed_lo), + ), + ( + max(device_column_lo, allowed_hi), + device_column_hi, + ), + ): + if ( + column_lo < column_hi + and torch.any( + value[ + row_slice, + column_lo - column_start: + column_hi - column_start, + ] + != 0 + ) + ): + raise ValueError( + "Non-neighbouring block contribution would be lost." + ) + + + for tab, window in lead_windows.items(): + boundary = 0 if tab == "lead_L" else len(blocks) - 1 + if block_index == boundary: + _window_add( + contacts[tab], + value, + row_start, + column_start, + row_window, + window, + ) + + else: + column_lo = max(column_start, window[0]) + column_hi = min( + column_start + value.shape[1], window[1] + ) + + if ( + column_lo < column_hi + and torch.any( + value[ + row_slice, + column_lo - column_start: + column_hi - column_start, + ] + != 0 + ) + ): + raise ValueError( + "Device-lead coupling extends beyond its " + "boundary block." + ) + + for column_block, output in outputs.items(): + suffix = ( + "d" + if column_block == block_index + else ("u" if column_block > block_index else "l") + ) + name = prefix + suffix + cache_index = min(block_index, column_block) + output.mul_(factor) + _write_complex( + cache[name], + f"{name}_k{k_index}_b{cache_index}", + output, + ) + del outputs, output + + if not orthogonal: + + for tab, window in lead_windows.items(): + first_atom = ( + np.searchsorted(local.offsets, window[0], side="right") - 1 + ) + last_atom = np.searchsorted( + local.offsets, window[1], side="left" + ) + + for source, target, value in local.pairs( + kpoint, range(first_atom, last_atom), phases=phases + ): + _window_add( + principal[tab], + value, + local.offsets[source], + local.offsets[target], + window, + window, + ) + return contacts, principal + + +@torch.no_grad() +def initialize_direct(owner, kpoints, structure_leads): + """Predict local features, partition exact H/S, and stream caches.""" + stages = [] + + data = AtomicData.from_ase(owner.structase, **owner.AtomicData_options) + data = owner.model.idp( + AtomicData.to_AtomicDataDict(data.to(owner.torch_device)) + ) + data[AtomicDataDict.KPOINT_KEY] = torch.as_tensor( + kpoints, dtype=owner.model.dtype, device=owner.torch_device + ) + data = owner.model(data) + + owner.overlap = AtomicDataDict.EDGE_OVERLAP_KEY in data + owner.remove_bonds_nonpbc(data, owner.pbc_negf, owner.overlap) + + local_h = LocalBlocks(data, owner.model.idp) + local_s = ( + LocalBlocks(data, owner.model.idp, overlap=True) + if owner.overlap + else None + ) + del data + + offsets = local_h.offsets + device_window = offsets[owner.device_id] + device_size = int(device_window[1] - device_window[0]) + + lead_windows = {} + for tab, atom_range in owner.lead_ids.items(): + start, end = offsets[atom_range] + lead_windows[tab] = (start, start + (end - start) // 2) + + edge, reverse_edge, contact_extents = _build_profiles( + (local_h, local_s), + kpoints, + offsets, + device_window, + lead_windows, + ) + left = contact_extents["lead_L"] + right = contact_extents["lead_R"] + + fixed_edges = owner._self_energy_cache_edge_sizes + if fixed_edges: + if fixed_edges[0] < left or fixed_edges[1] < right: + raise ValueError( + "Cached boundary cannot contain the full H/S device-lead " + "coupling." + ) + left, right = fixed_edges + log.warning( + "Reusing legacy self-energy cache: dimensions do not verify " + "model/contact/k/E compatibility; caller must ensure identical " + "physical inputs." + ) + + blocks = partition_profiles( + edge, + reverse_edge, + left, + right, + fixed=bool(fixed_edges), + ) + + elements = sum(size**2 for size in blocks) + 2 * sum( + left_size * right_size + for left_size, right_size in zip(blocks[:-1], blocks[1:]) + ) + estimated_mib = 2 * 16 * elements / 2**20 + log.info( + "Direct BTD: orbitals=%s blocks=%s max_block=%s " + "H+S complex128=%.2f MiB", + device_size, + len(blocks), + max(blocks), + estimated_mib, + ) + + if ( + owner.direct_btd_max_mib is not None + and estimated_mib > owner.direct_btd_max_mib + ): + raise MemoryError( + f"BTD H/S requires {estimated_mib:.2f} MiB, exceeding " + f"direct_btd_max_mib; largest block {max(blocks)}. " + "No matrices allocated." + ) + + log.info( + "Direct BTD block-row workspace: %.3f MiB; total resident BTD H/S " + "at downstream read: %.3f MiB.", + 3 + * max(blocks) ** 2 + * (8 if owner.model.dtype == torch.float32 else 16) + / 2**20, + estimated_mib, + ) + owner.subblocks = blocks + + from dptb.nn.hr2hk import HR2HK + + overlap_transform = ( + HR2HK( + idp=owner.model.idp, + overlap=True, + edge_field=AtomicDataDict.EDGE_OVERLAP_KEY, + node_field=AtomicDataDict.NODE_OVERLAP_KEY, + out_field=AtomicDataDict.OVERLAP_KEY, + dtype=owner.model.dtype, + device=owner.torch_device, + ) + if owner.overlap + else None + ) + + device_path = os.path.join(owner.results_path, "HS_device.h5") + with h5py.File(device_path, "w") as cache: + cache.attrs["layout_version"] = 2 + cache.attrs["initializer"] = "direct" + cache.attrs["cache_id"] = uuid.uuid4().hex + cache.attrs["complete"] = False + + cache.create_dataset("kpoints", data=kpoints) + cache.create_dataset("subblocks", data=blocks) + cache.create_dataset("block_tridiagonal", data=True) + + for name in ("hd", "hu", "hl", "sd", "su", "sl"): + cache.create_group(name) + + for k_index, kpoint in enumerate(kpoints): + h_contacts, h_principal = write_block_rows( + cache, + "h", + local_h, + kpoint, + k_index, + device_window, + blocks, + lead_windows, + factor=owner.h_factor, + ) + s_contacts, s_principal = write_block_rows( + cache, + "s", + local_s if local_s is not None else local_h, + kpoint, + k_index, + device_window, + blocks, + lead_windows, + orthogonal=local_s is None, + ) + + for tab, structure in structure_leads.items(): + lead_data = AtomicData.from_ase( + structure, **owner.AtomicData_options + ) + lead_data = owner.model.idp( + AtomicData.to_AtomicDataDict( + lead_data.to(owner.torch_device) + ) + ) + lead_data[AtomicDataDict.KPOINT_KEY] = torch.as_tensor( + np.asarray(kpoint).reshape(1, 3), + dtype=owner.model.dtype, + device=owner.torch_device, + ) + lead_data = owner.model(lead_data) + + owner.remove_bonds_nonpbc( + lead_data, owner.pbc_negf, owner.overlap + ) + + lead_h = owner.h2k(lead_data)[ + AtomicDataDict.HAMILTONIAN_KEY + ][0] + lead_s = ( + overlap_transform(lead_data)[ + AtomicDataDict.OVERLAP_KEY + ][0] + if owner.overlap + else torch.eye( + lead_h.shape[0], + dtype=owner.model.dtype, + device=owner.torch_device, + ) + ) + + principal_size = lead_h.shape[0] // 2 + principal_h = lead_h[:principal_size, :principal_size] + principal_s = lead_s[:principal_size, :principal_size] + + for name, actual, expected in ( + ("H", principal_h, h_principal[tab]), + ("S", principal_s, s_principal[tab]), + ): + rmse = ( + (actual - expected).abs().square().mean() + ).sqrt().item() + if rmse >= 1e-4: + raise ValueError( + f"{tab} independent/device {name} " + f"principal-layer mismatch: RMSE {rmse}." + ) + + hopping_h = lead_h[:principal_size, principal_size:].clone() + principal_h = principal_h.clone() + principal_h[principal_h.abs() < 1e-6] = 0 + hopping_h[hopping_h.abs() < 1e-6] = 0 + + values = { + "HL": principal_h.cdouble() * owner.h_factor, + "HLL": hopping_h.cdouble() * owner.h_factor, + "SL": principal_s, + "SLL": lead_s[:principal_size, principal_size:], + "HDL": h_contacts[tab].cdouble() * owner.h_factor, + "SDL": s_contacts[tab], + } + + lead_path = os.path.join( + owner.results_path, f"HS_{tab}.h5" + ) + with h5py.File( + lead_path, "w" if k_index == 0 else "a" + ) as lead_cache: + if k_index == 0: + boundary = 0 if tab == "lead_L" else len(blocks) - 1 + lead_cache.attrs.update( + layout_version=2, + contact_reduced=True, + lead=tab, + cache_id=cache.attrs["cache_id"], + device_block=boundary, + device_orbital_offset=sum(blocks[:boundary]), + device_orbitals=device_size, + ) + + lead_cache.create_dataset("kpoints", data=kpoints) + lead_cache.create_dataset("useBloch", data=False) + + for name in ("kpoints_bloch", "bloch_factor"): + lead_cache.create_dataset( + name, + data=np.array( + "None", dtype=h5py.string_dtype() + ), + ) + + for name, value in values.items(): + for component in ("real", "imag"): + lead_cache.create_dataset( + name + "_" + component, + shape=(len(kpoints),) + tuple(value.shape), + dtype=( + np.float32 + if owner.model.dtype == torch.float32 + else np.float64 + ), + ) + + for name, value in values.items(): + lead_cache[name + "_real"][k_index] = ( + value.real.cpu().numpy() + ) + lead_cache[name + "_imag"][k_index] = ( + value.imag.cpu().numpy() + if value.is_complex() + else 0 + ) + + del lead_data, lead_h, lead_s + + del h_contacts, s_contacts, h_principal, s_principal + + cache.attrs["complete"] = True + + owner.direct_btd_stages = stages \ No newline at end of file From cb924986682d054ac3cd26dfa15abaeb0c5d39e4 Mon Sep 17 00:00:00 2001 From: AsymmetryChou <181240085@smail.nju.edu.cn> Date: Thu, 1 Oct 2026 19:21:33 +0800 Subject: [PATCH 2/3] Refactor: update negf hamiltonian for better in hdf5 data flow --- dpnegf/negf/lead_property.py | 61 +++++- dpnegf/negf/negf_hamiltonian_init.py | 315 +++++++++++++++------------ dpnegf/utils/hs_cache.py | 92 ++++++++ 3 files changed, 325 insertions(+), 143 deletions(-) create mode 100644 dpnegf/utils/hs_cache.py diff --git a/dpnegf/negf/lead_property.py b/dpnegf/negf/lead_property.py index bf1513c..ccb08d0 100644 --- a/dpnegf/negf/lead_property.py +++ b/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, diff --git a/dpnegf/negf/negf_hamiltonian_init.py b/dpnegf/negf/negf_hamiltonian_init.py index aad9861..ca3e61b 100644 --- a/dpnegf/negf/negf_hamiltonian_init.py +++ b/dpnegf/negf/negf_hamiltonian_init.py @@ -1,5 +1,4 @@ import os -import re import h5py import logging from typing import Optional, Union, List @@ -18,7 +17,6 @@ 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 ''' @@ -65,6 +63,8 @@ def __init__(self, self_energy_save_path: Optional[str]=None, results_path:Optional[str]=None, torch_device: Union[str, torch.device]=torch.device('cpu'), + btd_initialization: str="dense", + direct_btd_max_mib: Optional[float]=None, ) -> None: # TODO: add dtype and device setting to the model @@ -75,6 +75,17 @@ def __init__(self, if isinstance(torch_device, str): torch_device = torch.device(torch_device) self.torch_device = torch_device + if btd_initialization not in ["dense", "direct"]: + raise ValueError( + f"Invalid btd_initialization value: {btd_initialization}. " + "Must be 'dense' or 'direct'.") + if btd_initialization == "direct" and not block_tridiagonal: + raise ValueError("Direct BTD initialization requires block_tridiagonal=True.") + if direct_btd_max_mib is not None and direct_btd_max_mib <= 0: + raise ValueError("direct_btd_max_mib must be a positive number.") + self.btd_initialization = btd_initialization + self.direct_btd_max_mib = direct_btd_max_mib + self.contact_metadata = {} self.model = model self.model.to(self.torch_device) self.AtomicData_options = AtomicData_options @@ -215,6 +226,17 @@ def initialize(self, kpoints, block_tridiagnal=False, plot_blocks=False, \ ''' + if self.btd_initialization == "direct": + if useBloch or not block_tridiagnal: + raise ValueError("Direct BTD initialization requires block_tridiagonal=True " + "and don't support Bloch expansion for lead self energy calculations.") + if plot_blocks: + raise ValueError("Direct BTD initialization does not support plot_blocks " + "because it avoids the full Hamiltonian construction.") + if set(self.lead_ids) != {"lead_L", "lead_R"}: + raise ValueError("Direct BTD initialization requires exactly two leads: 'lead_L' and 'lead_R'.") + + # structure initialization self.structase.set_pbc(self.pbc_negf) self.structase.pbc[2] = True @@ -256,7 +278,11 @@ def initialize(self, kpoints, block_tridiagnal=False, plot_blocks=False, \ log.info(msg="=="*40) else: self.saved_HS_path = self.results_path - self.Hamiltonian_initialized(kpoints,useBloch,bloch_factor,block_tridiagnal,plot_blocks,\ + if self.btd_initialization == "direct": + from dpnegf.utils.direct_btd import initialize_direct + initialize_direct(self, kpoints, structure_leads) + else: + 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.") @@ -869,88 +895,58 @@ def get_hs_device(self, kpoint=[0,0,0], V=None, block_tridiagonal=False, only_su HS_device_path_pth = os.path.join(self.saved_HS_path, "HS_device.pth") HS_device_path_h5 = os.path.join(self.saved_HS_path, "HS_device.h5") - torch_dtype = torch.float32 if torch.get_default_dtype() == torch.float32 else torch.float64 - #in this version, we only support complex128 for complex dtype to ensure the accuracy. - #TODO: check other complex dtype - complex_dtype = torch.complex128 if torch_dtype == torch.float32 else torch.complex128 - if os.path.exists(HS_device_path_h5): - # log.info(msg="The HS_device.h5 exists in the saved path {}.".format(self.saved_HS_path)) - HS_device_path = HS_device_path_h5 - HS_device = {} - with h5py.File(HS_device_path, "r") as f: - for key in f.keys(): - if isinstance(f[key], h5py.Dataset): - if key in ["kpoints", "subblocks"]: - HS_device[key] = np.array(f[key]) - elif key == "block_tridiagonal": - HS_device[key] = f[key][()] - else: - assert isinstance(f[key][()], np.ndarray),f"Expected np.ndarray, but got {type(f[key][()])}" - # read NumPy array: HD, SD - HS_device[key] = torch.tensor(f[key][()]) - else: - group = f[key] - items = [];sublist = [] - sub_keys = natsorted(group.keys()) # ensure the order of subblocks - current_idx_k = -1 - for sub_key in sub_keys: # sub_key format: f"{key}_k{idx_k}_b{idx_block}_real" - sub_key_type = sub_key.split("_")[-1] - if sub_key_type == "real": - parts = sub_key.split("_k") - if len(parts) < 2: - raise ValueError(f"Unexpected dataset format: {sub_key}") - match = re.search(r'_k(\d+)_', sub_key) - if not match: - raise ValueError(f"Unexpected dataset format: {sub_key}") - idx_k = int(match.group(1)) - if idx_k != current_idx_k: - if sublist: - items.append(sublist) # store the previous sublist - sublist = [] # begin a new sublist - current_idx_k = idx_k - real_part = torch.tensor(group[sub_key][()], dtype=torch_dtype) - imag_part = torch.tensor(group[re.sub(r"(_real|_imag)$", "", sub_key)+"_imag"][()], dtype=torch_dtype) - sublist.append(torch.complex(real_part, imag_part).to(complex_dtype)) - # sublist.append(torch.tensor(group[sub_key][()])) - if sublist: - items.append(sublist) # store the last sublist - HS_device[key] = items # idx_k, idx_block - + with h5py.File(HS_device_path_h5, "r") as f: + if ( + f.attrs.get("layout_version", 1) == 2 + and not f.attrs.get("complete", False) + ): + raise ValueError( + "Incomplete direct BTD cache; initialization did not " + "finish." + ) + if only_subblocks: + if "subblocks" not in f: + log.warning(msg=" 'subblocks' might not be saved in the HS_device.h5 for old version.") + log.error(msg="The subblocks are not saved in the HS_device.h5.") + raise ValueError + return f["subblocks"][()] + + if block_tridiagonal: + from dpnegf.utils.hs_cache import read_btd_kpoint + return read_btd_kpoint( + f, kpoint, V, self.torch_device + ) + from dpnegf.utils.hs_cache import find_kpoint_index + ik = find_kpoint_index(f["kpoints"][()], kpoint) + HD_k = torch.from_numpy(f["HD"][ik]) + SD_k = torch.from_numpy(f["SD"][ik]) + elif os.path.exists(HS_device_path_pth): # log.info(msg="The HS_device.pth exists in the saved path {}.".format(self.saved_HS_path)) - HS_device_path = HS_device_path_pth - HS_device = torch.load(HS_device_path) + HS_device = torch.load(HS_device_path_pth) + if only_subblocks: + if "subblocks" not in HS_device: + log.warning(msg=" 'subblocks' might not be saved in the HS_device.pth for old version.") + log.error(msg="The subblocks are not saved in the HS_device.pth.") + raise ValueError + return HS_device["subblocks"] + + ik = self._find_kpoint_index(HS_device["kpoints"], kpoint) + if block_tridiagonal: + hd_k, sd_k, hl_k, su_k, sl_k, hu_k = ( + HS_device["hd"][ik], HS_device["sd"][ik], + HS_device["hl"][ik], HS_device["su"][ik], + HS_device["sl"][ik], HS_device["hu"][ik], + ) + else: + HD_k, SD_k = HS_device["HD"][ik], HS_device["SD"][ik] else: raise FileNotFoundError(f"Neither HS_device.pth nor HS_device.h5 found in {self.saved_HS_path}. " ) - - - if only_subblocks: - if "subblocks" not in HS_device: - log.warning(msg=" 'subblocks' might not be saved in the HS_device.pth for old version.") - log.error(msg="The subblocks are not saved in the HS_device.pth.") - - raise ValueError - subblocks = HS_device["subblocks"] - return subblocks - - kpoints = HS_device["kpoints"] - ik = None - for i, k in enumerate(kpoints): - if np.abs(np.array(k) - np.array(kpoint)).sum() < 1e-8: - ik = i - break - - assert ik is not None if block_tridiagonal: - # hd format: ( k_index,block_index, orb, orb) - hd_k, sd_k, hl_k, su_k, sl_k, hu_k = HS_device["hd"][ik], HS_device["sd"][ik],\ - HS_device["hl"][ik], HS_device["su"][ik], \ - HS_device["sl"][ik], HS_device["hu"][ik] - # move blocks onto torch_device so downstream RGF runs on the chosen device. hd_k = [b.to(self.torch_device) for b in hd_k] sd_k = [b.to(self.torch_device) for b in sd_k] @@ -977,12 +973,16 @@ def get_hs_device(self, kpoint=[0,0,0], V=None, block_tridiagonal=False, only_su return hd_k , sd_k, hl_k , su_k, sl_k, hu_k else: - HD_k, SD_k = HS_device["HD"][ik], HS_device["SD"][ik] HD_k = HD_k.to(self.torch_device) SD_k = SD_k.to(self.torch_device) V = torch.as_tensor(V).to(self.torch_device) return HD_k - V*SD_k, SD_k, [], [], [], [] + @staticmethod + def _find_kpoint_index(kpoints, kpoint): + from dpnegf.utils.hs_cache import find_kpoint_index + return find_kpoint_index(kpoints, kpoint) + def get_hs_lead(self, kpoint, tab, v): """get the lead Hamiltonian and overlap matrix at a specific kpoint @@ -1006,77 +1006,116 @@ def get_hs_lead(self, kpoint, tab, v): HS_lead_path_pth = os.path.join(self.saved_HS_path, "HS_{0}.pth".format(tab)) HS_lead_path_h5 = os.path.join(self.saved_HS_path, "HS_{0}.h5".format(tab)) - torch_dtype = torch.float32 if torch.get_default_dtype() == torch.float32 else torch.float64 - #in this version, we only support complex128 for complex dtype to ensure the accuracy. - complex_dtype = torch.complex128 if torch_dtype == torch.float32 else torch.complex128 - - HS_leads = {} if os.path.exists(HS_lead_path_h5): - # log.info(msg="The HS_{0}.h5 exists in the saved path {1}.".format(tab, self.saved_HS_path)) - HS_lead_path = HS_lead_path_h5 - with h5py.File(HS_lead_path, "r") as f: - for key in f.keys(): - dataset = f[key] - if key == "useBloch": - HS_leads[key] = bool(dataset[()]) - elif key == "kpoints": - HS_leads[key] = dataset[()] - elif key in ["kpoints_bloch", "bloch_factor"]: - - if isinstance(dataset[()], np.ndarray): - HS_leads[key] = dataset[()] - else: - try : - assert dataset[()].decode() == "None" - HS_leads[key] = None - except: - raise ValueError(f"Unsupported value {dataset[()]} for key {key}") - else: - HSkey = key.split("_") - if HSkey[1] == "real": - HS_leads[HSkey[0]] = torch.tensor(dataset[()], dtype=torch_dtype) \ - + 1j * torch.tensor(f[HSkey[0]+"_imag"][()], dtype=torch_dtype) - HS_leads[HSkey[0]] = HS_leads[HSkey[0]].to(complex_dtype) + with h5py.File(HS_lead_path_h5, "r") as f: + from dpnegf.utils.hs_cache import read_complex + + metadata = dict(f.attrs) + if metadata.get("contact_reduced", False): + if ( + metadata.get("layout_version") != 2 + or metadata.get("lead") != tab + ): + raise ValueError( + "Invalid reduced-contact cache metadata." + ) + device_path = os.path.join( + self.saved_HS_path, "HS_device.h5" + ) + with h5py.File(device_path, "r") as device_cache: + if ( + not device_cache.attrs.get("complete", False) + or metadata.get("cache_id") + != device_cache.attrs.get("cache_id") + ): + raise ValueError( + "Reduced contact and device cache identities " + "do not match." + ) + subblocks = self.get_hs_device(only_subblocks=True) + boundary = ( + 0 if tab == "lead_L" else len(subblocks) - 1 + ) + if ( + metadata["device_block"] != boundary + or metadata["device_orbital_offset"] + != sum(subblocks[:boundary]) + or metadata["device_orbitals"] != sum(subblocks) + or f["HDL_real"].shape[1] != subblocks[boundary] + ): + raise ValueError( + "Reduced contact and device cache partitions do " + "not match." + ) + if not hasattr(self, "contact_metadata"): + self.contact_metadata = {} + self.contact_metadata[tab] = metadata + kpoints = f["kpoints"][()] + kpoints_bloch = self._read_optional_hdf5_metadata(f["kpoints_bloch"]) + bloch_factor = self._read_optional_hdf5_metadata(f["bloch_factor"]) + + if kpoints_bloch is None: + ik = self._find_kpoint_index(kpoints, kpoint) + assert len(kpoints) == f["HL_real"].shape[0], "The number of kpoints in the lead Hamiltonian file does not match the number of kpoints." + lead_index = ik + coupling_index = ik + else: + multi_k_num = int(bloch_factor[0]*bloch_factor[1]) + ik_bloch = self._find_kpoint_index(kpoints_bloch, kpoint) + coupling_index = int(ik_bloch/multi_k_num) + assert len(kpoints_bloch) == f["HL_real"].shape[0], "The number of kpoints in the lead Hamiltonian file does not match the number of kpoints." + assert len(kpoints) == f["HDL_real"].shape[0], "The number of kpoints in the lead Hamiltonian file does not match the number of kpoints." + lead_index = ik_bloch + + hL = read_complex(f, "HL", index=lead_index) + hLL = read_complex(f, "HLL", index=lead_index) + sL = read_complex(f, "SL", index=lead_index) + sLL = read_complex(f, "SLL", index=lead_index) + hDL = read_complex(f, "HDL", index=coupling_index) + sDL = read_complex(f, "SDL", index=coupling_index) + + elif os.path.exists(HS_lead_path_pth): - # log.info(msg="The HS_{0}.pth exists in the saved path {1}.".format(tab, self.saved_HS_path)) HS_leads = torch.load(HS_lead_path_pth) + if hasattr(self, "contact_metadata"): + self.contact_metadata.pop(tab, None) + kpoints = HS_leads["kpoints"] + kpoints_bloch = HS_leads["kpoints_bloch"] + bloch_factor = HS_leads["bloch_factor"] + + if kpoints_bloch is None: + ik = self._find_kpoint_index(kpoints, kpoint) + assert len(kpoints) == HS_leads['HL'].shape[0], "The number of kpoints in the lead Hamiltonian file does not match the number of kpoints." + hL, hLL, sL, sLL = HS_leads["HL"][ik], HS_leads["HLL"][ik],HS_leads["SL"][ik], HS_leads["SLL"][ik] + hDL,sDL = HS_leads["HDL"][ik], HS_leads["SDL"][ik] + else: + multi_k_num = int(bloch_factor[0]*bloch_factor[1]) + ik_bloch = self._find_kpoint_index(kpoints_bloch, kpoint) + ik = int(ik_bloch/multi_k_num) + assert len(kpoints_bloch) == HS_leads['HL'].shape[0], "The number of kpoints in the lead Hamiltonian file does not match the number of kpoints." + assert len(kpoints) == HS_leads['HDL'].shape[0], "The number of kpoints in the lead Hamiltonian file does not match the number of kpoints." + hL, hLL, sL, sLL = HS_leads["HL"][ik_bloch], HS_leads["HLL"][ik_bloch],\ + HS_leads["SL"][ik_bloch], HS_leads["SLL"][ik_bloch] + hDL,sDL = HS_leads["HDL"][ik], HS_leads["SDL"][ik] else: log.error(msg="The HS_{0}.pth or HS_{0}.h5 does not exist in the saved path {1}.".format(tab, self.saved_HS_path)) raise ValueError - - kpoints = HS_leads["kpoints"] - kpoints_bloch = HS_leads["kpoints_bloch"] - bloch_factor = HS_leads["bloch_factor"] - - if kpoints_bloch is None: - ik = None - for i, k in enumerate(kpoints): - if np.abs(np.array(k) - np.array(kpoint)).sum() < 1e-8: - ik = i - break - - assert ik is not None - assert len(kpoints) == HS_leads['HL'].shape[0], "The number of kpoints in the lead Hamiltonian file does not match the number of kpoints." - hL, hLL, sL, sLL = HS_leads["HL"][ik], HS_leads["HLL"][ik],HS_leads["SL"][ik], HS_leads["SLL"][ik] - hDL,sDL = HS_leads["HDL"][ik], HS_leads["SDL"][ik] - else: - multi_k_num = int(bloch_factor[0]*bloch_factor[1]) - ik = None; ik_bloch = None - for i, k in enumerate(kpoints_bloch): - if np.abs(np.array(k) - np.array(kpoint)).sum() < 1e-8: - ik_bloch = i - ik = int(i/multi_k_num) - break - assert ik is not None - assert len(kpoints_bloch) == HS_leads['HL'].shape[0], "The number of kpoints in the lead Hamiltonian file does not match the number of kpoints." - assert len(kpoints) == HS_leads['HDL'].shape[0], "The number of kpoints in the lead Hamiltonian file does not match the number of kpoints." - hL, hLL, sL, sLL = HS_leads["HL"][ik_bloch], HS_leads["HLL"][ik_bloch],\ - HS_leads["SL"][ik_bloch], HS_leads["SLL"][ik_bloch] - hDL,sDL = HS_leads["HDL"][ik], HS_leads["SDL"][ik] return hL-v*sL, hLL-v*sLL, hDL, sL, sLL, sDL + + @staticmethod + def _read_optional_hdf5_metadata(dataset): + value = dataset[()] + if isinstance(value, np.ndarray): + return value + if isinstance(value, bytes) and value.decode() == "None": + return None + if isinstance(value, str) and value == "None": + return None + raise ValueError(f"Unsupported value {value} in HDF5 dataset {dataset.name}") + @property def device_norbs(self): """ diff --git a/dpnegf/utils/hs_cache.py b/dpnegf/utils/hs_cache.py new file mode 100644 index 0000000..13f5f9c --- /dev/null +++ b/dpnegf/utils/hs_cache.py @@ -0,0 +1,92 @@ +"""Selective HDF5 readers shared by dense and direct BTD caches.""" + +import numpy as np +import torch + + +def find_kpoint_index(kpoints, kpoint): + """Return the cache index matching ``kpoint`` within the legacy tolerance.""" + + matches = np.flatnonzero( + np.abs(np.asarray(kpoints) - np.asarray(kpoint)).sum(axis=1) < 1e-8 + ) + if not len(matches): + raise AssertionError( + f"The requested k-point {np.asarray(kpoint)} is not in the H/S cache." + ) + + return int(matches[0]) + + +def read_complex(group, name, device=None, index=None): + """Read one real/imaginary dataset pair without dtype down-casting.""" + + selection = () if index is None else index + real = torch.from_numpy(np.asarray(group[f"{name}_real"][selection])) + imag = torch.from_numpy(np.asarray(group[f"{name}_imag"][selection])) + value = torch.complex(real, imag).to(torch.complex128) + + return value.to(device) if device is not None else value + + +def read_optional_metadata(dataset): + """Decode an optional array stored as an array or the string ``None``.""" + + value = dataset[()] + if isinstance(value, np.ndarray): + return value + if isinstance(value, bytes) and value.decode() == "None": + return None + if isinstance(value, str) and value == "None": + return None + raise ValueError(f"Unsupported value {value} for key {dataset.name}") + + +def read_btd_kpoint(cache, kpoint, voltage, device): + """Read only one k point, retaining the existing RGF return convention.""" + + + k_index = find_kpoint_index(cache["kpoints"][()], kpoint) + blocks = cache["subblocks"][()] + parts = {} + + for name in ("hd", "sd", "hl", "su", "sl", "hu"): + + count = len(blocks) if name.endswith("d") else len(blocks) - 1 + parts[name] = [ + read_complex( + cache[name], + f"{name}_k{k_index}_b{block_index}", + device=device, + ) + for block_index in range(count) + ] + voltage = torch.as_tensor( + 0 if voltage is None else voltage, device=device + ) + if voltage.ndim == 0: + voltage = voltage.expand(int(sum(blocks))) + if voltage.numel() != sum(blocks): + raise ValueError( + "Device voltage must be scalar or have one value per orbital." + ) + start = 0 + + for block_index, size in enumerate(blocks): + local_voltage = voltage[start:start + size].reshape(-1, 1) + parts["hd"][block_index] -= ( + local_voltage * parts["sd"][block_index] + ) + if block_index < len(blocks) - 1: + parts["hu"][block_index] -= ( + local_voltage * parts["su"][block_index] + ) + if block_index > 0: + parts["hl"][block_index - 1] -= ( + local_voltage * parts["sl"][block_index - 1] + ) + start += size + + return tuple( + parts[name] for name in ("hd", "sd", "hl", "su", "sl", "hu") + ) \ No newline at end of file From 5c3873162c470cdd187321a4e13c49008eb33745 Mon Sep 17 00:00:00 2001 From: AsymmetryChou <181240085@smail.nju.edu.cn> Date: Thu, 1 Oct 2026 19:21:47 +0800 Subject: [PATCH 3/3] Add use_jit =False --- dpnegf/negf/split_btd.py | 9 +++++++-- dpnegf/tests/test_split_btd_compute_blocks.py | 12 ++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/dpnegf/negf/split_btd.py b/dpnegf/negf/split_btd.py index 179f245..ddf436c 100644 --- a/dpnegf/negf/split_btd.py +++ b/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() diff --git a/dpnegf/tests/test_split_btd_compute_blocks.py b/dpnegf/tests/test_split_btd_compute_blocks.py index bdadb1f..dd30a23 100644 --- a/dpnegf/tests/test_split_btd_compute_blocks.py +++ b/dpnegf/tests/test_split_btd_compute_blocks.py @@ -25,6 +25,18 @@ split_into_subblocks_optimized, ) +@pytest.mark.parametrize("seed", range(20)) +def test_nonjit_partition_matches_jit(seed): + random = np.random.default_rng(seed) + mask = random.random((80, 80)) > .95 + mask |= mask.T.copy() + edge, reverse = compute_edge(mask) + assert compute_blocks( + 3, 7, edge, reverse, use_jit=False + ) == compute_blocks( + 3, 7, edge, reverse, use_jit=True + ) + # --------------------------------------------------------------------------- # Legacy oracles — verbatim copies of the pre-refactor implementations.