Coverage for gpaw/test/linalg/test_magma.py: 30%

1"""Tests for MAGMA eigensolver wrappers""" 

2import numpy as np 

3import pytest 

4from gpaw.new.magma import eigh_magma_cpu, eigh_magma_gpu 

5from gpaw.cgpaw import have_magma 

6from gpaw.gpu import cupy as cp 

7from gpaw.gpu import cupy_is_fake 

8 

9 

10def fix_eigenvector_phase(inout_arr): 

11 """Helper function for comparing eigenvector output from different 

12 solvers. Rotates eigenvectors in the input matrix so that the first 

13 element of each vector is real and non-negative. 

14 Input is modified in-place. 

15 NB: eigenvectors are on columns. 

16 """ 

17 assert inout_arr.ndim == 2 

18 # Works for cupy arrays too because the dtypes are compatible 

19 

20 if np.issubdtype(inout_arr.dtype, np.complexfloating): 

21 # Complex matrices 

22 for i in range(inout_arr.shape[1]): 

23 phase = np.angle(inout_arr[0, i]) 

24 if phase != 0: 

25 rotation = np.exp(phase * (-1j)) 

26 inout_arr[:, i] *= rotation 

27 

28 elif np.issubdtype(inout_arr.dtype, np.floating): 

29 # Real matrices 

30 for i in range(inout_arr.shape[1]): 

31 if inout_arr[0, i] < 0: 

32 inout_arr[:, i] *= -1 

33 

34 return inout_arr 

35 

36 

37@pytest.fixture 

38def eigh_test_matrix(): 

39 def _generate(n: int, type: str = 'symmetric', 

40 backend: str = 'numpy', seed: int = 42): 

41 

42 assert type in ['symmetric', 'hermitian'] 

43 assert backend in ['numpy', 'cupy'] 

44 

45 if backend == 'cupy': 

46 xp = cp 

47 else: 

48 xp = np 

49 

50 rng = xp.random.default_rng(seed) 

51 if type == 'symmetric': 

52 A = rng.random((n, n)) 

53 return (A + A.T) / 2 

54 

55 else: 

56 # Create Hermitian matrix 

57 A = rng.random((n, n)) + 1j * rng.random((n, n)) 

58 return (A + A.T.conj()) / 2 

59 

60 return _generate 

61 

62 

63@pytest.mark.skipif(not have_magma, reason="No MAGMA") 

64@pytest.mark.parametrize("matrix_size, matrix_type, uplo", 

65 [(2, 'symmetric', 'L'), (4, 'hermitian', 'U')]) 

66def test_eigh_magma_cpu(eigh_test_matrix: np.ndarray, 

67 matrix_size: int, 

68 matrix_type: str, 

69 uplo: str) -> None: 

70 """Compare eigh output of Numpy and MAGMA""" 

71 

72 arr = eigh_test_matrix(matrix_size, type=matrix_type, backend='numpy') 

73 eigvals, eigvects = eigh_magma_cpu(arr, uplo) 

74 

75 eigvals_np, eigvects_np = np.linalg.eigh(arr, UPLO=uplo) 

76 

77 fix_eigenvector_phase(eigvects) 

78 fix_eigenvector_phase(eigvects_np) 

79 

80 np.testing.assert_allclose(eigvals, eigvals_np, atol=1e-12) 

81 np.testing.assert_allclose(eigvects, eigvects_np, atol=1e-12) 

82 

83 

84# MAGMA seems to do small matrices (N <= 128) on the CPU. 

85# So need a large matrix for honest GPU tests 

86@pytest.mark.skipif(not have_magma, reason="No MAGMA") 

87@pytest.mark.skipif(cupy_is_fake, 

88 reason="MAGMA GPU tests disabled for fake cupy") 

89@pytest.mark.gpu 

90@pytest.mark.parametrize("matrix_size, matrix_type, uplo", 

91 [(16, 'symmetric', 'L'), 

92 (150, 'hermitian', 'L'), 

93 (256, 'symmetric', 'U')]) 

94def test_eigh_magma_gpu(eigh_test_matrix: cp.ndarray, 

95 matrix_size: int, 

96 matrix_type: str, 

97 uplo: str): 

98 """Compare eigh output of CUPY and MAGMA""" 

99 

100 arr = eigh_test_matrix(matrix_size, type=matrix_type, backend='cupy') 

101 eigvals, eigvects = eigh_magma_gpu(arr, uplo) 

102 

103 eigvals_cp, eigvects_cp = cp.linalg.eigh(arr, UPLO=uplo) 

104 

105 fix_eigenvector_phase(eigvects) 

106 fix_eigenvector_phase(eigvects_cp) 

107 

108 cp.testing.assert_allclose(eigvals, eigvals_cp, atol=1e-12) 

109 cp.testing.assert_allclose(eigvects, eigvects_cp, atol=1e-12)