Coverage for gpaw/xas.py: 70%

1import pickle 

2from math import log, pi, sqrt, ceil 

3from typing import List, Tuple, Union 

4 

5import numpy as np 

6 

7from ase.units import Hartree 

8 

9from gpaw.overlap import Overlap 

10from gpaw.utilities.cg import CG 

11from gpaw.gaunt import gaunt 

12from gpaw.typing import Array1D, Array2D, Array3D, ArrayND 

13import gpaw.mpi as mpi 

14 

15 

16def dipole_matrix_elements(setup): 

17 """calculate length form dipole matrix elements of setup-states 

18 with the core-state""" 

19 l_core = setup.data.lcorehole 

20 lmax = max(setup.lmax, l_core) # include the f states 

21 G_LLL = gaunt(lmax) 

22 

23 # map m, l quantum numbers to L 

24 M = {0: [0]} 

25 for l in range(1, lmax + 1): 

26 M[l] = range(M[l - 1][-1] + 1, M[l - 1][-1] + (l * 2) + 2) 

27 

28 phi_jg = setup.data.phi_jg 

29 nj = len(phi_jg) 

30 

31 A_cmi = np.zeros((3, len(M[l_core]), setup.ni)) 

32 

33 i = 0 

34 for j in range(nj): 

35 l = setup.l_j[j] 

36 a = setup.rgd.integrate(phi_jg[j] * setup.data.phicorehole_g, 

37 n=1) / (4 * pi) 

38 

39 for L2 in M[l]: 

40 for L0 in M[1]: 

41 for m, L1 in enumerate(M[l_core]): 

42 G = sqrt(4 * pi / 3) * G_LLL[L2, L0, L1] 

43 

44 c = L0 % 3 

45 A_cmi[c, m, i] = G * a 

46 

47 i += 1 

48 assert i == setup.ni 

49 

50 return A_cmi 

51 

52 

53def logger(txt, mode, spin, nocc, center, setup): 

54 spin_txt = 'up' 

55 if spin == 1: 

56 spin_txt = 'down' 

57 

58 txt('\nXAS - Calculating matrix elements\n') 

59 txt('Mode: ', mode) 

60 txt('Spin: ', spin_txt, f'({spin})') 

61 txt('Occupied states: ', nocc) 

62 txt('Center: ', center) 

63 txt('Element: ', setup.symbol) 

64 txt('Setup:') 

65 setup.print_info(txt) 

66 

67 

68class XAS: 

69 def __init__(self, paw=None, *args, **kwargs): 

70 if paw is not None: 

71 self.__full_init__(paw, *args, **kwargs) 

72 

73 def __full_init__(self, paw, mode='xas', center=None, 

74 spin=0, relative_index_lumo=0): 

75 """_summary_ 

76 

77 Args: 

78 paw (_type_): GPAW calculator object, with core-hole 

79 mode (str, optional): xas, xes or all . Defaults to 'xas'. 

80 center (int, optional): index of atome with corehole. 

81 Defaults to None. 

82 spin (int, optional): spinprogjection. Defaults to 0. 

83 nocc_cor (int, optional): correction for number of occupied states 

84 used in e.g. XCH XAS simulations. Defaults to 0. 

85 """ 

86 

87 self.log = paw.log 

88 wfs = paw.wfs 

89 self.world = paw.world 

90 kd = wfs.kd 

91 bd = wfs.bd 

92 gd = wfs.gd 

93 self.orthogonal = wfs.gd.orthogonal 

94 self.cell_cv = np.array(wfs.gd.cell_cv) 

95 

96 my_atom_indices = wfs.atom_partition.my_indices 

97 

98 # to allow spin polarized calclulation 

99 nkpts = len(wfs.kd.ibzk_kc) 

100 

101 # the following lines are to stop the user to make mistakes 

102 # if mode is not 'xes' and spin == 1: 

103 # raise RuntimeError( 

104 # 'The core hole is always in spin 0: please use spin=0') 

105 kd_rank = kd.comm.rank 

106 kd_size = kd.comm.size 

107 

108 if wfs.nspins == 1: 

109 if spin != 0: 

110 raise RuntimeError( 

111 'use spin=0 for a spin paired calculation') 

112 nocc = wfs.setups.nvalence // 2 

