Coverage for gpaw/lcao/projected_wannier.py: 40%
222 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1import numpy as np
2import numpy.linalg as la
3from ase.units import Hartree
5from gpaw.lcao.overlap import NewTwoCenterIntegrals as TwoCenterIntegrals
6from gpaw.lfc import BasisFunctions
7from gpaw.utilities import unpack_hermitian
8from gpaw.utilities.tools import dagger, lowdin, tri2full
9from gpaw.lcao.tools import basis_subset2
12def dots(*args):
13 x = args[0]
14 for M in args[1:]:
15 x = np.dot(x, M)
16 return x
19def normalize(U, U2=None, norms=None):
20 if norms is None:
21 norms2 = np.ones(U.shape[1])
22 else:
23 norms2 = np.sqrt(norms)
25 if U2 is None:
26 for n, col in zip(norms2, U.T):
27 col *= n / la.norm(col)
28 else:
29 for n, col1, col2 in zip(norms2, U.T, U2.T):
30 norm = np.sqrt(np.vdot(col1, col1) + np.vdot(col2, col2))
31 col1 *= n / norm
32 col2 *= n / norm
35def normalize2(C, S):
36 C /= np.sqrt(dots(dagger(C), S, C).diagonal())
39def get_rot(F_MM, V_oM, L):
40 eps_M, U_MM = la.eigh(F_MM)
41 indices = eps_M.real.argsort()[-L:]
42 U_Ml = U_MM[:, indices]
43 U_Ml /= np.sqrt(dots(U_Ml.T.conj(), F_MM, U_Ml).diagonal())
45 U_ow = V_oM.copy()
46 U_lw = np.dot(U_Ml.T.conj(), F_MM)
47 for col1, col2 in zip(U_ow.T, U_lw.T):
48 norm = np.sqrt(np.vdot(col1, col1) + np.vdot(col2, col2))
49 col1 /= norm
50 col2 /= norm
51 return U_ow, U_lw, U_Ml
54def condition_number(S):
55 eps = la.eigvalsh(S).real
56 return eps.max() / eps.min()
59def eigvals(H, S):
60 return np.sort(la.eigvals(la.solve(S, H)).real)
63def get_bfs(calc):
64 wfs = calc.wfs
65 bfs = BasisFunctions(wfs.gd,
66 [setup.basis_functions_J for setup in wfs.setups],
67 wfs.kd, cut=True)
68 bfs.set_positions(wfs.spos_ac)
69 return bfs
72def get_lcao_projections_HSP(calc, bfs=None, spin=0, projectionsonly=True):
73 """Some title.
75 if projectionsonly is True, return the projections::
77 V_qnM = <psi_qn | Phi_qM>
79 else, also return the Hamiltonian, overlap, and projector overlaps::
81 H_qMM = <Phi_qM| H |Phi_qM'>
82 S_qMM = <Phi_qM|Phi_qM'>
83 P_aqMi = <pt^a_qi|Phi_qM>
84 """
85 if calc.wfs.kd.comm.size != 1:
86 raise NotImplementedError('Parallelization over spin/kpt not '
87 'implemented yet.')
88 comm = calc.wfs.gd.comm
89 nq = len(calc.wfs.kd.ibzk_qc)
90 nao = calc.wfs.setups.nao
91 dtype = calc.wfs.dtype
92 if bfs is None:
93 bfs = get_bfs(calc)
94 tci = TwoCenterIntegrals(calc.wfs.gd.cell_cv,
95 calc.wfs.gd.pbc_c,
96 calc.wfs.setups,
97 calc.wfs.kd.ibzk_qc,
98 calc.wfs.kd.gamma)
100 # Calculate projector overlaps, and (lower triangle of-) S and T matrices
101 S_qMM = np.zeros((nq, nao, nao), dtype)
102 T_qMM = np.zeros((nq, nao, nao), dtype)
103 P_aqMi = {}
105 for a in bfs.my_atom_indices:
106 ni = calc.wfs.setups[a].ni
107 P_aqMi[a] = np.zeros((nq, nao, ni), dtype)
108 tci.calculate(calc.spos_ac, S_qMM, T_qMM, P_aqMi)
110 from gpaw.lcao.atomic_correction import DenseAtomicCorrection
111 dS_aii = {a: calc.wfs.setups[a].dO_ii for a in bfs.my_atom_indices}
112 atomic_correction = DenseAtomicCorrection(P_aqMi, dS_aii, 0, nao)
113 atomic_correction.add_overlap_correction(S_qMM)
114 # calc.wfs.P_aqMi = P_aqMi
115 # for q, S_MM in enumerate(S_qMM):
116 # atomic_correction.calculate_manual(wfs, q,
117 # atomic_correction.add_overlap_correction(calc.wfs, S_qMM)
118 calc.wfs.gd.comm.sum(S_qMM)
119 calc.wfs.gd.comm.sum(T_qMM)
121 # Calculate projections
122 V_qnM = np.zeros((nq, calc.wfs.bd.nbands, nao), dtype)
123 for kpt in calc.wfs.kpt_u:
124 if kpt.s != spin:
125 continue
126 V_nM = V_qnM[kpt.q]
127 # bfs.integrate2(kpt.psit_nG[:], V_nM, kpt.q) # all bands to save time
128 for n, V_M in enumerate(V_nM): # band-by-band to save memory
129 bfs.integrate2(kpt.psit_nG[n][:], V_M, kpt.q)
130 for a, P_ni in kpt.projections.items():
131 dS_ii = calc.wfs.setups[a].dO_ii
132 P_Mi = P_aqMi[a][kpt.q]
133 V_nM += np.dot(P_ni, np.inner(dS_ii, P_Mi).conj())
134 comm.sum(V_qnM)
135 if projectionsonly:
136 return V_qnM
138 # Determine potential matrix
139 vt_G = calc.hamiltonian.vt_sG[spin]
140 H_qMM = np.zeros((nq, nao, nao), dtype)
141 for q, H_MM in enumerate(H_qMM):
142 bfs.calculate_potential_matrix(vt_G, H_MM, q)
144 # Make Hamiltonian as sum of kinetic (T) and potential (V) matrices
145 # and add atomic corrections
146 for a, P_qMi in P_aqMi.items():
147 dH_ii = unpack_hermitian(calc.hamiltonian.dH_asp[a][spin])
148 for P_Mi, H_MM in zip(P_qMi, H_qMM):
149 H_MM += np.dot(P_Mi, np.inner(dH_ii, P_Mi).conj())
150 comm.sum(H_qMM)
151 H_qMM += T_qMM
153 # Fill in the upper triangles of H and S
154 for H_MM, S_MM in zip(H_qMM, S_qMM):
155 tri2full(H_MM)
156 tri2full(S_MM)
157 H_qMM *= Hartree
159 return V_qnM, H_qMM, S_qMM, P_aqMi
162def convert_projection_data(symbols, V_qnM, H_qMM, S_qMM, P_aqMi,
163 originaltype='dzp', newtype='sz'):
164 Nq = len(H_qMM)
165 mask = basis_subset2(symbols, originaltype, newtype)
166 mask2 = np.outer(mask, mask)
167 NMnew = len(mask.nonzero()[0])
168 V_qnM = V_qnM[..., mask]
169 H_qMM = H_qMM[:, mask2].reshape(Nq, NMnew, NMnew)
170 S_qMM = S_qMM[:, mask2].reshape(Nq, NMnew, NMnew)
171 P2_aqMi = {}
172 for a, P_qMi in P_aqMi.items():
173 P2_aqMi[a] = P_qMi[:, mask, :]
174 return V_qnM, H_qMM, S_qMM, P2_aqMi
177def get_phs(calc, s=0):
178 return get_lcao_projections_HSP(calc, bfs=None,
179 spin=s, projectionsonly=False)
182class ProjectedWannierFunctions:
183 def __init__(self, projections, h_lcao, s_lcao, eigenvalues, kpoints,
184 L_k=None, M_k=None, N=None, fixedenergy=None):
185 """projections[n,i] = <psi_n|f_i>
186 h_lcao[i1, i2] = <f_i1|h|f_i2>
187 s_lcao[[i1, i2] = <f_i1|f_i2>
188 eps_n: Exact eigenvalues
189 L: Number of extra degrees of freedom
190 M: Number of states to exactly span
191 N: Total number of bands in the calculation
193 Methods:
194 -- get_hamiltonian_and_overlap_matrix --
195 will return the hamiltonian and identity operator
196 in the projected wannier function basis.
197 The following steps are performed:
199 1) calculate_edf -> self.b_il
200 2) calculate_rotations -> self.Uo_mi and self.Uu_li
201 3) calculate_overlaps -> self.S_ii
202 4) calculate_hamiltonian_matrix -> self.H_ii
204 -- get_eigenvalues --
205 gives the eigenvalues of of the hamiltonian in the
206 projected wannier function basis.
208 -- indices --
209 i localized function index
210 n eigenstate index
211 l edf index
212 m fixed eigenstate index
213 k k-point index
214 """
216 self.eps_kn = eigenvalues
217 self.ibzk_kc = kpoints
218 self.nk = len(self.ibzk_kc)
219 self.V_kni = projections # <psi_n1|f_i1>
220 self.dtype = self.V_kni.dtype
221 self.Nw = self.V_kni.shape[2]
222 self.s_lcao_kii = s_lcao # F_ii[i1,i2] = <f_i1|f_i2>
223 self.h_lcao_kii = h_lcao
225 if N is None:
226 N = self.V_kni.shape[1]
227 self.N = N
229 if fixedenergy is None:
230 raise NotImplementedError(
231 'Only fixedenergy is implemented for now')
232 else:
233 self.fixedenergy = fixedenergy
234 self.M_k = [sum(eps_n <= fixedenergy) for eps_n in self.eps_kn]
235 self.L_k = [self.Nw - M for M in self.M_k]
236 print("fixedenergy =", self.fixedenergy)
238 print('N =', self.N)
239 print('skpt_kc = ')
240 print(self.ibzk_kc)
241 print('M_k =', self.M_k)
242 print('L_k =', self.L_k)
243 print('Nw =', self.Nw)
245 def get_hamiltonian_and_overlap_matrix(self, useibl=True):
246 self.calculate_edf(useibl=useibl)
247 self.calculate_rotations()
248 self.calculate_overlaps()
249 self.calculate_hamiltonian_matrix(useibl=useibl)
250 return self.H_kii, self.S_kii
252 def calculate_edf(self, useibl=True):
253 """Calculate the coefficients b_il in the expansion of the EDF.
255 ``|phi_l> = sum_i b_il |f^u_i>``, in terms of ``|f^u_i> = P^u|f_i>``.
257 To use the infinite band limit set useibl=True.
258 N is the total number of bands to use.
259 """
261 for k, L in enumerate(self.L_k):
262 if L == 0:
263 assert L != 0, 'L_k=0 for k=%i. Not implemented' % k
265 self.Vo_kni = [V_ni[:M] for V_ni, M in zip(self.V_kni, self.M_k)]
267 self.Fo_kii = np.asarray([np.dot(dagger(Vo_ni), Vo_ni)
268 for Vo_ni in self.Vo_kni])
270 if useibl:
271 self.Fu_kii = self.s_lcao_kii - self.Fo_kii
272 else:
273 self.Vu_kni = [V_ni[M:self.N]
274 for V_ni, M in zip(self.V_kni, self.M_k)]
275 self.Fu_kii = np.asarray([np.dot(dagger(Vu_ni), Vu_ni)
276 for Vu_ni in self.Vu_kni])
277 self.b_kil = []
278 for Fu_ii, L in zip(self.Fu_kii, self.L_k):
279 b_i, b_ii = la.eigh(Fu_ii)
280 ls = b_i.real.argsort()[-L:]
281 b_il = b_ii[:, ls] # pick out the eigenvec with largest eigenvals.
282 normalize2(b_il, Fu_ii) # normalize the EDF: <phi_l|phi_l> = 1
283 self.b_kil.append(b_il)
285 def calculate_rotations(self):
286 Uo_kni = [Vo_ni.copy() for Vo_ni in self.Vo_kni]
287 Uu_kli = [np.dot(dagger(b_il), Fu_ii)
288 for b_il, Fu_ii in zip(self.b_kil, self.Fu_kii)]
289 # Normalize such that <omega_i|omega_i> = <f_i|f_i>
290 for Uo_ni, Uu_li, s_ii in zip(Uo_kni, Uu_kli, self.s_lcao_kii):
291 normalize(Uo_ni, Uu_li, s_ii.diagonal())
292 self.Uo_kni = Uo_kni
293 self.Uu_kli = Uu_kli
295 def calculate_overlaps(self):
296 Wo_kii = [np.dot(dagger(Uo_ni), Uo_ni) for Uo_ni in self.Uo_kni]
297 Wu_kii = [np.dot(dagger(Uu_li), Uu_li) for Uu_li in self.Uu_kli]
298 # Wu_kii = [dots(dagger(Uu_li), dagger(b_il), Fu_ii, b_il, Uu_li)
299 # for Uu_li, b_il, Fu_ii in zip(self.Uu_kli, self.b_kil, self.Fu_kii)]
300 self.S_kii = np.asarray(Wo_kii) + np.asarray(Wu_kii)
302 def get_condition_number(self):
303 return np.asarray([condition_number(S) for S in self.S_kii])
305 def calculate_hamiltonian_matrix(self, useibl=True):
306 """Calculate H_kij = H^o_i(k)j(k) + H^u_i(k)j(k)
307 i(k): Bloch sum of omega_i
308 """
310 epso_kn = [eps_n[:M] for eps_n, M in zip(self.eps_kn, self.M_k)]
311 self.Ho_kii = np.asarray([np.dot(dagger(Uo_ni) * epso_n, Uo_ni)
312 for Uo_ni, epso_n in zip(self.Uo_kni,
313 epso_kn)])
315 if self.h_lcao_kii is not None and useibl:
316 print("Using h_lcao and infinite band limit")
317 Huf_kii = [h_lcao_ii - np.dot(dagger(Vo_ni) * epso_n, Vo_ni)
318 for h_lcao_ii, Vo_ni, epso_n in zip(self.h_lcao_kii,
319 self.Vo_kni,
320 epso_kn)]
321 self.Huf_kii = np.asarray(Huf_kii)
322 else:
323 print("Using finite band limit (not using h_lcao)")
324 epsu_kn = [eps_n[M:self.N]
325 for eps_n, M in zip(self.eps_kn, self.M_k)]
326 Huf_kii = [np.dot(dagger(Vu_ni) * epsu_n, Vu_ni)
327 for Vu_ni, epsu_n in zip(self.Vu_kni, epsu_kn)]
328 self.Huf_kii = np.asarray(Huf_kii)
330 Hu_kii = [dots(dagger(Uu_li), dagger(b_il), Huf_ii, b_il, Uu_li)
331 for Uu_li, b_il, Huf_ii in zip(self.Uu_kli, self.b_kil,
332 self.Huf_kii)]
333 self.Hu_kii = np.asarray(Hu_kii)
334 self.H_kii = self.Ho_kii + self.Hu_kii
336 def get_eigenvalues(self):
337 return np.asarray([eigvals(H, S)
338 for H, S in zip(self.H_kii, self.S_kii)])
340 def get_lcao_eigenvalues(self):
341 return np.asarray([eigvals(H, S)
342 for H, S in zip(self.h_lcao_kii, self.s_lcao_kii)])
344 def get_norm_of_projection(self):
345 Sinv_kii = np.asarray([la.inv(S_ii) for S_ii in self.S_kii])
347 normo_kn = np.asarray([dots(Uo_ni, Sinv_ii, dagger(Uo_ni)).diagonal()
348 for Uo_ni, Sinv_ii in
349 zip(self.Uo_kni, Sinv_kii)])
351 Vu_kni = np.asarray([V_ni[M:self.N]
352 for V_ni, M in zip(self.V_kni, self.M_k)])
354 Pu_kni = [dots(Vu_ni, b_il, Uu_li)
355 for Vu_ni, b_il, Uu_li in
356 zip(Vu_kni, self.b_kil, self.Uu_kli)]
358 normu_kn = np.asarray([dots(Pu_ni, Sinv_ii, dagger(Pu_ni)).diagonal()
359 for Pu_ni, Sinv_ii in zip(Pu_kni, Sinv_kii)])
361 return np.hstack((normo_kn, normu_kn))
363 def calculate_functions(self, calc, basis, k=0):
364 from gpaw.io import FileReference
365 psit_nG = calc.wfs.kpt_u[k].psit_nG
366 Uo_ni = self.Uo_kni[k]
367 tarinstance = isinstance(psit_nG, FileReference)
368 if tarinstance:
369 psit_nG = np.asarray([psit_nG[i] for i in range(self.M_k[k])])
371 basis_functions = get_bfs(calc)
372 b_il, Uu_li, Vo_ni = self.b_kil[k], self.Uu_kli[k], self.Vo_kni[k]
373 a_iG = -np.tensordot(Vo_ni, psit_nG, axes=[0, 0]) # a_iG = -fo_iG
374 # self.fo_iG = -a_iG.copy()#f projected onto the occupied subspace
375 C_iM = np.identity(self.Nw, dtype=self.dtype)
376 basis_functions.lcao_to_grid(C_iM, a_iG, q=-1) # a_iG=fu_iG=f_iG-fo_iG
377 # self.f_iG = self.fo_iG + a_iG # check
378 a_iG = np.tensordot(b_il, a_iG, axes=[0, 0]) # a_iG = EDF
379 a_iG = np.tensordot(Uu_li, a_iG, axes=[0, 0]) # a_iG = wu_iG
380 a_iG += np.tensordot(Uo_ni, psit_nG, axes=[0, 0]) # ai_G = wu_iG+wo_iG
381 self.w_iG = a_iG
383 def get_mlwf_initial_guess(self):
384 """calculate initial guess for maximally localized
385 wannier functions. Does not work for the infinite band limit.
386 cu_nl: rotation coefficents of unoccupied states
387 U_ii: rotation matrix of eigenstates and edf.
388 """
389 Vu_ni = self.Vu_ni[self.M: self.N]
390 cu_nl = np.dot(Vu_ni, self.b_il)
391 U_ii = np.vstack((self.Uo_ni, self.Uu_li))
392 lowdin(U_ii)
393 return U_ii, cu_nl