1"""Omega matrix for functionals with Hartree-Fock exchange.
2
3"""
4from math import sqrt
5
6from ase.units import Hartree
7from ase.utils.timing import Timer
8import numpy as np
9from numpy.linalg import inv
10from scipy.linalg import eigh
11
12from gpaw import debug
13import gpaw.mpi as mpi
14from gpaw.lrtddft.omega_matrix import OmegaMatrix
15from gpaw.pair_density import PairDensity
16from gpaw.helmholtz import HelmholtzSolver
17from gpaw.utilities.blas import gemm
18
19
20class ApmB(OmegaMatrix):
21
22    """Omega matrix for functionals with Hartree-Fock exchange.
23
24    """
25
26    def get_full(self):
27
28        hybrid = ((self.xc is not None) and
29                  hasattr(self.xc, 'hybrid') and
30                  (self.xc.hybrid > 0.0))
31        if self.fullkss.npspins < 2 and hybrid:
32            raise RuntimeError('Does not work spin-unpolarized ' +
33                               'with hybrids (use nspins=2)')
34
35        if hasattr(self.xc, 'rsf') and (self.xc.rsf == 'Yukawa'):
36            self.screened_poissonsolver = HelmholtzSolver(
37                k2=-self.xc.omega**2, eps=1e-11, nn=3)
38            self.screened_poissonsolver.set_grid_descriptor(self.gd)
39        self.paw.timer.start('ApmB RPA')
40        self.ApB = self.Om
41        self.AmB = self.get_rpa()
42        self.paw.timer.stop()
43
44        if self.xc is not None:
45            self.paw.timer.start('ApmB XC')
46            self.get_xc()  # inherited from OmegaMatrix
47            self.paw.timer.stop()
48
49    def get_rpa(self):
50        """Calculate RPA and Hartree-fock part of the A+-B matrices."""
51
52        # shorthands
53        kss = self.fullkss
54        finegrid = self.finegrid
55        yukawa = hasattr(self.xc, 'rsf') and (self.xc.rsf == 'Yukawa')
56
57        # calculate omega matrix
58        nij = len(kss)
59        self.log('RPAhyb', nij, 'transitions')
60
61        AmB = np.zeros((nij, nij))
62        ApB = self.ApB
63
64        # storage place for Coulomb integrals
65        integrals = {}
66        if yukawa:
67            rsf_integrals = {}
68        # setup things for IVOs
69        if (hasattr(self.xc, 'excitation') and
70            (self.xc.excitation is not None or self.xc.excited != 0)):
71            sin_tri_weight = 1
72            if self.xc.excitation is not None:
73                ex_type = self.xc.excitation.lower()
74                if ex_type == 'singlet':
75                    sin_tri_weight = 2
76                elif ex_type == 'triplet':
77                    sin_tri_weight = 0
78            homo = int(self.paw.get_number_of_electrons() // 2)
79            ivo_l = homo - self.xc.excited - 1
80        else:
81            ivo_l = None
82
83        for ij in range(nij):
84            self.log('RPAhyb kss[' + '%d' % ij + ']=', kss[ij])
85
86            timer = Timer()
87            timer.start('init')
88            timer2 = Timer()
89
90            # smooth density including compensation charges
91            timer2.start('with_compensation_charges 0')
92            rhot_p = kss[ij].with_compensation_charges(
93                finegrid != 0)
94            timer2.stop()
95
96            # integrate with 1/|r_1-r_2|
97            timer2.start('poisson')
98            phit_p = np.zeros(rhot_p.shape, rhot_p.dtype)
99            self.poisson.solve(phit_p, rhot_p, charge=None)
100            timer2.stop()
101
102            timer.stop()
103            t0 = timer.get_time('init')
104            timer.start(ij)
105
106            if finegrid == 1:
107                rhot = kss[ij].with_compensation_charges()
108                phit = self.gd.zeros()
109                self.restrict(phit_p, phit)
110            else:
111                phit = phit_p
112                rhot = rhot_p
113
114            for kq in range(ij, nij):
115                if kq != ij:
116                    # smooth density including compensation charges
117                    timer2.start('kq with_compensation_charges')
118                    rhot = kss[kq].with_compensation_charges(
119                        finegrid == 2)
120                    timer2.stop()
121                pre = self.weight_Kijkq(ij, kq)
122
123                timer2.start('integrate')
124                I = self.Coulomb_integral_kss(kss[ij], kss[kq], phit, rhot)
125                if kss[ij].spin == kss[kq].spin:
126                    name = self.Coulomb_integral_name(kss[ij].i, kss[ij].j,
127                                                      kss[kq].i, kss[kq].j,
128                                                      kss[ij].spin)
129                    integrals[name] = I
130                ApB[ij, kq] = pre * I
131                timer2.stop()
132
133                if ij == kq:
134                    epsij = kss[ij].get_energy() / kss[ij].get_weight()
135                    AmB[ij, kq] += epsij
136                    ApB[ij, kq] += epsij
137
138            timer.stop()
139# timer2.write()
140            if ij < (nij - 1):
141                self.log('RPAhyb estimated time left',
142                         self.time_left(timer, t0, ij, nij))
143
144        # add HF parts and apply symmetry
145        if hasattr(self.xc, 'hybrid'):
146            weight = self.xc.hybrid
147        else:
148            weight = 0.0
149        for ij in range(nij):
150            self.log('HF kss[' + '%d' % ij + ']')
151            timer = Timer()
152            timer.start('init')
153            timer.stop()
154            t0 = timer.get_time('init')
155            timer.start(ij)
156
157            i = kss[ij].i
158            j = kss[ij].j
159            s = kss[ij].spin
160            for kq in range(ij, nij):
161                if kss[ij].pspin == kss[kq].pspin:
162                    k = kss[kq].i
163                    q = kss[kq].j
164                    ikjq = self.Coulomb_integral_ijkq(i, k, j, q, s, integrals)
165                    iqkj = self.Coulomb_integral_ijkq(i, q, k, j, s, integrals)
166                    if yukawa:  # Yukawa integrals might be caches
167                        ikjq -= self.Coulomb_integral_ijkq(
168                            i, k, j, q, s, rsf_integrals, yukawa)
169                        iqkj -= self.Coulomb_integral_ijkq(
170                            i, q, k, j, s, rsf_integrals, yukawa)
171                    ApB[ij, kq] -= weight * (ikjq + iqkj)
172                    AmB[ij, kq] -= weight * (ikjq - iqkj)
173
174                ApB[kq, ij] = ApB[ij, kq]
175                AmB[kq, ij] = AmB[ij, kq]
176
177            timer.stop()
178            if ij < (nij - 1):
179                self.log('HF estimated time left',
180                         self.time_left(timer, t0, ij, nij))
181
182        if ivo_l is not None:
183            # IVO RPA after Berman, Kaldor, Chem. Phys. 43 (3) 1979
184            # doi: 10.1016/0301-0104(79)85205-2
185            l = ivo_l
186            for ij in range(nij):
187                i = kss[ij].i
188                j = kss[ij].j
189                s = kss[ij].spin
190                for kq in range(ij, nij):
191                    if kss[kq].i == i and kss[ij].pspin == kss[kq].pspin:
192                        k = kss[kq].i
193                        q = kss[kq].j
194                        jqll = self.Coulomb_integral_ijkq(j, q, l, l, s,
195                                                          integrals)
196                        jllq = self.Coulomb_integral_ijkq(j, l, l, q, s,
197                                                          integrals)
198                        if yukawa:
199                            jqll -= self.Coulomb_integral_ijkq(j, q, l, l, s,
200                                                               rsf_integrals,
201                                                               yukawa)
202                            jllq -= self.Coulomb_integral_ijkq(j, l, l, q, s,
203                                                               rsf_integrals,
204                                                               yukawa)
205                        jllq *= sin_tri_weight
206                        ApB[ij, kq] += weight * (jqll - jllq)
207                        AmB[ij, kq] += weight * (jqll - jllq)
208                        ApB[kq, ij] = ApB[ij, kq]
209                        AmB[kq, ij] = AmB[ij, kq]
210        return AmB
211
212    def Coulomb_integral_name(self, i, j, k, l, spin):
213        """return a unique name considering the Coulomb integral
214        symmetry"""
215        def ij_name(i, j):
216            return str(max(i, j)) + ' ' + str(min(i, j))
217
218        # maximal gives the first
219        if max(i, j) >= max(k, l):
220            base = ij_name(i, j) + ' ' + ij_name(k, l)
221        else:
222            base = ij_name(k, l) + ' ' + ij_name(i, j)
223        return base + ' ' + str(spin)
224
225    def Coulomb_integral_ijkq(self, i, j, k, q, spin, integrals,
226                              yukawa=False):
227        name = self.Coulomb_integral_name(i, j, k, q, spin)
228        if name in integrals:
229            return integrals[name]
230
231        # create the Kohn-Sham singles
232        kss_ij = PairDensity(self.paw)
233        kss_ij.initialize(self.paw.wfs.kpt_u[spin], i, j)
234        kss_kq = PairDensity(self.paw)
235        kss_kq.initialize(self.paw.wfs.kpt_u[spin], k, q)
236
237        rhot_p = kss_ij.with_compensation_charges(
238            self.finegrid != 0)
239        phit_p = np.zeros(rhot_p.shape, rhot_p.dtype)
240        if yukawa:
241            self.screened_poissonsolver.solve(phit_p, rhot_p, charge=None)
242        else:
243            self.poisson.solve(phit_p, rhot_p, charge=None)
244
245        if self.finegrid == 1:
246            phit = self.gd.zeros()
247            self.restrict(phit_p, phit)
248        else:
249            phit = phit_p
250
251        rhot = kss_kq.with_compensation_charges(
252            self.finegrid == 2)
253
254        integrals[name] = self.Coulomb_integral_kss(kss_ij, kss_kq,
255                                                    phit, rhot,
256                                                    yukawa=yukawa)
257        return integrals[name]
258
259    def timestring(self, t):
260        ti = int(t + .5)
261        td = int(ti // 86400)
262        st = ''
263        if td > 0:
264            st += '%d' % td + 'd'
265            ti -= td * 86400
266        th = int(ti // 3600)
267        if th > 0:
268            st += '%d' % th + 'h'
269            ti -= th * 3600
270        tm = int(ti // 60)
271        if tm > 0:
272            st += '%d' % tm + 'm'
273            ti -= tm * 60
274        st += '%d' % ti + 's'
275        return st
276
277    def mapAB(self, restrict={}):
278        """Map A+B, A-B matrices according to constraints."""
279
280        map, self.kss = self.get_map(restrict)
281        if map is None:
282            ApB = self.ApB.copy()
283            AmB = self.AmB.copy()
284        else:
285            nij = len(self.kss)
286            ApB = np.empty((nij, nij))
287            AmB = np.empty((nij, nij))
288            for ij in range(nij):
289                for kq in range(nij):
290                    ApB[ij, kq] = self.ApB[map[ij], map[kq]]
291                    AmB[ij, kq] = self.AmB[map[ij], map[kq]]
292
293        return ApB, AmB
294
295    def diagonalize(self, restrict={}, TDA=False):
296        """Evaluate Eigenvectors and Eigenvalues"""
297
298        ApB, AmB = self.mapAB(restrict)
299        nij = len(self.kss)
300
301        if TDA:
302            # Tamm-Dancoff approximation (B=0)
303            eigenvalues, evecs = eigh(0.5 * (ApB + AmB))
304            self.eigenvectors = evecs.T
305            self.eigenvalues = eigenvalues ** 2
306        else:
307            # the occupation matrix
308            C = np.empty((nij,))
309            for ij in range(nij):
310                C[ij] = 1. / self.kss[ij].fij
311
312            S = C * inv(AmB) * C
313            S = sqrt_matrix(inv(S).copy())
314
315            # get Omega matrix
316            M = np.zeros(ApB.shape)
317            gemm(1.0, ApB, S, 0.0, M)
318            self.eigenvectors = np.zeros(ApB.shape)
319            gemm(1.0, S, M, 0.0, self.eigenvectors)
320
321            self.eigenvalues, self.eigenvectors.T[:] = eigh(self.eigenvectors)
322
323    def read(self, filename=None, fh=None):
324        """Read myself from a file"""
325        if mpi.rank == 0:
326            if fh is None:
327                f = open(filename, 'r')
328            else:
329                f = fh
330
331            f.readline()
332            nij = int(f.readline())
333            ApB = np.zeros((nij, nij))
334            for ij in range(nij):
335                l = f.readline().split()
336                for kq in range(ij, nij):
337                    ApB[ij, kq] = float(l[kq - ij])
338                    ApB[kq, ij] = ApB[ij, kq]
339            self.ApB = ApB
340
341            f.readline()
342            nij = int(f.readline())
343            AmB = np.zeros((nij, nij))
344            for ij in range(nij):
345                l = f.readline().split()
346                for kq in range(ij, nij):
347                    AmB[ij, kq] = float(l[kq - ij])
348                    AmB[kq, ij] = AmB[ij, kq]
349            self.AmB = AmB
350
351            if fh is None:
352                f.close()
353
354    def weight_Kijkq(self, ij, kq):
355        """weight for the coupling matrix terms"""
356        return 2.
357
358    def write(self, filename=None, fh=None):
359        """Write current state to a file."""
360        if mpi.rank == 0:
361            if fh is None:
362                f = open(filename, 'w')
363            else:
364                f = fh
365
366            f.write('# A+B\n')
367            nij = len(self.fullkss)
368            f.write('%d\n' % nij)
369            for ij in range(nij):
370                for kq in range(ij, nij):
371                    f.write(' %g' % self.ApB[ij, kq])
372                f.write('\n')
373
374            f.write('# A-B\n')
375            nij = len(self.fullkss)
376            f.write('%d\n' % nij)
377            for ij in range(nij):
378                for kq in range(ij, nij):
379                    f.write(' %g' % self.AmB[ij, kq])
380                f.write('\n')
381
382            if fh is None:
383                f.close()
384
385    def __str__(self):
386        string = '<ApmB> '
387        if hasattr(self, 'eigenvalues'):
388            string += 'dimension ' + ('%d' % len(self.eigenvalues))
389            string += '\neigenvalues: '
390            for ev in self.eigenvalues:
391                string += ' ' + ('%f' % (sqrt(ev) * Hartree))
392        return string
393
394
395def sqrt_matrix(a, preserve=False):
396    """Get the sqrt of a symmetric matrix a (diagonalize is used).
397    The matrix is kept if preserve=True, a=sqrt(a) otherwise."""
398    n = len(a)
399    if debug:
400        assert a.flags.contiguous
401        assert a.dtype == float
402        assert a.shape == (n, n)
403    if preserve:
404        b = a.copy()
405    else:
406        b = a
407
408    # diagonalize to get the form b = Z * D * Z^T
409    # where D is diagonal
410    D = np.empty((n,))
411    D, b.T[:] = eigh(b, lower=True)
412    ZT = b.copy()
413    Z = np.transpose(b)
414
415    # c = Z * sqrt(D)
416    c = Z * np.sqrt(D)
417
418    # sqrt(b) = c * Z^T
419    gemm(1., ZT, np.ascontiguousarray(c), 0., b)
420    return b
421