1from math import pi
2
3import numpy as np
4try:
5    from scipy.special import spherical_jn
6
7    def sphj(n, z):
8        return spherical_jn(range(n), z)
9
10except ImportError:
11    from scipy.special import sph_jn
12
13    def sphj(n, z):
14        return sph_jn(n - 1, z)[0]
15
16
17from gpaw.gaunt import nabla, gaunt
18from gpaw.spherical_harmonics import Y
19
20
21def two_phi_planewave_integrals(k_Gv, setup=None, Gstart=0, Gend=None,
22                                rgd=None, phi_jg=None,
23                                phit_jg=None, l_j=None):
24    r"""Calculate PAW-correction matrix elements with planewaves.
25
26    ::
27
28      /  _       _   ik.r     _     ~   _   ik.r ~   _
29      | dr [phi (r) e    phi (r) - phi (r) e    phi (r)]
30      /        1            2         1            2
31
32                     ll   -  /     2                      ~       ~
33      = 4pi \sum_lm i  Y (k) | dr r  [ phi (r) phi (r) - phi (r) phi (r) j (kr)
34                        lm   /            1       2         1       2     ll
35
36          /
37        * | d\Omega Y     Y     Y
38          /          l1m1  l2m2  lm
39
40    """
41
42    if Gend is None:
43        Gend = len(k_Gv)
44
45    if setup is not None:
46        rgd = setup.rgd
47        l_j = setup.l_j
48        # Obtain the phi_j and phit_j
49        phi_jg = []
50        phit_jg = []
51        rcut2 = 2 * max(setup.rcut_j)
52        gcut2 = rgd.ceil(rcut2)
53        for phi_g, phit_g in zip(setup.data.phi_jg, setup.data.phit_jg):
54            phi_g = phi_g.copy()
55            phit_g = phit_g.copy()
56            phi_g[gcut2:] = phit_g[gcut2:] = 0.
57            phi_jg.append(phi_g)
58            phit_jg.append(phit_g)
59    else:
60        assert rgd is not None
61        assert phi_jg is not None
62        assert l_j is not None
63
64    ng = rgd.N
65    r_g = rgd.r_g
66    dr_g = rgd.dr_g
67
68    # Construct L (l**2 + m) and j (nl) index
69    L_i = []
70    j_i = []
71    for j, l in enumerate(l_j):
72        for m in range(2 * l + 1):
73            L_i.append(l**2 + m)
74            j_i.append(j)
75    ni = len(L_i)
76    nj = len(l_j)
77    lmax = max(l_j) * 2 + 1
78
79    if setup is not None:
80        assert ni == setup.ni and nj == setup.nj
81
82    # Initialize
83    npw = k_Gv.shape[0]
84    R_jj = np.zeros((nj, nj))
85    R_ii = np.zeros((ni, ni))
86    phi_Gii = np.zeros((npw, ni, ni), dtype=complex)
87    j_lg = np.zeros((lmax, ng))
88
89    # Store (phi_j1 * phi_j2 - phit_j1 * phit_j2 ) for further use
90    tmp_jjg = np.zeros((nj, nj, ng))
91    for j1 in range(nj):
92        for j2 in range(nj):
93            tmp_jjg[j1, j2] = (phi_jg[j1] * phi_jg[j2] -
94                               phit_jg[j1] * phit_jg[j2])
95
96    G_LLL = gaunt(max(l_j))
97
98    # Loop over G vectors
99    for iG in range(Gstart, Gend):
100        kk = k_Gv[iG]
101        k = np.sqrt(np.dot(kk, kk))  # calculate length of q+G
102
103        # Calculating spherical bessel function
104        for g, r in enumerate(r_g):
105            j_lg[:, g] = sphj(lmax, k * r)
106
107        for li in range(lmax):
108            # Radial part
109            for j1 in range(nj):
110                for j2 in range(nj):
111                    R_jj[j1, j2] = np.dot(r_g**2 * dr_g,
112                                          tmp_jjg[j1, j2] * j_lg[li])
113
114            for mi in range(2 * li + 1):
115                # Angular part
116                for i1 in range(ni):
117                    L1 = L_i[i1]
118                    j1 = j_i[i1]
119                    for i2 in range(ni):
120                        L2 = L_i[i2]
121                        j2 = j_i[i2]
122                        R_ii[i1, i2] = G_LLL[L1, L2, li**2 + mi] * R_jj[j1, j2]
123
124                if np.allclose(k, 0.):  # Avoid division by zero
125                    k = 1.
126                phi_Gii[iG] += R_ii * Y(li**2 + mi,
127                                        kk[0] / k,
128                                        kk[1] / k,
129                                        kk[2] / k) * (-1j)**li
130
131    phi_Gii *= 4 * pi
132
133    return phi_Gii.reshape(npw, ni * ni)
134
135
136def two_phi_nabla_planewave_integrals(k_Gv, setup=None, Gstart=0, Gend=None,
137                                      rgd=None, phi_jg=None,
138                                      phit_jg=None, l_j=None):
139    """Calculate PAW-correction matrix elements with planewaves and gradient.
140
141    ::
142
143      /  _       _   ik.r d       _     ~   _   ik.r d   ~   _
144      | dr [phi (r) e     -- phi (r) - phi (r) e     -- phi (r)]
145      /        1          dx    2         1          dx    2
146
147    and similar for y and z."""
148
149    if Gend is None:
150        Gend = len(k_Gv)
151
152    if setup is not None:
153        rgd = setup.rgd
154        l_j = setup.l_j
155        # Obtain the phi_j and phit_j
156        phi_jg = []
157        phit_jg = []
158        rcut2 = 2 * max(setup.rcut_j)
159        gcut2 = rgd.ceil(rcut2)
160        for phi_g, phit_g in zip(setup.data.phi_jg, setup.data.phit_jg):
161            phi_g = phi_g.copy()
162            phit_g = phit_g.copy()
163            phi_g[gcut2:] = phit_g[gcut2:] = 0.
164            phi_jg.append(phi_g)
165            phit_jg.append(phit_g)
166    else:
167        assert rgd is not None
168        assert phi_jg is not None
169        assert l_j is not None
170
171    ng = rgd.N
172    r_g = rgd.r_g
173    dr_g = rgd.dr_g
174
175    # Construct L (l**2 + m) and j (nl) index
176    L_i = []
177    j_i = []
178    for j, l in enumerate(l_j):
179        for m in range(2 * l + 1):
180            L_i.append(l**2 + m)
181            j_i.append(j)
182    ni = len(L_i)
183    nj = len(l_j)
184    lmax = max(l_j) * 2 + 1
185
186    ljdef = 3
187    l2max = max(l_j) * (max(l_j) > ljdef) + ljdef * (max(l_j) <= ljdef)
188    G_LLL = gaunt(l2max)
189    Y_LLv = nabla(2 * l2max)
190
191    if setup is not None:
192        assert ni == setup.ni and nj == setup.nj
193
194    # Initialize
195    npw = k_Gv.shape[0]
196    R1_jj = np.zeros((nj, nj))
197    R2_jj = np.zeros((nj, nj))
198    R_vii = np.zeros((3, ni, ni))
199    phi_vGii = np.zeros((3, npw, ni, ni), dtype=complex)
200    j_lg = np.zeros((lmax, ng))
201
202    # Store (phi_j1 * dphidr_j2 - phit_j1 * dphitdr_j2) for further use
203    tmp_jjg = np.zeros((nj, nj, ng))
204    tmpder_jjg = np.zeros((nj, nj, ng))
205    for j1 in range(nj):
206        for j2 in range(nj):
207            dphidr_g = np.empty_like(phi_jg[j2])
208            rgd.derivative_spline(phi_jg[j2], dphidr_g)
209            dphitdr_g = np.empty_like(phit_jg[j2])
210            rgd.derivative_spline(phit_jg[j2], dphitdr_g)
211
212            tmpder_jjg[j1, j2] = (phi_jg[j1] * dphidr_g -
213                                  phit_jg[j1] * dphitdr_g)
214            tmp_jjg[j1, j2] = (phi_jg[j1] * phi_jg[j2] -
215                               phit_jg[j1] * phit_jg[j2])
216
217    # Loop over G vectors
218    for iG in range(Gstart, Gend):
219        kk = k_Gv[iG]
220        k = np.sqrt(np.dot(kk, kk))  # calculate length of q+G
221
222        # Calculating spherical bessel function
223        for g, r in enumerate(r_g):
224            j_lg[:, g] = sphj(lmax, k * r)
225
226        for li in range(lmax):
227            # Radial part
228            for j1 in range(nj):
229                for j2 in range(nj):
230                    R1_jj[j1, j2] = np.dot(r_g**2 * dr_g,
231                                           tmpder_jjg[j1, j2] * j_lg[li])
232                    R2_jj[j1, j2] = np.dot(r_g * dr_g,
233                                           tmp_jjg[j1, j2] * j_lg[li])
234
235            for mi in range(2 * li + 1):
236                if k == 0:
237                    Ytmp = Y(li**2 + mi, 1.0, 0, 0)
238                else:
239                    # Note the spherical bessel gives
240                    # zero when k == 0 for li != 0
241                    Ytmp = Y(li**2 + mi, kk[0] / k, kk[1] / k, kk[2] / k)
242
243                for v in range(3):
244                    Lv = 1 + (v + 2) % 3
245                    # Angular part
246                    for i1 in range(ni):
247                        L1 = L_i[i1]
248                        j1 = j_i[i1]
249                        for i2 in range(ni):
250                            L2 = L_i[i2]
251                            j2 = j_i[i2]
252                            l2 = l_j[j2]
253
254                            R_vii[v, i1, i2] = (
255                                (4 * pi / 3)**0.5 *
256                                np.dot(G_LLL[L1, L2],
257                                       G_LLL[Lv, li**2 + mi]) *
258                                (R1_jj[j1, j2] - l2 *
259                                 R2_jj[j1, j2]))
260
261                            R_vii[v, i1, i2] += (R2_jj[j1, j2] *
262                                                 np.dot(G_LLL[L1, li**2 + mi],
263                                                        Y_LLv[:, L2, v]))
264
265                phi_vGii[:, iG] += (R_vii * Ytmp * (-1j)**li)
266
267    phi_vGii *= 4 * pi
268
269    return phi_vGii.reshape(3, npw, ni * ni)
270