 ### Merge branch 'sparse-modes' into 'master'

```Clean up linear algebra (related to modes)

See merge request kwant/kwant!389```
parents a8947abf 5832d028
Pipeline #74710 passed with stages
in 9 minutes and 44 seconds
 # Copyright 2011-2013 Kwant authors. # Copyright 2011-2021 Kwant authors. # # This file is part of Kwant. It is subject to the license terms in the file # LICENSE.rst found in the top-level directory of this distribution and at ... ... @@ -10,10 +10,6 @@ __all__ = ['lapack'] from . import lapack # Merge the public interface of the other submodules. from .decomp_lu import * from .decomp_schur import * from .decomp_ev import * __all__.extend([decomp_lu.__all__, decomp_ev.__all__, decomp_schur.__all__]) __all__.extend([decomp_schur.__all__])
 # Copyright 2011-2013 Kwant authors. # # This file is part of Kwant. It is subject to the license terms in the file # LICENSE.rst found in the top-level directory of this distribution and at # https://kwant-project.org/license. A list of Kwant authors can be found in # the file AUTHORS.rst at the top-level directory of this distribution and at # https://kwant-project.org/authors. __all__ = ['gen_eig'] from . import lapack def gen_eig(a, b, left=False, right=True, overwrite_ab=False): """Compute the eigenvalues and -vectors of the matrix pencil (a,b), i.e. of the generalized (unsymmetric) eigenproblem a v = lambda b v where a and b are square (unsymmetric) matrices, v the eigenvector and lambda the eigenvalues. The eigenvalues are returned as numerator alpha and denominator beta, i.e. lambda = alpha/beta. This is advantageous, as lambda can be infinity which is well-defined in this case as beta = 0. Parameters ---------- a : array, shape (M, M) b : array, shape (M, M) `a` and `b` are the two matrices defining the generalized eigenproblem left : boolean Whether to calculate and return left eigenvectors right : boolean Whether to calculate and return right eigenvectors overwrite_ab : boolean Whether to overwrite data in `a` and `b` (may improve performance) Returns ------- alpha : complex array, shape (M,) beta : real or complex array, shape (M,) The eigenvalues in the form ``alpha/beta`` (if left == True) vl : double or complex array, shape (M, M) The left eigenvector corresponding to the eigenvalue ``alpha[i]/beta[i]`` is the column ``vl[:,i]``. (if right == True) vr : double or complex array, shape (M, M) The right eigenvector corresponding to the eigenvalue ``alpha[i]/beta[i]`` is the column ``vr[:,i]``. """ a, b = lapack.prepare_for_lapack(overwrite_ab, a, b) return lapack.ggev(a, b, left, right)
 # Copyright 2011-2013 Kwant authors. # # This file is part of Kwant. It is subject to the license terms in the file # LICENSE.rst found in the top-level directory of this distribution and at # https://kwant-project.org/license. A list of Kwant authors can be found in # the file AUTHORS.rst at the top-level directory of this distribution and at # https://kwant-project.org/authors. __all__ = ['lu_factor', 'lu_solve', 'rcond_from_lu'] import numpy as np from . import lapack def lu_factor(a, overwrite_a=False): """Compute the LU factorization of a matrix A = P * L * U. The function returns a tuple (lu, p, singular), where lu contains the LU factorization storing the unit lower triangular matrix L in the strictly lower triangle (the unit diagonal is not stored) and the upper triangular matrix U in the upper triangle. p is a vector of pivot indices, and singular a Boolean value indicating whether the matrix A is singular up to machine precision. NOTE: This function mimics the behavior of scipy.linalg.lu_factor (except that it has in addition the flag singular). The main reason is that lu_factor in SciPy has a bug that depending on the type of NumPy matrix passed to it, it would not return what was descirbed in the documentation. This bug will be (probably) fixed in 0.10.0 but until this is standard, this version is better to use. Parameters ---------- a : array, shape (M, M) Matrix to factorize overwrite_a : boolean Whether to overwrite data in a (may increase performance) Returns ------- lu : array, shape (N, N) Matrix containing U in its upper triangle, and L in its lower triangle. The unit diagonal elements of L are not stored. piv : array, shape (N,) Pivot indices representing the permutation matrix P: row i of matrix was interchanged with row piv[i]. singular : boolean Whether the matrix a is singular (up to machine precision) """ a = lapack.prepare_for_lapack(overwrite_a, a) return lapack.getrf(a) def lu_solve(matrix_factorization, b): """Solve a linear system of equations, a x = b, given the LU factorization of a Parameters ---------- matrix_factorization Factorization of the coefficient matrix a, as given by lu_factor b : array (vector or matrix) Right-hand side Returns ------- x : array (vector or matrix) Solution to the system """ (lu, ipiv, singular) = matrix_factorization if singular: raise RuntimeWarning("In lu_solve: the flag singular indicates " "a singular matrix. Result of solve step " "are probably unreliable") lu, b = lapack.prepare_for_lapack(False, lu, b) ipiv = np.ascontiguousarray(np.asanyarray(ipiv), dtype=lapack.int_dtype) return lapack.getrs(lu, ipiv, b) def rcond_from_lu(matrix_factorization, norm_a, norm="1"): """Compute the reciprocal condition number from the LU decomposition as returned from lu_factor(), given additionally the norm of the matrix a in norm_a. The reciprocal condition number is given as 1/(||A||*||A^-1||), where ||...|| is a matrix norm. Parameters ---------- matrix_factorization Factorization of the matrix a, as given by lu_factor norm_a : float or complex norm of the original matrix a (type of norm is specified in norm) norm : {'1', 'I'}, optional type of matrix norm which should be used to compute the condition number ("1": 1-norm, "I": infinity norm). Default: '1'. Returns ------- rcond : float or complex reciprocal condition number of a with respect to the type of matrix norm specified in norm """ (lu, ipiv, singular) = matrix_factorization norm = norm.encode('utf8') # lapack expects bytes lu = lapack.prepare_for_lapack(False, lu) return lapack.gecon(lu, norm_a, norm)
 ... ... @@ -302,8 +302,7 @@ def gen_schur(a, b, calc_q=True, calc_z=True, calc_ev=True, problem (the entries of the diagonal of the complex Schur form are the eigenvalues of the matrix, for example), and the routine can optionally also return the generalized eigenvalues in the form (alpha, beta), such that alpha/beta is a generalized eigenvalue of the pencil (a, b) (see also gen_eig()). that alpha/beta is a generalized eigenvalue of the pencil (a, b). Parameters ---------- ... ...
 ... ... @@ -8,9 +8,7 @@ """Low-level access to LAPACK functions. """ __all__ = ['getrf', 'getrs', 'gecon', __all__ = ['gecon', 'ggev', 'gees', 'trsen', ... ... @@ -92,85 +90,6 @@ cdef l_int lwork_from_qwork(scalar qwork): return qwork.real def getrf(np.ndarray[scalar, ndim=2] A): cdef l_int M, N, info cdef np.ndarray[l_int] ipiv assert_fortran_mat(A) M = A.shape N = A.shape ipiv = np.empty(min(M,N), dtype = int_dtype) if scalar is float: lapack.sgetrf(&M, &N, A.data, &M, ipiv.data, &info) elif scalar is double: lapack.dgetrf(&M, &N, A.data, &M, ipiv.data, &info) elif scalar is float_complex: lapack.cgetrf(&M, &N, A.data, &M, ipiv.data, &info) elif scalar is double_complex: lapack.zgetrf(&M, &N, A.data, &M, ipiv.data, &info) assert info >= 0, "Argument error in getrf" return (A, ipiv, info > 0 or M != N) def getrs(np.ndarray[scalar, ndim=2] LU, np.ndarray[l_int] IPIV, np.ndarray B): cdef l_int N, NRHS, info assert_fortran_mat(LU) # Consistency checks for LU and B if B.descr.type_num != LU.descr.type_num: raise TypeError('B must have same dtype as LU') # Workaround for 1x1-Fortran bug in NumPy < v2.0 if ((B.ndim == 2 and (B.shape > 1 or B.shape > 1) and not B.flags["F_CONTIGUOUS"])): raise ValueError("B must be Fortran ordered") if B.ndim > 2: raise ValueError("B must be a vector or matrix") if LU.shape != B.shape: raise ValueError('LU and B have incompatible shapes') N = LU.shape if B.ndim == 1: NRHS = 1 elif B.ndim == 2: NRHS = B.shape if scalar is float: lapack.sgetrs("N", &N, &NRHS, LU.data, &N, IPIV.data, B.data, &N, &info) elif scalar is double: lapack.dgetrs("N", &N, &NRHS, LU.data, &N, IPIV.data, B.data, &N, &info) elif scalar is float_complex: lapack.cgetrs("N", &N, &NRHS, LU.data, &N, IPIV.data, B.data, &N, &info) elif scalar is double_complex: lapack.zgetrs("N", &N, &NRHS, LU.data, &N, IPIV.data, B.data, &N, &info) assert info == 0, "Argument error in getrs" return B def gecon(np.ndarray[scalar, ndim=2] LU, double normA, char *norm = b"1"): cdef l_int N, info cdef float srcond, snormA ... ... @@ -264,155 +183,6 @@ def ggev_postprocess(dtype, alphar, alphai, vl_r=None, vr_r=None): return (alpha, vl, vr) def ggev(np.ndarray[scalar, ndim=2] A, np.ndarray[scalar, ndim=2] B, left=False, right=True): cdef l_int N, info, lwork # Parameter checks assert_fortran_mat(A, B) if A.ndim != 2 or A.ndim != 2: raise ValueError("gen_eig requires both a and be to be matrices") if A.shape != A.shape: raise ValueError("gen_eig requires square matrix input") if A.shape != B.shape or A.shape != B.shape: raise ValueError("A and B do not have the same shape") # Allocate workspaces N = A.shape cdef np.ndarray[scalar] alphar, alphai if scalar in cmplx: alphar = np.empty(N, dtype=A.dtype) alphai = None else: alphar = np.empty(N, dtype=A.dtype) alphai = np.empty(N, dtype=A.dtype) cdef np.ndarray[scalar] beta = np.empty(N, dtype=A.dtype) cdef np.ndarray rwork = None if scalar is float_complex: rwork = np.empty(8 * N, dtype=np.float32) elif scalar is double_complex: rwork = np.empty(8 * N, dtype=np.float64) cdef np.ndarray vl cdef scalar *vl_ptr cdef char *jobvl if left: vl = np.empty((N,N), dtype=A.dtype, order='F') vl_ptr = vl.data jobvl = "V" else: vl = None vl_ptr = NULL jobvl = "N" cdef np.ndarray vr cdef scalar *vr_ptr cdef char *jobvr if right: vr = np.empty((N,N), dtype=A.dtype, order='F') vr_ptr = vr.data jobvr = "V" else: vr = None vr_ptr = NULL jobvr = "N" # Workspace query # Xggev expects &qwork as a (even though it's an integer) lwork = -1 cdef scalar qwork if scalar is float: lapack.sggev(jobvl, jobvr, &N, A.data, &N, B.data, &N, alphar.data, alphai.data, beta.data, vl_ptr, &N, vr_ptr, &N, &qwork, &lwork, &info) elif scalar is double: lapack.dggev(jobvl, jobvr, &N, A.data, &N, B.data, &N, alphar.data, alphai.data, beta.data, vl_ptr, &N, vr_ptr, &N, &qwork, &lwork, &info) elif scalar is float_complex: lapack.cggev(jobvl, jobvr, &N, A.data, &N, B.data, &N, alphar.data, beta.data, vl_ptr, &N, vr_ptr, &N, &qwork, &lwork, rwork.data, &info) elif scalar is double_complex: lapack.zggev(jobvl, jobvr, &N, A.data, &N, B.data, &N, alphar.data, beta.data, vl_ptr, &N, vr_ptr, &N, &qwork, &lwork, rwork.data, &info) assert info == 0, "Argument error in ggev" lwork = lwork_from_qwork(qwork) cdef np.ndarray[scalar] work = np.empty(lwork, dtype=A.dtype) # The actual calculation if scalar is float: lapack.sggev(jobvl, jobvr, &N, A.data, &N, B.data, &N, alphar.data, alphai.data, beta.data, vl_ptr, &N, vr_ptr, &N, work.data, &lwork, &info) elif scalar is double: lapack.dggev(jobvl, jobvr, &N, A.data, &N, B.data, &N, alphar.data, alphai.data, beta.data, vl_ptr, &N, vr_ptr, &N, work.data, &lwork, &info) elif scalar is float_complex: lapack.cggev(jobvl, jobvr, &N, A.data, &N, B.data, &N, alphar.data, beta.data, vl_ptr, &N, vr_ptr, &N, work.data, &lwork, rwork.data, &info) elif scalar is double_complex: lapack.zggev(jobvl, jobvr, &N, A.data, &N, B.data, &N, alphar.data, beta.data, vl_ptr, &N, vr_ptr, &N, work.data, &lwork, rwork.data, &info) if info > 0: raise LinAlgError("QZ iteration failed to converge in sggev") assert info == 0, "Argument error in ggev" if scalar is float: post_dtype = np.complex64 elif scalar is double: post_dtype = np.complex128 cdef np.ndarray alpha alpha = alphar if scalar in floating: alpha, vl, vr = ggev_postprocess(post_dtype, alphar, alphai, vl, vr) return filter_args((True, True, left, right), (alpha, beta, vl, vr)) def gees(np.ndarray[scalar, ndim=2] A, calc_q=True, calc_ev=True): cdef l_int N, lwork, sdim, info ... ...
This diff is collapsed.
 ... ... @@ -9,6 +9,7 @@ from math import sin, cos, sqrt, pi, copysign from collections import namedtuple import warnings from itertools import combinations_with_replacement import numpy as np ... ... @@ -19,11 +20,34 @@ from scipy.linalg import block_diag from scipy.sparse import (identity as sp_identity, hstack as sp_hstack, csr_matrix) dot = np.dot __all__ = ['selfenergy', 'modes', 'PropagatingModes', 'StabilizedModes'] def lu_factor_rcond(mat: np.ndarray): """Perform LU factorization and check condition number. Parameters ---------- mat : numpy array Matrix to be factorized Returns ------- sol : LU factorization rcond : float Condition number """ with warnings.catch_warnings(): warnings.simplefilter("ignore", la.LinAlgWarning) sol = la.lu_factor(mat) lu = np.asfortranarray(sol) gecon = la.lapack.get_lapack_funcs("gecon", [lu]) return sol, gecon(lu, npl.norm(mat, 1)) def nonzero_symm_projection(matrix): """Check whether a discrete symmetry relation between two blocks of the Hamiltonian vanishes or not. ... ... @@ -169,7 +193,7 @@ class StabilizedModes: outgoing = slice(self.nmodes, None) vecs = self.vecs[:, outgoing] vecslmbdainv = self.vecslmbdainv[:, outgoing] return dot(v, dot(vecs, la.solve(vecslmbdainv, v.T.conj()))) return v @ vecs @ la.solve(vecslmbdainv, v.T.conj()) # Auxiliary functions that perform different parts of the calculation. ... ... @@ -180,7 +204,7 @@ def setup_linsys(h_cell, h_hop, tol=1e6, stabilization=None): ---------- h_cell : numpy array with shape (n, n) Hamiltonian of a single lead unit cell. h_hop : numpy array with shape (n, m), m <= n h_hop : numpy array with shape (n, n) Hopping Hamiltonian from a cell to the next one. tol : float Numbers are considered zero when they are smaller than `tol` times ... ... @@ -214,7 +238,6 @@ def setup_linsys(h_cell, h_hop, tol=1e6, stabilization=None): """ n = h_cell.shape m = h_hop.shape if stabilization is not None: stabilization = list(stabilization) ... ... @@ -235,7 +258,6 @@ def setup_linsys(h_cell, h_hop, tol=1e6, stabilization=None): # s is the largest singular value.) u, s, vh = la.svd(h_hop) assert m == vh.shape, "Corrupt output of svd." n_nonsing = np.sum(s > eps * s) if (n_nonsing == n and stabilization is None): ... ... @@ -245,18 +267,18 @@ def setup_linsys(h_cell, h_hop, tol=1e6, stabilization=None): A = h_hop / h_hop_sqrt B = h_hop_sqrt B_H_inv = 1.0 / B # just a real scalar here A_inv = la.inv(A) mA_inv = -la.inv(A) lhs = np.zeros((2*n, 2*n), dtype=np.common_type(h_cell, h_hop)) lhs[:n, :n] = -dot(A_inv, h_cell) * B_H_inv lhs[:n, n:] = -A_inv * B lhs[:n, :n] = (mA_inv @ h_cell) * B_H_inv lhs[:n, n:] = mA_inv * B lhs[n:, :n] = A.T.conj() * B_H_inv def extract_wf(psi, lmbdainv): return B_H_inv * np.copy(psi[:n]) return B_H_inv * psi[:n] matrices = (lhs, None) v_out = h_hop_sqrt * np.eye(n) v = h_hop_sqrt * np.eye(n) else: if stabilization is None: stabilization = [None, False] ... ... @@ -269,10 +291,7 @@ def setup_linsys(h_cell, h_hop, tol=1e6, stabilization=None): u = u[:, :n_nonsing] s = s[:n_nonsing] u = u * np.sqrt(s) # pad v with zeros if necessary v = np.zeros((n, n_nonsing), dtype=vh.dtype) v[:vh.shape] = vh[:n_nonsing].T.conj() v = v * np.sqrt(s) v = vh[:n_nonsing].T.conj() * np.sqrt(s) # Eliminating the zero eigenvalues requires inverting the on-site # Hamiltonian, possibly including a self-energy-like term. The ... ... @@ -289,8 +308,7 @@ def setup_linsys(h_cell, h_hop, tol=1e6, stabilization=None): # Check if there is a chance we will not need to add an imaginary term. if not add_imaginary: h = h_cell sol = kla.lu_factor(h) rcond = kla.rcond_from_lu(sol, npl.norm(h, 1)) sol, rcond = lu_factor_rcond(h) if rcond < eps: need_to_stabilize = True ... ... @@ -301,11 +319,10 @@ def setup_linsys(h_cell, h_hop, tol=1e6, stabilization=None): need_to_stabilize = True # Matrices are complex or need self-energy-like term to be # stabilized. temp = dot(u, u.T.conj()) + dot(v, v.T.conj()) temp = u @ u.T.conj() + v @ v.T.conj() h = h_cell + 1j * temp sol = kla.lu_factor(h) rcond = kla.rcond_from_lu(sol, npl.norm(h, 1)) sol, rcond = lu_factor_rcond(h) # If the condition number of the stabilized h is # still bad, there is nothing we can do. ... ... @@ -317,11 +334,11 @@ def setup_linsys(h_cell, h_hop, tol=1e6, stabilization=None): # the projected one (v^dagger psi lambda^-1, s u^dagger psi). def extract_wf(psi, lmbdainv): wf = -dot(u, psi[: n_nonsing] * lmbdainv) - dot(v, psi[n_nonsing:]) wf = -u @ (psi[: n_nonsing] * lmbdainv) - v @ psi[n_nonsing:] if need_to_stabilize: wf += 1j * (dot(v, psi[: n_nonsing]) + dot(u, psi[n_nonsing:] * lmbdainv)) return kla.lu_solve(sol, wf) wf += 1j * (v @ psi[: n_nonsing] + u @ (psi[n_nonsing:] * lmbdainv)) return la.lu_solve(sol, wf) # Setup the generalized eigenvalue problem. ... ... @@ -331,29 +348,27 @@ def setup_linsys(h_cell, h_hop, tol=1e6, stabilization=None): begin, end = slice(n_nonsing), slice(n_nonsing, None) A[end, begin] = np.identity(n_nonsing) temp = kla.lu_solve(sol, v) temp2 = dot(u.T.conj(), temp) temp = la.lu_solve(sol, v) temp2 = u.T.conj() @ temp if need_to_stabilize: