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

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 

5 

6 

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. 

11 

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'. 

19 

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 """ 

28 

29 assert cgpaw.have_magma, "Must compile with MAGMA support" 

30 

31 if matrix.dtype == np.complex128: 

32 eigvals, eigvects = cgpaw.eigh_magma_zheevd(matrix, UPLO) 

33 

34 elif matrix.dtype == np.float64: 

35 eigvals, eigvects = cgpaw.eigh_magma_dsyevd(matrix, UPLO) 

36 

37 else: 

38 raise TypeError("Unsupported matrix dtype") 

39 

40 # MAGMA eigenvectors are on rows, numpy/cupy has them on columns 

41 return eigvals, np.conjugate(eigvects).T 

42 

43 

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. 

48 

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'. 

56 

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" 

66 

67 assert matrix.ndim == 2 and matrix.shape[0] == matrix.shape[1] 

68 

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) 

72 

73 # Alloc output arrays with CUPY. 

74 # Necessary because the C code has no easy access to CUPY array creation 

75 

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) 

80 

81 if matrix.dtype == np.complex128: 

82 cgpaw.eigh_magma_zheevd_gpu(matrix, 

83 UPLO, 

84 eigvals, 

85 eigvects) 

86 

87 elif matrix.dtype == np.float64: 

88 cgpaw.eigh_magma_dsyevd_gpu(matrix, 

89 UPLO, 

90 eigvals, 

91 eigvects) 

92 

93 else: 

94 raise TypeError("Unsupported matrix dtype") 

95 

96 # MAGMA eigenvectors are on rows, numpy/cupy has them on columns 

97 return eigvals, cp.conjugate(eigvects).T