/
eigenvalue.py
125 lines (92 loc) · 3.57 KB
/
eigenvalue.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
import cupy
from cupy import cuda
from cupy.cuda import cublas
from cupy.cuda import device
if cuda.cusolver_enabled:
from cupy.cuda import cusolver
def _syevd(a, UPLO, with_eigen_vector):
if UPLO not in ('L', 'U'):
raise ValueError("UPLO argument must be 'L' or 'U'")
if a.dtype == 'f' or a.dtype == 'e':
dtype = 'f'
ret_type = a.dtype
else:
# NumPy uses float64 when an input is not floating point number.
dtype = 'd'
ret_type = 'd'
# Note that cuSolver assumes fortran array
v = a.astype(dtype, order='F', copy=True)
m, lda = a.shape
w = cupy.empty(m, dtype)
dev_info = cupy.empty((), 'i')
handle = device.Device().cusolver_handle
if with_eigen_vector:
jobz = cusolver.CUSOLVER_EIG_MODE_VECTOR
else:
jobz = cusolver.CUSOLVER_EIG_MODE_NOVECTOR
if UPLO == 'L':
uplo = cublas.CUBLAS_FILL_MODE_LOWER
else: # UPLO == 'U'
uplo = cublas.CUBLAS_FILL_MODE_UPPER
if dtype == 'f':
buffer_size = cupy.cuda.cusolver.ssyevd_bufferSize
syevd = cupy.cuda.cusolver.ssyevd
elif dtype == 'd':
buffer_size = cupy.cuda.cusolver.dsyevd_bufferSize
syevd = cupy.cuda.cusolver.dsyevd
else:
raise RuntimeError('Only float and double are supported')
work_size = buffer_size(
handle, jobz, uplo, m, v.data.ptr, lda, w.data.ptr)
work = cupy.empty(work_size, dtype)
syevd(
handle, jobz, uplo, m, v.data.ptr, lda,
w.data.ptr, work.data.ptr, work_size, dev_info.data.ptr)
return w.astype(ret_type, copy=False), v.astype(ret_type, copy=False)
# TODO(okuta): Implement eig
def eigh(a, UPLO='L'):
"""Eigenvalues and eigenvectors of a symmetric matrix.
This method calculates eigenvalues and eigenvectors of a given
symmetric matrix.
.. note::
Currenlty only 2-D matrix is supported.
.. note::
CUDA >=8.0 is required.
Args:
a (cupy.ndarray): A symmetric 2-D square matrix.
UPLO (str): Select from ``'L'`` or ``'U'``. It specifies which
part of ``a`` is used. ``'L'`` uses the lower triangular part of
``a``, and ``'U'`` uses the upper triangular part of ``a``.
Returns:
tuple of :class:`~cupy.ndarray`:
Returns a tuple ``(w, v)``. ``w`` contains eigenvalues and
``v`` contains eigenvectors. ``v[:, i]`` is an eigenvector
corresponding to an eigenvalue ``w[i]``.
.. seealso:: :func:`numpy.linalg.eigh`
"""
if not cuda.cusolver_enabled:
raise RuntimeError('Current cupy only supports cusolver in CUDA 8.0')
return _syevd(a, UPLO, True)
# TODO(okuta): Implement eigvals
def eigvalsh(a, UPLO='L'):
"""Calculates eigenvalues of a symmetric matrix.
This method calculates eigenvalues a given symmetric matrix.
Note that :func:`cupy.linalg.eigh` calculates both eigenvalues and
eigenvectors.
.. note::
Currenlty only 2-D matrix is supported.
.. note::
CUDA >=8.0 is required.
Args:
a (cupy.ndarray): A symmetric 2-D square matrix.
UPLO (str): Select from ``'L'`` or ``'U'``. It specifies which
part of ``a`` is used. ``'L'`` uses the lower triangular part of
``a``, and ``'U'`` uses the upper triangular part of ``a``.
Returns:
cupy.ndarray:
Returns eigenvalues as a vector.
.. seealso:: :func:`numpy.linalg.eigvalsh`
"""
if not cuda.cusolver_enabled:
raise RuntimeError('Current cupy only supports cusolver in CUDA 8.0')
return _syevd(a, UPLO, False)[0]