1import pickle
2import numpy as np
3from os.path import isfile
4from gpaw.mpi import world
5from gpaw.utilities import unpack2
6from gpaw.lcao.projected_wannier import dots
7from gpaw.utilities.tools import tri2full
8from gpaw.lfc import LocalizedFunctionsCollection as LFC
9from ase.units import Bohr, Ha
10from gpaw.utilities.timing import StepTimer, nulltimer
11
12"""This module is used to calculate the electron-phonon coupling matrix,
13    expressed in terms of GPAW LCAO orbitals."""
14
15
16class ElectronPhononCouplingMatrix:
17    r"""Class for calculating the electron-phonon coupling matrix, defined
18    by the electron phonon interaction.
19
20    ::
21
22                  __                   _____
23                  \     l   cc        /  h             cc
24        H      =   )   M   c   c     /------   ( b  + b   ),
25         el-ph    /_    ij  i   j  \/   2 W       l    l
26                    l,ij                     l
27
28    where the electron phonon coupling matrix is given by::
29
30            l           ___
31            M   = < i | \ /  V   * v  |j>
32             ij          'u   eff   l
33
34    """
35
36    def __init__(self, atoms, indices=None, name='v', delta=0.005, nfree=2,
37                 derivativemethod='tci'):
38        assert nfree in [2, 4]
39        self.nfree = nfree
40        self.delta = delta
41
42        if indices is None:
43            indices = range(len(self.atoms))
44        self.calc = atoms.calc
45        self.atoms = atoms
46        self.indices = np.asarray(indices)
47        self.name = name
48        self.p0 = self.atoms.positions.copy()
49
50        self.derivativemethod = derivativemethod
51        if derivativemethod == 'grid':
52            self.get_dP_aMix = get_grid_dP_aMix
53        elif derivativemethod == 'grid2':
54            self.get_dP_aMix = get_grid2_dP_aMix
55        elif derivativemethod == 'tci':
56            self.get_dP_aMix = get_tci_dP_aMix
57        else:
58            raise ValueError('derivativemethod must be grid, grid2, or tci')
59
60    def run(self):
61        if not isfile(self.name + '.eq.pckl'):
62
63            self.calc.calculate(self.atoms)
64            Vt_G = self.calc.get_effective_potential(pad=False)
65            Vt_G = self.calc.wfs.gd.collect(Vt_G, broadcast=True) / Ha
66            dH_asp = self.calc.hamiltonian.dH_asp
67            setups = self.calc.wfs.setups
68            nspins = self.calc.wfs.nspins
69            gd_comm = self.calc.wfs.gd.comm
70
71            alldH_asp = {}
72            for a, setup in enumerate(setups):
73                ni = setup.ni
74                nii = ni * (ni + 1) // 2
75                tmpdH_sp = np.zeros((nspins, nii))
76                if a in dH_asp:
77                    tmpdH_sp[:] = dH_asp[a]
78                gd_comm.sum(tmpdH_sp)
79                alldH_asp[a] = tmpdH_sp
80
81            forces = self.atoms.get_forces()
82            self.calc.write('eq.gpw')
83
84            world.barrier()
85            if world.rank == 0:
86                vd = open(self.name + '.eq.pckl', 'wb')
87                fd = open('vib.eq.pckl', 'wb')
88                pickle.dump((Vt_G, alldH_asp), vd, 2)
89                pickle.dump(forces, fd)
90                vd.close()
91                fd.close()
92            world.barrier()
93
94        p = self.atoms.positions.copy()
95        for a in self.indices:
96            for j in range(3):
97                for sign in [-1, 1]:
98                    for ndis in range(1, self.nfree / 2 + 1):
99                        name = '.%d%s%s.pckl' % (a, 'xyz'[j],
100                                                 ndis * ' +-'[sign])
101                        if isfile(self.name + name):
102                            continue
103                        self.atoms.positions[a, j] = (p[a, j] +
104                                                      sign * ndis * self.delta)
105                        self.calc.calculate(self.atoms)
106                        Vt_G = self.calc.get_effective_potential(pad=False)
107                        Vt_G = self.calc.wfs.gd.collect(Vt_G,
108                                                        broadcast=True) / Ha
109                        dH_asp = self.calc.hamiltonian.dH_asp
110
111                        alldH_asp = {}
112                        for a2, setup in enumerate(setups):
113                            ni = setup.ni
114                            nii = ni * (ni + 1) // 2
115                            tmpdH_sp = np.zeros((nspins, nii))
116                            if a2 in dH_asp:
117                                tmpdH_sp[:] = dH_asp[a2]
118                            gd_comm.sum(tmpdH_sp)
119                            alldH_asp[a2] = tmpdH_sp
120
121                        forces = self.atoms.get_forces()
122                        world.barrier()
123                        if world.rank == 0:
124                            vd = open(self.name + name, 'wb')
125                            fd = open('vib' + name, 'wb')
126                            pickle.dump((Vt_G, alldH_asp), vd)
127                            pickle.dump(forces, fd)
128                            vd.close()
129                            fd.close()
130                        world.barrier()
131                        self.atoms.positions[a, j] = p[a, j]
132        self.atoms.set_positions(p)
133
134    def get_gradient(self):
135        """Calculates gradient"""
136        nx = len(self.indices) * 3
137        veqt_G, dHeq_asp = pickle.load(open(self.name + '.eq.pckl', 'rb'))
138        gpts = veqt_G.shape
139        dvt_Gx = np.zeros(gpts + (nx, ))
140        ddH_aspx = {}
141        for a, dH_sp in dHeq_asp.items():
142            ddH_aspx[a] = np.empty(dH_sp.shape + (nx,))
143
144        x = 0
145        for a in self.indices:
146            for i in range(3):
147                name = '%s.%d%s' % (self.name, a, 'xyz'[i])
148                vtm_G, dHm_asp = pickle.load(open(name + '-.pckl', 'rb'))
149                vtp_G, dHp_asp = pickle.load(open(name + '+.pckl', 'rb'))
150
151                if self.nfree == 4:
152                    vtmm_G, dHmm_asp = pickle.load(
153                        open(name + '--.pckl', 'rb'))
154                    vtpp_G, dHpp_asp = pickle.load(
155                        open(name + '++.pckl', 'rb'))
156                    dvtdx_G = (-vtpp_G + 8.0 * vtp_G -
157                               8 * vtm_G + vtmm_G) / (12 * self.delta / Bohr)
158                    dvt_Gx[:, :, :, x] = dvtdx_G
159                    for atom, ddH_spx in ddH_aspx.items():
160                        ddH_aspx[atom][:, :, x] = (
161                            -dHpp_asp[atom]
162                            + 8.0 * dHp_asp[atom]
163                            - 8.0 * dHm_asp[atom]
164                            + dHmm_asp[atom]) / (12 * self.delta / Bohr)
165                else:  # nfree = 2
166                    dvtdx_G = (vtp_G - vtm_G) / (2 * self.delta / Bohr)
167                    dvt_Gx[..., x] = dvtdx_G
168                    for atom, ddH_spx in ddH_aspx.items():
169                        ddH_aspx[atom][:, :, x] = (
170                            dHp_asp[atom]
171                            - dHm_asp[atom]) / (2 * self.delta / Bohr)
172                x += 1
173        return dvt_Gx, ddH_aspx
174
175    def get_M(self, modes, log=None, q=0):
176        r"""Calculate el-ph coupling matrix for given modes(s).
177
178        XXX:
179        kwarg "q=0" is an ugly hack for k-points.
180        There shuold be loops over q!
181
182        Note that modes must be given as a dictionary with mode
183        frequencies in eV and corresponding mode vectors in units
184        of 1/sqrt(amu), where amu = 1.6605402e-27 Kg is an atomic mass unit.
185        In short frequencies and mode vectors must be given in ase units.
186
187        ::
188
189                  d                   d  ~
190            < w | -- v | w' > = < w | -- v | w'>
191                  dP                  dP
192
193                               _
194                              \        ~a     d   .       ~a
195                            +  ) < w | p  >   -- /_\H   < p | w' >
196                              /_        i     dP     ij    j
197                              a,ij
198
199                               _
200                              \        d  ~a     .        ~a
201                            +  ) < w | -- p  >  /_\H    < p | w' >
202                              /_       dP  i        ij     j
203                              a,ij
204
205                               _
206                              \        ~a     .        d  ~a
207                            +  ) < w | p  >  /_\H    < -- p  | w' >
208                              /_        i        ij    dP  j
209                              a,ij
210
211        """
212        if log is None:
213            timer = nulltimer
214        elif log == '-':
215            timer = StepTimer(name='EPCM')
216        else:
217            timer = StepTimer(name='EPCM', out=open(log, 'w'))
218
219        modes1 = modes.copy()
220        # convert to atomic units
221        amu = 1.6605402e-27  # atomic unit mass [Kg]
222        me = 9.1093897e-31  # electron mass    [Kg]
223        modes = {}
224        for k in modes1.keys():
225            modes[k / Ha] = modes1[k] / np.sqrt(amu / me)
226
227        dvt_Gx, ddH_aspx = self.get_gradient()
228
229        from gpaw import restart
230        atoms, calc = restart('eq.gpw', txt=None)
231        if calc.wfs.S_qMM is None:
232            calc.initialize(atoms)
233            calc.initialize_positions(atoms)
234
235        wfs = calc.wfs
236        nao = wfs.setups.nao
237        bfs = wfs.basis_functions
238        dtype = wfs.dtype
239        spin = 0  # XXX
240
241        M_lii = {}
242        timer.write_now('Starting gradient of pseudo part')
243        for f, mode in modes.items():
244            mo = []
245            M_ii = np.zeros((nao, nao), dtype)
246            for a in self.indices:
247                mo.append(mode[a])
248            mode = np.asarray(mo).flatten()
249            dvtdP_G = np.dot(dvt_Gx, mode)
250            bfs.calculate_potential_matrix(dvtdP_G, M_ii, q=q)
251            tri2full(M_ii, 'L')
252            M_lii[f] = M_ii
253        timer.write_now('Finished gradient of pseudo part')
254
255        P_aqMi = calc.wfs.P_aqMi
256        # Add the term
257        #  _
258        # \        ~a     d   .       ~a
259        #  ) < w | p  >   -- /_\H   < p | w' >
260        # /_        i     dP     ij    j
261        # a,ij
262
263        Ma_lii = {}
264        for f, mode in modes.items():
265            Ma_lii[f] = np.zeros_like(M_lii.values()[0])
266
267        timer.write_now('Starting gradient of dH^a part')
268        for f, mode in modes.items():
269            mo = []
270            for a in self.indices:
271                mo.append(mode[a])
272            mode = np.asarray(mo).flatten()
273
274            for a, ddH_spx in ddH_aspx.items():
275                ddHdP_sp = np.dot(ddH_spx, mode)
276                ddHdP_ii = unpack2(ddHdP_sp[spin])
277                Ma_lii[f] += dots(P_aqMi[a][q], ddHdP_ii, P_aqMi[a][q].T)
278        timer.write_now('Finished gradient of dH^a part')
279
280        timer.write_now('Starting gradient of projectors part')
281        dP_aMix = self.get_dP_aMix(calc.spos_ac, wfs, q, timer)
282        timer.write_now('Finished gradient of projectors part')
283
284        dH_asp = pickle.load(open('v.eq.pckl', 'rb'))[1]
285
286        Mb_lii = {}
287        for f, mode in modes.items():
288            Mb_lii[f] = np.zeros_like(M_lii.values()[0])
289
290        for f, mode in modes.items():
291            for a, dP_Mix in dP_aMix.items():
292                dPdP_Mi = np.dot(dP_Mix, mode[a])
293                dH_ii = unpack2(dH_asp[a][spin])
294                dPdP_MM = dots(dPdP_Mi, dH_ii, P_aqMi[a][q].T)
295                Mb_lii[f] -= dPdP_MM + dPdP_MM.T
296                # XXX The minus sign here is quite subtle.
297                # It is related to how the derivative of projector
298                # functions in GPAW is calculated.
299                # More thorough explanations, anyone...?
300
301        # Units of M_lii are Hartree/(Bohr * sqrt(m_e))
302        for mode in M_lii.keys():
303            M_lii[mode] += Ma_lii[mode] + Mb_lii[mode]
304
305        # conversion to eV. The prefactor 1 / sqrt(hb^2 / 2 * hb * f)
306        # has units Bohr * sqrt(me)
307        M_lii_1 = M_lii.copy()
308        M_lii = {}
309
310        for f in M_lii_1.keys():
311            M_lii[f * Ha] = M_lii_1[f] * Ha / np.sqrt(2 * f)
312
313        return M_lii
314
315
316#####################################################
317# XXX grid and grid 2 sometimes gives random numbers,
318# XXX sometimes even nan!
319#####################################################
320
321def get_grid_dP_aMix(spos_ac, wfs, q, timer=nulltimer):  # XXXXXX q
322    nao = wfs.setups.nao
323    C_MM = np.identity(nao, dtype=wfs.dtype)
324    # XXX In the future use the New Two-Center integrals
325    # to evaluate this
326    dP_aMix = {}
327    for a, setup in enumerate(wfs.setups):
328        ni = 0
329        dP_Mix = np.zeros((nao, setup.ni, 3))
330        pt = LFC(wfs.gd, [setup.pt_j],
331                 wfs.kd.comm, dtype=wfs.dtype, forces=True)
332        spos1_ac = [spos_ac[a]]
333        pt.set_k_points(wfs.ibzk_qc)
334        pt.set_positions(spos1_ac)
335        for b, setup_b in enumerate(wfs.setups):
336            nao = setup_b.nao
337            phi_MG = wfs.gd.zeros(nao, wfs.dtype)
338            phi_MG = wfs.gd.collect(phi_MG, broadcast=False)
339            wfs.basis_functions.lcao_to_grid(C_MM[ni:ni + nao], phi_MG, q)
340            dP_bMix = pt.dict(len(phi_MG), derivative=True)
341            pt.derivative(phi_MG, dP_bMix, q=q)
342            dP_Mix[ni:ni + nao] = dP_bMix[0]
343            ni += nao
344            timer.write_now('projector grad. doing atoms (%s, %s) ' %
345                            (a, b))
346
347        dP_aMix[a] = dP_Mix
348    return dP_aMix
349
350
351def get_grid2_dP_aMix(spos_ac, wfs, q, *args, **kwargs):  # XXXXXX q
352    nao = wfs.setups.nao
353    C_MM = np.identity(nao, dtype=wfs.dtype)
354    bfs = wfs.basis_functions
355    phi_MG = wfs.gd.zeros(nao, wfs.dtype)
356    bfs.lcao_to_grid(C_MM, phi_MG, q)
357    setups = wfs.setups
358    pt = LFC(wfs.gd, [setup.pt_j for setup in setups],
359             wfs.kd.comm, dtype=wfs.dtype, forces=True)
360    pt.set_k_points(wfs.ibzk_qc)
361    pt.set_positions(spos_ac)
362    dP_aMix = pt.dict(len(phi_MG), derivative=True)
363    pt.derivative(phi_MG, dP_aMix, q=q)
364    return dP_aMix
365
366
367def get_tci_dP_aMix(spos_ac, wfs, q, *args, **kwargs):
368    # container for spline expansions of basis function-projector pairs
369    # (note to self: remember to conjugate/negate because of that)
370    from gpaw.lcao.overlap import ManySiteDictionaryWrapper,\
371        TwoCenterIntegralCalculator, NewTwoCenterIntegrals
372
373    if not isinstance(wfs.tci, NewTwoCenterIntegrals):
374        raise RuntimeError('Please remember --gpaw=usenewtci=True')
375
376    dP_aqxMi = {}
377    nao = wfs.setups.nao
378    nq = len(wfs.ibzk_qc)
379    for a, setup in enumerate(wfs.setups):
380        dP_aqxMi[a] = np.zeros((nq, 3, nao, setup.ni), wfs.dtype)
381
382    calc = TwoCenterIntegralCalculator(wfs.ibzk_qc, derivative=True)
383    expansions = ManySiteDictionaryWrapper(wfs.tci.P_expansions, dP_aqxMi)
384    calc.calculate(wfs.tci.atompairs, [expansions], [dP_aqxMi])
385
386    dP_aMix = {}
387    for a in dP_aqxMi:
388        dP_aMix[a] = dP_aqxMi[a].transpose(0, 2, 3, 1).copy()[q]  # XXX q
389    return dP_aMix
390