1# Copyright (C) 2003  CAMP
2# Please see the accompanying LICENSE file for further information.
3
4from math import pi, sqrt
5
6import numpy as np
7
8from gpaw.gaunt import gaunt
9# load points and weights for the angular integration
10from gpaw.sphere.lebedev import R_nv, Y_nL, weight_n
11from gpaw.spherical_harmonics import nablarlYL
12
13r"""
14                           3
15             __   dn       __   __    dY
16   __  2    \       L  2  \    \        L  2
17  (\/n) = (  ) Y  --- )  + ) (  )  n  --- )
18            /__ L dr      /__  /__  L dr
19                                        v
20             L            v=1    L
21
22
23        dY
24          L
25  A   = --- r
26   Lv   dr
27          v
28
29"""
30# A_nvL is defined as above, n is an expansion point index (50 Lebedev points).
31rnablaY_nLv = np.empty((len(R_nv), 25, 3))
32for rnablaY_Lv, Y_L, R_v in zip(rnablaY_nLv, Y_nL, R_nv):
33    for l in range(5):
34        for L in range(l**2, (l + 1)**2):
35            rnablaY_Lv[L] = nablarlYL(L, R_v) - l * R_v * Y_L[L]
36
37
38class PAWXCCorrection:
39    def __init__(self,
40                 w_jg,    # all-lectron partial waves
41                 wt_jg,   # pseudo partial waves
42                 nc_g,    # core density
43                 nct_g,   # smooth core density
44                 rgd,     # radial grid descriptor
45                 jl,      # ?
46                 lmax,    # maximal angular momentum to consider
47                 e_xc0,    # xc energy of reference atom
48                 phicorehole_g,  # ?
49                 fcorehole,      # ?
50                 tauc_g,  # kinetic core energy array
51                 tauct_g):  # pseudo kinetic core energy array
52        self.e_xc0 = e_xc0
53        self.Lmax = (lmax + 1)**2
54        self.rgd = rgd
55        self.dv_g = rgd.dv_g
56        self.Y_nL = Y_nL[:, :self.Lmax]
57        self.rnablaY_nLv = rnablaY_nLv[:, :self.Lmax]
58        self.ng = ng = len(nc_g)
59        self.phi_jg = w_jg
60        self.phit_jg = wt_jg
61
62        self.jlL = [(j, l, l**2 + m) for j, l in jl for m in range(2 * l + 1)]
63        self.ni = ni = len(self.jlL)
64        self.nj = nj = len(jl)
65        self.nii = nii = ni * (ni + 1) // 2
66        njj = nj * (nj + 1) // 2
67
68        self.tauc_g = tauc_g
69        self.tauct_g = tauct_g
70        self.tau_npg = None
71        self.taut_npg = None
72
73        G_LLL = gaunt(max(l for j, l in jl))[:, :, :self.Lmax]
74        LGcut = G_LLL.shape[2]
75        B_Lqp = np.zeros((self.Lmax, njj, nii))
76        p = 0
77        i1 = 0
78        for j1, l1, L1 in self.jlL:
79            for j2, l2, L2 in self.jlL[i1:]:
80                if j1 < j2:
81                    q = j2 + j1 * nj - j1 * (j1 + 1) // 2
82                else:
83                    q = j1 + j2 * nj - j2 * (j2 + 1) // 2
84                B_Lqp[:LGcut, q, p] = G_LLL[L1, L2]
85                p += 1
86            i1 += 1
87        self.B_pqL = B_Lqp.T.copy()
88
89        #
90        self.n_qg = np.zeros((njj, ng))
91        self.nt_qg = np.zeros((njj, ng))
92        q = 0
93        for j1, l1 in jl:
94            for j2, l2 in jl[j1:]:
95                self.n_qg[q] = w_jg[j1] * w_jg[j2]
96                self.nt_qg[q] = wt_jg[j1] * wt_jg[j2]
97                q += 1
98
99        self.nc_g = nc_g
100        self.nct_g = nct_g
101
102        if fcorehole != 0.0:
103            self.nc_corehole_g = fcorehole * phicorehole_g**2 / (4 * pi)
104        else:
105            self.nc_corehole_g = None
106
107    def four_phi_integrals(self, D_sp, fxc):
108        """Calculate four-phi integrals.
109
110        The density is given by the density matrix ``D_sp`` in packed(pack)
111        form, and the resulting rank-four tensor is also returned in
112        packed format. ``fxc`` is a radial object???
113        """
114
115        ns, nii = D_sp.shape
116
117        assert ns == 1
118
119        D_p = D_sp[0]
120        D_Lq = np.dot(self.B_pqL.T, D_p)
121
122        # Expand all-electron density in spherical harmonics:
123        n_qg = self.n_qg
124        n_Lg = np.dot(D_Lq, n_qg)
125        n_Lg[0] += self.nc_g * sqrt(4 * pi)
126
127        # Expand pseudo electron density in spherical harmonics:
128        nt_qg = self.nt_qg
129        nt_Lg = np.dot(D_Lq, nt_qg)
130        nt_Lg[0] += self.nct_g * sqrt(4 * pi)
131
132        # Allocate array for result:
133        J_pp = np.zeros((nii, nii))
134
135        # Loop over 50 points on the sphere surface:
136        for w, Y_L in zip(weight_n, self.Y_nL):
137            B_pq = np.dot(self.B_pqL, Y_L)
138
139            fxcdv = fxc(np.dot(Y_L, n_Lg)) * self.dv_g
140            dn2_qq = np.inner(n_qg * fxcdv, n_qg)
141
142            fxctdv = fxc(np.dot(Y_L, nt_Lg)) * self.dv_g
143            dn2_qq -= np.inner(nt_qg * fxctdv, nt_qg)
144
145            J_pp += w * np.dot(B_pq, np.inner(dn2_qq, B_pq))
146
147        return J_pp
148