#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf import ao2mo
from pyscf.fci import cistring
from pyscf.fci import direct_spin1
[docs]
def contract_2e(eri, civec_strs, norb, nelec, link_index=None):
'''Compute E_{pq}E_{rs}|CI>'''
if isinstance(nelec, (int, numpy.integer)):
nelecb = nelec//2
neleca = nelec - nelecb
else:
neleca, nelecb = nelec
eri = ao2mo.restore(1, eri, norb)
h_ps = numpy.einsum('pqqs->ps', eri)
ci_coeff, ci_strs = civec_strs
strsa, strsb = ci_strs
strsa = numpy.asarray(strsa)
strsb = numpy.asarray(strsb)
if link_index is None:
cd_indexa = cre_des_linkstr(strsa, norb, neleca)
dd_indexa = des_des_linkstr(strsa, norb, neleca)
cd_indexb = cre_des_linkstr(strsb, norb, nelecb)
dd_indexb = des_des_linkstr(strsb, norb, nelecb)
else:
cd_indexa, dd_indexa, cd_indexb, dd_indexb = link_index
ma = len(dd_indexa)
mb = len(dd_indexb)
na = len(strsa)
nb = len(strsb)
# Adding Eq (*) below to fcinew because contract_2e function computes the
# contraction "E_{pq}E_{rs} V_{pqrs} |CI>" (~ p^+ q r^+ s |CI>) while
# the actual contraction for (aa|aa) and (bb|bb) part is
# "p^+ r^+ s q V_{pqrs} |CI>". To make (aa|aa) and (bb|bb) code reproduce
# "p^+ q r^+ s |CI>", we employ the identity
# p^+ q r^+ s = p^+ r^+ s q + delta(qr) p^+ s
# the second term is the source of Eq (*)
fcivec = ci_coeff.reshape(na,nb)
fcinew = numpy.zeros_like(fcivec)
# (bb|aa)
t1 = numpy.zeros((norb,norb,na,nb))
for str1, tab in enumerate(cd_indexa):
for a, i, str0, sign in tab:
if a >= 0:
t1[a,i,str1] += sign * fcivec[str0]
fcinew += numpy.einsum('ps,psab->ab', h_ps, t1) # (*)
t1 = numpy.dot(eri.reshape(norb*norb,-1), t1.reshape(norb*norb,-1))
t1 = t1.reshape(norb,norb,na,nb)
for str1, tab in enumerate(cd_indexb):
for a, i, str0, sign in tab:
if a >= 0:
fcinew[:,str0] += sign * t1[a,i,:,str1]
# (aa|bb)
t1 = numpy.zeros((norb,norb,na,nb))
for str1, tab in enumerate(cd_indexb):
for a, i, str0, sign in tab:
if a >= 0:
t1[a,i,:,str1] += sign * fcivec[:,str0]
fcinew += numpy.einsum('ps,psab->ab', h_ps, t1) # (*)
t1 = numpy.dot(eri.reshape(norb*norb,-1), t1.reshape(norb*norb,-1))
t1 = t1.reshape(norb,norb,na,nb)
for str1, tab in enumerate(cd_indexa):
for a, i, str0, sign in tab:
if a >= 0:
fcinew[str0] += sign * t1[a,i,str1]
eri1 = eri.transpose(0,2,1,3)
# (aa|aa)
t1 = numpy.zeros((norb,norb,ma,nb))
for str1, tab in enumerate(dd_indexa):
for i, j, str0, sign in tab:
if i >= 0:
t1[i,j,str1] += sign * fcivec[str0]
t1 = numpy.dot(eri1.reshape(norb*norb,-1), t1.reshape(norb*norb,-1))
t1 = t1.reshape(norb,norb,ma,nb)
for str1, tab in enumerate(dd_indexa):
for i, j, str0, sign in tab:
if i >= 0:
fcinew[str0] += sign * t1[i,j,str1]
# (bb|bb)
t1 = numpy.zeros((norb,norb,na,mb))
for str1, tab in enumerate(dd_indexb):
for i, j, str0, sign in tab:
if i >= 0:
t1[i,j,:,str1] += sign * fcivec[:,str0]
t1 = numpy.dot(eri1.reshape(norb*norb,-1), t1.reshape(norb*norb,-1))
t1 = t1.reshape(norb,norb,na,mb)
for str1, tab in enumerate(dd_indexb):
for i, j, str0, sign in tab:
if i >= 0:
fcinew[:,str0] += sign * t1[i,j,:,str1]
return fcinew.reshape(ci_coeff.shape)
[docs]
def enlarge_space(myci, civec_strs, eri, norb, nelec):
if isinstance(nelec, (int, numpy.integer)):
nelecb = nelec//2
neleca = nelec - nelecb
else:
neleca, nelecb = nelec
ci_coeff, ci_strs = civec_strs
strsa, strsb = ci_strs
na = len(strsa)
nb = len(strsb)
ci_coeff = ci_coeff.reshape(na,nb)
eri = ao2mo.restore(1, eri, norb)
eri_pq_max = abs(eri.reshape(norb**2,-1)).max(axis=1).reshape(norb,norb)
civec_a_max = abs(ci_coeff).max(axis=1)
civec_b_max = abs(ci_coeff).max(axis=0)
aidx = civec_a_max > myci.ci_coeff_cutoff
bidx = civec_b_max > myci.ci_coeff_cutoff
ci_coeff = ci_coeff[aidx][:,bidx]
civec_a_max = civec_a_max[aidx]
civec_b_max = civec_b_max[bidx]
strsa = numpy.asarray(strsa)[aidx]
strsb = numpy.asarray(strsb)[bidx]
def select_strs(civec_max, strs, nelec):
strs_add = []
for ia, str0 in enumerate(strs):
occ = []
vir = []
for i in range(norb):
if str0 & (1 << i):
occ.append(i)
else:
vir.append(i)
ca = civec_max[ia]
for i1, i in enumerate(occ):
for a1, a in enumerate(vir):
if eri_pq_max[a,i]*ca > myci.select_cutoff:
str1 = str0 ^ (1 << i) | (1 << a)
strs_add.append(str1)
if i < nelec and a >= nelec:
for j in occ[:i1]:
for b in vir[a1+1:]:
if abs(eri[a,i,b,j])*ca > myci.select_cutoff:
strs_add.append(str1 ^ (1 << j) | (1 << b))
strs_add = sorted(set(strs_add) - set(strs))
return numpy.asarray(strs_add, dtype=int)
strsa_add = select_strs(civec_a_max, strsa, neleca)
strsb_add = select_strs(civec_b_max, strsb, nelecb)
strsa = numpy.append(strsa, strsa_add)
strsb = numpy.append(strsb, strsb_add)
aidx = numpy.argsort(strsa)
bidx = numpy.argsort(strsb)
ci_strs = (strsa[aidx], strsb[bidx])
aidx = numpy.where(aidx < ci_coeff.shape[0])[0]
bidx = numpy.where(bidx < ci_coeff.shape[1])[0]
ci_coeff1 = numpy.zeros((len(strsa),len(strsb)))
lib.takebak_2d(ci_coeff1, ci_coeff, aidx, bidx)
return ci_coeff1, ci_strs
[docs]
def cre_des_linkstr(strs, norb, nelec):
addrs = dict(zip(strs, range(len(strs))))
nvir = norb - nelec
link_index = numpy.zeros((len(addrs),nelec+nelec*nvir,4), dtype=int)
link_index[:,:,0] = -1
for i0, str1 in enumerate(strs):
occ = []
vir = []
for i in range(norb):
if str1 & (1 << i):
occ.append(i)
else:
vir.append(i)
k = 0
for i in occ:
link_index[i0,k] = (i, i, i0, 1)
k += 1
for a in vir:
for i in occ:
str0 = str1 ^ (1 << i) | (1 << a)
if str0 in addrs:
# [cre, des, targetddress, parity]
link_index[i0,k] = (a, i, addrs[str0], cistring.cre_des_sign(a, i, str1))
k += 1
return link_index
[docs]
def des_des_linkstr(strs, norb, nelec):
inter = []
for str0 in strs:
occ = [i for i in range(norb) if str0 & (1 << i)]
for i1, i in enumerate(occ):
for j in occ[:i1]:
inter.append(str0 ^ (1 << i) ^ (1 << j))
inter = sorted(set(inter))
addrs = dict(zip(strs, range(len(strs))))
nvir = norb - nelec + 2
link_index = numpy.zeros((len(inter),nvir*nvir,4), dtype=int)
link_index[:,:,0] = -1
for i1, str1 in enumerate(inter):
vir = [i for i in range(norb) if not str1 & (1 << i)]
k = 0
for i in vir:
for j in vir:
str0 = str1 | (1 << i) | (1 << j)
if i != j and str0 in addrs:
# from intermediate str1, create i, create j -> str0
# (str1 = des_i des_j str0)
# [cre_j, cre_i, targetddress, parity]
sign = cistring.cre_sign(i, str1)
sign*= cistring.cre_sign(j, str1 | (1 << i))
link_index[i1,k] = (i, j, addrs[str0], sign)
k += 1
return link_index
[docs]
def make_hdiag(h1e, g2e, ci_strs, norb, nelec):
if isinstance(nelec, (int, numpy.integer)):
nelecb = nelec//2
neleca = nelec - nelecb
else:
neleca, nelecb = nelec
strsa, strsb = ci_strs
strsa = numpy.asarray(strsa)
strsb = numpy.asarray(strsb)
occslista = [[i for i in range(norb) if str0 & (1 << i)] for str0 in strsa]
occslistb = [[i for i in range(norb) if str0 & (1 << i)] for str0 in strsb]
g2e = ao2mo.restore(1, g2e, norb)
diagj = numpy.einsum('iijj->ij',g2e)
diagk = numpy.einsum('ijji->ij',g2e)
hdiag = []
for aocc in occslista:
for bocc in occslistb:
e1 = h1e[aocc,aocc].sum() + h1e[bocc,bocc].sum()
e2 = diagj[aocc][:,aocc].sum() + diagj[aocc][:,bocc].sum() \
+ diagj[bocc][:,aocc].sum() + diagj[bocc][:,bocc].sum() \
- diagk[aocc][:,aocc].sum() - diagk[bocc][:,bocc].sum()
hdiag.append(e1 + e2*.5)
return numpy.array(hdiag)
[docs]
def kernel(h1e, eri, norb, nelec, ecore=0, verbose=logger.NOTE):
if isinstance(nelec, (int, numpy.integer)):
nelecb = nelec//2
neleca = nelec - nelecb
else:
neleca, nelecb = nelec
h2e = direct_spin1.absorb_h1e(h1e, eri, norb, nelec, .5)
namax = cistring.num_strings(norb, neleca)
nbmax = cistring.num_strings(norb, nelecb)
myci = SelectedCI()
strsa = [int('1'*neleca, 2)]
strsb = [int('1'*nelecb, 2)]
ci_strs = (strsa, strsb)
ci0 = numpy.ones((1,1))
ci0, ci_strs = enlarge_space(myci, (ci0, ci_strs), h2e, norb, nelec)
def all_linkstr_index(ci_strs):
cd_indexa = cre_des_linkstr(ci_strs[0], norb, neleca)
dd_indexa = des_des_linkstr(ci_strs[0], norb, neleca)
cd_indexb = cre_des_linkstr(ci_strs[1], norb, nelecb)
dd_indexb = des_des_linkstr(ci_strs[1], norb, nelecb)
return cd_indexa, dd_indexa, cd_indexb, dd_indexb
def hop(c):
hc = contract_2e(h2e, (c, ci_strs), norb, nelec, link_index)
return hc.reshape(-1)
precond = lambda x, e, *args: x/(hdiag-e+1e-4)
e_last = 0
tol = 1e-2
# conv = False
for icycle in range(norb):
tol = max(tol*1e-2, myci.float_tol)
link_index = all_linkstr_index(ci_strs)
hdiag = make_hdiag(h1e, eri, ci_strs, norb, nelec)
e, ci0 = lib.davidson(hop, ci0.reshape(-1), precond, tol=tol,
verbose=verbose)
print('icycle %d ci.shape %s E = %.15g' %
(icycle, (len(ci_strs[0]), len(ci_strs[1])), e))
if ci0.shape == (namax,nbmax) or abs(e-e_last) < myci.float_tol*10:
# conv = True
break
ci1, ci_strs = enlarge_space(myci, (ci0, ci_strs), h2e, norb, nelec)
if ci1.size < ci0.size*1.02:
# conv = True
break
e_last = e
ci0 = ci1
link_index = all_linkstr_index(ci_strs)
hdiag = make_hdiag(h1e, eri, ci_strs, norb, nelec)
e, ci0 = lib.davidson(hop, ci0.reshape(-1), precond, tol=myci.conv_tol,
verbose=verbose)
na = len(ci_strs[0])
nb = len(ci_strs[1])
return e+ecore, (ci0.reshape(na,nb), ci_strs)
# dm_pq = <|p^+ q|>
[docs]
def make_rdm1(civec_strs, norb, nelec):
if isinstance(nelec, (int, numpy.integer)):
nelecb = nelec//2
neleca = nelec - nelecb
else:
neleca, nelecb = nelec
ci_coeff, ci_strs = civec_strs
strsa, strsb = ci_strs
strsa = numpy.asarray(strsa)
strsb = numpy.asarray(strsb)
cd_indexa = cre_des_linkstr(strsa, norb, neleca)
cd_indexb = cre_des_linkstr(strsb, norb, nelecb)
na = len(strsa)
nb = len(strsb)
fcivec = ci_coeff.reshape(na,nb)
rdm1 = numpy.zeros((norb,norb))
for str1, tab in enumerate(cd_indexa):
for a, i, str0, sign in tab:
if a >= 0:
rdm1[a,i] += sign * numpy.dot(fcivec[str1], fcivec[str0])
for str1, tab in enumerate(cd_indexb):
for a, i, str0, sign in tab:
if a >= 0:
rdm1[a,i] += sign * numpy.dot(fcivec[:,str1], fcivec[:,str0])
return rdm1.T
# dm_pq,rs = <|p^+ q r^+ s|>
[docs]
def make_rdm2(civec_strs, norb, nelec):
if isinstance(nelec, (int, numpy.integer)):
nelecb = nelec//2
neleca = nelec - nelecb
else:
neleca, nelecb = nelec
ci_coeff, ci_strs = civec_strs
strsa, strsb = ci_strs
strsa = numpy.asarray(strsa)
strsb = numpy.asarray(strsb)
cd_indexa = cre_des_linkstr(strsa, norb, neleca)
dd_indexa = des_des_linkstr(strsa, norb, neleca)
cd_indexb = cre_des_linkstr(strsb, norb, nelecb)
dd_indexb = des_des_linkstr(strsb, norb, nelecb)
ma = len(dd_indexa)
mb = len(dd_indexb)
na = len(strsa)
nb = len(strsb)
fcivec = ci_coeff.reshape(na,nb)
rdm2 = numpy.zeros((norb,norb,norb,norb))
# (bb|aa) and (aa|bb)
t1a = numpy.zeros((norb,norb,na,nb))
t1b = numpy.zeros((norb,norb,na,nb))
for str1, tab in enumerate(cd_indexa):
for a, i, str0, sign in tab:
if a >= 0:
t1a[a,i,str1] += sign * fcivec[str0]
for str1, tab in enumerate(cd_indexb):
for a, i, str0, sign in tab:
if a >= 0:
t1b[a,i,:,str1] += sign * fcivec[:,str0]
tmp = numpy.dot(t1a.reshape(norb**2,-1), t1b.reshape(norb**2,-1).T)
tmp = tmp.reshape([norb]*4).transpose(1,0,2,3)
rdm2 += tmp
rdm2 += tmp.transpose(2,3,0,1)
# (aa|aa)
t1a = numpy.zeros((norb,norb,ma,nb))
for str1, tab in enumerate(dd_indexa):
for i, j, str0, sign in tab:
if i >= 0:
t1a[i,j,str1] += sign * fcivec[str0]
tmp = numpy.dot(t1a.reshape(norb**2,-1), t1a.reshape(norb**2,-1).T)
rdm2 -= tmp.reshape([norb]*4).transpose(0,3,1,2)
# (bb|bb)
t1b = numpy.zeros((norb,norb,na,mb))
for str1, tab in enumerate(dd_indexb):
for i, j, str0, sign in tab:
if i >= 0:
t1b[i,j,:,str1] += sign * fcivec[:,str0]
tmp = numpy.dot(t1b.reshape(norb**2,-1), t1b.reshape(norb**2,-1).T)
rdm2 -= tmp.reshape([norb]*4).transpose(0,3,1,2)
return rdm2
[docs]
class SelectedCI:
def __init__(self):
self.ci_coeff_cutoff = 1e-3
self.select_cutoff = 1e-3
self.float_tol = 1e-6
self.conv_tol = 1e-10
if __name__ == '__main__':
from functools import reduce
from pyscf import gto
from pyscf import scf
mol = gto.Mole()
mol.verbose = 0
mol.output = None
mol.atom = [
['H', ( 1.,-1. , 0. )],
['H', ( 0.,-1. ,-1. )],
['H', ( 1.,-0.5 ,-1. )],
['H', ( 0.,-0. ,-1. )],
['H', ( 1.,-0.5 , 0. )],
['H', ( 0., 1. , 1. )],
['H', ( 1., 2. , 3. )],
['H', ( 1., 2. , 4. )],
]
mol.basis = 'sto-3g'
mol.build()
m = scf.RHF(mol)
m.kernel()
norb = m.mo_coeff.shape[1]
nelec = mol.nelectron
h1e = reduce(numpy.dot, (m.mo_coeff.T, m.get_hcore(), m.mo_coeff))
eri = ao2mo.kernel(m._eri, m.mo_coeff, compact=False)
eri = eri.reshape(norb,norb,norb,norb)
e1, c1 = kernel(h1e, eri, norb, nelec)
e2, c2 = direct_spin1.kernel(h1e, eri, norb, nelec)
print(e1, e1 - -11.894559902235565, 'diff to FCI', e1-e2)
print(c1[0].shape, c2.shape)
dm1_1 = make_rdm1(c1, norb, nelec)
dm1_2 = direct_spin1.make_rdm1(c2, norb, nelec)
print(abs(dm1_1 - dm1_2).sum())
dm2_1 = make_rdm2(c1, norb, nelec)
dm2_2 = direct_spin1.make_rdm12(c2, norb, nelec)[1]
print(abs(dm2_1 - dm2_2).sum())
