Skip to content
Snippets Groups Projects
Commit be339a22 authored by Joseph Weston's avatar Joseph Weston
Browse files

switch Xtgevc to fused types

parent 28b93b16
Branches
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
kwant/linalg/decomp_schur.py
+ 1
− 16
View file @ be339a22
......@@ -613,27 +613,12 @@ def evecs_from_gen_schur(s, t, q=None, z=None, select=None,
ltype, s, t, q, z = lapack.prepare_for_lapack(overwrite_qz, s, t, q, z)
if (s.ndim != 2 or t.ndim != 2 or
(q is not None and q.ndim != 2) or
(z is not None and z.ndim != 2)):
raise ValueError("Expect matrices as input")
if ((s.shape[0] != s.shape[1] or t.shape[0] != t.shape[1] or
s.shape[0] != t.shape[0]) or
(q is not None and (q.shape[0] != q.shape[1] or
s.shape[0] != q.shape[0])) or
(z is not None and (z.shape[0] != z.shape[1] or
s.shape[0] != z.shape[0]))):
raise ValueError("Invalid Schur decomposition as input")
if left and q is None:
raise ValueError("Matrix q must be provided for left eigenvectors")
if right and z is None:
raise ValueError("Matrix z must be provided for right eigenvectors")
tgevc = getattr(lapack, ltype + "tgevc")
# Check if select is a function or an array.
if select is not None:
isfun = isarray = True
......@@ -659,4 +644,4 @@ def evecs_from_gen_schur(s, t, q=None, z=None, select=None,
else:
selectarr = None
return tgevc(s, t, q, z, selectarr, left, right)
return lapack.tgevc(s, t, q, z, selectarr, left, right)
kwant/linalg/lapack.pyx
+ 80
− 324
View file @ be339a22
......@@ -17,7 +17,7 @@ __all__ = ['getrf',
'trevc',
'gges',
'tgsen',
'stgevc', 'dtgevc', 'ctgevc', 'ztgevc',
'tgevc',
'prepare_for_lapack']
import numpy as np
......@@ -1108,139 +1108,43 @@ def tgsen(np.ndarray[l_logical] select,
(S, T, Q, Z, alpha, beta))
# xTGEVC
def stgevc(np.ndarray[np.float32_t, ndim=2] S,
np.ndarray[np.float32_t, ndim=2] T,
np.ndarray[np.float32_t, ndim=2] Q=None,
np.ndarray[np.float32_t, ndim=2] Z=None,
np.ndarray[l_logical] select=None,
left=False, right=True):
def tgevc(np.ndarray[scalar, ndim=2] S,
np.ndarray[scalar, ndim=2] T,
np.ndarray[scalar, ndim=2] Q,
np.ndarray[scalar, ndim=2] Z,
np.ndarray[l_logical] select,
left=False, right=True):
cdef l_int N, info, M, MM
cdef char *side
cdef char *howmny
cdef np.ndarray[np.float32_t, ndim=2] vl_r, vr_r
cdef float *vl_r_ptr
cdef float *vr_r_ptr
cdef np.ndarray[l_logical] select_cpy
cdef l_logical *select_ptr
cdef np.ndarray[np.float32_t] work
assert_fortran_mat(S, T, Q, Z)
N = S.shape[0]
work = np.empty(6*N, dtype = np.float32)
if left and right:
side = "B"
elif left:
side = "L"
elif right:
side = "R"
else:
return
backtr = False
if select is not None:
howmny = "S"
MM = select.nonzero()[0].size
# Correct for possible additional storage if a single complex
# eigenvalue is selected.
# For that: Figure out the positions of the 2x2 blocks.
cmplxindx = np.diagonal(S, -1).nonzero()[0]
for i in cmplxindx:
if bool(select[i]) != bool(select[i+1]):
MM += 1
# select is overwritten in stgevc
select_cpy = np.array(select, dtype = logical_dtype,
order = 'F')
select_ptr = <l_logical *>select_cpy.data
else:
MM = N
select_ptr = NULL
if ((left and right and Q is not None and Z is not None) or
(left and not right and Q is not None) or
(right and not left and Z is not None)):
howmny = "B"
backtr = True
else:
howmny = "A"
if left:
if backtr:
vl_r = Q
else:
vl_r = np.empty((N, MM), dtype = np.float32, order='F')
vl_r_ptr = <float *>vl_r.data
else:
vl_r_ptr = NULL
if right:
if backtr:
vr_r = Z
else:
vr_r = np.empty((N, MM), dtype = np.float32, order='F')
vr_r_ptr = <float *>vr_r.data
else:
vr_r_ptr = NULL
# Check parameters
lapack.stgevc(side, howmny, select_ptr,
&N, <float *>S.data, &N,
<float *>T.data, &N,
vl_r_ptr, &N, vr_r_ptr, &N, &MM, &M,
<float *>work.data, &info)
if ((S.shape[0] != S.shape[1] or T.shape[0] != T.shape[1] or
S.shape[0] != T.shape[0]) or
(Q is not None and (Q.shape[0] != Q.shape[1] or
S.shape[0] != Q.shape[0])) or
(Z is not None and (Z.shape[0] != Z.shape[1] or
S.shape[0] != Z.shape[0]))):
raise ValueError("Invalid Schur decomposition as input")
assert info == 0, "Argument error in stgevc"
assert MM == M, "Unexpected number of eigenvectors returned in stgevc"
assert_fortran_mat(S, T, Q, Z)
if not backtr:
if left:
vl_r = np.asfortranarray(np.dot(Q, vl_r))
if right:
vr_r = np.asfortranarray(np.dot(Z, vr_r))
# Allocate workspaces
# If there are complex eigenvalues, we need to postprocess the eigenvectors
if np.diagonal(S, -1).nonzero()[0].size:
if left:
vl = txevc_postprocess(np.complex64, S, vl_r, select)
if right:
vr = txevc_postprocess(np.complex64, S, vr_r, select)
else:
if left:
vl = vl_r
if right:
vr = vr_r
N = S.shape[0]
if left and right:
return (vl, vr)
elif left:
return vl
cdef np.ndarray[scalar] work
if scalar in floating:
work = np.empty(6 * N, dtype=S.dtype)
else:
return vr
work = np.empty(2 * N, dtype=S.dtype)
cdef np.ndarray rwork = None
if scalar is float_complex:
rwork = np.empty(2 * N, dtype=np.float32)
elif scalar is double_complex:
rwork = np.empty(2 * N, dtype=np.float64)
def dtgevc(np.ndarray[np.float64_t, ndim=2] S,
np.ndarray[np.float64_t, ndim=2] T,
np.ndarray[np.float64_t, ndim=2] Q=None,
np.ndarray[np.float64_t, ndim=2] Z=None,
np.ndarray[l_logical] select=None,
left=False, right=True):
cdef l_int N, info, M, MM
cdef char *side
cdef char *howmny
cdef np.ndarray[np.float64_t, ndim=2] vl_r, vr_r
cdef double *vl_r_ptr
cdef double *vr_r_ptr
cdef np.ndarray[l_logical] select_cpy
cdef l_logical *select_ptr
cdef np.ndarray[np.float64_t] work
assert_fortran_mat(S, T, Q, Z)
N = S.shape[0]
work = np.empty(6*N, dtype = np.float64)
if left and right:
side = "B"
elif left:
......@@ -1250,8 +1154,11 @@ def dtgevc(np.ndarray[np.float64_t, ndim=2] S,
else:
return
backtr = False
cdef l_logical backtr = False
cdef char *howmny
cdef np.ndarray[l_logical] select_cpy = None
cdef l_logical *select_ptr
if select is not None:
howmny = "S"
MM = select.nonzero()[0].size
......@@ -1263,8 +1170,8 @@ def dtgevc(np.ndarray[np.float64_t, ndim=2] S,
if bool(select[i]) != bool(select[i+1]):
MM += 1
# select is overwritten in dtgevc
select_cpy = np.array(select, dtype = logical_dtype,
# select is overwritten in tgevc
select_cpy = np.array(select, dtype=logical_dtype,
order = 'F')
select_ptr = <l_logical *>select_cpy.data
else:
......@@ -1278,32 +1185,55 @@ def dtgevc(np.ndarray[np.float64_t, ndim=2] S,
else:
howmny = "A"
cdef np.ndarray[scalar, ndim=2] vl_r
cdef scalar *vl_r_ptr
if left:
if backtr:
vl_r = Q
else:
vl_r = np.empty((N, MM), dtype = np.float64, order='F')
vl_r_ptr = <double *>vl_r.data
vl_r = np.empty((N, MM), dtype=S.dtype, order='F')
vl_r_ptr = <scalar *>vl_r.data
else:
vl_r_ptr = NULL
cdef np.ndarray[scalar, ndim=2] vr_r
cdef scalar *vr_r_ptr
if right:
if backtr:
vr_r = Z
else:
vr_r = np.empty((N, MM), dtype = np.float64, order='F')
vr_r_ptr = <double *>vr_r.data
vr_r = np.empty((N, MM), dtype=S.dtype, order='F')
vr_r_ptr = <scalar *>vr_r.data
else:
vr_r_ptr = NULL
lapack.dtgevc(side, howmny, select_ptr,
&N, <double *>S.data, &N,
<double *>T.data, &N,
vl_r_ptr, &N, vr_r_ptr, &N, &MM, &M,
<double *>work.data, &info)
if scalar is float:
lapack.stgevc(side, howmny, select_ptr,
&N, <float *>S.data, &N,
<float *>T.data, &N,
vl_r_ptr, &N, vr_r_ptr, &N, &MM, &M,
<float *>work.data, &info)
elif scalar is double:
lapack.dtgevc(side, howmny, select_ptr,
&N, <double *>S.data, &N,
<double *>T.data, &N,
vl_r_ptr, &N, vr_r_ptr, &N, &MM, &M,
<double *>work.data, &info)
elif scalar is float_complex:
lapack.ctgevc(side, howmny, select_ptr,
&N, <float complex *>S.data, &N,
<float complex *>T.data, &N,
vl_r_ptr, &N, vr_r_ptr, &N, &MM, &M,
<float complex *>work.data, <float *>rwork.data, &info)
elif scalar is double_complex:
lapack.ztgevc(side, howmny, select_ptr,
&N, <double complex *>S.data, &N,
<double complex *>T.data, &N,
vl_r_ptr, &N, vr_r_ptr, &N, &MM, &M,
<double complex *>work.data, <double *>rwork.data, &info)
assert info == 0, "Argument error in dtgevc"
assert MM == M, "Unexpected number of eigenvectors returned in dtgevc"
assert info == 0, "Argument error in tgevc"
assert MM == M, "Unexpected number of eigenvectors returned in tgevc"
if not backtr:
if left:
......@@ -1311,196 +1241,22 @@ def dtgevc(np.ndarray[np.float64_t, ndim=2] S,
if right:
vr_r = np.asfortranarray(np.dot(Z, vr_r))
# If there are complex eigenvalues, we need to postprocess the
# eigenvectors.
if np.diagonal(S, -1).nonzero()[0].size:
if left:
vl = txevc_postprocess(np.complex128, S, vl_r, select)
if right:
vr = txevc_postprocess(np.complex128, S, vr_r, select)
else:
if left:
vl = vl_r
if right:
vr = vr_r
if left and right:
return (vl, vr)
elif left:
return vl
else:
return vr
def ctgevc(np.ndarray[np.complex64_t, ndim=2] S,
np.ndarray[np.complex64_t, ndim=2] T,
np.ndarray[np.complex64_t, ndim=2] Q=None,
np.ndarray[np.complex64_t, ndim=2] Z=None,
np.ndarray[l_logical] select=None,
left=False, right=True):
cdef l_int N, info, M, MM
cdef char *side
cdef char *howmny
cdef np.ndarray[np.complex64_t, ndim=2] vl, vr
cdef float complex *vl_ptr
cdef float complex *vr_ptr
cdef l_logical *select_ptr
cdef np.ndarray[np.complex64_t] work
cdef np.ndarray[np.float32_t] rwork
assert_fortran_mat(S, T, Q, Z)
N = S.shape[0]
work = np.empty(2*N, dtype = np.complex64)
rwork = np.empty(2*N, dtype = np.float32)
if left and right:
side = "B"
elif left:
side = "L"
elif right:
side = "R"
else:
return
backtr = False
if select is not None:
howmny = "S"
MM = select.nonzero()[0].size
select_ptr = <l_logical *>select.data
else:
MM = N
select_ptr = NULL
if ((left and right and Q is not None and Z is not None) or
(left and not right and Q is not None) or
(right and not left and Z is not None)):
howmny = "B"
backtr = True
else:
howmny = "A"
if left:
if backtr:
vl = Q
else:
vl = np.empty((N, MM), dtype = np.complex64, order='F')
vl_ptr = <float complex *>vl.data
else:
vl_ptr = NULL
if right:
if backtr:
vr = Z
else:
vr = np.empty((N, MM), dtype = np.complex64, order='F')
vr_ptr = <float complex *>vr.data
else:
vr_ptr = NULL
lapack.ctgevc(side, howmny, select_ptr,
&N, <float complex *>S.data, &N,
<float complex *>T.data, &N,
vl_ptr, &N, vr_ptr, &N, &MM, &M,
<float complex *>work.data, <float *>rwork.data, &info)
assert info == 0, "Argument error in ctgevc"
assert MM == M, "Unexpected number of eigenvectors returned in ctgevc"
if not backtr:
if left:
vl = np.asfortranarray(np.dot(Q, vl))
if right:
vr = np.asfortranarray(np.dot(Z, vr))
if left and right:
return (vl, vr)
elif left:
return vl
else:
return vr
def ztgevc(np.ndarray[np.complex128_t, ndim=2] S,
np.ndarray[np.complex128_t, ndim=2] T,
np.ndarray[np.complex128_t, ndim=2] Q=None,
np.ndarray[np.complex128_t, ndim=2] Z=None,
np.ndarray[l_logical] select=None,
left=False, right=True):
cdef l_int N, info, M, MM
cdef char *side
cdef char *howmny
cdef np.ndarray[np.complex128_t, ndim=2] vl, vr
cdef double complex *vl_ptr
cdef double complex *vr_ptr
cdef l_logical *select_ptr
cdef np.ndarray[np.complex128_t] work
cdef np.ndarray[np.float64_t] rwork
assert_fortran_mat(S, T, Q, Z)
N = S.shape[0]
work = np.empty(2*N, dtype = np.complex128)
rwork = np.empty(2*N, dtype = np.float64)
if left and right:
side = "B"
elif left:
side = "L"
elif right:
side = "R"
else:
return
backtr = False
if select is not None:
howmny = "S"
MM = select.nonzero()[0].size
select_ptr = <l_logical *>select.data
else:
MM = N
select_ptr = NULL
if ((left and right and Q is not None and Z is not None) or
(left and not right and Q is not None) or
(right and not left and Z is not None)):
howmny = "B"
backtr = True
else:
howmny = "A"
# If there are complex eigenvalues, we need to postprocess the eigenvectors
cdef np.ndarray vl, vr
if left:
if backtr:
vl = Q
else:
vl = np.empty((N, MM), dtype = np.complex128, order='F')
vl_ptr = <double complex *>vl.data
else:
vl_ptr = NULL
vl = vl_r
if right:
if backtr:
vr = Z
vr = vr_r
if scalar in floating:
if scalar is float:
dtype = np.complex64
else:
vr = np.empty((N, MM), dtype = np.complex128, order='F')
vr_ptr = <double complex *>vr.data
else:
vr_ptr = NULL
lapack.ztgevc(side, howmny, select_ptr,
&N, <double complex *>S.data, &N,
<double complex *>T.data, &N,
vl_ptr, &N, vr_ptr, &N, &MM, &M,
<double complex *>work.data, <double *>rwork.data, &info)
assert info == 0, "Argument error in ztgevc"
assert MM == M, "Unexpected number of eigenvectors returned in ztgevc"
if not backtr:
if left:
vl = np.asfortranarray(np.dot(Q, vl))
if right:
vr = np.asfortranarray(np.dot(Z, vr))
dtype = np.complex128
if np.diagonal(S, -1).nonzero()[0].size:
if left:
vl = txevc_postprocess(dtype, S, vl_r, select)
if right:
vr = txevc_postprocess(dtype, S, vr_r, select)
if left and right:
return (vl, vr)
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment