Coverage for gpaw/lrtddft2/lr_layouts.py: 10%

1import numpy as np 

2 

3from gpaw.blacs import BlacsGrid, Redistributor 

4import gpaw.cgpaw as cgpaw 

5 

6 

7class LrDiagonalizeLayout: 

8 """BLACS layout for distributed Omega matrix in linear response 

9 time-dependet DFT calculations""" 

10 

11 def __init__(self, sl_lrtddft, nrows, lr_comms): 

12 self.mprocs, self.nprocs, self.block_size = tuple(sl_lrtddft) 

13 

14 self.lr_comms = lr_comms 

15 

16 # original grid, ie, how matrix is stored 

17 self.matrix_grid = BlacsGrid(self.lr_comms.parent_comm, 

18 self.lr_comms.dd_comm.size, 

19 self.lr_comms.eh_comm.size) 

20 

21 # diagonalization grid 

22 self.diag_grid = BlacsGrid(self.lr_comms.parent_comm, self.mprocs, 

23 self.nprocs) 

24 

25 # ----------------------------------------------------------------- 

26 # for SCALAPACK we need TRANSPOSED MATRIX (and vector) 

27 # 

28 # M = rows, N = cols 

29 M = nrows 

30 N = nrows 

31 mb = 1 

32 nb = 1 

33 self.matrix_descr = self.matrix_grid.new_descriptor(N, M, nb, mb) 

34 

35 bs = self.block_size 

36 self.diag_descr = self.diag_grid.new_descriptor(N, M, bs, bs) 

37 

38 self.diag_in_redist = Redistributor(self.lr_comms.parent_comm, 

39 self.matrix_descr, self.diag_descr) 

40 

41 self.diag_out_redist = Redistributor(self.lr_comms.parent_comm, 

42 self.diag_descr, 

43 self.matrix_descr) 

44 

45 def diagonalize(self, eigenvectors, eigenvalues): 

46 """Diagonalize symmetric distributed Casida matrix using Scalapack. 

47 Parameters: 

48 

49 eigenvectors 

50 distributed Casida matrix on input, distributed eigenvectors on 

51 output 

52 

53 eigenvalues 

54 zero array on input, eigenvalues on output 

55 """ 

56 O_diag = self.diag_descr.empty(dtype=float) 

57 if self.matrix_descr.blacsgrid.is_active(): 

58 O_orig = eigenvectors 

59 else: 

60 O_orig = np.empty((0, 0), dtype=float) 

61 

62 self.diag_in_redist.redistribute(O_orig, O_diag) 

63 

64 # print O_diag 

65 

66 self.diag_descr.diagonalize_dc(O_diag.copy(), O_diag, eigenvalues, 'L') 

67 

68 self.diag_out_redist.redistribute(O_diag, O_orig) 

69 

70 self.lr_comms.parent_comm.broadcast(eigenvalues, 0) 

71 

72 

73class LrTDDFPTSolveLayout: 

74 """BLACS layouts for distributed TD-DFPT""" 

75 

76 def __init__(self, sl_lrtddft, nrows, lr_comms): 

77 self.mprocs, self.nprocs, self.block_size = tuple(sl_lrtddft) 

78 

79 self.lr_comms = lr_comms 

80 

81 # for SCALAPACK we need TRANSPOSED MATRIX (and vector) 

82 # 

83 # ----------------------------------------------------------------- 

84 # matrix 

85 

86 # original grid, ie, how matrix is stored 

87 self.orig_matrix_grid = BlacsGrid(self.lr_comms.parent_comm, 

88 self.lr_comms.dd_comm.size, 

89 self.lr_comms.eh_comm.size) 

90 

91 # solve grid 

92 self.solve_matrix_grid = BlacsGrid(self.lr_comms.parent_comm, 

93 self.mprocs, self.nprocs) 

94 

95 # M = rows, N = cols 

96 M = nrows * 4 

97 N = nrows * 4 

98 mb = 4 

99 nb = 4 

100 self.orig_matrix_descr = self.orig_matrix_grid.new_descriptor( 

101 N, M, nb, mb) 

102 

103 bs = self.block_size 

