Coverage for gpaw/wavefunctions/lcao.py: 89%

1import numpy as np 

2from ase.units import Bohr 

3from ase.utils.timing import timer 

4 

5# from gpaw import debug 

6from gpaw.directmin.etdm_lcao import LCAOETDM 

7from gpaw.directmin.tools import loewdin_lcao, gramschmidt_lcao 

8from gpaw.lcao.atomic_correction import (DenseAtomicCorrection, 

9 SparseAtomicCorrection) 

10# from gpaw.lcao.overlap import NewTwoCenterIntegrals as NewTCI 

11from gpaw.lcao.tci import TCIExpansions 

12from gpaw.lfc import BasisFunctions 

13from gpaw.utilities import unpack_hermitian 

14from gpaw.utilities.blas import mmm, gemmdot 

15from gpaw.utilities.tools import tri2full 

16from gpaw.wavefunctions.base import WaveFunctions 

17from gpaw.wavefunctions.mode import Mode 

18 

19 

20class LCAO(Mode): 

21 name = 'lcao' 

22 

23 def __init__(self, atomic_correction=None, interpolation=3, 

24 force_complex_dtype=False): 

25 self.atomic_correction = atomic_correction 

26 self.interpolation = interpolation 

27 Mode.__init__(self, force_complex_dtype) 

28 

29 def __call__(self, *args, **kwargs): 

30 return LCAOWaveFunctions(*args, 

31 atomic_correction=self.atomic_correction, 

32 **kwargs) 

33 

34 def __repr__(self): 

35 return f'LCAO({self.todict()})' 

36 

37 def todict(self): 

38 dct = Mode.todict(self) 

39 dct['interpolation'] = self.interpolation 

40 return dct 

41 

42 

43def update_phases(C_unM, q_u, ibzk_qc, spos_ac, oldspos_ac, setups, Mstart): 

44 """Complex-rotate coefficients compensating discontinuous phase shift. 

45 

46 This changes the coefficients to counteract the phase discontinuity 

47 of overlaps when atoms move across a cell boundary.""" 

48 

49 # We don't want to apply any phase shift unless we crossed a cell 

50 # boundary. So we round the shift to either 0 or 1. 

51 # 

52 # Example: spos_ac goes from 0.01 to 0.99 -- this rounds to 1 and 

53 # we apply the phase. If someone moves an atom by half a cell 

54 # without crossing a boundary, then we are out of luck. But they 

55 # should have reinitialized from LCAO anyway. 

56 phase_qa = np.exp(2j * np.pi * 

57 np.dot(ibzk_qc, (spos_ac - oldspos_ac).T.round())) 

58 

59 for q, C_nM in zip(q_u, C_unM): 

60 if C_nM is None: 

61 continue 

62 for a in range(len(spos_ac)): 

63 M1 = setups.M_a[a] - Mstart 

64 M2 = M1 + setups[a].nao 

65 M1 = max(0, M1) 

66 C_nM[:, M1:M2] *= phase_qa[q, a] # (may truncate M2) 

67 

68 

69# replace by class to make data structure perhaps a bit less confusing 

70def get_r_and_offsets(nl, spos_ac, cell_cv): 

71 r_and_offset_aao = {} 

72 

73 def add(a1, a2, R_c, offset): 

74 if (a1, a2) not in r_and_offset_aao: 

75 r_and_offset_aao[(a1, a2)] = [] 

76 r_and_offset_aao[(a1, a2)].append((R_c, offset)) 

77 

78 for a1, spos1_c in enumerate(spos_ac): 

79 a2_a, offsets = nl.get_neighbors(a1) 

80 for a2, offset in zip(a2_a, offsets): 

81 spos2_c = spos_ac[a2] + offset 

82 

83 R_c = np.dot(spos2_c - spos1_c, cell_cv) 

84 add(a1, a2, R_c, offset) 

85 if a1 != a2 or offset.any(): 

86 add(a2, a1, -R_c, -offset) 

87 

88 return r_and_offset_aao 

89 

90 

91class LCAOWaveFunctions(WaveFunctions): 

92 mode = 'lcao' 

93 

94 def __init__(self, ksl, gd, nvalence, setups, bd, 

95 dtype, world, kd, kptband_comm, timer, 

96 atomic_correction=None, collinear=True): 