113 self.list_kpts = range(nkpts) 

114 

115 else: 

116 self.list_kpts = [] 

117 

118 if spin != 0 and spin != 1: 

119 print('spin', spin) 

120 raise RuntimeError( 

121 'use either spin=0 or spin=1') 

122 

123 # find kpoints with correct spin 

124 for i, kpt in enumerate(wfs.kpt_u): 

125 if kpt.s == spin: 

126 self.list_kpts.append(i) 

127 

128 # find number of occupied orbitals, if no fermi smearing 

129 nocc = 0.0 

130 for i in self.list_kpts: 

131 nocc += sum(wfs.kpt_u[i].f_n) 

132 

133 nocc = kd.comm.sum_scalar(nocc) 

134 nocc = int(nocc + 0.5) 

135 

136 nocc += relative_index_lumo 

137 self.nocc = nocc 

138 

139 # look for the center with the corehole 

140 if center is not None: 

141 setup = wfs.setups[center] 

142 a = center 

143 else: 

144 for a, setup in enumerate(wfs.setups): 

145 if setup.phicorehole_g is not None: 

146 break 

147 

148 assert setup.phicorehole_g is not None, 'There is no corehole' 

149 

150 A_cmi = dipole_matrix_elements(setup) 

151 bd_rank = bd.comm.rank 

152 bd_size = bd.comm.size 

153 

154 # xas, xes or all modes 

155 if mode == 'xas': 

156 if bd_rank == 0: 

157 n_start = nocc 

158 else: 

159 n_start = 0 

160 n_end = ceil(wfs.bd.nbands / bd_size) 

161 n = wfs.bd.nbands - nocc 

162 n_diff0 = n_end - nocc 

163 assert n_diff0 > 0, 'Over band parellaised' 

164 n_diff = n_end - n_start 

165 i_n = n_diff0 + n_diff * (bd_size - 1) - n 

166 elif mode == 'xes': # FIX XES later 

167 assert bd_size == 1, "'xes' does not suport band paralisation" 

168 n_start = 0 

169 n_end = nocc 

170 n = n_diff = nocc 

171 

172 elif mode == 'all': 

173 n_start = 0 

174 n_end = ceil(wfs.bd.nbands / bd_size) 

175 n_diff = n_diff0 = n_end - n_start 

176 n = wfs.bd.nbands 

177 i_n = n_diff * bd_size - n 

178 else: 

179 raise RuntimeError( 

180 "wrong keyword for 'mode', use 'xas', 'xes' or 'all'") 

181 

182 self.n = n 

183 l_core = setup.data.lcorehole 

184 self.eps_kn = np.zeros((nkpts, n)) 

185 self.sigma_cmkn = np.zeros((3, l_core * 2 + 1, nkpts, n), complex) 

186 

187 n1 = 0 

188 if bd_rank != 0: 

189 n1 += n_diff0 + n_diff * (bd_rank - 1) 

190 k = kd_rank * (nkpts // kd_size) 

191 

192 for kpt in wfs.kpt_u: 

193 if kpt.s != spin: 

194 continue 

195 

196 n2 = n1 + n_diff 

197 if bd_size != 1 and bd_rank == bd_size - 1: 

198 n2 -= i_n 

199 self.eps_kn[k, n1:n2] = kpt.eps_n[n_start:n_end] * Hartree 

200 

201 if a in my_atom_indices: 

202 P_ni = kpt.P_ani[a][n_start:n_end] 

203 a_cmn = np.inner(A_cmi, P_ni) 

204 weight = kpt.weight * wfs.nspins / 2 

205 self.sigma_cmkn[:, :, k, n1:n2] = weight**0.5 * a_cmn # .real 

206 

207 k += 1 

208 

209 kd.comm.sum(self.sigma_cmkn) 

210 kd.comm.sum(self.eps_kn) 

211 

212 bd.comm.sum(self.sigma_cmkn) 

213 bd.comm.sum(self.eps_kn) 

214 

215 gd.comm.sum(self.sigma_cmkn) 

216 

217 self.symmetry = wfs.kd.symmetry 

218 

219 logger(self.log, mode, spin, nocc, a, setup) 

220 

221 def write(self, fname: str): 

222 """Write matrix elements out to a file""" 

223 if self.world.rank == 0: 

224 self.log(f'Writing to {fname}\n') 

225 with open(fname, mode='wb') as f: 

226 np.savez_compressed( 

227 f, eps_kn=self.eps_kn, sigma_cmkn=self.sigma_cmkn, 

228 orthogonal=self.orthogonal) 

229 self.world.barrier() 

230 

231 @classmethod 

232 def restart(cls, fname: str): 

233 """Read from a matrix elements file""" 

234 self = XAS() 

235 with open(fname, mode='rb') as f: 

236 data = dict(np.load(f)).values() 

237 self.eps_kn, self.sigma_cmkn, self.orthogonal = data 

238 return self 

239 

240 def projection(self, proj=None, proj_xyz=True): 

241 if proj_xyz: 

242 proj_3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], float) 

243 else: 

244 proj_3 = np.array([], float) 

245 

246 if proj is not None: 

247 assert self.orthogonal 

248 proj_2 = np.array(proj, float) 

249 if len(proj_2.shape) == 1: 

250 proj_2 = np.array([proj], float) 

251 

252 for i, p in enumerate(proj_2): 

253 if sum(p**2)**0.5 != 1.0: 

254 print('proj_2 %s not normalized' % i) 

255 proj_2[i] /= sum(p**2)**0.5 

256 

257 proj_tmp = np.zeros((proj_3.shape[0] + proj_2.shape[0], 3), float) 

258 

259 for i, p in enumerate(proj_3): 

260 proj_tmp[i, :] = proj_3[i, :] 

261 

262 for i, p in enumerate(proj_2): 

263 proj_tmp[proj_3.shape[0] + i, :] = proj_2[i, :] 

264 

265 proj_3 = proj_tmp.copy() 

266 

267 return proj_3 

268 

269 def get_oscillator_strength( 

270 self, dks: Union[float, List], kpoint=None, 

271 proj=None, proj_xyz: bool = True, 

272 w: Array1D = None, 

273 raw: bool = False) -> Tuple[Array1D, ArrayND]: 

274 """Calculate stick spectra. 

275 

276 Parameters: 

277 

278 dks: 

279 Energy of first transition. Can be a list for spin-orbit split 

280 spectra. 

281 kpoint: 

282 select a specific k-point to calculate spectrum for 

283 proj: 

284 a list of vectors to project the transition dipole on. Default 

285 is None then only x,y,z components are calculated. a_stick and 

286 a_c squares of the transition moments in resp. direction 

287 proj_xyz: 

288 if True keep projections in x, y and z. a_stck and a_c will have 

289 length 3 + len(proj). if False only those projections 

290 defined by proj keyword, a_stick and a_c will have length len(proj) 

291 

292 Symmtrization has been moved inside get_spectra because we want to 

293 symmtrice squares of transition dipoles. 

294 

295 Returns: 

296 energies: 1D array [n] 

297 oscillator strengths: 3D array [c, m, n] 

298 """ 

299 proj_3 = self.projection(proj=proj, proj_xyz=proj_xyz) 

300 

301 sigma2_cmkn = np.zeros((proj_3.shape[0], 

302 self.sigma_cmkn.shape[1], 

303 self.sigma_cmkn.shape[2], 

304 self.sigma_cmkn.shape[3]), float) 

305 

306 for i, p in enumerate(proj_3): 

307 for m in range(self.sigma_cmkn.shape[1]): 

308 for k in range(self.sigma_cmkn.shape[2]): 

309 s_tmp = np.dot(p, self.sigma_cmkn[:, m, k, :]) 

310 sigma2_cmkn[i, m, k, :] += (s_tmp * 

311 np.conjugate(s_tmp)).real 

312 

313 eps_kn0 = np.min(self.eps_kn) 

314 k_pts = sigma2_cmkn.shape[2] 

315 n = sigma2_cmkn.shape[3] 

316 

317 if isinstance(dks, float) or isinstance(dks, int): 

318 dks = [dks] 

319 

320 energy_kn = np.zeros((k_pts, n * len(dks))) 

321 f_cmkn = np.zeros((sigma2_cmkn.shape[0], 

322 sigma2_cmkn.shape[1], 

323 k_pts, n * len(dks))) 

324 

325 if w is None: 

326 w = np.ones(len(dks)) 

