Coverage for gpaw/response/modelinteraction.py: 95%
105 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
1import numpy as np
2from ase.units import Ha
3from gpaw.mpi import world
4from gpaw.response import ResponseContext
5from gpaw.response.coulomb_kernels import CoulombKernel
6from gpaw.response.screened_interaction import initialize_w_calculator
7from gpaw.response import timer
8from gpaw.response.pw_parallelization import Blocks1D
9from gpaw.response.pair import KPointPairFactory
10from gpaw.wannier90 import read_uwan
11from gpaw.response.screened_interaction import GammaIntegrationMode
14def ibz2bz_map(qd):
15 """ Maps each k in BZ to corresponding k in IBZ. """
16 out_map = [[] for _ in range(qd.nibzkpts)]
17 for iK in range(qd.nbzkpts):
18 ik = qd.bz2ibz_k[iK]
19 out_map[ik].append(iK)
20 return out_map
23def initialize_w_model(chi0calc, truncation=None,
24 integrate_gamma=GammaIntegrationMode('sphere'),
25 q0_correction=False, txt='w_model.out',
26 eta=None, world=world, timer=None):
27 """ Helper function to initialize ModelInteraction
29 Parameters
30 ----------
31 chi0calc: Chi0Calculator
32 truncation: str
33 Coulomb truncation scheme. Can be either 2D, 1D, 0D or None.
34 integrate_gamma: GammaIntegrationMode
35 q0_correction: bool
36 Analytic correction to the q=0 contribution applicable to 2D
37 systems.
38 txt: str
39 Filename of output files.
40 world: MPI comm
41 timer: ResponseContext timer
42 """
43 gs = chi0calc.gs
44 wcontext = ResponseContext(txt=txt,
45 comm=world, timer=timer)
46 coulomb = CoulombKernel.from_gs(gs, truncation=truncation)
47 if eta is not None:
48 eta /= Ha
49 wcalc = initialize_w_calculator(chi0calc,
50 wcontext,
51 coulomb=coulomb,
52 xc='RPA',
53 integrate_gamma=integrate_gamma,
54 q0_correction=q0_correction,
55 eta=eta)
56 return ModelInteraction(wcalc)
59class ModelInteraction:
61 def __init__(self, wcalc):
62 """Class to compute interaction matrix in Wannier basis
64 Parameters
65 ----------
66 wcalc: WCalculator
67 """
69 self.wcalc = wcalc
70 self.gs = wcalc.gs
71 self.context = self.wcalc.context
72 self.qd = self.wcalc.qd
74 @timer('Calculate W in Wannier')
75 def calc_in_Wannier(self, chi0calc, Uwan_mnk, bandrange, spin=0):
76 """Calculates the screened interaction matrix in Wannier basis.
77 NOTE: Does not work with SOC!
79 Parameters:
80 ----------
81 chi0calc: Chi0Calculator
82 Uwan_mnk: cmplx or str
83 Wannier transformation matrix. if str name of wannier90 output
84 m: Wannier index, n: bandindex, k: k-index
85 bandrange: int
86 range of bands that the wannier functions were constructed from
87 spin: int
88 spin index (0 or 1)
90 Documentation
91 -------------
93 W_n1,n2;n3,n4(R=0, w) =
94 <w^*_{n1,R=0} w_{n2, R=0} | W(r,r',w) |w^*_{n3,R=0} w_{n4, R=0} >
96 w_{n R} = V/(2pi)^3 int_{BZ} dk e^{-kR} psi^w_{nk}
97 psi^w_{nk} = sum_m U_nm(k) psi^{KS}_{mk}
99 In this implementation we compute W_n1,n2;n3,n4(R=0) efficiently
100 by expressing it in terms of the reduced Wannier density matrices
102 A_{n1,n2,q} = sum_{k} rhowan^{n1,k}_{n2, k+q}
104 where the Wannier density matrix is calculated from the usual
105 density matrix
106 rho^{m1,k}_{m2 k+q} = < psi_{m1,k} | e^{i(q+G)r} | psi_{m2, k+q} >
107 and the Wannier transformation matrices U_{nm}(k) as
109 rhowan^{n1, k}_{n2, k+q} =
110 = sum_{m1, m5} U^*_{n1,m1}(k) U_{n2, m2}(k+q) rho^{m1,k}_{m2 k+q}
111 """
113 ibz2bz = ibz2bz_map(self.gs.kd)
115 # Create pair_factory with nblocks = 1
116 pair_factory = KPointPairFactory(self.gs, self.context)
117 pair_calc = pair_factory.pair_calculator()
119 if isinstance(Uwan_mnk, str):
120 # read w90 transformation matrix from file
121 seed = Uwan_mnk
122 Uwan_mnk, nk, nwan, nband = read_uwan(seed,
123 self.gs.kd,
124 dis=False)
125 if bandrange[1] - bandrange[0] > nband:
126 Uwan_mnk, nk, nwan, nband = read_uwan(seed,
127 self.gs.kd,
128 dis=True)
129 else:
130 nk = Uwan_mnk.shape[2]
131 nband = Uwan_mnk.shape[1]
132 nwan = Uwan_mnk.shape[0]
134 assert nk == self.gs.kd.nbzkpts
135 assert bandrange[1] - bandrange[0] == nband
136 wd = chi0calc.wd
137 nfreq = len(wd)
139 # Find frequency range for block distributed arrays
140 blockdist = chi0calc.chi0_body_calc.get_blockdist()
141 blocks1d = Blocks1D(blockdist.blockcomm, nfreq)
142 mynfreq = blocks1d.nlocal
143 self.intrablockcomm = blockdist.intrablockcomm
145 # Variable that will store the screened interaction in Wannier basis
146 Wwan_wijkl = np.zeros([mynfreq, nwan, nwan, nwan, nwan], dtype=complex)
148 for iq, q_c in enumerate(self.gs.kd.ibzk_kc):
149 self.context.print('iq = ', iq, '/', self.gs.kd.nibzkpts)
150 # Calculate chi0 and W for IBZ k-point q
151 self.context.print('calculating chi0...')
152 self.context.timer.start('chi0')
153 chi0 = chi0calc.calculate(q_c)
154 self.context.timer.stop('chi0')
155 qpd = chi0.qpd
156 self.context.print('calculating W_wGG...')
157 W_wGG = self.wcalc.calculate_W_wGG(chi0,
158 fxc_mode='GW',
159 only_correlation=False)
160 pawcorr = chi0calc.chi0_body_calc.pawcorr
162 self.context.print('Projecting to localized Wannier basis...')
163 # Loop over all equivalent k-points
164 for iQ in ibz2bz[iq]:
165 Q_c = self.wcalc.qd.bzk_kc[iQ]
166 assert self.wcalc.qd.where_is_q(Q_c,
167 self.wcalc.qd.bzk_kc) == iQ
169 A_mnG = self.get_reduced_wannier_density_matrix(spin,
170 Q_c,
171 iq,
172 bandrange,
173 pawcorr,
174 qpd,
175 pair_calc,
176 pair_factory,
177 Uwan_mnk)
178 if self.qd.time_reversal_k[iQ]:
179 # TR corresponds to complex conjugation
180 A_mnG = A_mnG.conj()
182 Wwan_wijkl += np.einsum('ijk,lkm,pqm->lipjq',
183 A_mnG.conj(),
184 W_wGG,
185 A_mnG,
186 optimize='optimal')
188 # factor from three BZ summations and taking from Hartree to eV
189 # and volume factor from matrix element in PW basis
190 factor = Ha / self.gs.kd.nbzkpts**3 / self.gs.volume
191 Wwan_wijkl *= factor
192 self.context.write_timer()
194 return wd.omega_w * Ha, blocks1d.all_gather(Wwan_wijkl)
196 @timer('get_reduced_wannier_density_matrix')
197 def get_reduced_wannier_density_matrix(self, spin, Q_c, iq, bandrange,
198 pawcorr, qpd, pair_calc,
199 pair_factory, Uwan_mnk):
200 """
201 Returns sum_k sum_(m1,m2) U_{n1m1}* U_{n2m2} rho^{m1 k}_{m2 k-q}(G)
202 where rho is the usual density matrix and U are wannier tranformation
203 matrices.
204 """
205 nG = qpd.ngmax
206 nwan = Uwan_mnk.shape[0]
207 A_mnG = np.zeros([nwan, nwan, nG], dtype=complex)
209 # Parallell sum over k-points
210 for iK1 in self.myKrange():
211 kpt1 = pair_factory.get_k_point(spin, iK1,
212 bandrange[0],
213 bandrange[1])
214 iK2 = self.gs.kd.find_k_plus_q(Q_c, [kpt1.K])[0]
215 kpt2 = pair_factory.get_k_point(spin, iK2,
216 bandrange[0],
217 bandrange[1])
219 # Calculate density matrix
220 rholoc, iqloc, sign = self.get_density_matrix(kpt1,
221 kpt2,
222 pawcorr,
223 qpd,
224 pair_calc)
225 assert iqloc == iq
227 # Rotate to Wannier basis and sum to get reduced Wannier
228 # density matrix A
229 A_mnG += np.einsum('ia,jb,abG->ijG',
230 Uwan_mnk[:, :, iK1].conj(),
231 Uwan_mnk[:, :, iK2],
232 rholoc)
233 self.intrablockcomm.sum(A_mnG)
234 return A_mnG
236 def get_density_matrix(self, kpt1, kpt2, pawcorr, qpd, pair_calc):
237 from gpaw.response.g0w0 import QSymmetryOp, get_nmG
239 symop, iq = QSymmetryOp.get_symop_from_kpair(
240 self.gs.kd, self.qd, kpt1, kpt2)
241 nG = qpd.ngmax
242 mypawcorr, I_G = symop.apply_symop_q(
243 qpd, pawcorr, kpt1, kpt2)
245 rho_mnG = np.zeros((len(kpt1.eps_n), len(kpt2.eps_n), nG),
246 complex)
247 for m in range(len(rho_mnG)):
248 rho_mnG[m] = get_nmG(kpt1, kpt2, mypawcorr, m, qpd, I_G, pair_calc)
250 return rho_mnG, iq, symop.sign
252 def myKrange(self, rank=None):
253 if rank is None:
254 rank = self.intrablockcomm.rank
255 nK = self.gs.kd.nbzkpts
256 myKsize = -(-nK // self.intrablockcomm.size)
257 myKrange = range(rank * myKsize,
258 min((rank + 1) * myKsize, nK))
259 # myKsize = len(myKrange)
260 return myKrange