97 WaveFunctions.__init__(self, gd, nvalence, setups, bd, 

98 dtype, collinear, world, kd, 

99 kptband_comm, timer) 

100 self.ksl = ksl 

101 self.S_qMM = None 

102 self.T_qMM = None 

103 self.P_aqMi = None 

104 self.debug_tci = False 

105 

106 if atomic_correction is None: 

107 atomic_correction = 'sparse' if ksl.using_blacs else 'dense' 

108 

109 if atomic_correction == 'sparse': 

110 self.atomic_correction_cls = SparseAtomicCorrection 

111 else: 

112 assert atomic_correction == 'dense' 

113 self.atomic_correction_cls = DenseAtomicCorrection 

114 

115 # self.tci = NewTCI(gd.cell_cv, gd.pbc_c, setups, kd.ibzk_qc, kd.gamma) 

116 with self.timer('TCI: Evaluate splines'): 

117 self.tciexpansions = TCIExpansions.new_from_setups(setups) 

118 

119 self.basis_functions = BasisFunctions(gd, 

120 [setup.basis_functions_J 

121 for setup in setups], 

122 kd, 

123 dtype=dtype, 

124 cut=True) 

125 

126 self.coefficients_read_from_file = False 

127 self.set_orthonormalized(False) 

128 

129 def set_orthonormalized(self, flag): 

130 self.orthonormalized = flag 

131 

132 @timer('Orthonormalize') 

133 def orthonormalize(self, kpt=None, type='gramschmidt'): 

134 assert type == 'gramschmidt' or type == 'loewdin' 

135 if kpt is None: 

136 for kpt in self.kpt_u: 

137 self.orthonormalize(kpt) 

138 self.orthonormalized = True 

139 return 

140 if type == 'loewdin': 

141 kpt.C_nM[:] = loewdin_lcao(kpt.C_nM, kpt.S_MM.conj()) 

142 elif type == 'gramschmidt': 

143 kpt.C_nM[:] = gramschmidt_lcao(kpt.C_nM, kpt.S_MM.conj()) 

144 

145 def empty(self, n=(), global_array=False, realspace=False): 

146 if realspace: 

147 return self.gd.empty(n, self.dtype, global_array) 

148 else: 

149 if isinstance(n, int): 

150 n = (n,) 

151 nao = self.setups.nao 

152 return np.empty(n + (nao,), self.dtype) 

153 

154 def __str__(self): 

155 s = 'Wave functions: LCAO\n' 

156 s += ' Diagonalizer: %s\n' % self.ksl.get_description() 

157 s += (' Atomic Correction: %s\n' 

158 % self.atomic_correction_cls.description) 

159 s += ' Data-type: %s\n' % self.dtype.__name__ 

160 return s 

161 

162 def set_eigensolver(self, eigensolver): 

163 WaveFunctions.set_eigensolver(self, eigensolver) 

164 if eigensolver: 

165 if isinstance(eigensolver, LCAOETDM): 

166 eigensolver.initialize(self.gd, self.dtype, self.bd.nbands, 

167 self.kd.nibzkpts, self.setups.nao, 

168 self.ksl.using_blacs, 

169 self.bd.comm.size, self.kpt_u) 

170 else: 

171 eigensolver.initialize(self.gd, self.dtype, self.setups.nao, 

172 self.ksl) 

173 

174 def set_positions(self, spos_ac, atom_partition=None, move_wfs=False): 

175 oldspos_ac = self.spos_ac 

176 with self.timer('Basic WFS set positions'): 

177 WaveFunctions.set_positions(self, spos_ac, atom_partition) 

178 

179 with self.timer('Basis functions set positions'): 

180 self.basis_functions.set_positions(spos_ac) 

181 

182 if self.ksl is not None: 

183 self.basis_functions.set_matrix_distribution(self.ksl.Mstart, 

184 self.ksl.Mstop) 

185 

186 Mstop = self.ksl.Mstop 

187 Mstart = self.ksl.Mstart 

188 

189 # if self.ksl.using_blacs: # XXX 

190 # S and T have been distributed to a layout with blacs, so 

191 # discard them to force reallocation from scratch. 

192 # 

