1# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14#
15# Localization based on UNO from UHF/UKS check files
16#
17from functools import reduce
18import numpy
19import scipy.linalg
20import h5py
21from pyscf import tools,gto,scf,dft
22from pyscf.tools import molden
23import pmloc
24import ulocal
25
26def sqrtm(s):
27   e, v = numpy.linalg.eigh(s)
28   return numpy.dot(v*numpy.sqrt(e), v.T.conj())
29
30def lowdin(s):
31   e, v = numpy.linalg.eigh(s)
32   return numpy.dot(v/numpy.sqrt(e), v.T.conj())
33
34def dumpLUNO(fname,thresh=0.01):
35   chkfile = fname+'.chk'
36   outfile = fname+'_cmo.molden'
37   tools.molden.from_chkfile(outfile, chkfile)
38   #=============================
39   # Natural orbitals
40   # Lowdin basis X=S{-1/2}
41   # psi = chi * C
42   #     = chi' * C'
43   #     = chi*X*(X{-1}C')
44   #=============================
45   mol,mf = scf.chkfile.load_scf(chkfile)
46   mo_coeff = mf["mo_coeff"]
47   ova=mol.intor_symmetric("cint1e_ovlp_sph")
48   nb = mo_coeff.shape[1]
49   # Check overlap
50   diff = reduce(numpy.dot,(mo_coeff[0].T,ova,mo_coeff[0])) - numpy.identity(nb)
51   print numpy.linalg.norm(diff)
52   diff = reduce(numpy.dot,(mo_coeff[1].T,ova,mo_coeff[1])) - numpy.identity(nb)
53   print numpy.linalg.norm(diff)
54   # UHF-alpha/beta
55   ma = mo_coeff[0]
56   mb = mo_coeff[1]
57   nalpha = (mol.nelectron+mol.spin)/2
58   nbeta  = (mol.nelectron-mol.spin)/2
59   # Spin-averaged DM
60   pTa = numpy.dot(ma[:,:nalpha],ma[:,:nalpha].T)
61   pTb = numpy.dot(mb[:,:nbeta],mb[:,:nbeta].T)
62   pT = 0.5*(pTa+pTb)
63   # Lowdin basis
64   s12 = sqrtm(ova)
65   s12inv = lowdin(ova)
66   pTOAO = reduce(numpy.dot,(s12,pT,s12))
67   eig,coeff = scipy.linalg.eigh(-pTOAO)
68   eig = -2.0*eig
69   eig[eig<0.0]=0.0
70   eig[abs(eig)<1.e-14]=0.0
71   ifplot = False #True
72   if ifplot:
73      import matplotlib.pyplot as plt
74      plt.plot(range(nb),eig,'ro')
75      plt.show()
76   # Back to AO basis
77   coeff = numpy.dot(s12inv,coeff)
78   diff = reduce(numpy.dot,(coeff.T,ova,coeff)) - numpy.identity(nb)
79   print 'CtSC-I',numpy.linalg.norm(diff)
80   #
81   # Averaged Fock
82   #
83   enorb = mf["mo_energy"]
84   fa = reduce(numpy.dot,(ma,numpy.diag(enorb[0]),ma.T))
85   fb = reduce(numpy.dot,(mb,numpy.diag(enorb[1]),mb.T))
86   # Non-orthogonal cases: FC=SCE
87   # Fao = SC*e*C{-1} = S*C*e*Ct*S
88   fav = 0.5*(fa+fb)
89   # Expectation value of natural orbitals <i|F|i>
90   fexpt = reduce(numpy.dot,(coeff.T,ova,fav,ova,coeff))
91   enorb = numpy.diag(fexpt)
92   nocc = eig.copy()
93   #
94   # Reordering and define active space according to thresh
95   #
96   idx = 0
97   active=[]
98   for i in range(nb):
99      if nocc[i]<=2.0-thresh and nocc[i]>=thresh:
100         active.append(True)
101      else:
102         active.append(False)
103   print '\nNatural orbitals:'
104   for i in range(nb):
105      print 'orb:',i,active[i],nocc[i],enorb[i]
106   active = numpy.array(active)
107   actIndices = list(numpy.argwhere(active==True).flatten())
108   cOrbs = coeff[:,:actIndices[0]]
109   aOrbs = coeff[:,actIndices]
110   vOrbs = coeff[:,actIndices[-1]+1:]
111   nb = cOrbs.shape[0]
112   nc = cOrbs.shape[1]
113   na = aOrbs.shape[1]
114   nv = vOrbs.shape[1]
115   print 'core orbs:',cOrbs.shape
116   print 'act  orbs:',aOrbs.shape
117   print 'vir  orbs:',vOrbs.shape
118   assert nc+na+nv == nb
119   # dump UNO
120   with open(fname+'_uno.molden','w') as thefile:
121       molden.header(mol,thefile)
122       molden.orbital_coeff(mol,thefile,coeff)
123   #=====================
124   # Population analysis
125   #=====================
126   from pyscf import lo
127   aux = lo.orth_ao(mol,method='meta_lowdin')
128   #clmo = ulocal.scdm(cOrbs,ova,aux)
129   #almo = ulocal.scdm(aOrbs,ova,aux)
130   clmo = cOrbs
131   almo = aOrbs
132   ierr,uc = pmloc.loc(mol,clmo)
133   ierr,ua = pmloc.loc(mol,almo)
134   clmo = clmo.dot(uc)
135   almo = almo.dot(ua)
136   vlmo = ulocal.scdm(vOrbs,ova,aux)
137   # P-SORT
138   mo_c,n_c,e_c = ulocal.psort(ova,fav,pT,clmo)
139   mo_o,n_o,e_o = ulocal.psort(ova,fav,pT,almo)
140   mo_v,n_v,e_v = ulocal.psort(ova,fav,pT,vlmo)
141   lmo = numpy.hstack((mo_c,mo_o,mo_v)).copy()
142   enorb = numpy.hstack([e_c,e_o,e_v])
143   occ = numpy.hstack([n_c,n_o,n_v])
144   # CHECK
145   diff = reduce(numpy.dot,(lmo.T,ova,lmo)) - numpy.identity(nb)
146   print 'diff=',numpy.linalg.norm(diff)
147   ulocal.lowdinPop(mol,lmo,ova,enorb,occ)
148   ulocal.dumpLMO(mol,fname,lmo)
149   print 'nalpha,nbeta,mol.spin,nb:',\
150          nalpha,nbeta,mol.spin,nb
151   return mol,ova,fav,pT,nb,nalpha,nbeta,nc,na,nv,lmo,enorb,occ
152
153def dumpAct(fname,info,actlst,base=1):
154   actlst2 = [i-base for i in actlst]
155   mol,ova,fav,pT,nb,nalpha,nbeta,nc,na,nv,lmo,enorb,occ = info
156   corb = set(range(nc))
157   aorb = set(range(nc,nc+na))
158   vorb = set(range(nc+na,nc+na+nv))
159   print '[dumpAct]'
160   print ' corb=',corb
161   print ' aorb=',aorb
162   print ' vorb=',vorb
163   sorb = set(actlst2)
164   rcorb = corb.difference(corb.intersection(sorb))
165   #assuming act in actlst
166   #raorb = aorb.difference(aorb.intersection(sorb))
167   rvorb = vorb.difference(vorb.intersection(sorb))
168   corb = list(rcorb)
169   aorb = list(sorb)
170   vorb = list(rvorb)
171   print ' corb=',corb
172   print ' aorb=',aorb
173   print ' vorb=',vorb
174   clmo = lmo[:,corb].copy()
175   almo = lmo[:,aorb].copy()
176   vlmo = lmo[:,vorb].copy()
177   ierr,ua = pmloc.loc(mol,almo)
178   almo = almo.dot(ua)
179   #>>> DUMP <<<#
180   # P-SORT
181   mo_c = clmo
182   mo_v = vlmo
183   e_c = enorb[corb].copy()
184   e_v = enorb[vorb].copy()
185   n_c = occ[corb].copy()
186   n_v = occ[vorb].copy()
187   mo_o,n_o,e_o = ulocal.psort(ova,fav,pT,almo)
188   lmo2 = numpy.hstack((mo_c,mo_o,mo_v))
189   enorb = numpy.hstack([e_c,e_o,e_v])
190   occ = numpy.hstack([n_c,n_o,n_v])
191   assert len(enorb)==nb
192   assert len(occ)==nb
193   # CHECK
194   diff = reduce(numpy.dot,(lmo2.T,ova,lmo2)) - numpy.identity(nb)
195   print 'diff=',numpy.linalg.norm(diff)
196   ulocal.lowdinPop(mol,lmo,ova,enorb,occ)
197   ulocal.dumpLMO(mol,fname+'_new',lmo2)
198   print 'nalpha,nbeta,mol.spin,nb:',\
199          nalpha,nbeta,mol.spin,nb
200   print 'diff(LMO2-LMO)=',numpy.linalg.norm(lmo2-lmo)
201   nc = len(e_c)
202   na = len(e_o)
203   nv = len(e_v)
204   assert na == len(actlst)
205   assert nc+na+nv == nb
206   print 'nc,na,nv,nb=',nc,na,nv,nb
207   return lmo2,nc,na,nv
208
209if __name__ == '__main__':
210   fname = 'hs_bp86'
211   info = dumpLUNO(fname)
212   actlst = [117,118,119,120,125,126]+range(127,137)
213   dumpAct(fname,info,actlst,base=1)
214