327 elif isinstance(w, float) or isinstance(w, int): 

328 w = [w] 

329 

330 for i in range(len(dks)): 

331 shift = dks[i] - eps_kn0 

332 ienergy_kn = self.eps_kn + shift 

333 

334 if_cmkn = w[i] * 2 * sigma2_cmkn[:, :, :, :] * ienergy_kn / Hartree 

335 

336 energy_kn[:, i * n:(1 + i) * n] = ienergy_kn 

337 f_cmkn[:, :, :, i * n:(1 + i) * n] = if_cmkn 

338 

339 if kpoint is not None: 

340 energy_n = energy_kn[kpoint, :] 

341 f_cmn = f_cmkn[:, :, kpoint, :] 

342 else: 

343 energy_n = np.zeros((k_pts * n * len(dks))) 

344 f_cmn = np.zeros((sigma2_cmkn.shape[0], 

345 sigma2_cmkn.shape[1], 

346 k_pts * n * len(dks))) 

347 

348 for k in range(k_pts): 

349 energy_n[n * k * len(dks): 

350 (k + 1) * n * len(dks)] = energy_kn[k, :] 

351 f_cmn[:, :, n * k * len(dks): 

352 (k + 1) * n * len(dks)] = f_cmkn[:, :, k, :] 

353 if raw: 

354 return energy_n, f_cmn 

355 

356 return energy_n, f_cmn.sum(axis=1) 

357 

358 def get_spectra(self, fwhm=0.5, E_in=None, linbroad=None, 

359 N=1000, kpoint=None, proj=None, proj_xyz=True, 

360 stick=False, dks: Union[float, List] = [0], 

361 w: Array1D = None): 

362 """Calculate spectra. 

363 

364 Parameters: 

365 

366 fwhm: 

367 the full width half maximum in eV for gaussian broadening 

368 linbroad: 

369 a list of three numbers, the first fwhm2, the second the value 

370 where the linear increase starts and the third the value where 

371 the broadening has reached fwhm2. example [0.5, 540, 550] 

372 E_in: 

373 a list of energy values where the spectrum is to be computed 

374 if None the orbital energies will be used to compute the energy 

375 range 

376 N: 

377 the number of bins in the broadened spectrum. If E_in is given N 

378 has no effect 

379 kpoint: 

380 select a specific k-point to calculate spectrum for 

381 proj: 

382 a list of vectors to project the transition dipole on. Default 

383 is None then only x,y,z components are calculated. a_stick and 

384 a_c squares of the transition moments in resp. direction 

385 proj_xyz: 

386 if True keep projections in x, y and z. a_stck and a_c will have 

387 length 3 + len(proj). if False only those projections 

388 defined by proj keyword, a_stick and a_c will have length len(proj) 

389 stick: 

390 if False return broadened spectrum, if True return stick spectrum 

391 

392 Symmtrization has been moved inside get_spectra because we want to 

393 symmtrice squares of transition dipoles. 

394 

395 Returns: 

396 energies: 1D array 

397 oscillator strengths: 3D array 

398 """ 

399 energy_n, f_cmn = self.get_oscillator_strength( 

400 kpoint=kpoint, proj=proj, proj_xyz=proj_xyz, dks=dks, w=w, 

401 raw=True) 

402 

403 if stick: 

404 return energy_n, f_cmn.sum(axis=1) 

405 

406 else: 

407 if E_in is not None: 

408 energy_i = np.array(E_in) 

409 else: 

410 emin = min(energy_n) - 2 * fwhm 

411 emax = max(energy_n) + 2 * fwhm 

412 energy_i = emin + np.arange(N + 1) * ((emax - emin) / N) 

413 

414 if linbroad is None: 

415 return self.constant_broadening( 

416 fwhm, energy_n, f_cmn, energy_i) 

417 

418 else: 

419 return self.variable_broadening( 

420 fwhm, linbroad, energy_n, f_cmn, energy_i) 

421 

422 def variable_broadening( 

423 self, fwhm: float, linbroad: List[float], 

424 eps_n: Array1D, f_cmn: Array3D, 

425 e: Array1D) -> Tuple[Array1D, Array2D]: 