193 # TODO: evaluate S and T when they *are* distributed, thus saving 

194 # memory and avoiding this problem 

195 for kpt in self.kpt_u: 

196 kpt.S_MM = None 

197 kpt.T_MM = None 

198 

199 # Free memory in case of old matrices: 

200 self.S_qMM = self.T_qMM = self.P_aqMi = None 

201 

202 if self.dtype == complex and oldspos_ac is not None: 

203 update_phases([kpt.C_nM for kpt in self.kpt_u], 

204 [kpt.q for kpt in self.kpt_u], 

205 self.kd.ibzk_qc, spos_ac, oldspos_ac, 

206 self.setups, Mstart) 

207 

208 self.timer.start('mktci') 

209 manytci = self.tciexpansions.get_manytci_calculator( 

210 self.setups, self.gd, spos_ac, self.kd.ibzk_qc, self.dtype, 

211 self.timer) 

212 self.timer.stop('mktci') 

213 self.manytci = manytci 

214 self.newtci = manytci.tci 

215 

216 my_atom_indices = self.basis_functions.my_atom_indices 

217 self.timer.start('ST tci') 

218 newS_qMM, newT_qMM = manytci.O_qMM_T_qMM(self.gd.comm, 

219 Mstart, Mstop, 

220 self.ksl.using_blacs) 

221 self.timer.stop('ST tci') 

222 self.timer.start('P tci') 

223 P_qIM = manytci.P_qIM(my_atom_indices) 

224 self.timer.stop('P tci') 

225 self.P_aqMi = manytci.P_aqMi(my_atom_indices) 

226 self.P_qIM = P_qIM # XXX atomic correction 

227 

228 self.atomic_correction = self.atomic_correction_cls.new_from_wfs(self) 

229 

230 # TODO 

231 # OK complex/conj, periodic images 

232 # OK scalapack 

233 # derivatives/forces 

234 # sparse 

235 # use symmetry/conj tricks to reduce calculations 

236 # enable caching of spherical harmonics 

237 

238 self.atomic_correction.add_overlap_correction(newS_qMM) 

239 self.allocate_arrays_for_projections(my_atom_indices) 

240 

241 newS_qMM = self.ksl.distribute_overlap_matrix(newS_qMM, root=-1) 

242 newT_qMM = self.ksl.distribute_overlap_matrix(newT_qMM, root=-1) 

243 

244 self.positions_set = True 

245 

246 for kpt in self.kpt_u: 

247 q = kpt.q 

248 kpt.S_MM = newS_qMM[q] 

249 kpt.T_MM = newT_qMM[q] 

250 self.S_qMM = newS_qMM 

251 self.T_qMM = newT_qMM 

252 

253 # Elpa wants to reuse the decomposed form of S_qMM. 

254 # We need to keep track of the existence of that object here, 

255 # since this is where we change S_qMM. Hence, expect this to 

256 # become arrays after the first diagonalization: 

257 self.decomposed_S_qMM = [None] * len(self.S_qMM) 

258 self.set_orthonormalized(False) 

259 

260 def initialize(self, density, hamiltonian, spos_ac): 

261 # Note: The above line exists also in set_positions. 

262 # This is guaranteed to be correct, but we can probably remove one. 

263 # Of course no human can understand the initialization process, 

264 # so this will be some other day. 

265 self.timer.start('LCAO WFS Initialize') 

266 if density.nt_sG is None: 

267 if self.kpt_u[0].f_n is None or self.kpt_u[0].C_nM is None: 

268 density.initialize_from_atomic_densities(self.basis_functions) 

269 else: 

270 # We have the info we need for a density matrix, so initialize 

271 # from that instead of from scratch. This will be the case 

272 # after set_positions() during a relaxation 

273 density.initialize_from_wavefunctions(self) 

274 # Initialize GLLB-potential from basis function orbitals 

275 if hamiltonian.xc.type == 'GLLB': 

276 hamiltonian.xc.initialize_from_atomic_orbitals( 

277 self.basis_functions) 

278 

279 else: 

280 # After a restart, nt_sg doesn't exist yet, so we'll have to 

281 # make sure it does. Of course, this should have been taken care 

282 # of already by this time, so we should improve the code elsewhere 

283 density.calculate_normalized_charges_and_mix() 

