Coverage for gpaw/test/parallel/test_scalapack.py: 96%

1"""Test of ScaLAPACK diagonalize and inverse cholesky. 

2 

3The test generates a random symmetric matrix H0 and 

4positive definite matrix S0 on a 1-by-1 BLACS grid. They 

5are redistributed to a mprocs-by-nprocs BLACS grid, 

6diagonalized in parallel, and eigenvalues are compared 

7against LAPACK. Eigenvectors are not compared. 

8""" 

9 

10import pytest 

11import numpy as np 

12from scipy.linalg import eigh, inv, cholesky 

13 

14from gpaw.mpi import world, rank 

15from gpaw.blacs import BlacsGrid, Redistributor 

16from gpaw.utilities.tools import tri2full 

17from gpaw.utilities import compiled_with_sl 

18from gpaw.utilities.blas import rk 

19from gpaw.utilities.scalapack import \ 

20 scalapack_general_diagonalize_dc, \ 

21 scalapack_diagonalize_dc, \ 

22 scalapack_inverse_cholesky, \ 

23 scalapack_inverse, \ 

24 scalapack_solve 

25 

26from .test_pblas import \ 

27 initialize_random, initialize_matrix, \ 

28 calculate_error, mnprocs_i 

29 

30 

31pytestmark = pytest.mark.skipif(not compiled_with_sl(), 

32 reason='not compiled with scalapack') 

33 

34 

35tol = 1.0e-8 

36 

37 

38@pytest.mark.parametrize('mprocs, nprocs', mnprocs_i) 

39@pytest.mark.parametrize('dtype', [float, complex]) 

40def test_scalapack_diagonalize_inverse(dtype, mprocs, nprocs, 

41 N=72, seed=42): 

42 gen = np.random.RandomState(seed) 

43 grid = BlacsGrid(world, mprocs, nprocs) 

44 

45 if dtype == complex: 

46 epsilon = 1.0j 

47 else: 

48 epsilon = 0.0 

49 

50 # Create descriptors for matrices on master: 

51 glob = grid.new_descriptor(N, N, N, N) 

52 

53 # print globA.asarray() 

54 # Populate matrices local to master: 

55 H0 = glob.zeros(dtype=dtype) + gen.rand(*glob.shape) 

56 S0 = glob.zeros(dtype=dtype) + gen.rand(*glob.shape) 

57 C0 = glob.empty(dtype=dtype) 

58 if rank == 0: 

59 # Complex case must have real numbers on the diagonal. 

60 # We make a simple complex Hermitian matrix below. 

61 H0 = H0 + epsilon * (0.1 * np.tri(N, N, k=-N // nprocs) + 

62 0.3 * np.tri(N, N, k=-1)) 

63 S0 = S0 + epsilon * (0.2 * np.tri(N, N, k=-N // nprocs) + 

64 0.4 * np.tri(N, N, k=-1)) 

65 # Make matrices symmetric 

66 rk(1.0, H0.copy(), 0.0, H0) 

67 rk(1.0, S0.copy(), 0.0, S0) 

68 # Overlap matrix must be semi-positive definite 

69 S0 = S0 + 50.0 * np.eye(N, N, 0) 

70 # Hamiltonian is usually diagonally dominant 

71 H0 = H0 + 75.0 * np.eye(N, N, 0) 

72 C0 = S0.copy() 

73 S0_inv = S0.copy() 

74 

75 # Local result matrices 

76 W0 = np.empty((N), dtype=float) 

77 W0_g = np.empty((N), dtype=float) 

78 

79 # Calculate eigenvalues / other serial results 

80 print(rank, C0.shape, C0.dtype) 

81 if rank == 0: 

82 W0 = eigh(H0, eigvals_only=True) 

83 W0_g = eigh(H0, S0, eigvals_only=True) 

84 C0 = inv(cholesky(C0, lower=True)).copy() # result in lower triangle 

85 tri2full(S0_inv, 'L') 

86 S0_inv = inv(S0_inv) 

87 # tri2full(C0) # symmetrize 

88 

89 print(rank, C0.shape, C0.dtype) 

90 assert glob.check(H0) 

91 assert glob.check(S0) 

92 assert glob.check(C0) 

93 

94 # Create distributed destriptors with various block sizes: 

95 dist = grid.new_descriptor(N, N, 8, 8) 

96 

97 # Distributed matrices: 

98 # We can use empty here, but end up with garbage on 

99 # on the other half of the triangle when we redistribute. 

100 # This is fine because ScaLAPACK does not care. 

101 

102 H = dist.empty(dtype=dtype) 

103 S = dist.empty(dtype=dtype) 

104 Sinv = dist.empty(dtype=dtype) 

105 Z = dist.empty(dtype=dtype) 

106 C = dist.empty(dtype=dtype) 

107 Sinv = dist.empty(dtype=dtype) 

108 

109 # Eigenvalues are non-BLACS matrices 

110 # W = np.empty((N), dtype=float) 

111 W_dc = np.empty((N), dtype=float) 

112 # W_mr3 = np.empty((N), dtype=float) 

113 # W_g = np.empty((N), dtype=float) 

114 W_g_dc = np.empty((N), dtype=float) 

115 # W_g_mr3 = np.empty((N), dtype=float) 

116 

117 Glob2dist = Redistributor(world, glob, dist) 

118 Glob2dist.redistribute(H0, H, uplo='L') 

119 Glob2dist.redistribute(S0, S, uplo='L') 

120 Glob2dist.redistribute(S0, C, uplo='L') # C0 was previously overwritten 

121 Glob2dist.redistribute(S0, Sinv, uplo='L') 

122 

123 # we don't test the expert drivers anymore since there 

124 # might be a buffer overflow error 

125 # scalapack_diagonalize_ex(dist, H.copy(), Z, W, 'L') 

126 scalapack_diagonalize_dc(dist, H.copy(), Z, W_dc, 'L') 

127 # scalapack_diagonalize_mr3(dist, H.copy(), Z, W_mr3, 'L') 

128 # scalapack_general_diagonalize_ex(dist, H.copy(), S.copy(), Z, W_g, 'L') 

129 scalapack_general_diagonalize_dc(dist, H.copy(), S.copy(), Z, W_g_dc, 'L') 

130 

131 scalapack_inverse_cholesky(dist, C, 'L') 

132 

133 if dtype == complex: # Only supported for complex for now 

134 scalapack_inverse(dist, Sinv, 'L') 

135 # Undo redistribute 

136 C_test = glob.empty(dtype=dtype) 

137 Sinv_test = glob.empty(dtype=dtype) 

138 Dist2glob = Redistributor(world, dist, glob) 

139 Dist2glob.redistribute(C, C_test) 

140 Dist2glob.redistribute(Sinv, Sinv_test) 

141 

142 if rank == 0: 

143 diag_dc_err = abs(W_dc - W0).max() 

144 # diag_mr3_err = abs(W_mr3 - W0).max() 

145 # general_diag_ex_err = abs(W_g - W0_g).max() 

146 general_diag_dc_err = abs(W_g_dc - W0_g).max() 

147 # general_diag_mr3_err = abs(W_g_mr3 - W0_g).max() 

148 print(C_test[:2, :2]) 

149 print(C0[:2, :2]) 

150 inverse_chol_err = abs(C_test - C0).max() 

151 

152 tri2full(Sinv_test, 'L') 

153 inverse_err = abs(Sinv_test - S0_inv).max() 

154 # print 'diagonalize ex err', diag_ex_err 

155 print('diagonalize dc err', diag_dc_err) 

156 # print 'diagonalize mr3 err', diag_mr3_err 

157 # print 'general diagonalize ex err', general_diag_ex_err 

158 print('general diagonalize dc err', general_diag_dc_err) 

159 # print 'general diagonalize mr3 err', general_diag_mr3_err 

160 print('inverse chol err', inverse_chol_err) 

161 if dtype == complex: 

162 print('inverse err', inverse_err) 

163 else: 

164 # diag_ex_err = 0.0 

165 diag_dc_err = 0.0 

166 # diag_mr3_err = 0.0 

167 # general_diag_ex_err = 0.0 

168 general_diag_dc_err = 0.0 

169 # general_diag_mr3_err = 0.0 

170 inverse_chol_err = 0.0 

171 inverse_err = 0.0 

172 

173 # We don't like exceptions on only one cpu 

174 # diag_ex_err = world.sum(diag_ex_err) 

175 diag_dc_err = world.sum_scalar(diag_dc_err) 

176 # diag_mr3_err = world.sum(diag_mr3_err) 

177 # general_diag_ex_err = world.sum(general_diag_ex_err) 

178 general_diag_dc_err = world.sum_scalar(general_diag_dc_err) 

179 # general_diag_mr3_err = world.sum(general_diag_mr3_err) 

180 inverse_chol_err = world.sum_scalar(inverse_chol_err) 

181 inverse_err = world.sum_scalar(inverse_err) 

182 # assert diag_ex_err < tol 

183 assert diag_dc_err < tol 

184 # assert diag_mr3_err < tol 

185 # assert general_diag_ex_err < tol 

186 assert general_diag_dc_err < tol 

187 # assert general_diag_mr3_err < tol 

188 assert inverse_chol_err < tol 

189 if dtype == complex: 

190 assert inverse_err < tol 

191 

192 

193@pytest.mark.parametrize('mprocs, nprocs', mnprocs_i) 

194@pytest.mark.parametrize('dtype', [float, complex]) 

195def test_scalapack_solve(dtype, mprocs, nprocs, 

196 M=160, K=120, seed=42): 

197 """Test scalapack_solve()""" 

198 random = initialize_random(seed, dtype) 

199 grid = BlacsGrid(world, mprocs, nprocs) 

200 

201 # Initialize matrices 

202 shapeA = (M, M) 

203 shapeB = (K, M) 

204 A0, A, descA = initialize_matrix(grid, *shapeA, 2, 2, random) 

205 B0, B, descB = initialize_matrix(grid, *shapeB, 3, 2, random) 

206 

207 if grid.comm.rank == 0: 

208 # Calculate reference with numpy 

209 ref_B0 = np.linalg.solve(A0.T, B0.T).T 

210 else: 

211 ref_B0 = None 

212 

213 # Calculate with scalapack 

214 scalapack_solve(descA, descB, A, B) 

215 

216 # Check error 

217 err = calculate_error(ref_B0, B, descB) 

218 tol = {float: 8e-12, complex: 2e-13}[dtype] 

219 assert err < tol