104 self.solve_matrix_descr = self.solve_matrix_grid.new_descriptor( 

105 N, M, bs, bs) 

106 

107 self.matrix_in_redist = Redistributor(self.lr_comms.parent_comm, 

108 self.orig_matrix_descr, 

109 self.solve_matrix_descr) 

110 

111 # ----------------------------------------------------------------- 

112 # vector 

113 

114 # original grid, ie, how vector is stored 

115 self.orig_vector_grid = BlacsGrid( 

116 self.lr_comms.parent_comm, 1, 

117 (self.lr_comms.dd_comm.size * self.lr_comms.eh_comm.size)) 

118 

119 # solve grid 

120 

121 # M = rows, N = cols 

122 M = nrows * 4 

123 Nrhs = 1 

124 mb = 4 

125 nb = 1 

126 self.orig_vector_descr = self.orig_vector_grid.new_descriptor( 

127 Nrhs, M, nb, mb) 

128 

129 bs = self.block_size 

130 self.solve_vector_descr = self.solve_matrix_grid.new_descriptor( 

131 Nrhs, M, 1, bs) 

132 

133 self.vector_in_redist = Redistributor(self.lr_comms.parent_comm, 

134 self.orig_vector_descr, 

135 self.solve_vector_descr) 

136 

137 self.vector_out_redist = Redistributor(self.lr_comms.parent_comm, 

138 self.solve_vector_descr, 

139 self.orig_vector_descr) 

140 

141 def solve(self, A_orig, b_orig): 

142 """Solve TD-DFPT equation using Scalapack. 

143 """ 

144 

145 A_solve = self.solve_matrix_descr.empty(dtype=float) 

146 if not self.orig_matrix_descr.blacsgrid.is_active(): 

147 A_orig = np.empty((0, 0), dtype=float) 

148 

149 self.matrix_in_redist.redistribute(A_orig, A_solve) 

150 

151 b_solve = self.solve_vector_descr.empty(dtype=float) 

152 if not self.orig_vector_descr.blacsgrid.is_active(): 

153 b_orig = np.empty((0, 0), dtype=float) 

154 

155 self.vector_in_redist.redistribute(b_orig, b_solve) 

156 

157 # if False: 

158 # np.set_printoptions(precision=5, suppress=True) 

159 # for i in range(self.lr_comms.parent_comm.size): 

160 # if ( self.lr_comms.parent_comm.rank == i ): 

161 # print 'rank ', i 

162 # print A_orig 

163 # print A_solve 

164 # print 

165 # print b_orig 

166 # print b_solve 

167 # print 

168 # print 

169 # print self.solve_matrix_descr.asarray() 

170 # print self.solve_vector_descr.asarray() 

171 # print 

172 # print '---' 

173 # print 

174 # self.lr_comms.parent_comm.barrier() 

175 

176 info = 0 

177 if self.solve_matrix_descr.blacsgrid.is_active(): 

178 cgpaw.scalapack_solve(A_solve, self.solve_matrix_descr.asarray(), 

179 b_solve, self.solve_vector_descr.asarray()) 

180 if info != 0: 

181 raise RuntimeError('scalapack_solve error: %d' % info) 

182 

183 self.vector_out_redist.redistribute(b_solve, b_orig) 

184 

185 # if False: 

186 # for i in range(self.lr_comms.parent_comm.size): 

187 # if ( self.lr_comms.parent_comm.rank == i ): 

188 # print 'rank ', i 

189 # print A_orig 

190 # print A_solve 

191 # print 

192 # print b_orig 

193 # print b_solve 

194 # print 

195 # print 

196 # self.lr_comms.parent_comm.barrier() 

197 

198 return b_orig 

199 

200 

201############################### 

202# BLACS layout for ScaLAPACK 

203class LrTDDFTLayouts: 

204 """BLACS layout for distributed Omega matrix in linear response 

205 time-dependet DFT calculations""" 

206 

207 def __init__(self, sl_lrtddft, nkq, dd_comm, eh_comm): 

208 mcpus, ncpus, blocksize = tuple(sl_lrtddft) 

209 self.world = eh_comm.parent 

210 self.dd_comm = dd_comm 

211 if self.world is None: 