284 

285 hamiltonian.update(density) 

286 self.timer.stop('LCAO WFS Initialize') 

287 

288 return 0, 0 

289 

290 def initialize_wave_functions_from_lcao(self): 

291 """Fill the calc.wfs.kpt_[u].psit_nG arrays with useful data. 

292 

293 Normally psit_nG is NOT used in lcao mode, but some extensions 

294 (like ase.dft.wannier) want to have it. 

295 This code is adapted from fd.py / initialize_from_lcao_coefficients() 

296 and fills psit_nG with data constructed from the current lcao 

297 coefficients (kpt.C_nM). 

298 

299 (This may or may not work in band-parallel case!) 

300 """ 

301 from gpaw.wavefunctions.arrays import UniformGridWaveFunctions 

302 bfs = self.basis_functions 

303 for kpt in self.kpt_u: 

304 kpt.psit = UniformGridWaveFunctions( 

305 self.bd.nbands, self.gd, self.dtype, kpt=kpt.q, dist=None, 

306 spin=kpt.s, collinear=True) 

307 kpt.psit_nG[:] = 0.0 

308 bfs.lcao_to_grid(kpt.C_nM, kpt.psit_nG[:self.bd.mynbands], kpt.q) 

309 

310 def initialize_wave_functions_from_restart_file(self): 

311 """Dummy function to ensure compatibility to fd mode""" 

312 self.initialize_wave_functions_from_lcao() 

313 

314 def add_orbital_density(self, nt_G, kpt, n): 

315 rank, q = self.kd.get_rank_and_index(kpt.k) 

316 u = q * self.nspins + kpt.s 

317 assert rank == self.kd.comm.rank 

318 assert self.kpt_u[u] is kpt 

319 psit_G = self._get_wave_function_array(u, n, realspace=True) 

320 self.add_realspace_orbital_to_density(nt_G, psit_G) 

321 

322 def calculate_density_matrix(self, f_n, C_nM, rho_MM=None): 

323 self.timer.start('Calculate density matrix') 

324 rho_MM = self.ksl.calculate_density_matrix(f_n, C_nM, rho_MM) 

325 self.timer.stop('Calculate density matrix') 

326 return rho_MM 

327 

328 if 1: 

329 # XXX Should not conjugate, but call gemm(..., 'c') 

330 # Although that requires knowing C_Mn and not C_nM. 

331 # that also conforms better to the usual conventions in literature 

332 Cf_Mn = C_nM.T.conj() * f_n 

333 self.timer.start('gemm') 

334 mmm(1.0, Cf_Mn, 'N', C_nM, 'N', 0.0, rho_MM) 

335 self.timer.stop('gemm') 

336 self.timer.start('band comm sum') 

337 self.bd.comm.sum(rho_MM) 

338 self.timer.stop('band comm sum') 

339 else: 

340 # Alternative suggestion. Might be faster. Someone should test this 

341 from gpaw.utilities.blas import r2k 

342 C_Mn = C_nM.T.copy() 

343 r2k(0.5, C_Mn, f_n * C_Mn, 0.0, rho_MM) 

344 tri2full(rho_MM) 

345 

346 def calculate_atomic_density_matrices_with_occupation(self, D_asp, f_un): 

347 # ac = self.atomic_correction 

348 # if ac.implements_distributed_projections(): 

349 # D2_asp = ac.redistribute(self, D_asp, type='asp', op='forth') 

350 # WaveFunctions.calculate_atomic_density_matrices_with_occupation( 

351 # self, D2_asp, f_un) 

352 # D3_asp = ac.redistribute(self, D2_asp, type='asp', op='back') 

353 # for a in D_asp: 

354 # D_asp[a][:] = D3_asp[a] 

355 # else: 

356 WaveFunctions.calculate_atomic_density_matrices_with_occupation( 

357 self, D_asp, f_un) 

358 

359 def calculate_density_matrix_delta(self, d_nn, C_nM, rho_MM=None): 

360 self.timer.start('Calculate density matrix') 

361 rho_MM = self.ksl.calculate_density_matrix_delta(d_nn, C_nM, rho_MM) 

362 self.timer.stop('Calculate density matrix') 

363 return rho_MM 

364 