426 """mpirun -n 6 python3 -m pytest test_xas_parallel.py 

427 fwhm: 

428 the full width half maximum in eV for gaussian broadening 

429 linbroad: 

430 a list of three numbers, the first fwhm2, the second the value 

431 where the linear increase starts and the third the value where 

432 the broadening has reached fwhm2. example [0.5, 540, 550] 

433 """ 

434 f_c = np.zeros((f_cmn.shape[0], len(e))) 

435 

436 # constant broadening fwhm until linbroad[1] and a 

437 # constant broadening over linbroad[2] with fwhm2= 

438 # linbroad[0] 

439 fwhm2 = linbroad[0] 

440 lin_e1 = linbroad[1] 

441 lin_e2 = linbroad[2] 

442 print('fwhm', fwhm, fwhm2, lin_e1, lin_e2) 

443 

444 f_cn = f_cmn.sum(axis=1) 

445 

446 # Fold 

447 for n, eps in enumerate(eps_n): 

448 if eps < lin_e1: 

449 alpha = 4 * log(2) / fwhm**2 

450 elif eps <= lin_e2: 

451 fwhm_lin = (fwhm + (eps - lin_e1) * 

452 (fwhm2 - fwhm) / (lin_e2 - lin_e1)) 

453 alpha = 4 * log(2) / fwhm_lin**2 

454 elif eps >= lin_e2: 

455 alpha = 4 * log(2) / fwhm2**2 

456 

457 x = -alpha * (e - eps)**2 

458 x = np.clip(x, -100.0, 100.0) 

459 f_c += np.outer(f_cn[:, n], 

460 (alpha / pi)**0.5 * np.exp(x)) 

461 

462 return e, f_c 

463 

464 def constant_broadening( 

465 self, fwhm: float, eps_n: Array1D, f_cmn, 

466 energy_i: Array1D) -> Tuple[Array1D, Array2D]: 

467 """ 

468 fwhm: 

469 the full width half maximum in eV for gaussian broadening 

470 """ 

471 

472 # constant broadening fwhm 

473 # alpha = 1 / (2 sigma^2) with fwhm = 2 sqrt{2 log 2} sigma 

474 alpha = 4 * log(2) / fwhm**2 

475 

476 f_cn = f_cmn.sum(axis=1) 

477 

478 # Fold 

479 f_ci = np.zeros((3, len(energy_i))) 

480 for n, eps in enumerate(eps_n): 

481 x = -alpha * (energy_i - eps) ** 2 

482 x = np.clip(x, -100.0, 100.0) 

483 f_ci += np.outer(f_cn[:, n], (alpha / pi)**0.5 * np.exp(x)) 

484 

485 return energy_i, f_ci 

486 

487 

488class RecursionMethod: 

489 """This class implements the Haydock recursion method. """ 

490 

491 def __init__(self, paw=None, filename=None, 

492 tol=1e-10, maxiter=100, proj=None, 

493 proj_xyz=True): 

494 

495 if paw is not None: 

496 wfs = paw.wfs 

497 assert wfs.gd.orthogonal 

498 

499 self.wfs = wfs 

500 self.hamiltonian = paw.hamiltonian 

501 self.nkpts = len(wfs.kd.ibzk_kc) * wfs.nspins 

502 self.nmykpts = len(wfs.kpt_u) 

503 

504 self.k1 = wfs.kd.comm.rank * self.nmykpts 

505 self.k2 = self.k1 + self.nmykpts 

506 

507 print('k1', self.k1, 'k2', self.k2) 

508 

509 # put spin and weight index in the columns corresponding 

510 # to this processors k-points 

511 self.spin_k = np.zeros(self.nkpts, int) 

512 self.weight_k = np.zeros(self.nkpts) 

513 

514 for n, i in enumerate(range(self.k1, self.k2)): 

515 self.spin_k[i] = wfs.kpt_u[n].s 

516 self.weight_k[i] = wfs.kpt_u[n].weight 

517 

518 self.op_scc = None 

519 if wfs.kd.symmetry is not None: 

520 self.op_scc = wfs.kd.symmetry.op_scc 

521 else: 

522 self.k1 = 0 

523 self.k2 = None 

524 self.wfs = None 

525 wfs = None 

526 

527 self.tol = tol 

528 self.maxiter = maxiter 

529 

530 if filename is not None: 

531 self.read(filename) 

532 if wfs is not None: 

