Coverage for gpaw/lcao/tools.py: 40%
425 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1from time import localtime
2import pickle
4import numpy as np
5from ase.units import Ha
6from ase.calculators.singlepoint import SinglePointCalculator
7from ase.utils import IOContext
9from gpaw.utilities import pack_density
10from gpaw.utilities.tools import tri2full
11from gpaw.utilities.blas import rk, mmm, mmmx
12from gpaw.basis_data import Basis
13from gpaw.setup import types2atomtypes
14from gpaw.coulomb import CoulombNEW as Coulomb
15from gpaw.mpi import world, rank, serial_comm
16from gpaw import GPAW
19def get_bf_centers(atoms, basis=None):
20 calc = atoms.calc
21 if calc is None or isinstance(calc, SinglePointCalculator):
22 symbols = atoms.get_chemical_symbols()
23 basis_a = types2atomtypes(symbols, basis, 'dzp')
24 nao_a = [Basis.find(symbol, type).nao
25 for symbol, type in zip(symbols, basis_a)]
26 else:
27 if not calc.initialized:
28 calc.initialize(atoms)
29 nao_a = [calc.wfs.setups[a].nao for a in range(len(atoms))]
30 pos_ic = []
31 for pos, nao in zip(atoms.get_positions(), nao_a):
32 pos_ic.extend(pos[None].repeat(nao, 0))
33 return np.array(pos_ic)
36def get_bfi(calc, a_list):
37 """basis function indices from a list of atom indices.
38 a_list: atom indices
39 Use: get_bfi(calc, [0, 4]) gives the functions indices
40 corresponding to atom 0 and 4"""
41 bfs_list = []
42 for a in a_list:
43 M = calc.wfs.basis_functions.M_a[a]
44 bfs_list += range(M, M + calc.wfs.setups[a].nao)
45 return bfs_list
48def get_mulliken(calc, a_list):
49 """Mulliken charges from a list of atom indices (a_list)."""
50 Q_a = {}
51 for a in a_list:
52 Q_a[a] = 0.0
53 for kpt in calc.wfs.kpt_u:
54 S_MM = calc.wfs.S_qMM[kpt.q]
55 nao = S_MM.shape[0]
56 rho_MM = np.empty((nao, nao), calc.wfs.dtype)
57 calc.wfs.calculate_density_matrix(kpt.f_n, kpt.C_nM, rho_MM)
58 Q_M = np.dot(rho_MM, S_MM).diagonal()
59 for a in a_list:
60 M1 = calc.wfs.basis_functions.M_a[a]
61 M2 = M1 + calc.wfs.setups[a].nao
62 Q_a[a] += np.sum(Q_M[M1:M2])
63 return Q_a
66def get_realspace_hs(h_skmm, s_kmm, bzk_kc, weight_k,
67 R_c=(0, 0, 0), direction='x',
68 symmetry={'enabled': False}):
70 from gpaw.symmetry import Symmetry
71 from ase.dft.kpoints import get_monkhorst_pack_size_and_offset, \
72 monkhorst_pack
74 if symmetry['point_group']:
75 raise NotImplementedError('Point group symmetry not implemented')
77 nspins, nk, nbf = h_skmm.shape[:3]
78 dir = 'xyz'.index(direction)
79 transverse_dirs = np.delete([0, 1, 2], [dir])
80 dtype = float
81 if len(bzk_kc) > 1 or np.any(bzk_kc[0] != [0, 0, 0]):
82 dtype = complex
84 kpts_grid, kpts_shift = get_monkhorst_pack_size_and_offset(bzk_kc)
86 # kpts in the transport direction
87 nkpts_p = kpts_grid[dir]
88 bzk_p_kc = monkhorst_pack((nkpts_p, 1, 1))[:, 0] + kpts_shift[dir]
89 weight_p_k = 1. / nkpts_p
91 # kpts in the transverse directions
92 offset = np.zeros((3,))
93 offset[:len(transverse_dirs)] = kpts_shift[transverse_dirs]
94 bzk_t_kc = monkhorst_pack(
95 tuple(kpts_grid[transverse_dirs]) + (1, )) + offset
96 if 'time_reversal' not in symmetry:
97 symmetry['time_reversal'] = True
98 if symmetry['time_reversal']:
99 # XXX a somewhat ugly hack:
100 # By default GPAW reduces inversion sym in the z direction. The steps
101 # below assure reduction in the transverse dirs.
102 # For now this part seems to do the job, but it may be written
103 # in a smarter way in the future.
104 symmetry = Symmetry([1], np.eye(3))
105 symmetry.prune_symmetries_atoms(np.zeros((1, 3)))
106 ibzk_kc, ibzweight_k = symmetry.reduce(bzk_kc)[:2]
107 ibzk_t_kc, weights_t_k = symmetry.reduce(bzk_t_kc)[:2]
108 ibzk_t_kc = ibzk_t_kc[:, :2]
109 nkpts_t = len(ibzk_t_kc)
110 else:
111 ibzk_kc = bzk_kc.copy()
112 ibzk_t_kc = bzk_t_kc
113 ibzk_t_kc = ibzk_t_kc[:, :2]
114 nkpts_t = len(bzk_t_kc)
115 weights_t_k = [1. / nkpts_t for k in range(nkpts_t)]
117 h_skii = np.zeros((nspins, nkpts_t, nbf, nbf), dtype)
118 if s_kmm is not None:
119 s_kii = np.zeros((nkpts_t, nbf, nbf), dtype)
121 tol = 7
122 for j, k_t in enumerate(ibzk_t_kc):
123 for k_p in bzk_p_kc:
124 k = np.zeros((3,))
125 k[dir] = k_p
126 k[transverse_dirs] = k_t
127 kpoint_list = [list(np.round(k_kc, tol)) for k_kc in ibzk_kc]
128 if list(np.round(k, tol)) not in kpoint_list:
129 k = -k # inversion
130 index = kpoint_list.index(list(np.round(k, tol)))
131 h = h_skmm[:, index].conjugate()
132 if s_kmm is not None:
133 s = s_kmm[index].conjugate()
134 k = -k
135 else: # kpoint in the ibz
136 index = kpoint_list.index(list(np.round(k, tol)))
137 h = h_skmm[:, index]
138 if s_kmm is not None:
139 s = s_kmm[index]
141 c_k = np.exp(2.j * np.pi * np.dot(k, R_c)) * weight_p_k
142 h_skii[:, j] += c_k * h
143 if s_kmm is not None:
144 s_kii[j] += c_k * s
146 if s_kmm is None:
147 return ibzk_t_kc, weights_t_k, h_skii
148 else:
149 return ibzk_t_kc, weights_t_k, h_skii, s_kii
152def remove_pbc(atoms, h, s=None, d=0, centers_ic=None, cutoff=None):
153 if h.ndim > 2:
154 raise KeyError('You have to run remove_pbc for each '
155 'spin/kpoint separately.')
156 L = atoms.cell[d, d]
157 nao = len(h)
158 dtype = h.dtype
159 if centers_ic is None:
160 centers_ic = get_bf_centers(atoms) # requires an attached LCAO calc
161 ni = len(centers_ic)
162 if nao != ni:
163 assert nao == 2 * ni
164 centers_ic = np.vstack((centers_ic, centers_ic))
165 centers_ic[ni:, d] += L
166 if cutoff is None:
167 cutoff = L - 1e-3
168 elif cutoff is None:
169 cutoff = 0.5 * L - 1e-3
170 pos_i = centers_ic[:, d]
171 for i in range(nao):
172 dpos_i = abs(pos_i - pos_i[i])
173 mask_i = (dpos_i < cutoff).astype(dtype)
174 h[i, :] *= mask_i
175 h[:, i] *= mask_i
176 if s is not None:
177 s[i, :] *= mask_i
178 s[:, i] *= mask_i
181def dump_hamiltonian(filename, atoms, direction=None, Ef=None):
182 h_skmm, s_kmm = get_lcao_hamiltonian(atoms.calc)
183 if direction is not None:
184 d = 'xyz'.index(direction)
185 for s in range(atoms.calc.nspins):
186 for k in range(atoms.calc.nkpts):
187 if s == 0:
188 remove_pbc(atoms, h_skmm[s, k], s_kmm[k], d)
189 else:
190 remove_pbc(atoms, h_skmm[s, k], None, d)
192 if atoms.calc.master:
193 with open(filename, 'wb') as fd:
194 pickle.dump((h_skmm, s_kmm), fd, 2)
195 atoms_data = {'cell': atoms.cell, 'positions': atoms.positions,
196 'numbers': atoms.numbers, 'pbc': atoms.pbc}
198 pickle.dump(atoms_data, fd, 2)
199 calc_data = {'weight_k': atoms.calc.weight_k,
200 'ibzk_kc': atoms.calc.ibzk_kc}
202 pickle.dump(calc_data, fd, 2)
204 world.barrier()
207def dump_hamiltonian_parallel(filename, atoms, direction=None, Ef=None):
208 raise DeprecationWarning(
209 'Please use dump_hamiltonian_and_overlap() instead.')
212def dump_hamiltonian_and_overlap(filename, atoms, direction=None):
213 """
214 Dump the lcao representation of H and S to file(s) beginning
215 with filename. If direction is x, y or z, the periodic boundary
216 conditions will be removed in the specified direction.
218 Note:
219 H and S are parallized over spin and k-points and
220 is for now dumped into a number of pickle files. This
221 may be changed into a dump to a single file in the future.
223 """
224 if direction is not None:
225 d = 'xyz'.index(direction)
227 calc = atoms.calc
228 wfs = calc.wfs
229 nao = wfs.setups.nao
230 nq = len(wfs.kpt_u) // wfs.nspins
231 H_qMM = np.empty((wfs.nspins, nq, nao, nao), wfs.dtype)
232 calc_data = {'k_q': {},
233 'skpt_qc': np.empty((nq, 3)),
234 'weight_q': np.empty(nq)}
236 S_qMM = wfs.S_qMM
238 for kpt in wfs.kpt_u:
239 calc_data['skpt_qc'][kpt.q] = calc.wfs.kd.ibzk_kc[kpt.k]
240 calc_data['weight_q'][kpt.q] = calc.wfs.kd.weight_k[kpt.k]
241 calc_data['k_q'][kpt.q] = kpt.k
242 H_MM = wfs.eigensolver.calculate_hamiltonian_matrix(calc.hamiltonian,
243 wfs,
244 kpt)
246 H_qMM[kpt.s, kpt.q] = H_MM
248 tri2full(H_qMM[kpt.s, kpt.q])
249 if kpt.s == 0:
250 tri2full(S_qMM[kpt.q])
251 if direction is not None:
252 remove_pbc(atoms, H_qMM[kpt.s, kpt.q], S_qMM[kpt.q], d)
253 else:
254 if direction is not None:
255 remove_pbc(atoms, H_qMM[kpt.s, kpt.q], None, d)
257 if wfs.gd.comm.rank == 0:
258 with open(filename + '%i.pckl' % wfs.kd.comm.rank, 'wb') as fd:
259 H_qMM *= Ha
260 pickle.dump((H_qMM, S_qMM), fd, 2)
261 pickle.dump(calc_data, fd, 2)
264def get_lcao_hamiltonian(calc):
265 """Return H_skMM, S_kMM on master, (None, None) on slaves. H is in eV."""
266 from gpaw.new.ase_interface import ASECalculator
267 if not isinstance(calc, ASECalculator):
268 return old_get_lcao_hamiltonian(calc)
269 dft = calc.dft
270 ibzwfs = dft.ibzwfs
271 ham = dft.scf_loop.hamiltonian
272 matcalc = ham.create_hamiltonian_matrix_calculator(dft.potential)
273 nM = ham.basis.Mmax
274 nK = len(ibzwfs.ibz)
275 H_skMM = np.zeros((ibzwfs.nspins, nK, nM, nM), ibzwfs.dtype)
276 S_kMM = np.zeros((nK, nM, nM), ibzwfs.dtype)
277 for wfs in ibzwfs:
278 H_MM = matcalc.calculate_matrix(wfs)
279 H_skMM[wfs.spin, wfs.k] = H_MM.data * Ha
280 S_MM = wfs.S_MM
281 if wfs.spin == 0:
282 S_kMM[wfs.k] = S_MM.data
283 ibzwfs.kpt_comm.sum(H_skMM)
284 ibzwfs.kpt_comm.sum(S_kMM)
285 if rank == 0:
286 return H_skMM, S_kMM
287 return None, None
290def old_get_lcao_hamiltonian(calc):
291 if calc.wfs.S_qMM is None:
292 calc.wfs.set_positions(calc.spos_ac)
293 dtype = calc.wfs.dtype
294 NM = calc.wfs.eigensolver.nao
295 Nk = calc.wfs.kd.nibzkpts
296 Ns = calc.wfs.nspins
298 S_kMM = np.zeros((Nk, NM, NM), dtype)
299 H_skMM = np.zeros((Ns, Nk, NM, NM), dtype)
300 for kpt in calc.wfs.kpt_u:
301 H_MM = calc.wfs.eigensolver.calculate_hamiltonian_matrix(
302 calc.hamiltonian, calc.wfs, kpt)
303 if kpt.s == 0:
304 S_kMM[kpt.k] = calc.wfs.S_qMM[kpt.q]
305 tri2full(S_kMM[kpt.k])
306 H_skMM[kpt.s, kpt.k] = H_MM * Ha
307 tri2full(H_skMM[kpt.s, kpt.k])
308 calc.wfs.kd.comm.sum(S_kMM, 0)
309 calc.wfs.kd.comm.sum(H_skMM, 0)
310 if rank == 0:
311 return H_skMM, S_kMM
312 else:
313 return None, None
316def get_lead_lcao_hamiltonian(calc, direction='x'):
317 H_skMM, S_kMM = get_lcao_hamiltonian(calc)
318 if rank == 0:
319 return lead_kspace2realspace(H_skMM, S_kMM,
320 bzk_kc=calc.wfs.kd.bzk_kc,
321 weight_k=calc.wfs.kd.weight_k,
322 direction=direction,
323 symmetry=calc.parameters.symmetry)
324 else:
325 return None, None, None, None
328def lead_kspace2realspace(h_skmm, s_kmm, bzk_kc, weight_k, direction='x',
329 symmetry={'point_group': False}):
330 """Convert a k-dependent Hamiltonian to tight-binding onsite and coupling.
332 For each transverse k-point:
333 Convert k-dependent (in transport direction) Hamiltonian
334 representing a lead to a real space tight-binding Hamiltonian
335 of double size representing two principal layers and the coupling between.
336 """
338 dir = 'xyz'.index(direction)
339 if symmetry['point_group']:
340 raise NotImplementedError
342 R_c = [0, 0, 0]
343 ibz_t_kc, weight_t_k, h_skii, s_kii = get_realspace_hs(
344 h_skmm, s_kmm, bzk_kc, weight_k, R_c, direction, symmetry)
346 R_c[dir] = 1
347 h_skij, s_kij = get_realspace_hs(
348 h_skmm, s_kmm, bzk_kc, weight_k, R_c, direction, symmetry)[-2:]
350 nspins, nk, nbf = h_skii.shape[:-1]
352 h_skmm = np.zeros((nspins, nk, 2 * nbf, 2 * nbf), h_skii.dtype)
353 s_kmm = np.zeros((nk, 2 * nbf, 2 * nbf), h_skii.dtype)
354 h_skmm[:, :, :nbf, :nbf] = h_skmm[:, :, nbf:, nbf:] = h_skii
355 h_skmm[:, :, :nbf, nbf:] = h_skij
356 h_skmm[:, :, nbf:, :nbf] = h_skij.swapaxes(2, 3).conj()
358 s_kmm[:, :nbf, :nbf] = s_kmm[:, nbf:, nbf:] = s_kii
359 s_kmm[:, :nbf, nbf:] = s_kij
360 s_kmm[:, nbf:, :nbf] = s_kij.swapaxes(1, 2).conj()
362 return ibz_t_kc, weight_t_k, h_skmm, s_kmm
365def zeta_pol(basis):
366 """Get number of zeta func. and polarization func. indices in Basis."""
367 zeta = []
368 pol = []
369 for bf in basis.bf_j:
370 if 'polarization' in bf.type:
371 pol.append(2 * bf.l + 1)
372 else:
373 zeta.append(2 * bf.l + 1)
374 zeta = sum(zeta)
375 pol = sum(pol)
376 assert zeta + pol == basis.nao
377 return zeta, pol
380def basis_subset(symbol, largebasis, smallbasis):
381 """Title.
383 Determine which basis function indices from ``largebasis`` are also
384 present in smallbasis.
385 """
386 blarge = Basis.find(symbol, largebasis)
387 zeta_large, pol_large = zeta_pol(blarge)
389 bsmall = Basis.find(symbol, smallbasis)
390 zeta_small, pol_small = zeta_pol(bsmall)
392 assert zeta_small <= zeta_large
393 assert pol_small <= pol_large
395 insmall = np.zeros(blarge.nao, bool)
396 insmall[:zeta_small] = True
397 insmall[zeta_large:zeta_large + pol_small] = True
398 return insmall
401def basis_subset2(symbols, largebasis='dzp', smallbasis='sz'):
402 """Same as basis_subset, but for an entire list of atoms."""
403 largebasis = types2atomtypes(symbols, largebasis, default='dzp')
404 smallbasis = types2atomtypes(symbols, smallbasis, default='sz')
405 mask = []
406 for symbol, large, small in zip(symbols, largebasis, smallbasis):
407 mask.extend(basis_subset(symbol, large, small))
408 return np.asarray(mask, bool)
411def collect_orbitals(a_xo, coords, root=0):
412 """Collect array distributed over orbitals to root-CPU.
414 Input matrix has last axis distributed amongst CPUs,
415 return is None on slaves, and the collected array on root.
417 The distribution can be uneven amongst CPUs. The list coords gives the
418 number of values for each CPU.
419 """
420 a_xo = np.ascontiguousarray(a_xo)
421 if world.size == 1:
422 return a_xo
424 # All slaves send their piece to ``root``:
425 # There can be several sends before the corresponding receives
426 # are posted, so use syncronous send here
427 if world.rank != root:
428 world.ssend(a_xo, root, 112)
429 return None
431 # On root, put the subdomains from the slaves into the big array
432 # for the whole domain on root:
433 xshape = a_xo.shape[:-1]
434 Norb2 = sum(coords) # total number of orbital indices
435 a_xO = np.empty(xshape + (Norb2,), a_xo.dtype)
436 o = 0
437 for rankx, norb in enumerate(coords):
438 if rankx != root:
439 tmp_xo = np.empty(xshape + (norb,), a_xo.dtype)
440 world.receive(tmp_xo, rankx, 112)
441 a_xO[..., o:o + norb] = tmp_xo
442 else:
443 a_xO[..., o:o + norb] = a_xo
444 o += norb
445 return a_xO
448def makeU(gpwfile='grid.gpw', orbitalfile='w_wG__P_awi.pckl',
449 rotationfile='eps_q__U_pq.pckl', tolerance=1e-5,
450 writeoptimizedpairs=False, dppname='D_pp.pckl', S_w=None):
452 # S_w: None or diagonal of overlap matrix. In the latter case
453 # the optimized and truncated pair orbitals are obtained from
454 # normalized (to 1) orbitals.
455 #
456 # Tolerance is used for truncation of optimized pairorbitals
457 # calc = GPAW(gpwfile, txt=None)
458 from gpaw import GPAW
459 from gpaw.mpi import world
460 calc = GPAW(gpwfile, txt='pairorb.txt') # XXX
461 gd = calc.wfs.gd
462 setups = calc.wfs.setups
463 myatoms = calc.density.D_asp.keys()
464 del calc
466 # Load orbitals on master and distribute to slaves
467 if world.rank == 0:
468 with open(orbitalfile, 'rb') as fd:
469 wglobal_wG, P_awi = pickle.load(fd)
470 Nw = len(wglobal_wG)
471 print('Estimated total (serial) mem usage: %0.3f GB' % (
472 np.prod(gd.N_c) * Nw**2 * 8 / 1024.**3))
473 else:
474 wglobal_wG = None
475 Nw = 0
476 Nw = gd.comm.sum_scalar(Nw) # distribute Nw to all nodes
477 w_wG = gd.empty(n=Nw)
478 gd.distribute(wglobal_wG, w_wG)
479 del wglobal_wG
481 # Make pairorbitals
482 f_pG = gd.zeros(n=Nw**2)
483 for p, (w1, w2) in enumerate(np.ndindex(Nw, Nw)):
484 np.multiply(w_wG[w1], w_wG[w2], f_pG[p])
485 del w_wG
486 assert f_pG.flags.contiguous
487 # Make pairorbital overlap (lower triangle only)
488 D_pp = np.zeros((Nw**2, Nw**2))
489 rk(gd.dv, f_pG, 0., D_pp)
491 # Add atomic corrections to pairorbital overlap
492 for a in myatoms:
493 if setups[a].type != 'ghost':
494 P_pp = np.array([pack_density(np.outer(P_awi[a][w1], P_awi[a][w2]))
495 for w1, w2 in np.ndindex(Nw, Nw)])
496 I4_pp = setups[a].four_phi_integrals()
497 A = np.zeros((len(I4_pp), len(P_pp)))
498 mmm(1.0, I4_pp, 'N', P_pp, 'T', 0.0, A)
499 mmm(1.0, P_pp, 'N', A, 'N', 1.0, D_pp)
500 # D_pp += np.dot(P_pp, np.dot(I4_pp, P_pp.T))
502 # Summ all contributions to master
503 gd.comm.sum(D_pp, 0)
505 if world.rank == 0:
506 if S_w is not None:
507 print('renormalizing pairorb overlap matrix (D_pp)')
508 S2 = np.sqrt(S_w)
509 for pa, (wa1, wa2) in enumerate(np.ndindex(Nw, Nw)):
510 for pb, (wb1, wb2) in enumerate(np.ndindex(Nw, Nw)):
511 D_pp[pa, pb] /= S2[wa1] * S2[wa2] * S2[wb1] * S2[wb2]
513 D_pp.dump(dppname)
514 # XXX if the diagonalization below (on MASTER only)
515 # fails, then one can always restart the stuff
516 # below using only the stored D_pp matrix
518 # Determine eigenvalues and vectors on master only
519 eps_q, U_pq = np.linalg.eigh(D_pp, UPLO='L')
520 del D_pp
521 indices = np.argsort(-eps_q.real)
522 eps_q = np.ascontiguousarray(eps_q.real[indices])
523 U_pq = np.ascontiguousarray(U_pq[:, indices])
525 # Truncate
526 indices = eps_q > tolerance
527 U_pq = np.ascontiguousarray(U_pq[:, indices])
528 eps_q = np.ascontiguousarray(eps_q[indices])
530 # Dump to file
531 with open(rotationfile, 'wb') as fd:
532 pickle.dump((eps_q, U_pq), fd, 2)
534 if writeoptimizedpairs is not False:
535 assert world.size == 1 # works in parallel if U and eps are broadcast
536 Uisq_qp = (U_pq / np.sqrt(eps_q)).T.copy()
537 g_qG = gd.zeros(n=len(eps_q))
538 mmmx(1.0, Uisq_qp, 'N', f_pG, 'N', 0.0, g_qG)
539 g_qG = gd.collect(g_qG)
540 if world.rank == 0:
541 P_app = {a: np.array([pack_density(np.outer(P_wi[w1], P_wi[w2]),
542 tolerance=1e3)
543 for w1, w2 in np.ndindex(Nw, Nw)])
544 for a, P_wi in P_awi.items()}
545 P_aqp = {a: np.dot(Uisq_qp, Px_pp) for a, Px_pp in P_app.items()}
546 with open(writeoptimizedpairs, 'wb') as fd:
547 pickle.dump((g_qG, P_aqp), fd, 2)
550def makeV(gpwfile='grid.gpw', orbitalfile='w_wG__P_awi.pckl',
551 rotationfile='eps_q__U_pq.pckl', coulombfile='V_qq.pckl',
552 log='V_qq.log', fft=False):
554 with IOContext() as io:
555 log = io.openfile(log, comm=serial_comm)
556 _makeV(gpwfile=gpwfile, orbitalfile=orbitalfile,
557 rotationfile=rotationfile, coulombfile=coulombfile,
558 log=log, fft=fft)
561def _makeV(gpwfile, orbitalfile, rotationfile, coulombfile, log, fft):
562 # Extract data from files
563 calc = GPAW(gpwfile, txt=None, communicator=serial_comm)
564 coulomb = Coulomb(calc.wfs.gd, calc.wfs.setups, calc.spos_ac, fft)
565 with open(orbitalfile, 'rb') as fd:
566 w_wG, P_awi = pickle.load(fd)
567 with open(rotationfile, 'rb') as fd:
568 eps_q, U_pq = pickle.load(fd)
569 del calc
571 # Make rotation matrix divided by sqrt of norm
572 Nq = len(eps_q)
573 Ni = len(w_wG)
574 Uisq_iqj = (U_pq / np.sqrt(eps_q)).reshape(
575 Ni, Ni, Nq).swapaxes(1, 2).copy()
576 del eps_q, U_pq
578 # Determine number of opt. pairorb on each cpu
579 Ncpu = world.size
580 nq, R = divmod(Nq, Ncpu)
581 nq_r = nq * np.ones(Ncpu, int)
582 if R > 0:
583 nq_r[-R:] += 1
585 # Determine number of opt. pairorb on this cpu
586 nq1 = nq_r[world.rank]
587 q1end = nq_r[:world.rank + 1].sum()
588 q1start = q1end - nq1
589 V_qq = np.zeros((Nq, nq1), float)
591 def make_optimized(qstart, qend):
592 g_qG = np.zeros((qend - qstart,) + w_wG.shape[1:], float)
593 P_aqp = {}
594 for a, P_wi in P_awi.items():
595 ni = P_wi.shape[1]
596 nii = ni * (ni + 1) // 2
597 P_aqp[a] = np.zeros((qend - qstart, nii), float)
598 for w1, w1_G in enumerate(w_wG):
599 U = Uisq_iqj[w1, qstart: qend].copy()
600 mmmx(1, U, 'N', w1_G * w_wG, 'N', 1, g_qG)
601 for a, P_wi in P_awi.items():
602 P_wp = np.array([pack_density(np.outer(P_wi[w1], P_wi[w2]))
603 for w2 in range(Ni)])
604 mmm(1., U, 'N', P_wp, 'N', 1.0, P_aqp[a])
605 return g_qG, P_aqp
607 g1_qG, P1_aqp = make_optimized(q1start, q1end)
608 for block, nq2 in enumerate(nq_r):
609 if block == world.rank:
610 g2_qG, P2_aqp = g1_qG, P1_aqp
611 q2start, q2end = q1start, q1end
612 else:
613 q2end = nq_r[:block + 1].sum()
614 q2start = q2end - nq2
615 g2_qG, P2_aqp = make_optimized(q2start, q2end)
617 for q1, q2 in np.ndindex(nq1, nq2):
618 P1_ap = {a: P_qp[q1] for a, P_qp in P1_aqp.items()}
619 P2_ap = {a: P_qp[q2] for a, P_qp in P2_aqp.items()}
620 V_qq[q2 + q2start, q1] = coulomb.calculate(g1_qG[q1], g2_qG[q2],
621 P1_ap, P2_ap)
622 if q2 == 0 and world.rank == 0:
623 T = localtime()
624 log.write(
625 'Block %i/%i is %4.1f percent done at %02i:%02i:%02i\n' %
626 (block + 1, world.size,
627 100.0 * q1 / nq1, T[3], T[4], T[5]))
628 log.flush()
630 # Collect V_qq array on master node
631 if world.rank == 0:
632 T = localtime()
633 log.write('Starting collect at %02i:%02i:%02i\n' % (
634 T[3], T[4], T[5]))
635 log.flush()
637 V_qq = collect_orbitals(V_qq, coords=nq_r, root=0)
638 if world.rank == 0:
639 # V can be slightly asymmetric due to numerics
640 V_qq = 0.5 * (V_qq + V_qq.T)
641 V_qq.dump(coulombfile)
643 T = localtime()
644 log.write('Finished at %02i:%02i:%02i\n' % (
645 T[3], T[4], T[5]))
646 log.flush()