365 def add_to_density_from_k_point_with_occupation(self, nt_sG, kpt, f_n): 

366 """Add contribution to pseudo electron-density. Do not use the standard 

367 occupation numbers, but ones given with argument f_n.""" 

368 # Custom occupations are used in calculation of response potential 

369 # with GLLB-potential 

370 if kpt.rho_MM is None: 

371 rho_MM = self.calculate_density_matrix(f_n, kpt.C_nM) 

372 if hasattr(kpt, 'c_on'): 

373 assert self.bd.comm.size == 1 

374 d_nn = np.zeros((self.bd.mynbands, self.bd.mynbands), 

375 dtype=kpt.C_nM.dtype) 

376 for ne, c_n in zip(kpt.ne_o, kpt.c_on): 

377 assert abs(c_n.imag).max() < 1e-14 

378 d_nn += ne * np.outer(c_n.conj(), c_n).real 

379 rho_MM += self.calculate_density_matrix_delta(d_nn, kpt.C_nM) 

380 else: 

381 rho_MM = kpt.rho_MM 

382 self.timer.start('Construct density') 

383 self.basis_functions.construct_density(rho_MM, nt_sG[kpt.s], kpt.q) 

384 self.timer.stop('Construct density') 

385 

386 def add_to_kinetic_density_from_k_point(self, taut_G, kpt): 

387 raise NotImplementedError('Kinetic density calculation for LCAO ' 

388 'wavefunctions is not implemented.') 

389 

390 def calculate_forces(self, hamiltonian, F_av): 

391 self.timer.start('LCAO forces') 

392 

393 Fref_av = np.zeros_like(F_av) 

394 self.forcecalc = LCAOforces(self.ksl, self.dtype, self.gd, 

395 self.bd, self.kd, self.kpt_u, self.nspins, 

396 self.basis_functions, self.newtci, 

397 self.P_aqMi, self.setups, 

398 self.manytci, hamiltonian, 

399 self, self.spos_ac, 

400 self.timer, Fref_av) 

401 

402 F_av[:, :] = self.forcecalc.get_forces_sum_GS() 

403 self.timer.stop('LCAO forces') 

404 

405 def _get_wave_function_array(self, u, n, realspace=True, periodic=False): 

406 # XXX Taking kpt is better than taking u 

407 kpt = self.kpt_u[u] 

408 C_M = kpt.C_nM[n] 

409 

410 if realspace: 

411 psit_G = self.gd.zeros(dtype=self.dtype) 

412 self.basis_functions.lcao_to_grid(C_M, psit_G, kpt.q) 

413 if periodic and self.dtype == complex: 

414 k_c = self.kd.ibzk_kc[kpt.k] 

415 return self.gd.plane_wave(-k_c) * psit_G 

416 return psit_G 

417 else: 

418 return C_M 

419 

420 def write(self, writer, write_wave_functions=False): 

421 WaveFunctions.write(self, writer) 

422 if write_wave_functions: 

423 self.write_wave_functions(writer) 

424 

425 def write_wave_functions(self, writer): 

426 writer.add_array( 

427 'coefficients', 

428 (self.nspins, self.kd.nibzkpts, self.bd.nbands, self.setups.nao), 

429 dtype=self.dtype) 

430 for s in range(self.nspins): 

431 for k in range(self.kd.nibzkpts): 

432 C_nM = self.collect_array('C_nM', k, s) 

433 writer.fill(C_nM * Bohr**-1.5) 

434 

435 def read(self, reader): 

436 WaveFunctions.read(self, reader) 

437 r = reader.wave_functions 

438 if 'coefficients' in r: 

439 self.read_wave_functions(r) 

440 

441 def read_wave_functions(self, reader): 

442 c = 1.0 if getattr(reader, 'version', 3) >= 4 else Bohr**1.5 

443 for kpt in self.kpt_u: 

444 C_nM = reader.proxy('coefficients', kpt.s, kpt.k) 

445 kpt.C_nM = self.bd.empty(self.setups.nao, dtype=self.dtype) 

446 for myn, C_M in enumerate(kpt.C_nM): 

447 n = self.bd.global_index(myn) 

448 # XXX number of bands could have been rounded up! 

449 if n >= len(C_nM): 

450 break 

451 C_M[:] = C_nM[n] * c 

