1#!/usr/bin/env python
2# Copyright 2014-2018,2021 The PySCF Developers. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16# Author: Timothy Berkelbach <tim.berkelbach@gmail.com>
17#
18
19'''PP with numeric integration.  See also pyscf/pbc/gto/pesudo/pp_int.py
20
21For GTH/HGH PPs, see:
22    Goedecker, Teter, Hutter, PRB 54, 1703 (1996)
23    Hartwigsen, Goedecker, and Hutter, PRB 58, 3641 (1998)
24'''
25
26import numpy as np
27import scipy.linalg
28import scipy.special
29from pyscf import lib
30from pyscf.gto import mole
31from pyscf.pbc.gto.pseudo import pp_int
32
33def get_alphas(cell):
34    '''alpha parameters from the non-divergent Hartree+Vloc G=0 term.
35
36    See ewald.pdf
37
38    Returns:
39        alphas : (natm,) ndarray
40    '''
41    return get_alphas_gth(cell)
42
43def get_alphas_gth(cell):
44    '''alpha parameters for the local GTH pseudopotential.'''
45    G0 = np.zeros((1,3))
46    return -get_gth_vlocG(cell, G0)
47
48def get_vlocG(cell, Gv=None):
49    '''Local PP kernel in G space: Vloc(G)
50
51    Returns:
52        (natm, ngrids) ndarray
53    '''
54    if Gv is None: Gv = cell.Gv
55    vlocG = get_gth_vlocG(cell, Gv)
56    return vlocG
57
58def get_gth_vlocG(cell, Gv):
59    '''Local part of the GTH pseudopotential.
60
61    See MH (4.79).
62
63    Args:
64        Gv : (ngrids,3) ndarray
65
66    Returns:
67         (natm, ngrids) ndarray
68    '''
69    vlocG = pp_int.get_gth_vlocG_part1(cell, Gv)
70
71    # Add the C1, C2, C3, C4 contributions
72    G2 = np.einsum('ix,ix->i', Gv, Gv)
73    for ia in range(cell.natm):
74        symb = cell.atom_symbol(ia)
75        if symb not in cell._pseudo:
76            continue
77
78        pp = cell._pseudo[symb]
79        rloc, nexp, cexp = pp[1:3+1]
80
81        G2_red = G2 * rloc**2
82        cfacs = 0
83        if nexp >= 1:
84            cfacs += cexp[0]
85        if nexp >= 2:
86            cfacs += cexp[1] * (3 - G2_red)
87        if nexp >= 3:
88            cfacs += cexp[2] * (15 - 10*G2_red + G2_red**2)
89        if nexp >= 4:
90            cfacs += cexp[3] * (105 - 105*G2_red + 21*G2_red**2 - G2_red**3)
91
92        vlocG[ia,:] -= (2*np.pi)**(3/2.)*rloc**3*np.exp(-0.5*G2_red) * cfacs
93
94    return vlocG
95
96def get_projG(cell, kpt=np.zeros(3)):
97    '''PP weight and projector for the nonlocal PP in G space.
98
99    Returns:
100        hs : list( list( np.array( , ) ) )
101         - hs[atm][l][i,j]
102        projs : list( list( list( list( np.array(ngrids) ) ) ) )
103         - projs[atm][l][m][i][ngrids]
104    '''
105    return get_gth_projG(cell, kpt+cell.Gv)
106
107def get_gth_projG(cell, Gvs):
108    r'''G space projectors from the FT of the real-space projectors.
109
110    \int e^{iGr} p_j^l(r) Y_{lm}^*(theta,phi)
111    = i^l p_j^l(G) Y_{lm}^*(thetaG, phiG)
112
113    See MH Eq.(4.80)
114    '''
115    Gs,thetas,phis = cart2polar(Gvs)
116
117    hs = []
118    projs = []
119    for ia in range(cell.natm):
120        symb = cell.atom_symbol(ia)
121        pp = cell._pseudo[symb]
122        # nproj_types = pp[4]
123        h_ia = []
124        proj_ia = []
125        for l,proj in enumerate(pp[5:]):
126            rl, nl, hl = proj
127            h_ia.append( np.array(hl) )
128            proj_ia_l = []
129            for m in range(-l,l+1):
130                projG_ang = Ylm(l,m,thetas,phis).conj()
131                proj_ia_lm = []
132                for i in range(nl):
133                    projG_radial = projG_li(Gs,l,i,rl)
134                    proj_ia_lm.append( (1j)**l * projG_radial*projG_ang )
135                proj_ia_l.append(proj_ia_lm)
136            proj_ia.append(proj_ia_l)
137        hs.append(h_ia)
138        projs.append(proj_ia)
139
140    return hs, projs
141
142def projG_li(G, l, i, rl):
143    G = np.array(G)
144    G_red = G*rl
145
146    # MH Eq. (4.81)
147    return (_qli(G_red,l,i) * np.pi**(5/4.) * G**l * np.sqrt(rl**(2*l+3))
148            / np.exp(0.5*G_red**2) )
149
150def _qli(x,l,i):
151    # MH Eqs. (4.82)-(4.93) :: beware typos!
152    # Mathematica formulas:
153    # p[l_, i_, r_] = Sqrt[2] r^(l + 2 (i - 1)) Exp[-r^2/(2 R^2)]/(R^(l + (4 i - 1)/2) Sqrt[Gamma[l + (4 i - 1)/2]])
154    # pG[l_, i_, G_] = Integrate[p[l, i, r] 4 Pi r^2 SphericalBesselJ[l, G r], {r, 0, Infinity}]
155    # qG[l_, i_, G_] := pG[l, i, G]/(Pi^(5/4) G^l Sqrt[R^(2 l + 3)]/Exp[(G R)^2/2])
156    # FullSimplify[qG[4, 3, G], R > 0 && G > 0]
157    sqrt = np.sqrt
158    if l==0 and i==0:
159        return 4*sqrt(2.)
160    elif l==0 and i==1:
161        return 8*sqrt(2/15.)*(3-x**2) # MH & GTH (right)
162        #return sqrt(8*2/15.)*(3-x**2) # HGH (wrong)
163    elif l==0 and i==2:
164        #return 16/3.*sqrt(2/105.)*(15-20*x**2+4*x**4) # MH (wrong)
165        return 16/3.*sqrt(2/105.)*(15-10*x**2+x**4) # HGH (right)
166    elif l==1 and i==0:
167        return 8*sqrt(1/3.)
168    elif l==1 and i==1:
169        return 16*sqrt(1/105.)*(5-x**2)
170    elif l==1 and i==2:
171        #return 32/3.*sqrt(1/1155.)*(35-28*x**2+4*x**4) # MH (wrong)
172        return 32/3.*sqrt(1/1155.)*(35-14*x**2+x**4) # HGH (right)
173    elif l==2 and i==0:
174        return 8*sqrt(2/15.)
175    elif l==2 and i==1:
176        return 16/3.*sqrt(2/105.)*(7-x**2)
177    elif l==2 and i==2:
178        #return 32/3.*sqrt(2/15015.)*(63-36*x**2+4*x**4) # MH (wrong I think)
179        return 32/3.*sqrt(2/15015.)*(63-18*x**2+x**4) # TCB
180    elif l==3 and i==0:
181        return 16*sqrt(1/105.)
182    elif l==3 and i==1:
183        return 32/3.*sqrt(1/1155.)*(9-x**2)
184    elif l==3 and i==2:
185        return 64/45.*sqrt(1/1001.)*(99-22*x**2+x**4)
186    elif l==4 and i==0:
187        return 16/3.*sqrt(2/105.)
188    elif l==4 and i==1:
189        return 32/3.*sqrt(2/15015.)*(11-x**2)
190    elif l==4 and i==2:
191        return 64/45.*sqrt(2/17017.)*(143-26*x**2+x**4)
192    else:
193        print("*** WARNING *** l =", l, ", i =", i, "not yet implemented for NL PP!")
194        return 0.
195
196def Ylm_real(l,m,theta,phi):
197    '''Real spherical harmonics, if desired.'''
198    Ylabsm = Ylm(l,np.abs(m),theta,phi)
199    if m < 0:
200        return np.sqrt(2.) * Ylabsm.imag
201    elif m > 0:
202        return np.sqrt(2.) * Ylabsm.real
203    else: # m == 0
204        return Ylabsm.real
205
206def Ylm(l,m,theta,phi):
207    '''
208    Spherical harmonics; returns a complex number
209
210    Note the "convention" for theta and phi:
211    http://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.special.sph_harm.html
212    '''
213    #return scipy.special.sph_harm(m=m,n=l,theta=phi,phi=theta)
214    return scipy.special.sph_harm(m,l,phi,theta)
215
216def cart2polar(rvec):
217    # The rows of rvec are the 3-component vectors
218    # i.e. rvec is N x 3
219    x,y,z = rvec.T
220    r = lib.norm(rvec, axis=1)
221    # theta is the polar angle, 0 < theta < pi
222    # catch possible 0/0
223    theta = np.arccos(z/(r+1e-200))
224    # phi is the azimuthal angle, 0 < phi < 2pi (or -pi < phi < pi)
225    phi = np.arctan2(y,x)
226    return r, theta, phi
227
228
229def get_pp(cell, kpt=np.zeros(3)):
230    '''Get the periodic pseudotential nuc-el AO matrix
231    '''
232    from pyscf.pbc import tools
233    coords = cell.get_uniform_grids()
234    aoR = cell.pbc_eval_gto('GTOval', coords, kpt=kpt)
235    nao = cell.nao_nr()
236
237    SI = cell.get_SI()
238    vlocG = get_vlocG(cell)
239    vpplocG = -np.sum(SI * vlocG, axis=0)
240    vpplocG[0] = np.sum(get_alphas(cell)) # from get_jvloc_G0 function
241
242    # vpploc evaluated in real-space
243    vpplocR = tools.ifft(vpplocG, cell.mesh).real
244    vpploc = np.dot(aoR.T.conj(), vpplocR.reshape(-1,1)*aoR)
245
246    # vppnonloc evaluated in reciprocal space
247    aokG = tools.fftk(np.asarray(aoR.T, order='C'),
248                      cell.mesh, np.exp(-1j*np.dot(coords, kpt))).T
249    ngrids = len(aokG)
250
251    fakemol = mole.Mole()
252    fakemol._atm = np.zeros((1,mole.ATM_SLOTS), dtype=np.int32)
253    fakemol._bas = np.zeros((1,mole.BAS_SLOTS), dtype=np.int32)
254    ptr = mole.PTR_ENV_START
255    fakemol._env = np.zeros(ptr+10)
256    fakemol._bas[0,mole.NPRIM_OF ] = 1
257    fakemol._bas[0,mole.NCTR_OF  ] = 1
258    fakemol._bas[0,mole.PTR_EXP  ] = ptr+3
259    fakemol._bas[0,mole.PTR_COEFF] = ptr+4
260    Gv = np.asarray(cell.Gv+kpt)
261    G_rad = lib.norm(Gv, axis=1)
262
263    vppnl = np.zeros((nao,nao), dtype=np.complex128)
264    for ia in range(cell.natm):
265        symb = cell.atom_symbol(ia)
266        if symb not in cell._pseudo:
267            continue
268        pp = cell._pseudo[symb]
269        for l, proj in enumerate(pp[5:]):
270            rl, nl, hl = proj
271            if nl > 0:
272                hl = np.asarray(hl)
273                fakemol._bas[0,mole.ANG_OF] = l
274                fakemol._env[ptr+3] = .5*rl**2
275                fakemol._env[ptr+4] = rl**(l+1.5)*np.pi**1.25
276                pYlm_part = fakemol.eval_gto('GTOval', Gv)
277
278                pYlm = np.empty((nl,l*2+1,ngrids))
279                for k in range(nl):
280                    qkl = _qli(G_rad*rl, l, k)
281                    pYlm[k] = pYlm_part.T * qkl
282                # pYlm is real
283                SPG_lmi = np.einsum('g,nmg->nmg', SI[ia].conj(), pYlm)
284                SPG_lm_aoG = np.einsum('nmg,gp->nmp', SPG_lmi, aokG)
285                tmp = np.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
286                vppnl += np.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp)
287    vppnl *= (1./ngrids**2)
288
289    if aoR.dtype == np.double:
290        return vpploc.real + vppnl.real
291    else:
292        return vpploc + vppnl
293
294
295def get_jvloc_G0(cell, kpt=np.zeros(3)):
296    '''Get the (separately divergent) Hartree + Vloc G=0 contribution.
297    '''
298    ovlp = cell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpt)
299    return 1./cell.vol * np.sum(get_alphas(cell)) * ovlp
300