#!/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 ctypes
import numpy
from pyscf import lib
from pyscf import ao2mo
from pyscf.fci import direct_spin1
from pyscf.fci import selected_ci
libfci = direct_spin1.libfci
[docs]
def contract_2e(eri, civec_strs, norb, nelec, link_index=None):
ci_coeff, nelec, ci_strs = selected_ci._unpack(civec_strs, nelec)
if link_index is None:
link_index = selected_ci._all_linkstr_index(ci_strs, norb, nelec)
cd_indexa, dd_indexa, cd_indexb, dd_indexb = link_index
na, nlinka = nb, nlinkb = cd_indexa.shape[:2]
eri = ao2mo.restore(1, eri, norb)
eri1 = eri.transpose(0,2,1,3) - eri.transpose(0,2,3,1)
idx,idy = numpy.tril_indices(norb, -1)
idx = idx * norb + idy
eri1 = lib.take_2d(eri1.reshape(norb**2,-1), idx, idx) * 2
lib.transpose_sum(eri1, inplace=True)
eri1 *= .5
fcivec = ci_coeff.reshape(na,nb)
# (aa|aa)
ci1 = numpy.zeros_like(fcivec)
if nelec[0] > 1:
ma, mlinka = mb, mlinkb = dd_indexa.shape[:2]
libfci.SCIcontract_2e_aaaa(eri1.ctypes.data_as(ctypes.c_void_p),
fcivec.ctypes.data_as(ctypes.c_void_p),
ci1.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(norb),
ctypes.c_int(na), ctypes.c_int(nb),
ctypes.c_int(ma), ctypes.c_int(mlinka),
dd_indexa.ctypes.data_as(ctypes.c_void_p))
h_ps = numpy.einsum('pqqs->ps', eri) * (.5/nelec[0])
eri1 = eri.copy()
for k in range(norb):
eri1[:,:,k,k] += h_ps
eri1[k,k,:,:] += h_ps
eri1 = ao2mo.restore(4, eri1, norb)
lib.transpose_sum(eri1, inplace=True)
eri1 *= .5
# (bb|aa)
libfci.SCIcontract_2e_bbaa(eri1.ctypes.data_as(ctypes.c_void_p),
fcivec.ctypes.data_as(ctypes.c_void_p),
ci1.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(norb),
ctypes.c_int(na), ctypes.c_int(nb),
ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
cd_indexa.ctypes.data_as(ctypes.c_void_p),
cd_indexb.ctypes.data_as(ctypes.c_void_p))
lib.transpose_sum(ci1, inplace=True)
return selected_ci._as_SCIvector(ci1.reshape(ci_coeff.shape), ci_strs)
[docs]
def enlarge_space(myci, civec_strs, eri, norb, nelec):
if isinstance(civec_strs, (tuple, list)):
nelec, (strsa, strsb) = selected_ci._unpack(civec_strs[0], nelec)[1:]
ci_coeff = lib.asarray(civec_strs)
else:
ci_coeff, nelec, (strsa, strsb) = selected_ci._unpack(civec_strs, nelec)
na = nb = len(strsa)
ci0 = ci_coeff.reshape(-1,na,nb)
abs_ci = abs(ci0).max(axis=0)
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.max(axis=1)
ci_aidx = numpy.where(civec_a_max > myci.ci_coeff_cutoff)[0]
civec_a_max = civec_a_max[ci_aidx]
strsa = strsa[ci_aidx]
strsa_add = selected_ci.select_strs(myci, eri, eri_pq_max, civec_a_max, strsa, norb, nelec[0])
strsa = numpy.append(strsa, strsa_add)
aidx = numpy.argsort(strsa)
strsa = strsa[aidx]
aidx = numpy.where(aidx < len(ci_aidx))[0]
ci_bidx = ci_aidx
strsb = strsa
bidx = aidx
ma = mb = len(strsa)
cs = []
for i in range(ci0.shape[0]):
ci1 = numpy.zeros((ma,mb))
tmp = lib.take_2d(ci0[i], ci_aidx, ci_bidx)
lib.takebak_2d(ci1, tmp, aidx, bidx)
cs.append(selected_ci._as_SCIvector(ci1, (strsa,strsb)))
if ci_coeff[0].ndim == 0 or ci_coeff[0].shape[-1] != nb:
cs = [c.ravel() for c in cs]
if (isinstance(ci_coeff, numpy.ndarray) and
ci_coeff.shape[0] == na or ci_coeff.shape[0] == na*nb):
cs = cs[0]
return cs
[docs]
def make_hdiag(h1e, eri, ci_strs, norb, nelec, compress=False):
hdiag = selected_ci.make_hdiag(h1e, eri, ci_strs, norb, nelec)
na = len(ci_strs[0])
lib.transpose_sum(hdiag.reshape(na,na), inplace=True)
hdiag *= .5
return hdiag
[docs]
def kernel(h1e, eri, norb, nelec, ci0=None, level_shift=1e-3, tol=1e-10,
lindep=1e-14, max_cycle=50, max_space=12, nroots=1,
davidson_only=False, pspace_size=400, orbsym=None, wfnsym=None,
select_cutoff=1e-3, ci_coeff_cutoff=1e-3, ecore=0, **kwargs):
return direct_spin1._kfactory(SelectedCI, h1e, eri, norb, nelec, ci0,
level_shift, tol, lindep, max_cycle,
max_space, nroots, davidson_only,
pspace_size, select_cutoff=select_cutoff,
ci_coeff_cutoff=ci_coeff_cutoff, ecore=ecore,
**kwargs)
make_rdm1s = selected_ci.make_rdm1s
make_rdm2s = selected_ci.make_rdm2s
make_rdm1 = selected_ci.make_rdm1
make_rdm2 = selected_ci.make_rdm2
trans_rdm1s = selected_ci.trans_rdm1s
trans_rdm1 = selected_ci.trans_rdm1
[docs]
class SelectedCI(selected_ci.SelectedCI):
[docs]
def contract_2e(self, eri, civec_strs, norb, nelec, link_index=None, **kwargs):
# The argument civec_strs is a CI vector in function FCISolver.contract_2e.
# Save and patch self._strs to make this contract_2e function compatible to
# FCISolver.contract_2e.
if getattr(civec_strs, '_strs', None) is not None:
self._strs = civec_strs._strs
else:
assert (civec_strs.size == len(self._strs[0])*len(self._strs[1]))
civec_strs = selected_ci._as_SCIvector(civec_strs, self._strs)
return contract_2e(eri, civec_strs, norb, nelec, link_index)
make_hdiag = staticmethod(make_hdiag)
enlarge_space = enlarge_space
SCI = SelectedCI
if __name__ == '__main__':
from functools import reduce
from pyscf import gto
from pyscf import scf
from pyscf.fci import direct_spin0
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_spin0.kernel(h1e, eri, norb, nelec)
print(e1, e1 - -11.894559902235565, 'diff to FCI', e1-e2)
