1#!/usr/bin/env python
2# Copyright 2014-2019 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'''
17Spherical harmonics
18'''
19
20import numpy
21import scipy.linalg
22from pyscf.symm.cg import cg_spin
23
24def real_sph_vec(r, lmax, reorder_p=False):
25    '''Computes (all) real spherical harmonics up to the angular momentum lmax'''
26    #:import scipy.special
27    #:ngrid = r.shape[0]
28    #:cosphi = r[:,2]
29    #:sinphi = (1-cosphi**2)**.5
30    #:costheta = numpy.ones(ngrid)
31    #:sintheta = numpy.zeros(ngrid)
32    #:costheta[sinphi!=0] = r[sinphi!=0,0] / sinphi[sinphi!=0]
33    #:sintheta[sinphi!=0] = r[sinphi!=0,1] / sinphi[sinphi!=0]
34    #:costheta[costheta> 1] = 1
35    #:costheta[costheta<-1] =-1
36    #:sintheta[sintheta> 1] = 1
37    #:sintheta[sintheta<-1] =-1
38    #:varphi = numpy.arccos(cosphi)
39    #:theta = numpy.arccos(costheta)
40    #:theta[sintheta<0] = 2*numpy.pi - theta[sintheta<0]
41    #:ylms = []
42    #:for l in range(lmax+1):
43    #:    ylm = numpy.empty((l*2+1,ngrid))
44    #:    ylm[l] = scipy.special.sph_harm(0, l, theta, varphi).real
45    #:    for m in range(1, l+1):
46    #:        f1 = scipy.special.sph_harm(-m, l, theta, varphi)
47    #:        f2 = scipy.special.sph_harm( m, l, theta, varphi)
48    #:        # complex to real spherical functions
49    #:        if m % 2 == 1:
50    #:            ylm[l-m] = (-f1.imag - f2.imag) / numpy.sqrt(2)
51    #:            ylm[l+m] = ( f1.real - f2.real) / numpy.sqrt(2)
52    #:        else:
53    #:            ylm[l-m] = (-f1.imag + f2.imag) / numpy.sqrt(2)
54    #:            ylm[l+m] = ( f1.real + f2.real) / numpy.sqrt(2)
55    #:    ylms.append(ylm)
56    #:return ylms
57
58    # When r is a normalized vector:
59    norm = 1./numpy.linalg.norm(r, axis=1)
60    return multipoles(r*norm.reshape(-1,1), lmax, reorder_p)
61
62
63def multipoles(r, lmax, reorder_dipole=True):
64    '''
65    Compute all multipoles upto lmax
66
67    rad = numpy.linalg.norm(r, axis=1)
68    ylms = real_ylm(r/rad.reshape(-1,1), lmax)
69    pol = [rad**l*y for l, y in enumerate(ylms)]
70
71    Kwargs:
72        reorder_p : bool
73            sort dipole to the order (x,y,z)
74    '''
75    from pyscf import gto
76
77# libcint cart2sph transformation provide the capability to compute
78# multipole directly.  cart2sph function is fast for low angular moment.
79    ngrid = r.shape[0]
80    xs = numpy.ones((lmax+1,ngrid))
81    ys = numpy.ones((lmax+1,ngrid))
82    zs = numpy.ones((lmax+1,ngrid))
83    for i in range(1,lmax+1):
84        xs[i] = xs[i-1] * r[:,0]
85        ys[i] = ys[i-1] * r[:,1]
86        zs[i] = zs[i-1] * r[:,2]
87    ylms = []
88    for l in range(lmax+1):
89        nd = (l+1)*(l+2)//2
90        c = numpy.empty((nd,ngrid))
91        k = 0
92        for lx in reversed(range(0, l+1)):
93            for ly in reversed(range(0, l-lx+1)):
94                lz = l - lx - ly
95                c[k] = xs[lx] * ys[ly] * zs[lz]
96                k += 1
97        ylm = gto.cart2sph(l, c.T).T
98        ylms.append(ylm)
99
100# when call libcint, p functions are ordered as px,py,pz
101# reorder px,py,pz to p(-1),p(0),p(1)
102    if (not reorder_dipole) and lmax >= 1:
103        ylms[1] = ylms[1][[1,2,0]]
104    return ylms
105
106def sph_pure2real(l, reorder_p=True):
107    r'''
108    Transformation matrix: from the pure spherical harmonic functions Y_m to
109    the real spherical harmonic functions O_m.
110          O_m = \sum Y_m' * U(m',m)
111    Y(-1) = 1/\sqrt(2){-iO(-1) + O(1)}; Y(1) = 1/\sqrt(2){-iO(-1) - O(1)}
112    Y(-2) = 1/\sqrt(2){-iO(-2) + O(2)}; Y(2) = 1/\sqrt(2){iO(-2) + O(2)}
113    O(-1) = i/\sqrt(2){Y(-1) + Y(1)};   O(1) = 1/\sqrt(2){Y(-1) - Y(1)}
114    O(-2) = i/\sqrt(2){Y(-2) - Y(2)};   O(2) = 1/\sqrt(2){Y(-2) + Y(2)}
115
116    Kwargs:
117        reorder_p (bool): Whether the p functions are in the (x,y,z) order.
118
119    Returns:
120        2D array U_{complex,real}
121    '''
122    n = 2 * l + 1
123    u = numpy.zeros((n,n), dtype=complex)
124    sqrthfr = numpy.sqrt(.5)
125    sqrthfi = numpy.sqrt(.5)*1j
126
127    if reorder_p and l == 1:
128        u[1,2] = 1
129        u[0,1] =  sqrthfi
130        u[2,1] =  sqrthfi
131        u[0,0] =  sqrthfr
132        u[2,0] = -sqrthfr
133    else:
134        u[l,l] = 1
135        for m in range(1, l+1, 2):
136            u[l-m,l-m] =  sqrthfi
137            u[l+m,l-m] =  sqrthfi
138            u[l-m,l+m] =  sqrthfr
139            u[l+m,l+m] = -sqrthfr
140        for m in range(2, l+1, 2):
141            u[l-m,l-m] =  sqrthfi
142            u[l+m,l-m] = -sqrthfi
143            u[l-m,l+m] =  sqrthfr
144            u[l+m,l+m] =  sqrthfr
145
146    return u
147
148def sph_real2pure(l, reorder_p=True):
149    '''
150    Transformation matrix: from real spherical harmonic functions to the pure
151    spherical harmonic functions.
152
153    Kwargs:
154        reorder_p (bool): Whether the real p functions are in the (x,y,z) order.
155    '''
156    # numpy.linalg.inv(sph_pure2real(l))
157    return sph_pure2real(l, reorder_p).conj().T
158
159# |spinor> = (|real_sph>, |real_sph>) * / u_alpha \
160#                                       \ u_beta  /
161# Return 2D array U_{sph,spinor}
162def sph2spinor(l, reorder_p=True):
163    if l == 0:
164        return numpy.array((0., 1.)).reshape(1,-1), \
165               numpy.array((1., 0.)).reshape(1,-1)
166    else:
167        u1 = sph_real2pure(l, reorder_p)
168        ua = numpy.zeros((2*l+1,4*l+2),dtype=complex)
169        ub = numpy.zeros((2*l+1,4*l+2),dtype=complex)
170        j = l * 2 - 1
171        mla = l + (-j-1)//2
172        mlb = l + (-j+1)//2
173        for k,mj in enumerate(range(-j, j+1, 2)):
174            ua[:,k] = u1[:,mla] * cg_spin(l, j, mj, 1)
175            ub[:,k] = u1[:,mlb] * cg_spin(l, j, mj,-1)
176            mla += 1
177            mlb += 1
178        j = l * 2 + 1
179        mla = l + (-j-1)//2
180        mlb = l + (-j+1)//2
181        for k,mj in enumerate(range(-j, j+1, 2)):
182            if mla < 0:
183                ua[:,l*2+k] = 0
184            else:
185                ua[:,l*2+k] = u1[:,mla] * cg_spin(l, j, mj, 1)
186            if mlb >= 2*l+1:
187                ub[:,l*2+k] = 0
188            else:
189                ub[:,l*2+k] = u1[:,mlb] * cg_spin(l, j, mj,-1)
190            mla += 1
191            mlb += 1
192    return ua, ub
193real2spinor = sph2spinor
194
195# Returns 2D array U_{sph,spinor}
196def sph2spinor_coeff(mol):
197    '''Transformation matrix that transforms real-spherical GTOs to spinor
198    GTOs for all basis functions
199
200    Examples::
201
202    >>> from pyscf import gto
203    >>> from pyscf.symm import sph
204    >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='ccpvtz')
205    >>> ca, cb = sph.sph2spinor_coeff(mol)
206    >>> s0 = mol.intor('int1e_ovlp_spinor')
207    >>> s1 = ca.conj().T.dot(mol.intor('int1e_ovlp_sph')).dot(ca)
208    >>> s1+= cb.conj().T.dot(mol.intor('int1e_ovlp_sph')).dot(cb)
209    >>> print(abs(s1-s0).max())
210    >>> 6.66133814775e-16
211    '''
212    lmax = max([mol.bas_angular(i) for i in range(mol.nbas)])
213    ualst = []
214    ublst = []
215    for l in range(lmax+1):
216        u1, u2 = sph2spinor(l, reorder_p=True)
217        ualst.append(u1)
218        ublst.append(u2)
219
220    ca = []
221    cb = []
222    for ib in range(mol.nbas):
223        l = mol.bas_angular(ib)
224        kappa = mol.bas_kappa(ib)
225        if kappa == 0:
226            ua = ualst[l]
227            ub = ublst[l]
228        elif kappa < 0:
229            ua = ualst[l][:,l*2:]
230            ub = ublst[l][:,l*2:]
231        else:
232            ua = ualst[l][:,:l*2]
233            ub = ublst[l][:,:l*2]
234        nctr = mol.bas_nctr(ib)
235        ca.extend([ua]*nctr)
236        cb.extend([ub]*nctr)
237    return numpy.stack([scipy.linalg.block_diag(*ca),
238                        scipy.linalg.block_diag(*cb)])
239real2spinor_whole = sph2spinor_coeff
240
241def cart2spinor(l):
242    '''Cartesian to spinor for angular moment l'''
243    from pyscf import gto
244    return gto.cart2spinor_l(l)
245
246
247if __name__ == '__main__':
248    for l in range(3):
249        print(sph_pure2real(l))
250        print(sph_real2pure(l))
251
252    for l in range(3):
253        print(sph2spinor(l)[0])
254        print(sph2spinor(l)[1])
255