533 self.allocate_tmp_arrays() 

534 else: 

535 self.initialize_start_vector(proj=proj, proj_xyz=proj_xyz) 

536 

537 def read(self, filename): 

538 with open(filename, 'rb') as fd: 

539 data = pickle.load(fd) 

540 self.nkpts = data['nkpts'] 

541 if 'swaps' in data: 

542 # This is an old file: 

543 self.op_scc = np.array([np.identity(3, int)[list(swap)] 

544 for swap in data['swaps']]) 

545 else: 

546 self.op_scc = data['symmetry operations'] 

547 self.weight_k = data['weight_k'] 

548 self.spin_k = data['spin_k'] 

549 self.dim = data['dim'] 

550 k1, k2 = self.k1, self.k2 

551 if k2 is None: 

552 k2 = self.nkpts 

553 a_kci, b_kci = data['ab'] 

554 self.a_uci = a_kci[k1:k2].copy() 

555 self.b_uci = b_kci[k1:k2].copy() 

556 

557 if self.wfs is not None and 'arrays' in data: 

558 print('reading arrays') 

559 w_kcG, wold_kcG, y_kcG = data['arrays'] 

560 i = [slice(k1, k2), slice(0, self.dim)] + self.wfs.gd.get_slice() 

561 self.w_ucG = w_kcG[i].copy() 

562 self.wold_ucG = wold_kcG[i].copy() 

563 self.y_ucG = y_kcG[i].copy() 

564 

565 def write(self, filename, mode=''): 

566 assert self.wfs is not None 

567 kpt_comm = self.wfs.kd.comm 

568 gd = self.wfs.gd 

569 

570 if gd.comm.rank == 0: 

571 if kpt_comm.rank == 0: 

572 nmyu, dim, ni = self.a_uci.shape 

573 a_kci = np.empty((kpt_comm.size, nmyu, dim, ni), 

574 self.wfs.dtype) 

575 b_kci = np.empty((kpt_comm.size, nmyu, dim, ni), 

576 self.wfs.dtype) 

577 

578 kpt_comm.gather(self.a_uci, 0, a_kci) 

579 kpt_comm.gather(self.b_uci, 0, b_kci) 

580 kpt_comm.sum(self.spin_k, 0) 

581 kpt_comm.sum(self.weight_k, 0) 

582 

583 a_kci.shape = (self.nkpts, dim, ni) 

584 b_kci.shape = (self.nkpts, dim, ni) 

585 data = {'ab': (a_kci, b_kci), 

586 'nkpts': self.nkpts, 

587 'symmetry operations': self.op_scc, 

588 'weight_k': self.weight_k, 

589 'spin_k': self.spin_k, 

590 'dim': dim} 

591 else: 

592 kpt_comm.gather(self.a_uci, 0) 

593 kpt_comm.gather(self.b_uci, 0) 

594 kpt_comm.sum(self.spin_k, 0) 

595 kpt_comm.sum(self.weight_k, 0) 

596 

597 if mode == 'all': 

598 w0_ucG = gd.collect(self.w_ucG) 

599 wold0_ucG = gd.collect(self.wold_ucG) 

600 y0_ucG = gd.collect(self.y_ucG) 

601 if gd.comm.rank == 0: 

602 if kpt_comm.rank == 0: 

603 w_kcG = gd.empty((self.nkpts, dim), self.wfs.dtype, 

604 global_array=True) 

605 wold_kcG = gd.empty((self.nkpts, dim), self.wfs.dtype, 

606 global_array=True) 

607 y_kcG = gd.empty((self.nkpts, dim), self.wfs.dtype, 

608 global_array=True) 

609 kpt_comm.gather(w0_ucG, 0, w_kcG) 

610 kpt_comm.gather(wold0_ucG, 0, wold_kcG) 

611 kpt_comm.gather(y0_ucG, 0, y_kcG) 

612 data['arrays'] = (w_kcG, wold_kcG, y_kcG) 

613 else: 

614 kpt_comm.gather(w0_ucG, 0) 

615 kpt_comm.gather(wold0_ucG, 0) 

616 kpt_comm.gather(y0_ucG, 0) 

617 

618 if self.wfs.world.rank == 0: 

619 with open(filename, 'wb') as fd: 

620 pickle.dump(data, fd) 

