diff --git a/dpnegf/negf/lead_property.py b/dpnegf/negf/lead_property.py index 1b5bd52..f3df706 100644 --- a/dpnegf/negf/lead_property.py +++ b/dpnegf/negf/lead_property.py @@ -815,7 +815,7 @@ def compute_all_self_energy(eta, lead_L, lead_R, kpoints_grid, energy_grid, # logging state and the WARNING default) can match it when they reinit. parent_log_level = logging.getLogger().getEffectiveLevel() if len(total_tasks) <= ek_batch_size: - Parallel(n_jobs=safe_n_jobs, backend="loky")( + Parallel(n_jobs=min(safe_n_jobs, len(total_tasks)), backend="loky")( delayed(_self_energy_worker_blas)(k, e, eta, leadL_pack, leadR_pack, self_energy_save_path, se_numba_jit, parent_log_level, blas_threads_per_worker) @@ -824,7 +824,7 @@ def compute_all_self_energy(eta, lead_L, lead_R, kpoints_grid, energy_grid, else: for i in range(0, len(total_tasks), ek_batch_size): batch = total_tasks[i:i+ek_batch_size] - Parallel(n_jobs=safe_n_jobs, backend="loky")( + Parallel(n_jobs=min(safe_n_jobs, len(batch)), backend="loky")( delayed(_self_energy_worker_blas)(k, e, eta, leadL_pack, leadR_pack, self_energy_save_path, se_numba_jit, parent_log_level, blas_threads_per_worker) @@ -888,11 +888,7 @@ def _precompute_lead_kdata(lead, kpoints_grid): HLk, HLLk, HDLk, SLk, SLLk, SDLk = lead.hamiltonian.get_hs_lead( k, tab=lead.tab, v=lead.voltage ) - pack["kdata"][key] = { - "subblocks": subblocks, - "HLk": HLk, "HLLk": HLLk, "HDLk": HDLk, - "SLk": SLk, "SLLk": SLLk, "SDLk": SDLk, - } + pack["kdata"][key] = _pack_lead_matrices(HLk, HLLk, HDLk, SLk, SLLk, SDLk, subblocks) else: kpoints_bloch = bloch_unfolder.unfold_points(list(np.asarray(k, dtype=float).reshape(3))) bloch_entries = [] @@ -901,11 +897,9 @@ def _precompute_lead_kdata(lead, kpoints_grid): HLk, HLLk, HDLk, SLk, SLLk, SDLk = lead.hamiltonian.get_hs_lead( kb_tensor, tab=lead.tab, v=lead.voltage ) - bloch_entries.append({ - "k_bloch": kb_tensor, - "HLk": HLk, "HLLk": HLLk, "HDLk": HDLk, - "SLk": SLk, "SLLk": SLLk, "SDLk": SDLk, - }) + entry = _pack_lead_matrices(HLk, HLLk, HDLk, SLk, SLLk, SDLk, subblocks=None) + entry["k_bloch"] = kb_tensor + bloch_entries.append(entry) pack["kdata"][key] = { "subblocks": subblocks, "bloch_entries": bloch_entries, @@ -914,6 +908,44 @@ def _precompute_lead_kdata(lead, kpoints_grid): return pack +def _to_numpy_c128(x): + """Detach torch tensors and cast to a C-contiguous complex128 numpy array. + + Workers receive these plain arrays so `selfEnergy` / `surface_green` skip the + per-energy torch→numpy round-trip, and the numba core can consume them + directly without the `@njit` boundary re-inferring layout. + """ + if isinstance(x, torch.Tensor): + arr = x.detach().numpy() + else: + arr = np.asarray(x) + return np.ascontiguousarray(arr, dtype=np.complex128) + + +def _pack_lead_matrices(HLk, HLLk, HDLk, SLk, SLLk, SDLk, subblocks): + """Build a per-k pack entry. + + HLk/HLLk/SLk/SLLk go into the surface-Green core as numpy — convert them + once here so workers don't repeat the torch→numpy round-trip per energy, + and precompute ``h10 = conj(HLLk.T)`` / ``s10 = conj(SLLk.T)`` for the same + reason. HDLk/SDLk stay as torch tensors because `LeadProperty.HDL_reduced` + and the Bloch-branch matmul below use torch ops on them. + """ + HLk_np = _to_numpy_c128(HLk) + HLLk_np = _to_numpy_c128(HLLk) + SLk_np = _to_numpy_c128(SLk) + SLLk_np = _to_numpy_c128(SLLk) + entry = { + "HLk": HLk_np, "HLLk": HLLk_np, "HDLk": HDLk, + "SLk": SLk_np, "SLLk": SLLk_np, "SDLk": SDLk, + "h10": np.ascontiguousarray(np.conj(HLLk_np.T)), + "s10": np.ascontiguousarray(np.conj(SLLk_np.T)), + } + if subblocks is not None: + entry["subblocks"] = subblocks + return entry + + def _compute_self_energy_from_pack(pack, k, e, eta_lead, method="Lopez-Sancho", se_numba_jit=None): """Pure-function port of LeadProperty.self_energy_cal that operates on the dict produced by _precompute_lead_kdata. Mirrors the math in @@ -942,7 +974,8 @@ def _compute_self_energy_from_pack(pack, k, e, eta_lead, method="Lopez-Sancho", E_ref=E_ref, etaLead=eta_lead, method=method, - numba_jit=se_numba_jit + numba_jit=se_numba_jit, + h10=entry["h10"], s10=entry["s10"], ) return se @@ -963,7 +996,8 @@ def _compute_self_energy_from_pack(pack, k, e, eta_lead, method="Lopez-Sancho", E_ref=E_ref, etaLead=eta_lead, method=method, - numba_jit=se_numba_jit + numba_jit=se_numba_jit, + h10=be["h10"], s10=be["s10"], ) phase_factor_m = torch.zeros([m_size, m_size], dtype=torch.complex128) bloch_R_list = pack["bloch_R_list"] diff --git a/dpnegf/negf/surface_green.py b/dpnegf/negf/surface_green.py index ba8de49..d14da04 100644 --- a/dpnegf/negf/surface_green.py +++ b/dpnegf/negf/surface_green.py @@ -2,14 +2,14 @@ import scipy.linalg as SLA import logging import torch -from numba import njit, float64, complex128, int64 log = logging.getLogger(__name__) -def selfEnergy(hL, hLL, sL, sLL, ee, hDL=None, sDL=None, etaLead=1e-8, Bulk=False, - E_ref=0.0, dtype=np.complex128, device='cpu', method='Lopez-Sancho', numba_jit=None): +def selfEnergy(hL, hLL, sL, sLL, ee, hDL=None, sDL=None, etaLead=1e-8, Bulk=False, + E_ref=0.0, dtype=np.complex128, device='cpu', method='Lopez-Sancho', numba_jit=None, + h10=None, s10=None): '''calculates the self-energy and surface Green's function for a given Hamiltonian and overlap matrix. - + Parameters ---------- hL @@ -23,7 +23,7 @@ def selfEnergy(hL, hLL, sL, sLL, ee, hDL=None, sDL=None, etaLead=1e-8, Bulk=Fals ee the given energy hDL - Hamiltonian matrix between the lead and the device. + Hamiltonian matrix between the lead and the device. sDL Overlap matrix between the lead and the device. etaLead @@ -33,22 +33,25 @@ def selfEnergy(hL, hLL, sL, sLL, ee, hDL=None, sDL=None, etaLead=1e-8, Bulk=Fals chemiPot the chemical potential of the lead. dtype - the data type of the tensors used in the calculations. + the data type of the tensors used in the calculations. device The "device" parameter specifies the device on which the calculations will be performed. It can be set to 'cpu' for CPU computation or 'cuda' for GPU computation. method - specify the method for calculating the surface Green's function.The available options + specify the method for calculating the surface Green's function.The available options are "Lopez-Sancho" and any other value will default to "Lopez-Sancho". numba_jit - A boolean flag that indicates whether to use Numba's Just-In-Time (JIT) compilation for the surface Green's function calculation. - + A boolean flag that indicates whether to use Numba's Just-In-Time (JIT) compilation for the surface Green's function calculation. + h10, s10 + Optional precomputed ``conj(hLL.T)`` and ``conj(sLL.T)``. Passed in as + C-contiguous complex128 numpy arrays from ``_precompute_lead_kdata`` so + the surface-Green core does not recompute them per energy. + Returns ------- two values: Sig and SGF. The former is self-energy and the latter is surface Green's function. - + ''' - # 确保输入是NumPy数组 hL = convert_to_numpy(hL) sL = convert_to_numpy(sL) hLL = convert_to_numpy(hLL) @@ -59,35 +62,37 @@ def selfEnergy(hL, hLL, sL, sLL, ee, hDL=None, sDL=None, etaLead=1e-8, Bulk=Fals if sDL is not None: sDL = convert_to_numpy(sDL) E_ref = convert_to_numpy(E_ref) + if h10 is not None: + h10 = convert_to_numpy(h10) + if s10 is not None: + s10 = convert_to_numpy(s10) - - if not isinstance(ee, np.ndarray): eeshifted = np.array(ee, dtype=dtype) + E_ref else: eeshifted = ee + E_ref - + eeshifted = eeshifted.item() - + if hDL is None: ESH = (eeshifted * sL - hL) - SGF = surface_green(hL, hLL, sL, sLL, eeshifted + 1j * etaLead , method, - numba_jit=numba_jit) - + SGF = surface_green(hL, hLL, sL, sLL, eeshifted + 1j * etaLead, method, + numba_jit=numba_jit, h10=h10, s10=s10) + if Bulk: Sig = np.linalg.inv(SGF) else: Sig = ESH - np.linalg.inv(SGF) else: a, b = hDL.shape - SGF = surface_green(hL, hLL, sL, sLL, eeshifted + 1j * etaLead , method, - numba_jit=numba_jit) - + SGF = surface_green(hL, hLL, sL, sLL, eeshifted + 1j * etaLead, method, + numba_jit=numba_jit, h10=h10, s10=s10) + Sig = (eeshifted*sDL-hDL) @ SGF[:b,:b] @ (eeshifted*sDL.conj().T-hDL.conj().T) - + Sig = torch.tensor(Sig, dtype=torch.complex128, device=device) SGF = torch.tensor(SGF, dtype=torch.complex128, device=device) - + return Sig, SGF @@ -95,68 +100,80 @@ def selfEnergy(hL, hLL, sL, sLL, ee, hDL=None, sDL=None, etaLead=1e-8, Bulk=Fals try: - from numba import njit, complex128, int64, float64 - from numba.types import Tuple - NumbaReturnType = Tuple((complex128[:,:], int64, float64, float64)) - - @njit(NumbaReturnType(complex128[:,:], complex128[:,:], complex128[:,:], complex128[:,:], complex128)) - def _surface_green_numba_core(H, h01, S, s01, ee): - - N = H.shape[0] - h10 = np.conj(h01.T) - s10 = np.conj(s01.T) - alpha = h10 - ee * s10 - beta = h01 - ee * s01 - eps = H.copy() - epss = H.copy() - - eS = ee * S # loop invariant - - # preallocated scratch buffer reused every iteration - RHS = np.empty((N, 2 * N), dtype=np.complex128) - - converged = False - iteration = 0 - while not converged: - iteration += 1 - # one LU solve with stacked RHS = [alpha | beta] - RHS[:, :N] = alpha - RHS[:, N:] = beta - sol = np.linalg.solve(eS - eps, RHS) - tmpa = np.ascontiguousarray(sol[:, :N]) - tmpb = np.ascontiguousarray(sol[:, N:]) - - # update eps/epss in place while alpha/beta still hold previous values; - # beta @ tmpa is reused in both updates, so compute once. - beta_tmpa = beta @ tmpa - eps += alpha @ tmpb - eps += beta_tmpa - epss += beta_tmpa - - alpha = alpha @ tmpa - beta = beta @ tmpb - - LopezConvTest = np.max(np.abs(alpha) + np.abs(beta)) - - if LopezConvTest < 1.0e-40: - gs = np.linalg.inv(eS - epss) - - test = eS - H - (ee * s01 - h01) @ gs @ (ee * s10 - h10) - myConvTest = np.max(np.abs((test @ gs) - np.eye(H.shape[0], dtype=h01.dtype))) - - if myConvTest < 3.0e-5: - converged = True - if myConvTest > 1.0e-8: # warning threshold - return gs, 1, myConvTest, ee.real + from numba import njit + + def _make_surface_green_numba_core(fastmath): + # Lazy signature: dispatcher compiles one specialization per input + # layout, so both writable arrays (parent process / tests) and + # joblib-memmapped readonly arrays (loky workers, see + # `_pack_lead_matrices` in lead_property.py) match without falling + # back to the scipy core. + @njit(fastmath=fastmath) + def _core(H, h01, S, s01, h10, s10, ee): + + N = H.shape[0] + alpha = h10 - ee * s10 + beta = h01 - ee * s01 + eps = H.copy() + epss = H.copy() + + eS = ee * S # loop invariant + + # preallocated scratch buffer reused every iteration + RHS = np.empty((N, 2 * N), dtype=np.complex128) + + converged = False + iteration = 0 + while not converged: + iteration += 1 + # one LU solve with stacked RHS = [alpha | beta] + RHS[:, :N] = alpha + RHS[:, N:] = beta + sol = np.linalg.solve(eS - eps, RHS) + + # Two batched GEMMs: AS = alpha @ [tmpa | tmpb], BS = beta @ [tmpa | tmpb]. + # Splits into the four (N,N)@(N,N) products the recursion needs, but the + # ZGEMM "A" panel is packed once per merged call instead of twice. + AS = alpha @ sol + BS = beta @ sol + + # next iter uses these as GEMM "A" operand; copy so they are C-contig + alpha = np.ascontiguousarray(AS[:, :N]) + beta = np.ascontiguousarray(BS[:, N:]) + + # AS[:, N:] = alpha_prev @ tmpb; BS[:, :N] = beta_prev @ tmpa + eps += AS[:, N:] + BS[:, :N] + epss += BS[:, :N] + + # Cheaper probe: max(max|alpha|, max|beta|) < 1e-40 is equivalent to the old + # max(|alpha|+|beta|) < 1e-40 up to a factor of 2, which is trivially covered. + LopezConvTest = max(np.max(np.abs(alpha)), np.max(np.abs(beta))) + + if LopezConvTest < 1.0e-40: + gs = np.linalg.inv(eS - epss) + + test = eS - H - (ee * s01 - h01) @ gs @ (ee * s10 - h10) + myConvTest = np.max(np.abs((test @ gs) - np.eye(H.shape[0], dtype=h01.dtype))) + + if myConvTest < 3.0e-5: + converged = True + if myConvTest > 1.0e-8: # warning threshold + return gs, 1, myConvTest, ee.real + else: + return gs, 0, 0, 0 else: - return gs, 0, 0, 0 - else: - raise ArithmeticError + raise ArithmeticError - if iteration >= 101: - raise RuntimeError + if iteration >= 101: + raise RuntimeError - return gs + return gs + return _core + + _surface_green_numba_core = _make_surface_green_numba_core(fastmath=True) + # Regression-test sibling: identical body, fastmath disabled. Not called by + # production code; kept so `tests/test_surface_green_fastmath.py` can compare. + _surface_green_numba_core_nofastmath = _make_surface_green_numba_core(fastmath=False) _numba_available = True log.info("Numba is available and JIT functions are compiled.") @@ -166,10 +183,8 @@ def _surface_green_numba_core(H, h01, S, s01, ee): _numba_available = False # Scipy-based implementation of the surface Green's function calculation -def _surface_green_scipy_core(H, h01, S, s01, ee): +def _surface_green_scipy_core(H, h01, S, s01, h10, s10, ee): N = H.shape[0] - h10 = np.conj(h01.T) - s10 = np.conj(s01.T) alpha = h10 - ee * s10 beta = h01 - ee * s01 @@ -190,20 +205,18 @@ def _surface_green_scipy_core(H, h01, S, s01, ee): RHS[:, N:] = beta sol = SLA.solve(eS - eps, RHS, overwrite_a=True, overwrite_b=True, check_finite=False) - tmpa = np.ascontiguousarray(sol[:, :N]) - tmpb = np.ascontiguousarray(sol[:, N:]) - # update eps/epss in place while alpha/beta still hold previous values; - # beta @ tmpa is reused in both updates, so compute once. - beta_tmpa = beta @ tmpa - eps += alpha @ tmpb - eps += beta_tmpa - epss += beta_tmpa + # Same batched-GEMM structure as the numba core. + AS = alpha @ sol + BS = beta @ sol - alpha = alpha @ tmpa - beta = beta @ tmpb + alpha = np.ascontiguousarray(AS[:, :N]) + beta = np.ascontiguousarray(BS[:, N:]) - LopezConvTest = np.max(np.abs(alpha) + np.abs(beta)) + eps += AS[:, N:] + BS[:, :N] + epss += BS[:, :N] + + LopezConvTest = max(np.max(np.abs(alpha)), np.max(np.abs(beta))) if LopezConvTest < 1.0e-40: gs = np.linalg.inv(eS - epss) @@ -228,55 +241,46 @@ def _surface_green_scipy_core(H, h01, S, s01, ee): def surface_green(H, h01, S, s01, ee, method='Lopez-Sancho', - numba_jit=None): + numba_jit=None, + h10=None, s10=None): '''calculate surface green function At this stage, we realized Lopez-Sancho scheme and GEP scheme. However, GEP scheme is not so stable, and we strongly recommended to implement the Lopez-Sancho scheme. + ``h10`` / ``s10`` (``conj(h01.T)`` / ``conj(s01.T)``) may be passed in from + the precomputed lead pack to avoid re-doing the transpose+conj per energy. + They are derived on the fly if omitted. + ''' # default: use numba whenever it compiled successfully if numba_jit is None: numba_jit = _numba_available + if h10 is None: + h10 = np.ascontiguousarray(np.conj(h01.T)) + if s10 is None: + s10 = np.ascontiguousarray(np.conj(s01.T)) + if method == 'GEP': gs = calcg0(ee, H, S, h01, s01) return gs else: # Lopez-Sancho scheme if numba_jit and _numba_available: try: - # check - # 1. type check - assert isinstance(H, np.ndarray), "H must be a NumPy array." - assert isinstance(h01, np.ndarray), "h01 must be a NumPy array." - assert isinstance(S, np.ndarray), "S must be a NumPy array." - assert isinstance(s01, np.ndarray), "s01 must be a NumPy array." - assert isinstance(ee, (complex, float, int)), "ee must be a complex, float, or integer scalar." - - # 2. dimension check - assert H.ndim == 2, "H must be a 2D array." - assert h01.ndim == 2, "h01 must be a 2D array." - assert S.ndim == 2, "S must be a 2D array." - assert s01.ndim == 2, "s01 must be a 2D array." - - # 3. complex type check - assert np.iscomplexobj(H), "H must be a complex array." - assert np.iscomplexobj(h01), "h01 must be a complex array." - assert np.iscomplexobj(S), "S must be a complex array." - assert np.iscomplexobj(s01), "s01 must be a complex array." # normalize ee to numpy complex128 to match the @njit signature ee = np.complex128(ee) log.debug("surface_green: using numba core") - gs, conv_flag, conv_test, e_real = _surface_green_numba_core(H, h01, S, s01, ee) + gs, conv_flag, conv_test, e_real = _surface_green_numba_core(H, h01, S, s01, h10, s10, ee) if conv_flag == 1: log.warning(f"Lopez-Sancho scheme not-so-well converged at E = {e_real:.4f} eV: {conv_test}") return gs except Exception as e: log.error(f"Numba JIT function failed at runtime. Falling back to NumPy. Error: {e}") - return _surface_green_scipy_core(H, h01, S, s01, ee) + return _surface_green_scipy_core(H, h01, S, s01, h10, s10, ee) else: log.debug("surface_green: using scipy core") - return _surface_green_scipy_core(H, h01, S, s01, ee) + return _surface_green_scipy_core(H, h01, S, s01, h10, s10, ee)