Coverage for gpaw/eigensolvers/davidson.py: 97%

1from functools import partial 

2 

3import numpy as np 

4from ase.utils.timing import timer 

5from gpaw import debug 

6from gpaw.eigensolvers.diagonalizerbackend import (ScalapackDiagonalizer, 

7 ElpaDiagonalizer, 

8 ScipyDiagonalizer) 

9from gpaw.eigensolvers.eigensolver import Eigensolver 

10from gpaw.hybrids import HybridXC 

11from gpaw.matrix import matrix_matrix_multiply as mmm 

12 

13 

14class DummyArray: 

15 def __getitem__(self, x): 

16 return np.empty((0, 0)) 

17 

18 

19class Davidson(Eigensolver): 

20 """Simple Davidson eigensolver 

21 

22 It is expected that the trial wave functions are orthonormal 

23 and the integrals of projector functions and wave functions 

24 ``nucleus.P_uni`` are already calculated. 

25 

26 Solution steps are: 

27 

28 * Subspace diagonalization 

29 * Calculate all residuals 

30 * Add preconditioned residuals to the subspace and diagonalize 

31 """ 

32 

33 def __init__( 

34 self, niter=2): 

35 super().__init__() 

36 self.niter = niter 

37 self.diagonalizer_backend = None 

38 

39 self.orthonormalization_required = False 

40 self.H_NN = DummyArray() 

41 self.S_NN = DummyArray() 

42 self.eps_N = DummyArray() 

43 

44 def __repr__(self): 

45 return f'Davidson(niter={self.niter})' 

46 

47 def todict(self): 

48 return {'name': 'dav', 'niter': self.niter} 

49 

50 def initialize(self, wfs, dist_backend='scalapack'): 

51 # dist_backend keyword exists due to this class having 

52 # no reference to the parallelization backend. As such, 

53 # the keyword must be specified by the user upon creation 

54 # of this object. Usually one would want this to be ELPA 

55 # if use_elpa in parallel is set to True. 

56 dist_diagonalizers = { 

57 'scalapack': ScalapackDiagonalizer, 

58 'elpa': ElpaDiagonalizer 

59 } 

60 

61 super().initialize(wfs) 

62 slcomm, nrows, ncols, slsize = wfs.scalapack_parameters 

63 

64 if wfs.gd.comm.rank == 0 and wfs.bd.comm.rank == 0: 

65 # Allocate arrays 

66 B = self.nbands 

67 self.H_NN = np.zeros((2 * B, 2 * B), self.dtype) 

68 self.S_NN = np.zeros((2 * B, 2 * B), self.dtype) 

69 self.eps_N = np.zeros(2 * B) 

70 

71 if slsize is not None: 

72 self.diagonalizer_backend = dist_diagonalizers[dist_backend]( 

73 arraysize=self.nbands * 2, 

74 grid_nrows=nrows, 

75 grid_ncols=ncols, 

76 scalapack_communicator=slcomm, 

77 dtype=self.dtype, 

78 blocksize=slsize) 

79 else: 

80 self.diagonalizer_backend = ScipyDiagonalizer(slcomm) 

81 

82 def estimate_memory(self, mem, wfs): 

83 super().estimate_memory(mem, wfs) 

84 nbands = wfs.bd.nbands 

85 mem.subnode('H_nn', nbands * nbands * mem.itemsize[wfs.dtype]) 

86 mem.subnode('S_nn', nbands * nbands * mem.itemsize[wfs.dtype]) 

87 mem.subnode('H_2n2n', 4 * nbands * nbands * mem.itemsize[wfs.dtype]) 

88 mem.subnode('S_2n2n', 4 * nbands * nbands * mem.itemsize[wfs.dtype]) 

89 mem.subnode('eps_2n', 2 * nbands * mem.floatsize) 

90 

91 @timer('Davidson') 

92 def iterate_one_k_point(self, ham, wfs, kpt, weights): 

93 """Do Davidson iterations for the kpoint""" 

94 if isinstance(ham.xc, HybridXC): 

95 niter = 1 

96 else: 

97 niter = self.niter 

98 

99 bd = wfs.bd 

100 B = bd.nbands 

101 

102 H_NN = self.H_NN 

103 S_NN = self.S_NN 

104 eps_N = self.eps_N 

105 

106 def integrate(a_G): 

107 if wfs.collinear: 

108 return np.real(wfs.integrate(a_G, a_G, global_integral=False)) 

109 return sum( 

110 np.real(wfs.integrate(b_G, b_G, global_integral=False)) 

111 for b_G in a_G) 

112 

113 self.subspace_diagonalize(ham, wfs, kpt) 

114 

115 psit = kpt.psit 

116 psit2 = psit.new(buf=wfs.work_array) 

117 P = kpt.projections 

118 P2 = P.new() 

119 P3 = P.new() 

120 M = wfs.work_matrix_nn 

121 dS = wfs.setups.dS 