452 

453 self.coefficients_read_from_file = True 

454 

455 def estimate_memory(self, mem): 

456 nq = len(self.kd.ibzk_qc) 

457 nao = self.setups.nao 

458 ni_total = sum([setup.ni for setup in self.setups]) 

459 itemsize = mem.itemsize[self.dtype] 

460 mem.subnode('C [qnM]', nq * self.bd.mynbands * nao * itemsize) 

461 nM1, nM2 = self.ksl.get_overlap_matrix_shape() 

462 mem.subnode('S, T [2 x qmm]', 2 * nq * nM1 * nM2 * itemsize) 

463 mem.subnode('P [aqMi]', nq * nao * ni_total // self.gd.comm.size) 

464 # self.tci.estimate_memory(mem.subnode('TCI')) 

465 self.basis_functions.estimate_memory(mem.subnode('BasisFunctions')) 

466 self.eigensolver.estimate_memory(mem.subnode('Eigensolver'), 

467 self.dtype) 

468 

469 

470class LCAOforces: 

471 

472 def __init__(self, ksl, dtype, gd, bd, kd, kpt_u, nspins, bfs, newtci, 

473 P_aqMi, setups, manytci, hamiltonian, wfs, spos_ac, 

474 timer, Fref_av): 

475 """ Object which calculates LCAO forces """ 

476 

477 self.ksl = ksl 

478 self.nao = ksl.nao 

479 self.mynao = ksl.mynao 

480 self.dtype = dtype 

481 self.newtci = newtci 

482 self.manytci = manytci 

483 self.P_aqMi = P_aqMi 

484 self.gd = gd 

485 self.bd = bd 

486 self.kd = kd 

487 self.kpt_u = kpt_u 

488 self.nspins = nspins 

489 self.bfs = bfs 

490 self.spos_ac = spos_ac 

491 self.Mstart = ksl.Mstart 

492 self.Mstop = ksl.Mstop 

493 self.setups = setups 

494 self.hamiltonian = hamiltonian 

495 self.wfs = wfs 

496 self.timer = timer 

497 self.Fref_av = Fref_av 

498 self.my_atom_indices = bfs.my_atom_indices 

499 self.atom_indices = bfs.atom_indices 

500 self.dH_asp = hamiltonian.dH_asp 

501 

502 from gpaw.kohnsham_layouts import BlacsOrbitalLayouts 

503 self.isblacs = isinstance(self.ksl, BlacsOrbitalLayouts) 

504 

505 if not self.isblacs: 

506 self.timer.start('TCI derivative') 

507 self.dThetadR_qvMM, self.dTdR_qvMM = self.manytci.O_qMM_T_qMM( 

508 self.gd.comm, self.Mstart, self.Mstop, False, derivative=True) 

509 self.dPdR_aqvMi = self.manytci.P_aqMi(self.bfs.my_atom_indices, 

510 derivative=True) 

511 

512 self.gd.comm.sum(self.dThetadR_qvMM) 

513 self.gd.comm.sum(self.dTdR_qvMM) 

514 self.timer.stop('TCI derivative') 

515 self.rhoT_uMM, self.ET_uMM = self.get_den_mat_and_E() 

516 

517 def get_forces_sum_GS(self): 

518 """ This function calculates ground state forces in LCAO mode """ 

519 if not self.isblacs: 

520 F_av = np.zeros_like(self.Fref_av) 

521 Fkin_av = self.get_kinetic_term() 

522 Fpot_av = self.get_pot_term() 

523 Ftheta_av = self.get_den_mat_term() 

524 Frho_av = self.get_den_mat_paw_term() 

525 Fatom_av = self.get_atomic_density_term() 

526 F_av += Fkin_av + Fpot_av + Ftheta_av + Frho_av + Fatom_av 

527 else: 

528 F_av = np.zeros_like(self.Fref_av) 

529 Fpot_av = self.get_pot_term_blacs() 

530 Fkin_av, Ftheta_av = self.get_kin_and_den_term_blacs() 

531 Fatom_av, Frho_av = self.get_at_den_and_den_paw_blacs() 

532 

533 F_av += Fkin_av + Fpot_av + Ftheta_av + Frho_av + Fatom_av 

534 

535 self.timer.start('Wait for sum') 

536 self.ksl.orbital_comm.sum(F_av) 

537 if self.bd.comm.rank == 0: 

538 self.kd.comm.sum(F_av, 0) 

539 self.timer.stop('Wait for sum') 

540 

541 return F_av 

542 

543 def _slices(self, indices): 

544 for a in indices: 

545 M1 = self.bfs.M_a[a] - self.Mstart 

546 M2 = M1 + self.setups[a].nao 

547 if M2 > 0: 

548 yield a, max(0, M1), M2 

549 

550 def slices(self): 

551 return self._slices(self.atom_indices) 

552 

553 def my_slices(self): 

554 return self._slices(self.my_atom_indices) 

555 

556 def get_den_mat_and_E(self): 

557 # 

558 # ----- ----- 

559 # \ -1 \ * 

560 # E = ) S H rho = ) c eps f c 

561 # mu nu / mu x x z z nu / n mu n n n nu 

562 # ----- ----- 

563 # x z n 

564 # 

565 # We use the transpose of that matrix. The first form is used 

566 # if rho is given, otherwise the coefficients are used. 

567 self.timer.start('Initial') 

568 if self.kpt_u[0].rho_MM is None: 

569 rhoT_uMM = [] 

570 ET_uMM = [] 

571 self.timer.start('Get density matrix') 

572 for kpt in self.kpt_u: 

573 rhoT_MM = self.ksl.get_transposed_density_matrix(kpt.f_n, 

574 kpt.C_nM) 

575 rhoT_uMM.append(rhoT_MM) 

576 ET_MM = self.ksl.get_transposed_density_matrix(kpt.f_n * 

577 kpt.eps_n, 

578 kpt.C_nM) 

579 ET_uMM.append(ET_MM) 

580 if hasattr(kpt, 'c_on'): 

581 # XXX does this work with BLACS/non-BLACS/etc.? 

582 assert self.bd.comm.size == 1 

583 d_nn = np.zeros((self.bd.mynbands, self.bd.mynbands), 

584 dtype=kpt.C_nM.dtype) 

585 for ne, c_n in zip(kpt.ne_o, kpt.c_on): 

586 d_nn += ne * np.outer(c_n.conj(), c_n) 

587 rhoT_MM += self.ksl.get_transposed_density_matrix_delta( 

588 d_nn, kpt.C_nM) 

589 ET_MM += self.ksl.get_transposed_density_matrix_delta( 

590 d_nn * kpt.eps_n, kpt.C_nM) 

591 self.timer.stop('Get density matrix') 

592 else: 

593 rhoT_uMM = [] 

594 ET_uMM = [] 

595 for kpt in self.kpt_u: 

596 H_MM = self.wfs.eigensolver.calculate_hamiltonian_matrix( 

597 self.hamiltonian, self.wfs, kpt) 

598 tri2full(H_MM) 

599 S_MM = kpt.S_MM.copy() 

600 tri2full(S_MM) 

601 ET_MM = np.linalg.solve(S_MM, gemmdot(H_MM, 

602 kpt.rho_MM)).T.copy() 

603 del S_MM, H_MM 

604 rhoT_MM = kpt.rho_MM.T.copy() 

605 rhoT_uMM.append(rhoT_MM) 

606 ET_uMM.append(ET_MM) 

607 self.timer.stop('Initial') 

608 return rhoT_uMM, ET_uMM 

609 

610 def get_kinetic_term(self): 

611 """Calculate Kinetic energy term in LCAO""" 

612 Fkin_av = np.zeros_like(self.Fref_av) 

613 self.timer.start('TCI derivative') 

614 # Kinetic energy contribution 

615 # 

616 # ----- d T 

617 # a \ mu nu 

618 # F += 2 Re ) -------- rho 

619 # / d R nu mu 

620 # ----- mu nu 

621 # mu in a; nu 

622 # 

623 Fkin_av = np.zeros_like(Fkin_av) 

624 for u, kpt in enumerate(self.kpt_u): 

625 dEdTrhoT_vMM = (self.dTdR_qvMM[kpt.q] * 

626 self.rhoT_uMM[u][np.newaxis]).real 

627 # XXX load distribution! 

628 for a, M1, M2 in self.my_slices(): 

629 Fkin_av[a, :] += \ 

630 2.0 * dEdTrhoT_vMM[:, M1:M2].sum(-1).sum(-1) 

631 self.timer.stop('TCI derivative') 

632 

633 return Fkin_av 

634 

635 def get_den_mat_term(self): 

636 """Calculate density matrix term in LCAO""" 

637 Ftheta_av = np.zeros_like(self.Fref_av) 

638 # Density matrix contribution due to basis overlap 

639 # 

640 # ----- d Theta 

641 # a \ mu nu 

642 # F += -2 Re ) ------------ E 

643 # / d R nu mu 

644 # ----- mu nu 

645 # mu in a; nu 

646 # 

647 Ftheta_av = np.zeros_like(Ftheta_av) 

648 for u, kpt in enumerate(self.kpt_u): 

649 dThetadRE_vMM = (self.dThetadR_qvMM[kpt.q] * 

650 self.ET_uMM[u][np.newaxis]).real 

651 for a, M1, M2 in self.my_slices(): 

652 Ftheta_av[a, :] += \ 

653 -2.0 * dThetadRE_vMM[:, M1:M2].sum(-1).sum(-1) 

654 

655 return Ftheta_av 

656 

657 def get_pot_term(self): 

658 """Calculate potential term""" 

659 Fpot_av = np.zeros_like(self.Fref_av) 

660 # Potential contribution 

661 # 

662 # ----- / d Phi (r) 

663 # a \ | mu ~ 

664 # F += -2 Re ) | ---------- v (r) Phi (r) dr rho 

665 # / | d R nu nu mu 

666 # ----- / a 

667 # mu in a; nu 

668 # 

669 self.timer.start('Potential') 

670 vt_sG = self.hamiltonian.vt_sG 

671 Fpot_av = np.zeros_like(Fpot_av) 

672 for u, kpt in enumerate(self.kpt_u): 

673 vt_G = vt_sG[kpt.s] 

674 Fpot_av += self.bfs.calculate_force_contribution(vt_G, 

675 self.rhoT_uMM[u], 

676 kpt.q) 

677 self.timer.stop('Potential') 

678 

679 return Fpot_av 

680 

681 def get_den_mat_paw_term(self): 

682 """Calcualte PAW correction""" 

683 # TO DO: split this function into 

684 # _get_den_mat_paw_term (which calculate Frho_av) and 

685 # get_paw_correction (which calculate ZE_MM) 

686 # Density matrix contribution from PAW correction 

687 # 

688 # ----- ----- 

689 # a \ a \ b 

690 # F += 2 Re ) Z E - 2 Re ) Z E 

691 # / mu nu nu mu / mu nu nu mu 

692 # ----- ----- 

693 # mu nu b; mu in a; nu 

694 # 

695 # with 

696 # b* 

697 # ----- dP 

698 # b \ i mu b b 

699 # Z = ) -------- dS P 

700 # mu nu / dR ij j nu 

701 # ----- b mu 

702 # ij 

703 # 

704 self.timer.start('Paw correction') 

705 Frho_av = np.zeros_like(self.Fref_av) 

706 for u, kpt in enumerate(self.kpt_u): 

707 work_MM = np.zeros((self.mynao, self.nao), self.dtype) 

708 ZE_MM = None 

709 for b in self.my_atom_indices: 

710 setup = self.setups[b] 

711 dO_ii = np.asarray(setup.dO_ii, self.dtype) 

712 dOP_iM = np.zeros((setup.ni, self.nao), self.dtype) 

713 mmm(1.0, dO_ii, 'N', self.P_aqMi[b][kpt.q], 'C', 0.0, dOP_iM) 

714 for v in range(3): 

715 mmm(1.0, 

716 self.dPdR_aqvMi[b][kpt.q][v][self.Mstart:self.Mstop], 

717 'N', 

718 dOP_iM, 'N', 

719 0.0, work_MM) 

720 ZE_MM = (work_MM * self.ET_uMM[u]).real 

721 for a, M1, M2 in self.slices(): 

722 dE = 2 * ZE_MM[M1:M2].sum() 

723 Frho_av[a, v] -= dE # the "b; mu in a; nu" term 

724 Frho_av[b, v] += dE # the "mu nu" term