621 

622 def allocate_tmp_arrays(self): 

623 

624 self.tmp1_cG = self.wfs.gd.zeros(self.dim, self.wfs.dtype) 

625 self.tmp2_cG = self.wfs.gd.zeros(self.dim, self.wfs.dtype) 

626 self.z_cG = self.wfs.gd.zeros(self.dim, self.wfs.dtype) 

627 

628 def initialize_start_vector(self, proj=None, proj_xyz=True): 

629 # proj is one list of vectors [[e1_x,e1_y,e1_z],[e2_x,e2_y,e2_z]] 

630 # ( or [ex,ey,ez] if only one projection ) 

631 # that the spectrum will be projected on 

632 # default is to only calculate the averaged spectrum 

633 # if proj_xyz is True, keep projection in x,y,z, if False 

634 # only calculate the projections in proj 

635 

636 # Create initial wave function: 

637 nmykpts = self.nmykpts 

638 

639 for a, setup in enumerate(self.wfs.setups): 

640 if setup.phicorehole_g is not None: 

641 break 

642 

643 A_cmi = dipole_matrix_elements(setup) 

644 A_ci = A_cmi[:, 0, :] 

645 

646 # 

647 # proj keyword 

648 # 

649 

650 # check normalization of incoming vectors 

651 if proj is not None: 

652 proj_2 = np.array(proj, float) 

653 if len(proj_2.shape) == 1: 

654 proj_2 = np.array([proj], float) 

655 

656 for i, p in enumerate(proj_2): 

657 if sum(p ** 2) ** 0.5 != 1.0: 

658 print('proj_2 %s not normalized' % i) 

659 proj_2[i] /= sum(p ** 2) ** 0.5 

660 

661 proj_tmp = [] 

662 for p in proj_2: 

663 proj_tmp.append(np.dot(p, A_ci)) 

664 proj_tmp = np.array(proj_tmp, float) 

665 

666 # if proj_xyz is True, append projections to A_ci 

667 if proj_xyz: 

668 A_ci_tmp = np.zeros((3 + proj_2.shape[0], A_ci.shape[1])) 

669 A_ci_tmp[0:3, :] = A_ci 

670 A_ci_tmp[3:, :] = proj_tmp 

671 

672 # otherwise, replace A_ci by projections 

673 else: 

674 A_ci_tmp = np.zeros((proj_2.shape[0], A_ci.shape[1])) 

675 A_ci_tmp = proj_tmp 

676 A_ci = A_ci_tmp 

677 

678 self.dim = len(A_ci) 

679 

680 self.allocate_tmp_arrays() 

681 

682 self.w_ucG = self.wfs.gd.zeros((nmykpts, self.dim), self.wfs.dtype) 

683 self.wold_ucG = self.wfs.gd.zeros((nmykpts, self.dim), self.wfs.dtype) 

684 self.y_ucG = self.wfs.gd.zeros((nmykpts, self.dim), self.wfs.dtype) 

685 

686 self.a_uci = np.zeros((nmykpts, self.dim, 0), self.wfs.dtype) 

687 self.b_uci = np.zeros((nmykpts, self.dim, 0), self.wfs.dtype) 

688 

689 A_aci = self.wfs.pt.dict(3, zero=True) 

690 if a in A_aci: 

691 A_aci[a] = A_ci.astype(self.wfs.dtype) 

692 for u in range(nmykpts): 

693 self.wfs.pt.add(self.w_ucG[u], A_aci, u) 

694 

695 def run(self, nsteps, inverse_overlap='exact'): 

696 

697 if inverse_overlap == 'exact': 

698 self.solver = self.solve 

699 elif inverse_overlap == 'approximate': 

700 self.solver = self.solve2 

701 elif inverse_overlap == 'noinverse': 

702 self.solver = self.solve3 

703 else: 

704 raise RuntimeError("Error, inverse_solver must be either 'exact' " 

705 "'approximate' or 'noinverse'") 

706 

707 self.overlap = Overlap() 

708 

709 ni = self.a_uci.shape[2] 

710 a_uci = np.empty((self.nmykpts, self.dim, ni + nsteps), self.wfs.dtype) 

711 b_uci = np.empty((self.nmykpts, self.dim, ni + nsteps), self.wfs.dtype)