Coverage for gpaw/utilities/dos.py: 70%
382 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
1from math import pi, sqrt
2import numpy as np
3from ase.parallel import paropen
4from gpaw.utilities import pack_density
5from gpaw.analyse.wignerseitz import wignerseitz
6from gpaw.setup_data import SetupData
7from gpaw.gauss import Gauss
8from gpaw.io.fmf import FMF
9from gpaw.utilities.blas import gemmdot
12def print_projectors(setup):
13 """Print information on the projectors of input nucleus object.
15 If nucleus is a string, treat this as an element name.
16 """
17 if isinstance(setup, str):
18 setup = SetupData(setup, 'LDA', 'paw')
19 n_j = setup.n_j
20 l_j = setup.l_j
21 else:
22 n_j = setup.n_j
23 l_j = setup.l_j
25 angular = [['1'],
26 ['y', 'z', 'x'],
27 ['xy', 'yz', '3z^2-r^2', 'xz', 'x^2-y^2'],
28 ['3x^2y-y^3', 'xyz', '5yz^2-yr^2', '5z^3-3zr^2',
29 '5xz^2-xr^2', 'x^2z-y^2z', 'x^3-3xy^2'],
30 ]
31 print(' i n l m')
32 print('--------')
33 i = 0
34 for n, l in zip(n_j, l_j):
35 for m in range(2 * l + 1):
36 if n == -1:
37 n = '*'
38 print('%2s %s %s_%s' % (i, n, 'spdf'[l], angular[l][m]))
39 i += 1
42def number_of_projectors(setup):
43 """Returns the number of the bound state projectors.
45 If setup is a string, treat this as an element name.
46 """
47 if isinstance(setup, str):
48 setup = SetupData(setup, 'LDA', 'paw')
49 n_j = setup.n_j
50 l_j = setup.l_j
51 else:
52 n_j = setup.n_j
53 l_j = setup.l_j
55 i = 0
56 for n, l in zip(n_j, l_j):
57 for m in range(2 * l + 1):
58 if n != -1:
59 i += 1
60 return i
63def get_angular_projectors(setup, angular, type='bound'):
64 """Determine the projector indices which have specified angula
65 quantum number.
67 angular can be s, p, d, f, or a list of these.
68 If type is 'bound', only bound state projectors are considered, otherwise
69 all projectors are included.
70 """
71 # Get the number of relevant j values
72 if type == 'bound':
73 nj = len([n for n in setup.n_j if n >= 0])
74 else:
75 nj = len(setup.n_j)
77 # Choose the relevant projectors
78 projectors = []
79 i = 0
80 j = 0
81 for j in range(nj):
82 m = 2 * setup.l_j[j] + 1
83 if 'spdf'[setup.l_j[j]] in angular:
84 projectors.extend(range(i, i + m))
85 j += 1
86 i += m
88 return projectors
91def delta(x, x0, width, mode='Gauss'):
92 """Return a gaussian of given width centered at x0."""
93 if mode == 'Gauss':
94 return np.exp(np.clip(-((x - x0) / width)**2,
95 -100.0, 100.0)) / (sqrt(pi) * width)
96 if mode == 'Lorentz':
97 return (2 / pi / width) / ((np.clip(((x - x0) / (width / 2))**2,
98 -100.0, 100.0)) + 1)
101def fold(energies, weights, npts, width, mode='Gauss'):
102 """Take a list of energies and weights, and sum a delta function
103 for each."""
104 emin = min(energies) - 5 * width
105 emax = max(energies) + 5 * width
106 e = np.linspace(emin, emax, npts)
107 dos_e = np.zeros(npts)
108 for e0, w in zip(energies, weights):
109 dos_e += w * delta(e, e0, width, mode=mode)
110 return e, dos_e
113def raw_orbital_LDOS(paw, a, spin, angular='spdf', nbands=None):
114 """Return a list of eigenvalues, and their weight on the specified atom.
116 angular can be s, p, d, f, or a list of these.
117 If angular is None, the raw weight for each projector is returned.
119 An integer value for ``angular`` can also be used to specify a specific
120 projector function.
122 Setting nbands limits the number of bands included. This speeds up the
123 calculation if one has many bands in the calculator but is only interested
124 in the DOS at low energies.
125 """
126 wfs = paw.wfs
127 w_k = wfs.kd.weight_k
128 nk = len(w_k)
129 if not nbands:
130 nb = wfs.bd.nbands
131 else:
132 nb = nbands
133 assert nb <= wfs.bd.nbands, ('nbands higher than available number' +
134 'of bands')
136 if a < 0:
137 # Allow list-style negative indices; we'll need the positive a for the
138 # dictionary lookup later
139 a = len(wfs.setups) + a
141 I1 = sum(setup.ni for setup in wfs.setups[:a])
142 setup = wfs.setups[a]
143 I2 = I1 + setup.ni
145 energies = np.empty(nb * nk)
146 weights_xi = np.empty((nb * nk, setup.ni))
147 x = 0
148 for k, w in enumerate(w_k):
149 eps_n = wfs.collect_eigenvalues(k=k, s=spin)
150 if len(eps_n) > 0:
151 energies[x:x + nb] = eps_n[:nb]
152 P_nI = wfs.collect_projections(k, spin)
153 if P_nI is not None:
154 weights_xi[x:x + nb, :] = w * abs(P_nI[:nb, I1:I2])**2
155 x += nb
157 wfs.world.broadcast(energies, 0)
158 wfs.world.broadcast(weights_xi, 0)
160 if angular is None:
161 return energies, weights_xi
162 elif isinstance(angular, int):
163 return energies, weights_xi[:, angular]
164 else:
165 projectors = get_angular_projectors(setup, angular, type='bound')
166 weights = np.sum(np.take(weights_xi,
167 indices=projectors, axis=1), axis=1)
168 return energies, weights
171def all_electron_LDOS(paw, mol, spin, lc=None, wf_k=None, P_aui=None):
172 """Returns a list of eigenvalues, and their weights on a given molecule
174 wf_k should be a list of pseudo_wavefunctons of a Kohn-Sham state,
175 corresponding to the different kpoints. It should be accompanied by a
176 list of projector overlaps: P_aui=[[kpt.P_ani[a][n] for kpt in
177 paw.wfs.kpt_u] for a in range(len(molecule))] for the band n. The
178 weights are then calculated as the overlap of all-electron
179 KS wavefunctions with wf_k
181 If wf_k is None, the weights are calculated as linear combinations of
182 atomic orbitals using P_uni. lc should then be a list of weights
183 for each atom. For example, the pure 2pi_x orbital of a CO or N2
184 molecule can be obtained with lc=[[0,0,0,1.0],[0,0,0,-1.0]]. mol
185 should be a list of atom numbers contributing to the molecule."""
187 w_k = paw.wfs.kd.weight_k
188 nk = len(w_k)
189 nb = paw.wfs.bd.nbands
191 P_kn = np.zeros((nk, nb), complex)
192 if wf_k is None:
193 if lc is None:
194 lc = [[1, 0, 0, 0] for a in mol]
195 for k, kpt in enumerate(kpt_s[spin] for kpt_s in paw.wfs.kpt_qs):
196 N = 0
197 for atom, w_a in zip(mol, lc):
198 i = 0
199 for w_o in w_a:
200 P_kn[k] += w_o * kpt.P_ani[atom][:, i]
201 N += abs(w_o)**2
202 i += 1
203 P_kn /= sqrt(N)
205 else:
206 P_aui = [np.array(P_ui).conj() for P_ui in P_aui]
207 for k, kpt in enumerate(kpt_s[spin] for kpt_s in paw.wfs.kpt_qs):
208 for n in range(nb):
209 P_kn[k][n] = paw.wfs.integrate(wf_k[k], kpt.psit_nG[n])
210 for a, b in zip(mol, range(len(mol))):
211 atom = paw.wfs.setups[a]
212 p_i = kpt.P_ani[a][n]
213 for i in range(len(p_i)):
214 for j in range(len(p_i)):
215 P_kn[k][n] += (P_aui[b][spin * nk + k][i] *
216 atom.dO_ii[i][j] * p_i[j])
217 print('# k', k, ' Sum_m |<m|n>|^2 =', sum(abs(P_kn[k])**2))
219 energies = np.empty(nb * nk)
220 weights = np.empty(nb * nk)
221 x = 0
222 for k, w in enumerate(w_k):
223 eps_n = paw.wfs.collect_eigenvalues(k=k, s=spin)
224 if len(eps_n) > 0:
225 energies[x:x + nb] = eps_n
226 weights[x:x + nb] = w * abs(P_kn[k])**2
227 x += nb
229 return energies, weights
232def get_all_electron_IPR(paw):
233 density = paw.density
234 wfs = paw.wfs
235 n_G = wfs.gd.empty()
236 n_g = density.finegd.empty()
237 print()
238 print('inverse participation function')
239 print('-' * 35)
240 print('%5s %5s %10s %10s' % ('k', 'band', 'eps', 'ipr'))
241 print('-' * 35)
242 for k, kpt in enumerate(paw.wfs.kpt_u):
243 for n, (eps, psit_G) in enumerate(zip(kpt.eps_n, kpt.psit_nG)):
244 n_G[:] = 0.0
245 wfs.add_orbital_density(n_G, kpt, n)
246 density.interpolator.apply(n_G, n_g)
247 norm = density.finegd.integrate(n_g)
248 n_g = n_g ** 2
249 ipr = density.finegd.integrate(n_g)
250 for a in kpt.P_ani:
251 # Get xccorr for atom a
252 setup = paw.density.setups[a]
253 xccorr = setup.xc_correction
255 # Get D_sp for atom a
256 D_sp = np.array(wfs.get_orbital_density_matrix(a, kpt, n))
258 # density a function of L and partial wave radial pair
259 # density coefficient
260 D_sLq = gemmdot(D_sp, xccorr.B_Lqp, trans='t')
262 # Create pseudo/ae density iterators for integration
263 n_iter = xccorr.expand_density(D_sLq, xccorr.n_qg, None)
264 nt_iter = xccorr.expand_density(D_sLq, xccorr.nt_qg, None)
266 # Take the spherical average of smooth and ae radial
267 # xc potentials
268 for n_sg, nt_sg, integrator in zip(n_iter,
269 nt_iter,
270 xccorr.get_integrator(
271 None)):
272 ipr += integrator.weight * np.sum((n_sg[0]**2 -
273 nt_sg[0]**2) *
274 xccorr.rgd.dv_g)
275 norm += integrator.weight * np.sum((n_sg[0] - nt_sg[0]) *
276 xccorr.rgd.dv_g)
278 print('%5i %5i %10.5f %10.5f' % (k, n, eps, ipr / norm**2))
279 print('-' * 35)
282def raw_wignerseitz_LDOS(paw, a, spin):
283 """Return a list of eigenvalues, and their weight on the specified atom.
284 Only supports real (gamma-point only) wavefunctions."""
285 wfs = paw.wfs
286 assert wfs.dtype == float
287 gd = wfs.gd
288 atom_index = wignerseitz(gd, paw.atoms)
290 w_k = wfs.kd.weight_k
291 nk = len(w_k)
292 nb = wfs.bd.nbands
294 energies = np.empty(nb * nk)
295 weights = np.empty(nb * nk)
296 x = 0
297 for k, w in enumerate(w_k):
298 u = spin * nk + k
299 energies[x:x + nb] = wfs.collect_eigenvalues(k=k, s=spin)
300 for n, psit_G in enumerate(wfs.kpt_u[u].psit_nG):
301 P_i = wfs.kpt_u[u].P_ani[a][n]
302 P_p = pack_density(np.outer(P_i, P_i))
303 Delta_p = sqrt(4 * pi) * wfs.setups[a].Delta_pL[:, 0]
304 weights[x + n] = w * (
305 gd.integrate(abs(np.where(atom_index == a, psit_G, 0.0))**2) +
306 np.dot(Delta_p, P_p))
307 x += nb
308 return energies, weights
311class RawLDOS:
312 """Class to get the unfolded LDOS"""
313 def __init__(self, calc):
314 self.paw = calc
315 for setup in calc.wfs.setups.setups.values():
316 if not hasattr(setup, 'l_i'):
317 # get the mapping
318 l_i = []
319 for l in setup.l_j:
320 for m in range(2 * l + 1):
321 l_i.append(l)
322 setup.l_i = l_i
324 def get(self, atom):
325 """Return the s,p,d weights for each state"""
326 wfs = self.paw.wfs
327 nibzkpts = len(wfs.kd.ibzk_kc)
328 spd = np.zeros((wfs.nspins, nibzkpts, wfs.bd.nbands, 3))
330 if hasattr(atom, '__iter__'):
331 # atom is a list of atom indices
332 for a in atom:
333 spd += self.get(a)
334 return spd
336 l_i = wfs.setups[atom].l_i
338 for kpt in self.paw.wfs.kpt_u:
339 if atom in kpt.P_ani:
340 for i, P_n in enumerate(kpt.P_ani[atom].T):
341 spd[kpt.s, kpt.k, :, l_i[i]] += np.abs(P_n)**2
343 wfs.gd.comm.sum(spd)
344 wfs.kd.comm.sum(spd)
345 return spd
347 def by_element(self, indices=None):
348 """Return a dict with elements as keys and LDOS as values.
350 Parameter:
351 -----------
352 indices: list or None
353 Atom indices to consider, default None
354 """
355 atoms = self.paw.atoms
356 if indices is None:
357 selected = range(len(atoms))
358 else:
359 selected = indices
361 elemi = {}
362 for i in selected:
363 symbol = atoms[i].symbol
364 if symbol in elemi:
365 elemi[symbol].append(i)
366 else:
367 elemi[symbol] = [i]
368 for key in elemi:
369 elemi[key] = self.get(elemi[key])
370 return elemi
372 def to_file(self,
373 ldbe,
374 filename,
375 width=None,
376 shift=True,
377 bound=False):
378 """Write the LDOS to a file.
380 Parameters:
381 -----------
382 ldbe: dict
383 weights
384 filename: string
385 width: float or None
386 If a width is given, the LDOS will be Gaussian folded
387 shift: bool
388 Set Fermi energy to 0 eV. Default True
389 bound: bool
390 If the calculation was performed spin-polarized with fixmagmom=True,
391 there are two Fermi-levels, one for each spin. False shifts
392 both spins by their own Fermi levels, True both
393 by the higher Fermi level. Default: False
394 """
396 f = paropen(filename, 'w')
398 def append_weight_strings(ldbe, data):
399 s = ''
400 for key in ldbe:
401 for l in 'spd':
402 data.append(key + '(' + l + ')-weight')
403 if len(key) == 1:
404 key = ' ' + key
405 s += ' ' + key + ':s p d '
406 return s
408 wfs = self.paw.wfs
410 if width is None:
411 # unfolded ldos
412 fmf = FMF(['Raw LDOS obtained from projector weights'])
413 print(fmf.header(), end=' ', file=f)
414 data = ['energy: energy [eV]',
415 'occupation number: occ',
416 'spin index: s',
417 'k-point index: k',
418 'band index: n',
419 'k-point weight: weight']
420 string = '# e_i[eV] occ s k n kptwght '
421 string += append_weight_strings(ldbe, data)
422 data.append('summed weights: sum')
423 string += ' sum'
424 print(fmf.data(data), end=' ', file=f)
425 print(string, file=f)
426 for k in range(wfs.kd.nibzkpts):
427 for s in range(wfs.nspins):
428 e_n = self.paw.get_eigenvalues(kpt=k, spin=s)
429 f_n = self.paw.get_occupation_numbers(kpt=k, spin=s)
430 if e_n is None:
431 continue
432 w = wfs.kd.weight_k[k]
433 for n in range(wfs.bd.nbands):
434 sum = 0.0
435 print('%10.5f %6.4f %2d %5d' % (e_n[n], f_n[n],
436 s, k), end=' ', file=f)
437 print('%6d %8.4f' % (n, w), end=' ', file=f)
438 for key in ldbe:
439 spd = ldbe[key][s, k, n]
440 for l in range(3):
441 sum += spd[l]
442 print('%8.4f' % spd[l], end=' ', file=f)
443 print('%8.4f' % sum, file=f)
444 else:
445 # folded ldos
446 fmf = FMF(['Folded raw LDOS obtained from projector weights'])
447 print(fmf.header(), end=' ', file=f)
449 gauss = Gauss(width)
450 print(fmf.field('folding',
451 ['function: Gauss',
452 'width: ' + str(width) + ' [eV]']),
453 end=' ', file=f)
455 data = ['energy: energy [eV]',
456 'spin index: s',
457 'k-point index: k',
458 'band index: n',
459 'k-point weight: weight']
461 # minimal and maximal energies
462 emin = 1.e32
463 emax = -1.e32
464 for k in range(wfs.kd.nibzkpts):
465 for s in range(wfs.nspins):
466 e_n = self.paw.get_eigenvalues(kpt=k, spin=s)
467 emin = min(e_n.min(), emin)
468 emax = max(e_n.max(), emax)
469 emin -= 4 * width
470 emax += 4 * width
472 # Fermi energies
473 efermis = self.paw.wfs.fermi_levels
475 eshift = 0.0
477 if shift and len(efermis) == 1:
478 eshift = -efermis[0]
480 # set de to sample 4 points in the width
481 de = width / 4
483 string = '# e[eV] s '
484 string += append_weight_strings(ldbe, data)
486 for s in range(wfs.nspins):
487 if len(efermis) == 2:
488 if not bound:
489 eshift = -efermis[s]
490 else:
491 eshift = -efermis.max()
493 print(fmf.data(data), end=' ', file=f)
495 print('# Gauss folded, width=%g [eV]' % width, file=f)
496 if shift:
497 print('# shifted to Fermi energy = 0', file=f)
498 print('# Fermi energy was', end=' ', file=f)
499 else:
500 print('# Fermi energy', end=' ', file=f)
501 print(efermis, 'eV', file=f)
502 print(string, file=f)
504 # loop over energies
505 emax = emax + 0.5 * de
506 e = emin
507 while e < emax:
508 val = {}
509 for key in ldbe:
510 val[key] = np.zeros(3)
511 for k in range(wfs.kd.nibzkpts):
512 w = wfs.kpt_u[k].weight
513 e_n = self.paw.get_eigenvalues(kpt=k, spin=s)
514 for n in range(wfs.bd.nbands):
515 w_i = w * gauss.get(e_n[n] - e)
516 for key in ldbe:
517 val[key] += w_i * ldbe[key][s, k, n]
519 print('%10.5f %2d' % (e + eshift, s), end=' ', file=f)
520 for key in val:
521 spd = val[key]
522 for l in range(3):
523 print('%8.4f' % spd[l], end=' ', file=f)
524 print(file=f)
525 e += de
527 f.close()
529 def by_element_to_file(self,
530 filename='ldos_by_element.dat',
531 width=None,
532 shift=True,
533 bound=False,
534 indices=None):
535 """
536 Write the LDOS by element to a file.
538 Parameters:
539 -----------
540 indices: list or None
541 Atom indices to consider, default None
542 """
543 ldbe = self.by_element(indices)
544 self.to_file(ldbe, filename, width, shift, bound)
547class LCAODOS:
548 """Class for calculating the projected subspace DOS.
550 The projected subspace density of states is defined only in LCAO
551 mode. The advantages to the PDOS based on projectors is that the
552 LCAODOS will properly take non-orthogonality and completeness into
553 account."""
554 def __init__(self, calc):
555 self.calc = calc
557 def get_orbital_pdos(self, M, ravel=True):
558 """Get projected DOS from LCAO basis function M."""
559 return self.get_subspace_pdos([M], ravel=ravel)
561 def get_atomic_subspace_pdos(self, a, ravel=True):
562 """Get projected subspace DOS from LCAO basis on atom(s) a.
564 a may be an atomic index or a list of indices."""
565 Mvalues = self.get_atom_indices(a)
566 return self.get_subspace_pdos(Mvalues, ravel=ravel)
568 def get_atom_indices(self, a):
569 """Get list of basis function indices of atom(s) a."""
570 if isinstance(a, int):
571 a = [a]
572 Mvalues = []
573 for a0 in a:
574 M = self.calc.wfs.basis_functions.M_a[a0]
575 Mvalues.extend(range(M, M + self.calc.wfs.setups[a0].nao))
576 return Mvalues
578 def get_subspace_pdos(self, Mvalues, ravel=True):
579 """Get projected subspace DOS from LCAO basis."""
580 wfs = self.calc.wfs
581 bd = wfs.bd
582 kd = wfs.kd
583 gd = wfs.gd
585 for kpt in wfs.kpt_u:
586 assert not np.isnan(kpt.eps_n).any()
588 w_skn = np.zeros((kd.nspins, kd.nibzkpts, bd.nbands))
589 eps_skn = np.zeros((kd.nspins, kd.nibzkpts, bd.nbands))
590 for u, kpt in enumerate(wfs.kpt_u):
591 C_nM = kpt.C_nM
592 from gpaw.kohnsham_layouts import BlacsOrbitalLayouts
593 if isinstance(wfs.ksl, BlacsOrbitalLayouts):
594 raise NotImplementedError('Something not quite working. '
595 'FIXME.')
596 S_MM = wfs.ksl.mmdescriptor.collect_on_master(kpt.S_MM)
597 if bd.rank != 0 or gd.rank != 0:
598 S_MM = np.empty((wfs.ksl.nao, wfs.ksl.nao),
599 dtype=wfs.dtype)
600 bd.comm.broadcast(S_MM, 0)
601 else:
602 S_MM = kpt.S_MM
604 if gd.comm.rank == 0:
605 S_mm = S_MM[Mvalues, :][:, Mvalues].copy()
606 iS_mm = np.linalg.inv(S_mm).copy()
607 CS_nm = np.dot(C_nM, S_MM[:, Mvalues])
608 CSiS_nm = np.dot(CS_nm, iS_mm)
609 w_n = (np.conj(CS_nm) * CSiS_nm).sum(axis=1) * kpt.weight
610 w_skn[kpt.s, kpt.k, bd.beg:bd.end:bd.step] = w_n.real
611 eps_skn[kpt.s, kpt.k, bd.beg:bd.end:bd.step] = kpt.eps_n
613 for arr in [eps_skn, w_skn]:
614 gd.comm.broadcast(arr, 0)
615 bd.comm.sum(arr)
616 kd.comm.sum(arr)
618 if ravel:
619 eps_n = eps_skn.ravel()
620 w_n = w_skn.ravel()
621 args = np.argsort(eps_n)
622 eps_n = eps_n[args]
623 w_n = w_n[args]
624 return eps_n, w_n
625 else:
626 return eps_skn, w_skn
629class RestartLCAODOS(LCAODOS):
630 """Class for calculating LCAO subspace PDOS from a restarted calculator.
632 Warning: This has side effects on the calculator. The
633 operation will allocate memory to diagonalize the Hamiltonian and
634 set coefficients plus positions."""
635 def __init__(self, calc):
636 LCAODOS.__init__(self, calc)
637 system = calc.get_atoms()
638 calc.set_positions(system)
639 calc.wfs.eigensolver.iterate(calc.hamiltonian, calc.wfs)