Coverage for gpaw/new/magma.py: 23%
26 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
1import gpaw.cgpaw as cgpaw
2from gpaw.gpu import cupy as cp, cupy_is_fake
3import numpy as np
4from gpaw.gpu.cpupy import asnumpy
7def eigh_magma_cpu(matrix: np.ndarray, UPLO: str) -> tuple[np.ndarray,
8 np.ndarray]:
9 """
10 Wrapper for MAGMA symmetric/Hermitian eigensolvers, CPU version.
12 Parameters
13 ----------
14 matrix : (N, N) numpy.ndarray
15 The matrix to diagonalize. Must be symmetric or Hermitian.
16 UPLO : str
17 Whether the upper or lower part of the matrix is stored.
18 Choose 'U' or 'L'.
20 Returns
21 -------
22 w : (N,) numpy.ndarray
23 Eigenvalues in ascending order
24 v : (N, N) numpy.ndarray
25 Matrix containing orthonormal eigenvectors.
26 Eigenvector corresponding to ``w[i]`` is in column ``v[:,i]``.
27 """
29 assert cgpaw.have_magma, "Must compile with MAGMA support"
31 if matrix.dtype == np.complex128:
32 eigvals, eigvects = cgpaw.eigh_magma_zheevd(matrix, UPLO)
34 elif matrix.dtype == np.float64:
35 eigvals, eigvects = cgpaw.eigh_magma_dsyevd(matrix, UPLO)
37 else:
38 raise TypeError("Unsupported matrix dtype")
40 # MAGMA eigenvectors are on rows, numpy/cupy has them on columns
41 return eigvals, np.conjugate(eigvects).T
44def eigh_magma_gpu(matrix: cp.ndarray, UPLO: str) -> tuple[cp.ndarray,
45 cp.ndarray]:
46 """
47 Wrapper for MAGMA symmetric/Hermitian eigensolvers, GPU version.
49 Parameters
50 ----------
51 matrix : (N, N) cupy.ndarray
52 The matrix to diagonalize. Must be symmetric or Hermitian.
53 UPLO : str
54 Whether the upper or lower part of the matrix is stored.
55 Choose 'U' or 'L'.
57 Returns
58 -------
59 w : (N,) cupy.ndarray
60 Eigenvalues in ascending order
61 v : (N, N) cupy.ndarray
62 Matrix containing orthonormal eigenvectors.
63 Eigenvector corresponding to ``w[i]`` is in column ``v[:,i]``.
64 """
65 assert cgpaw.have_magma, "Must compile with MAGMA support"
67 assert matrix.ndim == 2 and matrix.shape[0] == matrix.shape[1]
69 if cupy_is_fake:
70 eigval_np, eigvect_np = eigh_magma_cpu(asnumpy(matrix), UPLO)
71 return cp.asarray(eigval_np), cp.asarray(eigvect_np)
73 # Alloc output arrays with CUPY.
74 # Necessary because the C code has no easy access to CUPY array creation
76 eigvects = cp.empty_like(matrix)
77 # Only symmetric/Hermitian matrices supported for now,
78 # so eigenvalues are always real
79 eigvals = cp.empty((matrix.shape[0],), dtype=np.float64)
81 if matrix.dtype == np.complex128:
82 cgpaw.eigh_magma_zheevd_gpu(matrix,
83 UPLO,
84 eigvals,
85 eigvects)
87 elif matrix.dtype == np.float64:
88 cgpaw.eigh_magma_dsyevd_gpu(matrix,
89 UPLO,
90 eigvals,
91 eigvects)
93 else:
94 raise TypeError("Unsupported matrix dtype")
96 # MAGMA eigenvectors are on rows, numpy/cupy has them on columns
97 return eigvals, cp.conjugate(eigvects).T