212 self.world = self.dd_comm 

213 

214 # All the ranks within domain communicator contain the omega matrix 

215 # construct new communicator only on domain masters 

216 eh_ranks = np.arange(eh_comm.size) * dd_comm.size 

217 self.eh_comm2 = self.world.new_communicator(eh_ranks) 

218 

219 self.eh_grid = BlacsGrid(self.eh_comm2, eh_comm.size, 1) 

220 self.eh_descr = self.eh_grid.new_descriptor(nkq, nkq, 1, nkq) 

221 self.diag_grid = BlacsGrid(self.world, mcpus, ncpus) 

222 self.diag_descr = self.diag_grid.new_descriptor( 

223 nkq, nkq, blocksize, blocksize) 

224 

225 self.redistributor_in = Redistributor(self.world, self.eh_descr, 

226 self.diag_descr) 

227 self.redistributor_out = Redistributor(self.world, self.diag_descr, 

228 self.eh_descr) 

229 """ 

230 # ----------------------------------------------------------------- 

231 # for SCALAPACK we need TRANSPOSED MATRIX (and vector) 

232 # ----------------------------------------------------------------- 

233 # M = rows, N = cols 

234 M = nkq*4; N = nkq*4; mb = nkq*4; nb = 4; Nrhs = 1 

235 # Matrix, mp=1, np=eh_comm.size 

236 self.eh_grid2a = BlacsGrid(self.eh_comm2, eh_comm.size, 1) 

237 # Vector, mp=eh_comm.size, np=1 

238 self.eh_grid2b = BlacsGrid(self.eh_comm2, 1, eh_comm.size) 

239 self.eh_descr2a = self.eh_grid2a.new_descriptor(N, M, nb, mb) 

240 self.eh_descr2b = self.eh_grid2b.new_descriptor(Nrhs, N, 1, nb) 

241 

242 self.solve_descr2a =self.diag_grid.new_descriptor(N, M, 

243 blocksize, blocksize) 

244 self.solve_descr2b =self.diag_grid.new_descriptor(Nrhs, N, 

245 1, blocksize) 

246 

247 self.redist_solve_in_2a = Redistributor(self.world, 

248 self.eh_descr2a, 

249 self.solve_descr2a) 

250 self.redist_solve_in_2b = Redistributor(self.world, 

251 self.eh_descr2b, 

252 self.solve_descr2b) 

253 

254 self.redist_solve_out_2a = Redistributor(self.world, 

255 self.solve_descr2a, 

256 self.eh_descr2a) 

257 self.redist_solve_out_2b = Redistributor(self.world, 

258 self.solve_descr2b, 

259 self.eh_descr2b) 

260 """ 

261 

262 # ----------------------------------------------------------------- 

263 # for SCALAPACK we need TRANSPOSED MATRIX (and vector) 

264 # ----------------------------------------------------------------- 

265 # M = rows, N = cols 

266 M = nkq * 4 

267 N = nkq * 4 

268 mb = 4 

269 nb = 4 

270 Nrhs = 1 

271 # Matrix, mp=1, np=eh_comm.size 

272 self.eh_grid2a = BlacsGrid(self.world, dd_comm.size, eh_comm.size) 

273 # Vector, mp=eh_comm.size, np=1 

274 self.eh_grid2b = BlacsGrid(self.world, 1, dd_comm.size * eh_comm.size) 

275 self.eh_descr2a = self.eh_grid2a.new_descriptor(N, M, nb, mb) 

276 self.eh_descr2b = self.eh_grid2b.new_descriptor(Nrhs, N, Nrhs, nb) 

277 self.solve_descr2a = self.diag_grid.new_descriptor( 

278 N, M, blocksize, blocksize) 

279 self.solve_descr2b = self.diag_grid.new_descriptor( 

280 Nrhs, N, Nrhs, blocksize) 

281 

282 self.redist_solve_in_2a = Redistributor(self.world, self.eh_descr2a, 

283 self.solve_descr2a) 

284 self.redist_solve_in_2b = Redistributor(self.world, self.eh_descr2b, 

285 self.solve_descr2b) 

286 