122 comm = wfs.gd.comm 

123 

124 is_gridband_master = (comm.rank == 0) and (bd.comm.rank == 0) 

125 

126 if bd.comm.size > 1: 

127 M0 = M.new(dist=(bd.comm, 1, 1)) 

128 else: 

129 M0 = M 

130 

131 if comm.rank == 0: 

132 e_N = bd.collect(kpt.eps_n) 

133 if e_N is not None: 

134 eps_N[:B] = e_N 

135 

136 Ht = partial(wfs.apply_pseudo_hamiltonian, kpt, ham) 

137 

138 if self.keep_htpsit: 

139 R = psit.new(buf=self.Htpsit_nG) 

140 else: 

141 R = psit.apply(Ht) 

142 

143 self.calculate_residuals(kpt, wfs, ham, psit, P, kpt.eps_n, R, P2) 

144 

145 precond = self.preconditioner 

146 

147 for nit in range(niter): 

148 if nit == niter - 1: 

149 error = np.dot(weights, [integrate(R_G) for R_G in R.array]) 

150 

151 for psit_G, R_G, psit2_G in zip(psit.array, R.array, psit2.array): 

152 ekin = precond.calculate_kinetic_energy(psit_G, kpt) 

153 precond(R_G, kpt, ekin, out=psit2_G) 

154 

155 # Calculate projections 

156 psit2.matrix_elements(wfs.pt, out=P2) 

157 

158 def copy(M, C_nn): 

159 comm.sum(M.array, 0) 

160 if comm.rank == 0: 

161 M.complex_conjugate() 

162 M.redist(M0) 

163 if bd.comm.rank == 0: 

164 C_nn[:] = M0.array 

165 

166 with self.timer('calc. matrices'): 

167 # <psi2 | H | psi2> 

168 psit2.matrix_elements( 

169 operator=Ht, result=R, out=M, symmetric=True, cc=True) 

170 ham.dH(P2, out=P3) 

171 mmm(1.0, P2, 'N', P3, 'C', 1.0, M) # , symmetric=True) 

172 copy(M, H_NN[B:, B:]) 

173 

174 # <psi2 | H | psi> 

175 R.matrix_elements(psit, out=M, cc=True) 

176 mmm(1.0, P3, 'N', P, 'C', 1.0, M) 

177 copy(M, H_NN[B:, :B]) 

178 

179 # <psi2 | S | psi2> 

180 psit2.matrix_elements(out=M, symmetric=True, cc=True) 

181 dS.apply(P2, out=P3) 

182 mmm(1.0, P2, 'N', P3, 'C', 1.0, M) 

183 copy(M, S_NN[B:, B:]) 

184 

185 # <psi2 | S | psi> 

186 psit2.matrix_elements(psit, out=M, cc=True) 

187 mmm(1.0, P3, 'N', P, 'C', 1.0, M) 

188 copy(M, S_NN[B:, :B]) 

189 

190 if is_gridband_master: 

191 H_NN[:B, :B] = np.diag(eps_N[:B]) 

192 S_NN[:B, :B] = np.eye(B) 

193 

194 with self.timer('diagonalize'): 

195 if is_gridband_master and debug: 

196 H_NN[np.triu_indices(2 * B, 1)] = 42.0 

197 S_NN[np.triu_indices(2 * B, 1)] = 42.0 

198 

199 try: 

200 self.diagonalizer_backend.diagonalize( 

201 H_NN, S_NN, eps_N, debug=debug) 

202 except np.linalg.LinAlgError as ex: 

203 raise ValueError( 

204 'Too few plane waves or grid points') from ex 

205 if comm.rank == 0: 

206 bd.distribute(eps_N[:B], kpt.eps_n) 

207 comm.broadcast(kpt.eps_n, 0) 

208 

209 with self.timer('rotate_psi'): 

210 if comm.rank == 0: 

211 if bd.comm.rank == 0: 

212 M0.array[:] = H_NN[:B, :B].T 

213 M0.redist(M) 

214 comm.broadcast(M.array, 0) 

215 mmm(1.0, M, 'N', psit, 'N', 0.0, R) 

216 mmm(1.0, M, 'N', P, 'N', 0.0, P3) 

217 if comm.rank == 0: 

218 if bd.comm.rank == 0: 

219 M0.array[:] = H_NN[B:, :B].T 

220 M0.redist(M) 

221 comm.broadcast(M.array, 0) 

222 mmm(1.0, M, 'N', psit2, 'N', 1.0, R) 

223 mmm(1.0, M, 'N', P2, 'N', 1.0, P3) 

224 psit[:] = R 

225 P, P3 = P3, P 

226 kpt.projections = P 

227 

228 if nit < niter - 1: 

229 psit.apply(Ht, out=R) 

230 self.calculate_residuals( 

231 kpt, wfs, ham, psit, P, kpt.eps_n, R, P2) 

232 

233 error = wfs.gd.comm.sum_scalar(error) 

234 return error