Coverage for gpaw/lcao/pwf2.py: 92%
84 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-08 00:17 +0000
1import numpy as np
2from ase.units import Hartree
4from gpaw.lcao.projected_wannier import condition_number, dots, get_bfs
5from gpaw.utilities import unpack_hermitian
6from gpaw.utilities.tools import tri2full
9def get_lcao_xc(calc, P_aqMi, bfs=None, spin=0):
10 nq = len(calc.wfs.kd.ibzk_qc)
11 nao = calc.wfs.setups.nao
12 dtype = calc.wfs.dtype
13 if bfs is None:
14 bfs = get_bfs(calc)
16 if calc.density.nt_sg is None:
17 calc.density.interpolate_pseudo_density()
18 nt_sg = calc.density.nt_sg
19 vxct_sg = calc.density.finegd.zeros(calc.wfs.nspins)
20 calc.hamiltonian.xc.calculate(calc.density.finegd, nt_sg, vxct_sg)
21 vxct_G = calc.wfs.gd.zeros()
22 calc.hamiltonian.restrict_and_collect(vxct_sg[spin], vxct_G)
23 Vxc_qMM = np.zeros((nq, nao, nao), dtype)
24 for q, Vxc_MM in enumerate(Vxc_qMM):
25 bfs.calculate_potential_matrix(vxct_G, Vxc_MM, q)
26 tri2full(Vxc_MM, 'L')
28 # Add atomic PAW corrections
29 for a, P_qMi in P_aqMi.items():
30 D_sp = calc.density.D_asp[a][:]
31 H_sp = np.zeros_like(D_sp)
32 calc.hamiltonian.xc.calculate_paw_correction(calc.wfs.setups[a],
33 D_sp, H_sp)
34 H_ii = unpack_hermitian(H_sp[spin])
35 for Vxc_MM, P_Mi in zip(Vxc_qMM, P_qMi):
36 Vxc_MM += dots(P_Mi, H_ii, P_Mi.T.conj())
37 return Vxc_qMM * Hartree
40class LCAOwrap:
41 def __init__(self, calc, spin=0):
42 assert calc.wfs.gd.comm.size == 1
43 assert calc.wfs.kd.comm.size == 1
44 assert calc.wfs.bd.comm.size == 1
46 from gpaw.lcao.tools import get_lcao_hamiltonian
47 H_skMM, S_kMM = get_lcao_hamiltonian(calc)
49 self.calc = calc
50 self.dtype = calc.wfs.dtype
51 self.spin = spin
52 self.H_qww = H_skMM[spin]
53 self.S_qww = S_kMM
54 self.P_aqwi = calc.wfs.P_aqMi
55 self.Nw = self.S_qww.shape[-1]
57 for S in self.S_qww:
58 print('Condition number: %0.1e' % condition_number(S))
60 def get_hamiltonian(self, q=0, indices=None):
61 if indices is None:
62 return self.H_qww[q]
63 else:
64 return self.H_qww[q].take(indices, 0).take(indices, 1)
66 def get_overlap(self, q=0, indices=None):
67 if indices is None:
68 return self.S_qww[q]
69 else:
70 return self.S_qww[q].take(indices, 0).take(indices, 1)
72 def get_projections(self, q=0, indices=None):
73 if indices is None:
74 return {a: P_qwi[q] for a, P_qwi in self.P_aqwi.items()}
75 else:
76 return {a: P_qwi[q].take(indices, 0)
77 for a, P_qwi in self.P_aqwi.items()}
79 def get_orbitals(self, q=-1, indices=None):
80 assert q == -1
81 if indices is None:
82 indices = range(self.Nw)
83 Ni = len(indices)
84 C_wM = np.zeros((Ni, self.Nw), self.dtype)
85 for i, C_M in zip(indices, C_wM):
86 C_M[i] = 1.0
87 w_wG = self.calc.wfs.gd.zeros(Ni, dtype=self.dtype)
88 self.calc.wfs.basis_functions.lcao_to_grid(C_wM, w_wG, q=-1)
89 return w_wG
91 def get_Fcore(self, q=0, indices=None):
92 if indices is None:
93 Fcore_ww = np.zeros_like(self.H_qww[q])
94 else:
95 Fcore_ww = np.zeros((len(indices), len(indices)))
96 for a, P_wi in self.get_projections(q, indices).items():
97 if self.calc.wfs.setups[a].type != 'ghost':
98 X_ii = unpack_hermitian(self.calc.wfs.setups[a].X_p)
99 Fcore_ww -= dots(P_wi, X_ii, P_wi.T.conj())
100 return Fcore_ww * Hartree
102 def get_xc(self, q=0, indices=None):
103 if not hasattr(self, 'Vxc_qww'):
104 self.Vxc_qww = get_lcao_xc(self.calc, self.P_aqwi,
105 bfs=self.calc.wfs.basis_functions,
106 spin=self.spin)
107 if indices is None:
108 return self.Vxc_qww[q]
109 else:
110 return self.Vxc_qww[q].take(indices, 0).take(indices, 1)