287 self.redist_solve_out_2a = Redistributor(self.world, 

288 self.solve_descr2a, 

289 self.eh_descr2a) 

290 self.redist_solve_out_2b = Redistributor(self.world, 

291 self.solve_descr2b, 

292 self.eh_descr2b) 

293 

294 def solve(self, A, b): 

295 # if 0: 

296 # print 'edescr2a', rank, self.eh_descr2a.asarray() 

297 # print 'edescr2b', rank, self.eh_descr2b.asarray() 

298 # 

299 # sys.stdout.flush() 

300 # self.world.barrier() 

301 # 

302 # print 'sdescr2a', rank, self.solve_descr2a.asarray() 

303 # print 'sdescr2b', rank, self.solve_descr2b.asarray() 

304 # 

305 # sys.stdout.flush() 

306 # self.world.barrier() 

307 # 

308 # print 'A ', rank, A.shape 

309 # if b is not None: 

310 # print 'b ', rank, b.shape 

311 # 

312 # sys.stdout.flush() 

313 # self.world.barrier() 

314 

315 A_nn = self.solve_descr2a.empty(dtype=float) 

316 if self.eh_descr2a.blacsgrid.is_active(): 

317 A_Nn = A 

318 else: 

319 A_Nn = np.empty((0, 0), dtype=float) 

320 self.redist_solve_in_2a.redistribute(A_Nn, A_nn) 

321 

322 b_n = self.solve_descr2b.empty(dtype=float) 

323 if self.eh_descr2b.blacsgrid.is_active(): 

324 b_N = b.reshape(1, len(b)) 

325 else: 

326 b_N = np.empty((A_Nn.shape[0], 0), dtype=float) 

327 self.redist_solve_in_2b.redistribute(b_N, b_n) 

328 

329 # if 0: 

330 # print 'A_Nn ', rank, A_Nn.shape 

331 # print 'b_N ', rank, b_N.shape 

332 # sys.stdout.flush() 

333 # self.world.barrier() 

334 # print 'A_nn ', rank, A_nn.shape 

335 # print 'b_n ', rank, b_n.shape 

336 # sys.stdout.flush() 

337 # self.world.barrier() 

338 # 

339 # 

340 # print 'b_N ', rank, b_N 

341 # sys.stdout.flush() 

342 # self.world.barrier() 

343 # print 'b_n ', rank, b_n 

344 # sys.stdout.flush() 

345 # self.world.barrier() 

346 # 

347 # print 'A_Nn ', rank, A_Nn 

348 # sys.stdout.flush() 

349 # self.world.barrier() 

350 # print 'A_nn ', rank, A_nn 

351 # sys.stdout.flush() 

352 # self.world.barrier() 

353 

354 info = 0 

355 if self.solve_descr2a.blacsgrid.is_active(): 

356 cgpaw.scalapack_solve(A_nn, self.solve_descr2a.asarray(), b_n, 

357 self.solve_descr2b.asarray()) 

358 if info != 0: 

359 raise RuntimeError('scalapack_solve error: %d' % info) 

360 self.redist_solve_out_2b.redistribute(b_n, b_N) 

361 

362 if self.eh_descr2b.blacsgrid.is_active(): 

363 b_N = b_N.flatten() 

364 else: 

365 b_N = b 

366 

367 # self.dd_comm.broadcast(b_N, 0) 

368 

369 b[:] = b_N 

370 

371 def diagonalize(self, Om, eps_n): 

372 

373 O_nn = self.diag_descr.empty(dtype=float) 

374 if self.eh_descr.blacsgrid.is_active(): 

375 O_nN = Om 

376 else: 

377 O_nN = np.empty((0, 0), dtype=float) 

378 

379 self.redistributor_in.redistribute(O_nN, O_nn) 

380 self.diag_descr.diagonalize_dc(O_nn.copy(), O_nn, eps_n, 'L') 

381 self.redistributor_out.redistribute(O_nn, O_nN) 

382 self.world.broadcast(eps_n, 0) 

383 # Broadcast eigenvectors within domains 

384 if not self.eh_descr.blacsgrid.is_active(): 

385 O_nN = Om 

386 self.dd_comm.broadcast(O_nN, 0) 

387 

388 

389###############################################################################