def to_fci(cisdvec, norb, nelec):
if isinstance(nelec, (int, numpy.number)):
nelecb = nelec // 2
neleca = nelec - nelecb
else:
neleca, nelecb = nelec
nocc = neleca
nvir = norb - nocc
c0, c1, c2 = cisdvec_to_amplitudes(cisdvec, norb, nocc)
t1addr, t1sign = t1strs(norb, nocc)
na = cistring.num_strings(norb, nocc)
fcivec = numpy.zeros((na, na))
fcivec[0, 0] = c0
c1 = c1[::-1].T.ravel()
fcivec[0, t1addr] = fcivec[t1addr, 0] = c1 * t1sign
c2ab = c2[::-1, ::-1].transpose(2, 0, 3, 1).reshape(nocc * nvir, -1)
c2ab = numpy.einsum('i,j,ij->ij', t1sign, t1sign, c2ab)
lib.takebak_2d(fcivec, c2ab, t1addr, t1addr)
if nocc > 1 and nvir > 1:
hf_str = int('1' * nocc, 2)
for a in range(nocc, norb):
for b in range(nocc, a):
for i in reversed(range(1, nocc)):
for j in reversed(range(i)):
c2aa = c2[i, j, a - nocc,
b - nocc] - c2[j, i, a - nocc, b - nocc]
str1 = hf_str ^ (1 << j) | (1 << b)
c2aa *= cistring.cre_des_sign(b, j, hf_str)
c2aa *= cistring.cre_des_sign(a, i, str1)
str1 ^= (1 << i) | (1 << a)
addr = cistring.str2addr(norb, nocc, str1)
fcivec[0, addr] = fcivec[addr, 0] = c2aa
return fcivec
def trans(ci1, aindex, bindex, nea, neb):
if aindex is None or bindex is None:
return None
t1 = numpy.zeros(
(cistring.num_strings(norb, nea), cistring.num_strings(norb, neb)))
for i in range(norb):
signa = aindex[:, i, 1]
signb = bindex[:, i, 1]
maska = numpy.where(signa != 0)[0]
maskb = numpy.where(signb != 0)[0]
addra = aindex[maska, i, 0]
addrb = bindex[maskb, i, 0]
citmp = lib.take_2d(fcivec, maska, maskb)
citmp *= signa[maska].reshape(-1, 1)
citmp *= signb[maskb]
#: t1[addra.reshape(-1,1),addrb] += citmp
lib.takebak_2d(t1, citmp, addra, addrb)
for i in range(norb):
signa = aindex[:, i, 1]
signb = bindex[:, i, 1]
maska = numpy.where(signa != 0)[0]
maskb = numpy.where(signb != 0)[0]
addra = aindex[maska, i, 0]
addrb = bindex[maskb, i, 0]
citmp = lib.take_2d(t1, addra, addrb)
citmp *= signa[maska].reshape(-1, 1)
citmp *= signb[maskb]
#: ci1[maska.reshape(-1,1), maskb] += citmp
lib.takebak_2d(ci1, citmp, maska, maskb)
def trans(ci1, aindex, bindex):
if aindex is None or bindex is None:
return None
ma = len(aindex)
mb = len(bindex)
t1 = numpy.zeros((ma,mb))
for i in range(norb):
signa = aindex[:,i,1]
signb = bindex[:,i,1]
maska = numpy.where(signa!=0)[0]
maskb = numpy.where(signb!=0)[0]
addra = aindex[maska,i,0]
addrb = bindex[maskb,i,0]
citmp = lib.take_2d(ci_coeff, addra, addrb)
citmp *= signa[maska].reshape(-1,1)
citmp *= signb[maskb]
#: t1[addra.reshape(-1,1),addrb] += citmp
lib.takebak_2d(t1, citmp, maska, maskb)
for i in range(norb):
signa = aindex[:,i,1]
signb = bindex[:,i,1]
maska = numpy.where(signa!=0)[0]
maskb = numpy.where(signb!=0)[0]
addra = aindex[maska,i,0]
addrb = bindex[maskb,i,0]
citmp = lib.take_2d(t1, maska, maskb)
citmp *= signa[maska].reshape(-1,1)
citmp *= signb[maskb]
#: ci1[maska.reshape(-1,1), maskb] += citmp
lib.takebak_2d(ci1, citmp, addra, addrb)
def trans(ci1, aindex, bindex):
if aindex is None or bindex is None:
return None
ma = len(aindex)
mb = len(bindex)
t1 = numpy.zeros((ma, mb))
for i in range(norb):
signa = aindex[:, i, 1]
signb = bindex[:, i, 1]
maska = numpy.where(signa != 0)[0]
maskb = numpy.where(signb != 0)[0]
addra = aindex[maska, i, 0]
addrb = bindex[maskb, i, 0]
citmp = lib.take_2d(ci_coeff, addra, addrb)
citmp *= signa[maska].reshape(-1, 1)
citmp *= signb[maskb]
#: t1[addra.reshape(-1,1),addrb] += citmp
lib.takebak_2d(t1, citmp, maska, maskb)
for i in range(norb):
signa = aindex[:, i, 1]
signb = bindex[:, i, 1]
maska = numpy.where(signa != 0)[0]
maskb = numpy.where(signb != 0)[0]
addra = aindex[maska, i, 0]
addrb = bindex[maskb, i, 0]
citmp = lib.take_2d(t1, maska, maskb)
citmp *= signa[maska].reshape(-1, 1)
citmp *= signb[maskb]
#: ci1[maska.reshape(-1,1), maskb] += citmp
lib.takebak_2d(ci1, citmp, addra, addrb)
def trans(ci1, aindex, bindex, nea, neb):
if aindex is None or bindex is None:
return None
t1 = numpy.zeros((cistring.num_strings(norb,nea),
cistring.num_strings(norb,neb)))
for i in range(norb):
signa = aindex[:,i,1]
signb = bindex[:,i,1]
maska = numpy.where(signa!=0)[0]
maskb = numpy.where(signb!=0)[0]
addra = aindex[maska,i,0]
addrb = bindex[maskb,i,0]
citmp = lib.take_2d(fcivec, maska, maskb)
citmp *= signa[maska].reshape(-1,1)
citmp *= signb[maskb]
#: t1[addra.reshape(-1,1),addrb] += citmp
lib.takebak_2d(t1, citmp, addra, addrb)
for i in range(norb):
signa = aindex[:,i,1]
signb = bindex[:,i,1]
maska = numpy.where(signa!=0)[0]
maskb = numpy.where(signb!=0)[0]
addra = aindex[maska,i,0]
addrb = bindex[maskb,i,0]
citmp = lib.take_2d(t1, addra, addrb)
citmp *= signa[maska].reshape(-1,1)
citmp *= signb[maskb]
#: ci1[maska.reshape(-1,1), maskb] += citmp
lib.takebak_2d(ci1, citmp, maska, maskb)
def to_fci(cisdvec, norb, nelec):
nocc = nelec // 2
nvir = norb - nocc
c0 = cisdvec[0]
c1 = cisdvec[1:nocc*nvir+1].reshape(nocc,nvir)
c2 = cisdvec[nocc*nvir+1:].reshape(nocc,nocc,nvir,nvir)
t1addr, t1sign = t1strs(norb, nocc)
na = fci.cistring.num_strings(norb, nocc)
fcivec = numpy.zeros((na,na))
fcivec[0,0] = c0
c1 = c1[::-1].T.ravel()
fcivec[0,t1addr] = fcivec[t1addr,0] = c1 * t1sign
c2ab = c2[::-1,::-1].transpose(2,0,3,1).reshape(nocc*nvir,-1)
c2ab = numpy.einsum('i,j,ij->ij', t1sign, t1sign, c2ab)
lib.takebak_2d(fcivec, c2ab, t1addr, t1addr)
if nocc > 1 and nvir > 1:
t2addr, t2sign = t2strs(norb, nocc)
c2aa = []
for a in range(1,nvir):
for b in range(a):
for i in reversed(range(1,nocc)):
for j in reversed(range(i)):
c2aa.append((c2[i,j,a,b] - c2[j,i,a,b]))
c2aa = numpy.asarray(c2aa)
fcivec[0,t2addr] = fcivec[t2addr,0] = c2aa * t2sign
return fcivec
def to_fci(civec_strs, norb, nelec):
ci_coeff, nelec, ci_strs = _unpack(civec_strs, nelec)
addrsa = [cistring.str2addr(norb, nelec[0], x) for x in ci_strs[0]]
addrsb = [cistring.str2addr(norb, nelec[1], x) for x in ci_strs[1]]
na = cistring.num_strings(norb, nelec[0])
nb = cistring.num_strings(norb, nelec[1])
ci0 = numpy.zeros((na, nb))
lib.takebak_2d(ci0, ci_coeff, addrsa, addrsb)
return ci0
def convert(out, inp, idx1, idx2, idx3, idx4):
dim1 = inp.shape
dim2 = out.shape
assert (dim1 == (len(idx1), len(idx2), len(idx3), len(idx4)))
locx = (dim2[1] * idx1.reshape(-1, 1) + idx2).ravel()
locy = (dim2[3] * idx3.reshape(-1, 1) + idx4).ravel()
lib.takebak_2d(out.reshape(dim2[0] * dim2[1], dim2[2] * dim2[3]),
inp.reshape(dim1[0] * dim1[1], dim1[2] * dim1[3]), locx,
locy)
def to_fci(civec_strs, norb, nelec):
ci_coeff, nelec, ci_strs = _unpack(civec_strs, nelec)
addrsa = [cistring.str2addr(norb, nelec[0], x) for x in ci_strs[0]]
addrsb = [cistring.str2addr(norb, nelec[1], x) for x in ci_strs[1]]
na = cistring.num_strings(norb, nelec[0])
nb = cistring.num_strings(norb, nelec[1])
ci0 = numpy.zeros((na,nb))
lib.takebak_2d(ci0, ci_coeff, addrsa, addrsb)
return ci0
def contract_2e(eri,
fcivec,
norb,
nelec,
link_index=None,
orbsym=None,
wfnsym=0):
if orbsym is None:
return direct_spin0.contract_2e(eri, fcivec, norb, nelec, link_index)
eri = ao2mo.restore(4, eri, norb)
neleca, nelecb = direct_spin1._unpack_nelec(nelec)
assert (neleca == nelecb)
link_indexa = direct_spin0._unpack(norb, nelec, link_index)
na, nlinka = link_indexa.shape[:2]
eri_irs, rank_eri, irrep_eri = direct_spin1_symm.reorder_eri(
eri, norb, orbsym)
strsa = numpy.asarray(cistring.gen_strings4orblist(range(norb), neleca))
aidx, link_indexa = direct_spin1_symm.gen_str_irrep(
strsa, orbsym, link_indexa, rank_eri, irrep_eri)
Tirrep = ctypes.c_void_p * TOTIRREPS
linka_ptr = Tirrep(
*[x.ctypes.data_as(ctypes.c_void_p) for x in link_indexa])
eri_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in eri_irs])
dimirrep = (ctypes.c_int * TOTIRREPS)(*[x.shape[0] for x in eri_irs])
fcivec_shape = fcivec.shape
fcivec = fcivec.reshape((na, na), order='C')
ci1new = numpy.zeros_like(fcivec)
nas = (ctypes.c_int * TOTIRREPS)(*[x.size for x in aidx])
ci0 = []
ci1 = []
for ir in range(TOTIRREPS):
ma, mb = aidx[ir].size, aidx[wfnsym ^ ir].size
ci0.append(numpy.zeros((ma, mb)))
ci1.append(numpy.zeros((ma, mb)))
if ma > 0 and mb > 0:
lib.take_2d(fcivec, aidx[ir], aidx[wfnsym ^ ir], out=ci0[ir])
ci0_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in ci0])
ci1_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in ci1])
libfci.FCIcontract_2e_symm1(eri_ptrs, ci0_ptrs, ci1_ptrs,
ctypes.c_int(norb), nas, nas,
ctypes.c_int(nlinka), ctypes.c_int(nlinka),
linka_ptr, linka_ptr, dimirrep,
ctypes.c_int(wfnsym))
for ir in range(TOTIRREPS):
if ci0[ir].size > 0:
lib.takebak_2d(ci1new, ci1[ir], aidx[ir], aidx[wfnsym ^ ir])
return lib.transpose_sum(ci1new, inplace=True).reshape(fcivec_shape)
def contract_sladder(fcivec, norb, nelec, op=-1):
''' Contract spin ladder operator S+ or S- with fcivec.
Changes neleca - nelecb without altering <S2>
Obtained by modifying pyscf.fci.spin_op.contract_ss
'''
neleca, nelecb = _unpack_nelec(nelec)
na = cistring.num_strings(norb, neleca)
nb = cistring.num_strings(norb, nelecb)
fcivec = fcivec.reshape(na, nb)
assert (op in (-1, 1)), 'op = -1 or 1'
if ((op == -1 and (neleca == 0 or nelecb == norb))
or (op == 1 and (neleca == norb or nelecb == 0))):
return np.zeros((0, 0))
# ^ Annihilate vacuum state ^
def gen_map(fstr_index, nelec, des=True):
a_index = fstr_index(range(norb), nelec)
amap = np.zeros((a_index.shape[0], norb, 2), dtype=np.int32)
if des:
for k, tab in enumerate(a_index):
amap[k, tab[:, 1]] = tab[:, 2:]
else:
for k, tab in enumerate(a_index):
amap[k, tab[:, 0]] = tab[:, 2:]
return amap
if op == -1:
aindex = gen_map(cistring.gen_des_str_index, neleca)
bindex = gen_map(cistring.gen_cre_str_index, nelecb, False)
else:
aindex = gen_map(cistring.gen_cre_str_index, neleca, False)
bindex = gen_map(cistring.gen_des_str_index, nelecb)
ci1 = np.zeros((cistring.num_strings(norb, neleca + op),
cistring.num_strings(norb, nelecb - op)))
for i in range(norb):
signa = aindex[:, i, 1]
signb = bindex[:, i, 1]
maska = np.where(signa != 0)[0]
maskb = np.where(signb != 0)[0]
addra = aindex[maska, i, 0]
addrb = bindex[maskb, i, 0]
citmp = lib.take_2d(fcivec, maska, maskb)
citmp *= signa[maska].reshape(-1, 1)
citmp *= signb[maskb]
#: ci1[addra.reshape(-1,1),addrb] += citmp
lib.takebak_2d(ci1, citmp, addra, addrb)
ci1 /= linalg.norm(ci1) # ???
return ci1
def test_takebak_2d(self):
b = numpy.arange(9.).reshape((3,3))
a = numpy.arange(49.).reshape(7,7)
idx = numpy.array([3,0,5])
idy = numpy.array([5,4,1])
ref = a.copy()
ref[idx[:,None],idy] += b
lib.takebak_2d(a, b, idx, idy)
self.assertTrue(numpy.array_equal(ref, a))
b = numpy.arange(9, dtype=numpy.int32).reshape((3,3))
a = numpy.arange(49, dtype=numpy.int32).reshape(7,7)
ref = a.copy()
ref[idx[:,None],idy] += b
lib.takebak_2d(a, b, idx, idy)
self.assertTrue(numpy.array_equal(ref, a))
def test_takebak_2d(self):
b = numpy.arange(9.).reshape((3, 3))
a = numpy.arange(49.).reshape(7, 7)
idx = numpy.array([3, 0, 5])
idy = numpy.array([5, 4, 1])
ref = a.copy()
ref[idx[:, None], idy] += b
lib.takebak_2d(a, b, idx, idy)
self.assertTrue(numpy.array_equal(ref, a))
b = numpy.arange(9, dtype=numpy.int32).reshape((3, 3))
a = numpy.arange(49, dtype=numpy.int32).reshape(7, 7)
ref = a.copy()
ref[idx[:, None], idy] += b
lib.takebak_2d(a, b, idx, idy)
self.assertTrue(numpy.array_equal(ref, a))
def to_fci(cisdvec, nelec, orbspin):
from pyscf import fci
nocc = nelec
norb = len(orbspin)
c0, c1, c2 = cisdvec_to_amplitudes(cisdvec, nocc, norb - nocc)
oidxa = orbspin[:nocc] == 0
oidxb = orbspin[:nocc] == 1
vidxa = orbspin[nocc:] == 0
vidxb = orbspin[nocc:] == 1
c1a = c1[oidxa][:, vidxa]
c1b = c1[oidxb][:, vidxb]
c2aa = c2[oidxa][:, oidxa][:, :, vidxa][:, :, :, vidxa]
c2bb = c2[oidxb][:, oidxb][:, :, vidxb][:, :, :, vidxb]
c2ab = c2[oidxa][:, oidxb][:, :, vidxa][:, :, :, vidxb]
nocca = numpy.count_nonzero(oidxa)
noccb = numpy.count_nonzero(oidxb)
nvira = numpy.count_nonzero(vidxa)
nvirb = numpy.count_nonzero(vidxb)
norba = nocca + nvira
norbb = noccb + nvirb
t1addra, t1signa = t1strs(norba, nocca)
t1addrb, t1signb = t1strs(norbb, noccb)
na = fci.cistring.num_strings(norba, nocca)
nb = fci.cistring.num_strings(norbb, noccb)
fcivec = numpy.zeros((na, nb))
fcivec[0, 0] = c0
fcivec[t1addra, 0] = c1a[::-1].T.ravel() * t1signa
fcivec[0, t1addrb] = c1b[::-1].T.ravel() * t1signb
c2ab = c2ab[::-1, ::-1].transpose(2, 0, 3, 1).reshape(nocca * nvira, -1)
c2ab = numpy.einsum('i,j,ij->ij', t1signa, t1signb, c2ab)
lib.takebak_2d(fcivec, c2ab, t1addra, t1addrb)
if nocca > 1 and nvira > 1:
ooidx = numpy.tril_indices(nocca, -1)
vvidx = numpy.tril_indices(nvira, -1)
c2aa = c2aa[ooidx][:, vvidx[0], vvidx[1]]
t2addra, t2signa = t2strs(norba, nocca)
fcivec[t2addra, 0] = c2aa[::-1].T.ravel() * t2signa
if noccb > 1 and nvirb > 1:
ooidx = numpy.tril_indices(noccb, -1)
vvidx = numpy.tril_indices(nvirb, -1)
c2bb = c2bb[ooidx][:, vvidx[0], vvidx[1]]
t2addrb, t2signb = t2strs(norbb, noccb)
fcivec[0, t2addrb] = c2bb[::-1].T.ravel() * t2signb
return fcivec
def enlarge_space(myci, civec_strs, eri, norb, nelec):
if isinstance(civec_strs, (tuple, list)):
nelec, (strsa, strsb) = _unpack(civec_strs[0], nelec)[1:]
ci_coeff = lib.asarray(civec_strs)
else:
ci_coeff, nelec, (strsa, strsb) = _unpack(civec_strs, nelec)
na = len(strsa)
nb = len(strsb)
ci0 = ci_coeff.reshape(-1, na, nb)
civec_a_max = lib.norm(ci0, axis=2).max(axis=0)
civec_b_max = lib.norm(ci0, axis=1).max(axis=0)
ci_aidx = numpy.where(civec_a_max > myci.ci_coeff_cutoff)[0]
ci_bidx = numpy.where(civec_b_max > myci.ci_coeff_cutoff)[0]
civec_a_max = civec_a_max[ci_aidx]
civec_b_max = civec_b_max[ci_bidx]
strsa = strsa[ci_aidx]
strsb = strsb[ci_bidx]
eri = ao2mo.restore(1, eri, norb)
eri_pq_max = abs(eri.reshape(norb**2, -1)).max(axis=1).reshape(norb, norb)
strsa_add = select_strs(myci, eri, eri_pq_max, civec_a_max, strsa, norb,
nelec[0])
strsb_add = select_strs(myci, eri, eri_pq_max, civec_b_max, strsb, norb,
nelec[1])
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 < len(ci_aidx))[0]
bidx = numpy.where(bidx < len(ci_bidx))[0]
ma = len(strsa)
mb = len(strsb)
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(_as_SCIvector(ci1, ci_strs))
if not isinstance(civec_strs, (tuple, list)) and civec_strs.ndim < 3:
cs = cs[0]
return cs
def contract_2e(eri, fcivec, norb, nelec, link_index=None, orbsym=None, wfnsym=0):
if orbsym is None:
return direct_spin0.contract_2e(eri, fcivec, norb, nelec, link_index)
eri = ao2mo.restore(4, eri, norb)
neleca, nelecb = direct_spin1._unpack_nelec(nelec)
assert(neleca == nelecb)
link_indexa = direct_spin0._unpack(norb, nelec, link_index)
na, nlinka = link_indexa.shape[:2]
eri_irs, rank_eri, irrep_eri = direct_spin1_symm.reorder_eri(eri, norb, orbsym)
totirrep = len(eri_irs)
strsa = numpy.asarray(cistring.gen_strings4orblist(range(norb), neleca))
aidx, link_indexa = direct_spin1_symm.gen_str_irrep(strsa, orbsym, link_indexa,
rank_eri, irrep_eri, totirrep)
Tirrep = ctypes.c_void_p*totirrep
linka_ptr = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in link_indexa])
eri_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in eri_irs])
dimirrep = (ctypes.c_int*totirrep)(*[x.shape[0] for x in eri_irs])
fcivec_shape = fcivec.shape
fcivec = fcivec.reshape((na,na), order='C')
ci1new = numpy.zeros_like(fcivec)
nas = (ctypes.c_int*8)(*[x.size for x in aidx])
ci0 = []
ci1 = []
for ir in range(totirrep):
ma, mb = aidx[ir].size, aidx[wfnsym^ir].size
ci0.append(numpy.zeros((ma,mb)))
ci1.append(numpy.zeros((ma,mb)))
if ma > 0 and mb > 0:
lib.take_2d(fcivec, aidx[ir], aidx[wfnsym^ir], out=ci0[ir])
ci0_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in ci0])
ci1_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in ci1])
libfci.FCIcontract_2e_symm1(eri_ptrs, ci0_ptrs, ci1_ptrs,
ctypes.c_int(norb), nas, nas,
ctypes.c_int(nlinka), ctypes.c_int(nlinka),
linka_ptr, linka_ptr, dimirrep,
ctypes.c_int(totirrep), ctypes.c_int(wfnsym))
for ir in range(totirrep):
if ci0[ir].size > 0:
lib.takebak_2d(ci1new, ci1[ir], aidx[ir], aidx[wfnsym^ir])
return lib.transpose_sum(ci1new, inplace=True).reshape(fcivec_shape)
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
def enlarge_space(myci, civec_strs, eri, norb, nelec):
if isinstance(civec_strs, (tuple, list)):
nelec, (strsa, strsb) = _unpack(civec_strs[0], nelec)[1:]
ci_coeff = lib.asarray(civec_strs)
else:
ci_coeff, nelec, (strsa, strsb) = _unpack(civec_strs, nelec)
na = len(strsa)
nb = len(strsb)
ci0 = ci_coeff.reshape(-1,na,nb)
civec_a_max = lib.norm(ci0, axis=2).max(axis=0)
civec_b_max = lib.norm(ci0, axis=1).max(axis=0)
ci_aidx = numpy.where(civec_a_max > myci.ci_coeff_cutoff)[0]
ci_bidx = numpy.where(civec_b_max > myci.ci_coeff_cutoff)[0]
civec_a_max = civec_a_max[ci_aidx]
civec_b_max = civec_b_max[ci_bidx]
strsa = strsa[ci_aidx]
strsb = strsb[ci_bidx]
eri = ao2mo.restore(1, eri, norb)
eri_pq_max = abs(eri.reshape(norb**2,-1)).max(axis=1).reshape(norb,norb)
strsa_add = select_strs(myci, eri, eri_pq_max, civec_a_max, strsa, norb, nelec[0])
strsb_add = select_strs(myci, eri, eri_pq_max, civec_b_max, strsb, norb, nelec[1])
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 < len(ci_aidx))[0]
bidx = numpy.where(bidx < len(ci_bidx))[0]
ma = len(strsa)
mb = len(strsb)
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(_as_SCIvector(ci1, ci_strs))
if not isinstance(civec_strs, (tuple, list)) and civec_strs.ndim < 3:
cs = cs[0]
return cs
def vector_to_amplitudes(vector, nmo, nocc):
"""
vector to amps, with the same bahavior as pyscf gccsd.
"""
nvir = nmo - nocc
nov = nocc * nvir
nocc2 = nocc * (nocc - 1) // 2
otril = np.tril_indices(nocc, k=-1)
vlocs = [_task_location(nvir, task_id) for task_id in range(mpi.pool.size)]
vloc0, vloc1 = vlocs[rank]
nvir_seg = vloc1 - vloc0
vtril = get_vtril(nvir, vlocs[rank])
nvir2 = len(vtril[0])
if rank == 0:
t1T = vector[:nov].copy().reshape(nvir, nocc)
mpi.bcast(t1T)
t2tril = vector[nov:].reshape(nvir2, nocc2)
else:
t1T = mpi.bcast(None)
t2tril = vector.reshape(nvir2, nocc2)
t2T = np.zeros((nvir_seg * nvir, nocc**2), dtype=t2tril.dtype)
lib.takebak_2d(t2T, t2tril, vtril[0]*nvir+vtril[1], otril[0]*nocc+otril[1])
# anti-symmetry when exchanging two particle indices
lib.takebak_2d(t2T, -t2tril, vtril[0]*nvir+vtril[1], otril[1]*nocc+otril[0])
t2T = t2T.reshape(nvir_seg, nvir, nocc, nocc)
t2tmp = mpi.alltoall_new([t2T[:, p0:p1] for p0,p1 in vlocs], split_recvbuf=True)
for task_id, (p0, p1) in enumerate(vlocs):
if task_id < rank:
# do not need this part since it is already filled.
continue
elif task_id == rank:
# fill the trlu by -tril.
v_idx = get_vtril(nvir, vlocs[task_id], p0=p0, p1=p1)
tmp = t2tmp[task_id].reshape(p1-p0, nvir_seg, nocc, nocc)
t2T[v_idx[1]-p0, v_idx[0]+p0] = tmp[v_idx[0], v_idx[1]-p0].transpose(0, 2, 1)
else:
tmp = t2tmp[task_id].reshape(p1-p0, nvir_seg, nocc, nocc)
t2T[:, p0:p1] = tmp.transpose(1, 0, 3, 2)
return t1T.T, t2T.transpose(2, 3, 0, 1)
def enlarge_space(myci, civec_strs, eri, norb, nelec):
if isinstance(civec_strs, (tuple, list)):
nelec, (strsa, strsb) = select_ci._unpack(civec_strs[0], nelec)[1:]
ci_coeff = lib.asarray(civec_strs)
else:
ci_coeff, nelec, (strsa, strsb) = select_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 = select_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(select_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
def to_fci(cisdvec, nmoa_nmob, nocca_noccb):
from pyscf import fci
from pyscf.ci.cisd_slow import t2strs
norba, norbb = nmoa_nmob
nocca, noccb = nocca_noccb
nvira = norba - nocca
nvirb = norbb - noccb
c0, c1, c2 = cisdvec_to_amplitudes(cisdvec, nmoa_nmob, nocca_noccb)
c1a, c1b = c1
c2aa, c2ab, c2bb = c2
t1addra, t1signa = cisd.t1strs(norba, nocca)
t1addrb, t1signb = cisd.t1strs(norbb, noccb)
na = fci.cistring.num_strings(norba, nocca)
nb = fci.cistring.num_strings(norbb, noccb)
fcivec = numpy.zeros((na, nb))
fcivec[0, 0] = c0
fcivec[t1addra, 0] = c1a[::-1].T.ravel() * t1signa
fcivec[0, t1addrb] = c1b[::-1].T.ravel() * t1signb
c2ab = c2ab[::-1, ::-1].transpose(2, 0, 3, 1).reshape(nocca * nvira, -1)
c2ab = numpy.einsum('i,j,ij->ij', t1signa, t1signb, c2ab)
lib.takebak_2d(fcivec, c2ab, t1addra, t1addrb)
if nocca > 1 and nvira > 1:
ooidx = numpy.tril_indices(nocca, -1)
vvidx = numpy.tril_indices(nvira, -1)
c2aa = c2aa[ooidx][:, vvidx[0], vvidx[1]]
t2addra, t2signa = t2strs(norba, nocca)
fcivec[t2addra, 0] = c2aa[::-1].T.ravel() * t2signa
if noccb > 1 and nvirb > 1:
ooidx = numpy.tril_indices(noccb, -1)
vvidx = numpy.tril_indices(nvirb, -1)
c2bb = c2bb[ooidx][:, vvidx[0], vvidx[1]]
t2addrb, t2signb = t2strs(norbb, noccb)
fcivec[0, t2addrb] = c2bb[::-1].T.ravel() * t2signb
return fcivec
def to_fci(cisdvec, norb, nelec):
if isinstance(nelec, (int, numpy.number)):
nelecb = nelec//2
neleca = nelec - nelecb
else:
neleca, nelecb = nelec
nocc = neleca
nvir = norb - nocc
c0 = cisdvec[0]
c1 = cisdvec[1:nocc*nvir+1].reshape(nocc,nvir)
c2 = cisdvec[nocc*nvir+1:].reshape(nocc,nocc,nvir,nvir)
t1addr, t1sign = t1strs(norb, nocc)
na = cistring.num_strings(norb, nocc)
fcivec = numpy.zeros((na,na))
fcivec[0,0] = c0
c1 = c1[::-1].T.ravel()
fcivec[0,t1addr] = fcivec[t1addr,0] = c1 * t1sign
c2ab = c2[::-1,::-1].transpose(2,0,3,1).reshape(nocc*nvir,-1)
c2ab = numpy.einsum('i,j,ij->ij', t1sign, t1sign, c2ab)
lib.takebak_2d(fcivec, c2ab, t1addr, t1addr)
if nocc > 1 and nvir > 1:
hf_str = int('1'*nocc, 2)
for a in range(nocc, norb):
for b in range(nocc, a):
for i in reversed(range(1, nocc)):
for j in reversed(range(i)):
c2aa = c2[i,j,a-nocc,b-nocc] - c2[j,i,a-nocc,b-nocc]
str1 = hf_str ^ (1 << j) | (1 << b)
c2aa*= cistring.cre_des_sign(b, j, hf_str)
c2aa*= cistring.cre_des_sign(a, i, str1)
str1^= (1 << i) | (1 << a)
addr = cistring.str2addr(norb, nocc, str1)
fcivec[0,addr] = fcivec[addr,0] = c2aa
return fcivec
def _restore(mat, nvira, nvirb):
mat1 = numpy.zeros((nvira * nvira, nvirb * nvirb))
idxa = numpy.tril_indices(nvira)
idxb = numpy.tril_indices(nvirb)
lib.takebak_2d(mat1, mat, idxa[0] * nvira + idxa[1],
idxb[0] * nvirb + idxb[1])
lib.takebak_2d(mat1, mat, idxa[0] * nvira + idxa[1],
idxb[1] * nvirb + idxb[0])
lib.takebak_2d(mat1, mat, idxa[1] * nvira + idxa[0],
idxb[0] * nvirb + idxb[1])
lib.takebak_2d(mat1, mat, idxa[1] * nvira + idxa[0],
idxb[1] * nvirb + idxb[0])
mat1 = mat1.reshape(nvira, nvira, nvirb, nvirb)
idxa = numpy.arange(nvira)
idxb = numpy.arange(nvirb)
mat1[idxa, idxa] *= .5
mat1[:, :, idxb, idxb] *= .5
return mat1
def _get_vvVV(eris):
if eris.vvVV is None and getattr(eris, 'VVL', None) is not None: # DF eris
vvL = np.asarray(eris.vvL)
VVL = np.asarray(eris.VVL)
vvVV = lib.dot(vvL, VVL.T)
elif len(eris.vvVV.shape) == 2:
vvVV = np.asarray(eris.vvVV)
else:
return eris.vvVV
nvira = int(np.sqrt(vvVV.shape[0]*2))
nvirb = int(np.sqrt(vvVV.shape[1]*2))
vvVV1 = np.zeros((nvira**2,nvirb**2))
vtrila = np.tril_indices(nvira)
vtrilb = np.tril_indices(nvirb)
lib.takebak_2d(vvVV1, vvVV, vtrila[0]*nvira+vtrila[1], vtrilb[0]*nvirb+vtrilb[1])
lib.takebak_2d(vvVV1, vvVV, vtrila[1]*nvira+vtrila[0], vtrilb[1]*nvirb+vtrilb[0])
lib.takebak_2d(vvVV1, vvVV, vtrila[0]*nvira+vtrila[1], vtrilb[1]*nvirb+vtrilb[0])
lib.takebak_2d(vvVV1, vvVV, vtrila[1]*nvira+vtrila[0], vtrilb[0]*nvirb+vtrilb[1])
return vvVV1.reshape(nvira,nvira,nvirb,nvirb)
def _get_vvVV(eris):
if eris.vvVV is None and hasattr(eris, 'VVL'): # DF eris
vvL = np.asarray(eris.vvL)
VVL = np.asarray(eris.VVL)
vvVV = lib.dot(vvL, VVL.T)
elif len(eris.vvVV.shape) == 2:
vvVV = np.asarray(eris.vvVV)
else:
return eris.vvVV
nvira = int(np.sqrt(vvVV.shape[0]*2))
nvirb = int(np.sqrt(vvVV.shape[1]*2))
vvVV1 = np.zeros((nvira**2,nvirb**2))
vtrila = np.tril_indices(nvira)
vtrilb = np.tril_indices(nvirb)
lib.takebak_2d(vvVV1, vvVV, vtrila[0]*nvira+vtrila[1], vtrilb[0]*nvirb+vtrilb[1])
lib.takebak_2d(vvVV1, vvVV, vtrila[1]*nvira+vtrila[0], vtrilb[1]*nvirb+vtrilb[0])
lib.takebak_2d(vvVV1, vvVV, vtrila[0]*nvira+vtrila[1], vtrilb[1]*nvirb+vtrilb[0])
lib.takebak_2d(vvVV1, vvVV, vtrila[1]*nvira+vtrila[0], vtrilb[0]*nvirb+vtrilb[1])
return vvVV1.reshape(nvira,nvira,nvirb,nvirb)
def restore_t2_inplace(t2T):
tmp = numpy.zeros((nvir**2,nocc**2), dtype=t2T.dtype)
lib.takebak_2d(tmp, t2T.reshape(nvir**2,nocc**2), vv_idx, oo_idx)
t2 = lib.transpose(tmp.reshape(nvir**2,nocc**2), out=t2T)
return t2.reshape(nocc,nocc,nvir,nvir)
def update_amps(mycc, t1, t2, eris):
time0 = time.clock(), time.time()
log = logger.Logger(mycc.stdout, mycc.verbose)
nocc, nvir = t1.shape
nov = nocc*nvir
fock = eris.fock
t1t2new = numpy.zeros((nov+nov**2))
t1new = t1t2new[:nov].reshape(t1.shape)
t2new = t1t2new[nov:].reshape(t2.shape)
t2new_tril = numpy.zeros((nocc*(nocc+1)//2,nvir,nvir))
mycc.add_wvvVV_(t1, t2, eris, t2new_tril)
idxo = numpy.tril_indices(nocc)
lib.takebak_2d(t2new.reshape(nocc**2,nvir**2), t2new_tril.reshape(-1,nvir**2),
idxo[0]*nocc+idxo[1], numpy.arange(nvir**2))
idxo = numpy.arange(nocc)
t2new[idxo,idxo] *= .5
t2new_tril = None
time1 = log.timer_debug1('vvvv', *time0)
#** make_inter_F
fov = fock[:nocc,nocc:].copy()
t1new += fov
foo = fock[:nocc,:nocc].copy()
foo[range(nocc),range(nocc)] = 0
foo += .5 * numpy.einsum('ia,ja->ij', fock[:nocc,nocc:], t1)
fvv = fock[nocc:,nocc:].copy()
fvv[range(nvir),range(nvir)] = 0
fvv -= .5 * numpy.einsum('ia,ib->ab', t1, fock[:nocc,nocc:])
#: woooo = numpy.einsum('la,ikja->ikjl', t1, eris.ooov)
eris_ooov = _cp(eris.ooov)
foo += numpy.einsum('kc,jikc->ij', 2*t1, eris_ooov)
foo += numpy.einsum('kc,jkic->ij', -t1, eris_ooov)
woooo = lib.ddot(eris_ooov.reshape(-1,nvir), t1.T).reshape((nocc,)*4)
woooo = lib.transpose_sum(woooo.reshape(nocc**2,nocc**2), inplace=True)
woooo += _cp(eris.oooo).reshape(nocc**2,nocc**2)
woooo = _cp(woooo.reshape(nocc,nocc,nocc,nocc).transpose(0,2,1,3))
eris_ooov = None
time1 = log.timer_debug1('woooo', *time1)
unit = _memory_usage_inloop(nocc, nvir)
max_memory = max(2000, mycc.max_memory - lib.current_memory()[0])
blksize = min(nocc, max(BLKMIN, int(max_memory/unit)))
blknvir = int((max_memory*.9e6/8-blksize*nocc*nvir**2*6)/(blksize*nvir**2*2))
blknvir = min(nvir, max(BLKMIN, blknvir))
log.debug1('max_memory %d MB, nocc,nvir = %d,%d blksize = %d,%d',
max_memory, nocc, nvir, blksize, blknvir)
nvir_pair = nvir * (nvir+1) // 2
def prefect_ovvv(p0, p1, q0, q1, prefetch):
if q1 != nvir:
q0, q1 = q1, min(nvir, q1+blknvir)
readbuf = numpy.ndarray((p1-p0,q1-q0,nvir_pair), buffer=prefetch)
readbuf[:] = eris.ovvv[p0:p1,q0:q1]
def prefect_ovov(p0, p1, buf):
buf[:] = eris.ovov[p0:p1]
def prefect_oovv(p0, p1, buf):
buf[:] = eris.oovv[p0:p1]
buflen = max(nocc*nvir**2, nocc**3)
bufs = numpy.empty((5,blksize*buflen))
buf1, buf2, buf3, buf4, buf5 = bufs
for p0, p1 in prange(0, nocc, blksize):
#: wOoVv += numpy.einsum('iabc,jc->ijab', eris.ovvv, t1)
#: wOoVv -= numpy.einsum('jbik,ka->jiba', eris.ovoo, t1)
wOoVv = numpy.ndarray((nocc,p1-p0,nvir,nvir), buffer=buf3)
wooVV = numpy.ndarray((p1-p0,nocc,nvir,nvir), buffer=buf4)
handler = None
readbuf = numpy.empty((p1-p0,blknvir,nvir_pair))
prefetchbuf = numpy.empty((p1-p0,blknvir,nvir_pair))
ovvvbuf = numpy.empty((p1-p0,blknvir,nvir,nvir))
for q0, q1 in lib.prange(0, nvir, blknvir):
if q0 == 0:
readbuf[:] = eris.ovvv[p0:p1,q0:q1]
else:
readbuf, prefetchbuf = prefetchbuf, readbuf
handler = async_do(handler, prefect_ovvv, p0, p1, q0, q1, prefetchbuf)
eris_ovvv = numpy.ndarray(((p1-p0)*(q1-q0),nvir_pair), buffer=readbuf)
#:eris_ovvv = _cp(eris.ovvv[p0:p1,q0:q1])
eris_ovvv = lib.unpack_tril(eris_ovvv, out=ovvvbuf)
eris_ovvv = eris_ovvv.reshape(p1-p0,q1-q0,nvir,nvir)
#: tau = t2 + numpy.einsum('ia,jb->ijab', t1, t1)
#: tmp = numpy.einsum('ijcd,kcdb->ijbk', tau, eris.ovvv)
#: t2new += numpy.einsum('ka,ijbk->ijab', -t1, tmp)
if not mycc.direct:
eris_vovv = lib.transpose(eris_ovvv.reshape(-1,nvir))
eris_vovv = eris_vovv.reshape(nvir*(p1-p0),-1)
tmp = numpy.ndarray((nocc,nocc,nvir,p1-p0), buffer=buf1)
for j0, j1 in prange(0, nocc, blksize):
tau = numpy.ndarray((j1-j0,nocc,q1-q0,nvir), buffer=buf2)
tau = numpy.einsum('ia,jb->ijab', t1[j0:j1,q0:q1], t1, out=tau)
tau += t2[j0:j1,:,q0:q1]
lib.ddot(tau.reshape((j1-j0)*nocc,-1), eris_vovv.T, 1,
tmp[j0:j1].reshape((j1-j0)*nocc,-1), 0)
tmp1 = numpy.ndarray((nocc,nocc,nvir,p1-p0), buffer=buf2)
tmp1[:] = tmp.transpose(1,0,2,3)
lib.ddot(tmp1.reshape(-1,p1-p0), t1[p0:p1], -1, t2new.reshape(-1,nvir), 1)
eris_vovv = tau = tmp1 = tmp = None
fvv += numpy.einsum('kc,kcba->ab', 2*t1[p0:p1,q0:q1], eris_ovvv)
fvv[:,q0:q1] += numpy.einsum('kc,kbca->ab', -t1[p0:p1], eris_ovvv)
#: wooVV -= numpy.einsum('jc,icba->ijba', t1, eris_ovvv)
tmp = t1[:,q0:q1].copy()
for i in range(eris_ovvv.shape[0]):
lib.ddot(tmp, eris_ovvv[i].reshape(q1-q0,-1), -1,
wooVV[i].reshape(nocc,-1))
#: wOoVv += numpy.einsum('ibac,jc->jiba', eris_ovvv, t1)
tmp = numpy.ndarray((nocc,p1-p0,q1-q0,nvir), buffer=buf1)
lib.ddot(t1, eris_ovvv.reshape(-1,nvir).T, 1, tmp.reshape(nocc,-1))
wOoVv[:,:,q0:q1] = tmp
#: theta = t2.transpose(1,0,2,3) * 2 - t2
#: t1new += numpy.einsum('ijcb,jcba->ia', theta, eris.ovvv)
theta = tmp
theta[:] = t2[p0:p1,:,q0:q1,:].transpose(1,0,2,3)
theta *= 2
theta -= t2[:,p0:p1,q0:q1,:]
lib.ddot(theta.reshape(nocc,-1), eris_ovvv.reshape(-1,nvir), 1, t1new, 1)
theta = tmp = None
handler.join()
readbuf = prefetchbuf = ovvvbuf = eris_ovvv = None
time2 = log.timer_debug1('ovvv [%d:%d]'%(p0, p1), *time1)
tmp = numpy.ndarray((nocc,p1-p0,nvir,nocc), buffer=buf1)
tmp[:] = _cp(eris.ovoo[p0:p1]).transpose(2,0,1,3)
lib.ddot(tmp.reshape(-1,nocc), t1, -1, wOoVv.reshape(-1,nvir), 1)
eris_ooov = _cp(eris.ooov[p0:p1])
eris_oovv = numpy.empty((p1-p0,nocc,nvir,nvir))
handler = lib.background_thread(prefect_oovv, p0, p1, eris_oovv)
tmp = numpy.ndarray((p1-p0,nocc,nvir,nocc), buffer=buf1)
tmp[:] = eris_ooov.transpose(0,1,3,2)
#: wooVV = numpy.einsum('ka,ijkb->ijba', t1, eris.ooov[p0:p1])
lib.ddot(tmp.reshape(-1,nocc), t1, 1, wooVV.reshape(-1,nvir), 1)
t2new[p0:p1] += wOoVv.transpose(1,0,2,3)
#:eris_oovv = _cp(eris.oovv[p0:p1])
handler.join()
eris_ovov = numpy.empty((p1-p0,nvir,nocc,nvir))
handler = lib.background_thread(prefect_ovov, p0, p1, eris_ovov)
#: g2 = 2 * eris.oOVv - eris.oovv
#: t1new += numpy.einsum('jb,ijba->ia', t1, g2)
t1new[p0:p1] += numpy.einsum('jb,ijba->ia', -t1, eris_oovv)
wooVV -= eris_oovv
#tmp = numpy.einsum('ic,jkbc->jikb', t1, eris_oovv)
#t2new[p0:p1] += numpy.einsum('ka,jikb->ijba', -t1, tmp)
tmp1 = numpy.ndarray((nocc,nocc*nvir), buffer=buf1)
tmp2 = numpy.ndarray((nocc*nvir,nocc), buffer=buf2)
for j in range(p1-p0):
tmp = lib.ddot(t1, eris_oovv[j].reshape(-1,nvir).T, 1, tmp1)
lib.transpose(_cp(tmp).reshape(nocc,nocc,nvir), axes=(0,2,1), out=tmp2)
t2new[:,p0+j] -= lib.ddot(tmp2, t1).reshape(nocc,nvir,nvir)
eris_oovv = None
#:eris_ovov = _cp(eris.ovov[p0:p1])
handler.join()
for i in range(p1-p0):
t2new[p0+i] += eris_ovov[i].transpose(1,0,2) * .5
t1new[p0:p1] += numpy.einsum('jb,iajb->ia', 2*t1, eris_ovov)
#:tmp = numpy.einsum('ic,jbkc->jibk', t1, eris_ovov)
#:t2new[p0:p1] += numpy.einsum('ka,jibk->jiba', -t1, tmp)
for j in range(p1-p0):
lib.ddot(t1, eris_ovov[j].reshape(-1,nvir).T, 1, tmp1)
lib.ddot(tmp1.reshape(-1,nocc), t1, -1, t2new[p0+j].reshape(-1,nvir), 1)
tmp1 = tmp2 = tmp = None
fov[p0:p1] += numpy.einsum('kc,iakc->ia', t1, eris_ovov) * 2
fov[p0:p1] -= numpy.einsum('kc,icka->ia', t1, eris_ovov)
#: fvv -= numpy.einsum('ijca,ibjc->ab', theta, eris.ovov)
#: foo += numpy.einsum('iakb,jkba->ij', eris.ovov, theta)
tau = numpy.ndarray((nocc,nvir,nvir), buffer=buf1)
theta = numpy.ndarray((nocc,nvir,nvir), buffer=buf2)
for i in range(p1-p0):
tau = numpy.einsum('a,jb->jab', t1[p0+i]*.5, t1, out=tau)
tau += t2[p0+i]
theta = lib.transpose(tau, axes=(0,2,1), out=theta)
theta *= 2
theta -= tau
vov = lib.transpose(eris_ovov[i].reshape(nvir,-1), out=tau)
lib.ddot(vov.reshape(nocc,-1), theta.reshape(nocc,-1).T, 1, foo, 1)
lib.ddot(theta.reshape(-1,nvir).T, eris_ovov[i].reshape(nvir,-1).T, -1, fvv, 1)
tau = theta = vov = None
#: theta = t2.transpose(0,2,1,3) * 2 - t2.transpose(0,3,2,1)
#: t1new += numpy.einsum('jb,ijba->ia', fov, theta)
#: t1new -= numpy.einsum('kijb,kjba->ia', eris_ooov, theta)
theta = numpy.ndarray((p1-p0,nvir,nocc,nvir), buffer=buf1)
for i in range(p1-p0):
tmp = t2[p0+i].transpose(0,2,1) * 2
tmp-= t2[p0+i]
lib.ddot(eris_ooov[i].reshape(nocc,-1),
tmp.reshape(-1,nvir), -1, t1new, 1)
lib.transpose(_cp(tmp).reshape(-1,nvir), out=theta[i]) # theta[i] = tmp.transpose(2,0,1)
t1new += numpy.einsum('jb,jbia->ia', fov[p0:p1], theta)
eris_ooov = None
#: wOVov += eris.ovov
#: tau = theta - numpy.einsum('ic,kb->ikcb', t1, t1*2)
#: wOVov += .5 * numpy.einsum('jakc,ikcb->jiba', eris.ovov, tau)
#: wOVov -= .5 * numpy.einsum('jcka,ikcb->jiba', eris.ovov, t2)
#: t2new += numpy.einsum('ikca,kjbc->ijba', theta, wOVov)
for i in range(p1-p0):
wOoVv[:,i] += wooVV[i]*.5 #: jiba + ijba*.5
wOVov = lib.transpose(wOoVv.reshape(nocc,-1,nvir), axes=(0,2,1), out=buf5)
wOVov = wOVov.reshape(nocc,nvir,-1,nvir)
eris_OVov = lib.transpose(eris_ovov.reshape(-1,nov), out=buf3)
eris_OVov = eris_OVov.reshape(nocc,nvir,-1,nvir)
wOVov += eris_OVov
theta = theta.reshape(-1,nov)
for i in range(nocc): # OVov-OVov.transpose(0,3,2,1)*.5
eris_OVov[i] -= eris_OVov[i].transpose(2,1,0)*.5
for j0, j1 in prange(0, nocc, blksize):
tau = numpy.ndarray((j1-j0,nvir,nocc,nvir), buffer=buf2)
for i in range(j1-j0):
tau[i] = t2[j0+i].transpose(1,0,2) * 2
tau[i] -= t2[j0+i].transpose(2,0,1)
tau[i] -= numpy.einsum('a,jb->bja', t1[j0+i]*2, t1)
#: wOVov[j0:j1] += .5 * numpy.einsum('iakc,jbkc->jbai', eris_ovov, tau)
lib.ddot(tau.reshape(-1,nov), eris_OVov.reshape(nov,-1),
.5, wOVov[j0:j1].reshape((j1-j0)*nvir,-1), 1)
#theta = t2[p0:p1] * 2 - t2[p0:p1].transpose(0,1,3,2)
#: t2new[j0:j1] += numpy.einsum('iack,jbck->jiba', theta, wOVov[j0:j1])
tmp = lib.ddot(wOVov[j0:j1].reshape((j1-j0)*nvir,-1), theta, 1,
tau.reshape(-1,nov)).reshape(-1,nvir,nocc,nvir)
for i in range(j1-j0):
t2new[j0+i] += tmp[i].transpose(1,0,2)
theta = wOoVv = wOVov = eris_OVov = tmp = tau = None
time2 = log.timer_debug1('wOVov [%d:%d]'%(p0, p1), *time2)
#: tau = t2 + numpy.einsum('ia,jb->ijab', t1, t1)
#: woooo += numpy.einsum('ijba,klab->ijkl', eris.oOVv, tau)
#: tau = .5*t2 + numpy.einsum('ia,jb->ijab', t1, t1)
#: woVoV += numpy.einsum('jkca,ikbc->ijba', tau, eris.oOVv)
tmp = numpy.ndarray((p1-p0,nvir,nocc,nvir), buffer=buf1)
tmp[:] = wooVV.transpose(0,2,1,3)
woVoV = lib.transpose(_cp(tmp).reshape(-1,nov), out=buf4).reshape(nocc,nvir,p1-p0,nvir)
eris_oOvV = numpy.ndarray((p1-p0,nocc,nvir,nvir), buffer=buf3)
eris_oOvV[:] = eris_ovov.transpose(0,2,1,3)
eris_oVOv = lib.transpose(eris_oOvV.reshape(-1,nov,nvir), axes=(0,2,1), out=buf5)
eris_oVOv = eris_oVOv.reshape(-1,nvir,nocc,nvir)
for j0, j1 in prange(0, nocc, blksize):
tau = make_tau(t2[j0:j1], t1[j0:j1], t1, 1, out=buf2)
#: woooo[p0:p1,:,j0:j1] += numpy.einsum('ijab,klab->ijkl', eris_oOvV, tau)
_dgemm('N', 'T', (p1-p0)*nocc, (j1-j0)*nocc, nvir*nvir,
eris_oOvV.reshape(-1,nvir*nvir), tau.reshape(-1,nvir*nvir),
woooo[p0:p1].reshape(-1,nocc*nocc), 1, 1, 0, 0, j0*nocc)
for i in range(j1-j0):
tau[i] -= t2[j0+i] * .5
#: woVoV[j0:j1] += numpy.einsum('jkca,ickb->jiab', tau, eris_ovov)
lib.ddot(lib.transpose(tau.reshape(-1,nov,nvir), axes=(0,2,1)).reshape(-1,nov),
eris_oVOv.reshape(-1,nov).T,
1, woVoV[j0:j1].reshape((j1-j0)*nvir,-1), 1)
time2 = log.timer_debug1('woVoV [%d:%d]'%(p0, p1), *time2)
tau = make_tau(t2[p0:p1], t1[p0:p1], t1, 1, out=buf2)
#: t2new += .5 * numpy.einsum('klij,klab->ijab', woooo[p0:p1], tau)
lib.ddot(woooo[p0:p1].reshape(-1,nocc*nocc).T, tau.reshape(-1,nvir*nvir),
.5, t2new.reshape(nocc*nocc,-1), 1)
eris_ovov = eris_oVOv = eris_oOvV = wooVV = tau = tmp = None
t2ibja = lib.transpose(_cp(t2[p0:p1]).reshape(-1,nov,nvir), axes=(0,2,1),
out=buf1).reshape(-1,nvir,nocc,nvir)
tmp = numpy.ndarray((blksize,nvir,nocc,nvir), buffer=buf2)
for j0, j1 in prange(0, nocc, blksize):
#: t2new[j0:j1] += numpy.einsum('ibkc,kcja->ijab', woVoV[j0:j1], t2ibja)
lib.ddot(woVoV[j0:j1].reshape((j1-j0)*nvir,-1),
t2ibja.reshape(-1,nov), 1, tmp[:j1-j0].reshape(-1,nov))
for i in range(j1-j0):
t2new[j0+i] += tmp[i].transpose(1,2,0)
t2new[j0+i] += tmp[i].transpose(1,0,2) * .5
woVoV = t2ibja = tmp = None
time1 = log.timer_debug1('contract occ [%d:%d]'%(p0, p1), *time1)
buf1 = buf2 = buf3 = buf4 = buf5 = bufs = None
time1 = log.timer_debug1('contract loop', *time0)
woooo = None
ft_ij = foo + numpy.einsum('ja,ia->ij', .5*t1, fov)
ft_ab = fvv - numpy.einsum('ia,ib->ab', .5*t1, fov)
#: t2new += numpy.einsum('ijac,bc->ijab', t2, ft_ab)
#: t2new -= numpy.einsum('ki,kjab->ijab', ft_ij, t2)
lib.ddot(t2.reshape(-1,nvir), ft_ab.T, 1, t2new.reshape(-1,nvir), 1)
lib.ddot(ft_ij.T, t2.reshape(nocc,nocc*nvir**2),-1,
t2new.reshape(nocc,nocc*nvir**2), 1)
mo_e = fock.diagonal()
eia = mo_e[:nocc,None] - mo_e[None,nocc:]
t1new += numpy.einsum('ib,ab->ia', t1, fvv)
t1new -= numpy.einsum('ja,ji->ia', t1, foo)
t1new /= eia
#: t2new = t2new + t2new.transpose(1,0,3,2)
ij = 0
for i in range(nocc):
for j in range(i+1):
t2new[i,j] += t2new[j,i].T
t2new[i,j] /= lib.direct_sum('a,b->ab', eia[i], eia[j])
t2new[j,i] = t2new[i,j].T
ij += 1
time0 = log.timer_debug1('update t1 t2', *time0)
return t1new, t2new
def spatial2spin(tx, orbspin, kconserv):
'''Convert T1/T2 of spatial orbital representation to T1/T2 of
spin-orbital representation
'''
if isinstance(tx, numpy.ndarray) and tx.ndim == 3:
# KRCCSD t1 amplitudes
return spatial2spin((tx, tx), orbspin, kconserv)
elif isinstance(tx, numpy.ndarray) and tx.ndim == 7:
# KRCCSD t2 amplitudes
t2aa = tx - tx.transpose(0, 1, 2, 4, 3, 5, 6)
return spatial2spin((t2aa, tx, t2aa), orbspin, kconserv)
elif len(tx) == 2: # KUCCSD t1
t1a, t1b = tx
nocc_a, nvir_a = t1a.shape[1:]
nocc_b, nvir_b = t1b.shape[1:]
else: # KUCCSD t2
t2aa, t2ab, t2bb = tx
nocc_a, nocc_b, nvir_a, nvir_b = t2ab.shape[3:]
nkpts = len(orbspin)
nocc = nocc_a + nocc_b
nvir = nvir_a + nvir_b
idxoa = [numpy.where(orbspin[k][:nocc] == 0)[0] for k in range(nkpts)]
idxob = [numpy.where(orbspin[k][:nocc] == 1)[0] for k in range(nkpts)]
idxva = [numpy.where(orbspin[k][nocc:] == 0)[0] for k in range(nkpts)]
idxvb = [numpy.where(orbspin[k][nocc:] == 1)[0] for k in range(nkpts)]
if len(tx) == 2: # t1
t1 = numpy.zeros((nkpts, nocc, nvir), dtype=t1a.dtype)
for k in range(nkpts):
lib.takebak_2d(t1[k], t1a[k], idxoa[k], idxva[k])
lib.takebak_2d(t1[k], t1b[k], idxob[k], idxvb[k])
t1 = lib.tag_array(t1, orbspin=orbspin)
return t1
else:
t2 = numpy.zeros((nkpts, nkpts, nkpts, nocc**2, nvir**2),
dtype=t2aa.dtype)
for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
kb = kconserv[ki, ka, kj]
idxoaa = idxoa[ki][:, None] * nocc + idxoa[kj]
idxoab = idxoa[ki][:, None] * nocc + idxob[kj]
idxoba = idxob[kj][:, None] * nocc + idxoa[ki]
idxobb = idxob[ki][:, None] * nocc + idxob[kj]
idxvaa = idxva[ka][:, None] * nvir + idxva[kb]
idxvab = idxva[ka][:, None] * nvir + idxvb[kb]
idxvba = idxvb[kb][:, None] * nvir + idxva[ka]
idxvbb = idxvb[ka][:, None] * nvir + idxvb[kb]
tmp2aa = t2aa[ki, kj, ka].reshape(nocc_a * nocc_a, nvir_a * nvir_a)
tmp2bb = t2bb[ki, kj, ka].reshape(nocc_b * nocc_b, nvir_b * nvir_b)
tmp2ab = t2ab[ki, kj, ka].reshape(nocc_a * nocc_b, nvir_a * nvir_b)
lib.takebak_2d(t2[ki, kj, ka], tmp2aa, idxoaa.ravel(),
idxvaa.ravel())
lib.takebak_2d(t2[ki, kj, ka], tmp2bb, idxobb.ravel(),
idxvbb.ravel())
lib.takebak_2d(t2[ki, kj, ka], tmp2ab, idxoab.ravel(),
idxvab.ravel())
lib.takebak_2d(t2[kj, ki, kb], tmp2ab, idxoba.T.ravel(),
idxvba.T.ravel())
abba = -tmp2ab
lib.takebak_2d(t2[ki, kj, kb], abba, idxoab.ravel(),
idxvba.T.ravel())
lib.takebak_2d(t2[kj, ki, ka], abba, idxoba.T.ravel(),
idxvab.ravel())
t2 = t2.reshape(nkpts, nkpts, nkpts, nocc, nocc, nvir, nvir)
t2 = lib.tag_array(t2, orbspin=orbspin)
return t2
def to_fcivec(cisdvec, norb, nelec, frozen=0):
'''Convert CISD coefficients to FCI coefficients'''
if isinstance(nelec, (int, numpy.number)):
nelecb = nelec // 2
neleca = nelec - nelecb
else:
neleca, nelecb = nelec
assert (neleca == nelecb)
frozen_mask = numpy.zeros(norb, dtype=bool)
if isinstance(frozen, (int, numpy.integer)):
nfroz = frozen
frozen_mask[:frozen] = True
else:
nfroz = len(frozen)
frozen_mask[frozen] = True
nocc = numpy.count_nonzero(~frozen_mask[:neleca])
nmo = norb - nfroz
nvir = nmo - nocc
c0, c1, c2 = cisdvec_to_amplitudes(cisdvec, nmo, nocc)
t1addr, t1sign = t1strs(nmo, nocc)
na = cistring.num_strings(nmo, nocc)
fcivec = numpy.zeros((na, na))
fcivec[0, 0] = c0
c1 = c1[::-1].T.ravel()
fcivec[0, t1addr] = fcivec[t1addr, 0] = c1 * t1sign
c2ab = c2[::-1, ::-1].transpose(2, 0, 3, 1).reshape(nocc * nvir, -1)
c2ab = numpy.einsum('i,j,ij->ij', t1sign, t1sign, c2ab)
lib.takebak_2d(fcivec, c2ab, t1addr, t1addr)
if nocc > 1 and nvir > 1:
hf_str = int('1' * nocc, 2)
for a in range(nocc, nmo):
for b in range(nocc, a):
for i in reversed(range(1, nocc)):
for j in reversed(range(i)):
c2aa = c2[i, j, a - nocc,
b - nocc] - c2[j, i, a - nocc, b - nocc]
str1 = hf_str ^ (1 << j) | (1 << b)
c2aa *= cistring.cre_des_sign(b, j, hf_str)
c2aa *= cistring.cre_des_sign(a, i, str1)
str1 ^= (1 << i) | (1 << a)
addr = cistring.str2addr(nmo, nocc, str1)
fcivec[0, addr] = fcivec[addr, 0] = c2aa
if nfroz == 0:
return fcivec
assert (norb < 63)
strs = cistring.gen_strings4orblist(range(norb), neleca)
na = len(strs)
count = numpy.zeros(na, dtype=int)
parity = numpy.zeros(na, dtype=bool)
core_mask = numpy.ones(na, dtype=bool)
# During the loop, count saves the number of occupied orbitals that
# lower (with small orbital ID) than the present orbital i.
# Moving all the frozen orbitals to the beginning of the orbital list
# (before the occupied orbitals) leads to parity odd (= True, with
# negative sign) or even (= False, with positive sign).
for i in range(norb):
if frozen_mask[i]:
if i < neleca:
# frozen occupied orbital should be occupied
core_mask &= (strs & (1 << i)) != 0
parity ^= (count & 1) == 1
else:
# frozen virtual orbital should not be occupied.
# parity is not needed since it's unoccupied
core_mask &= (strs & (1 << i)) == 0
else:
count += (strs & (1 << i)) != 0
sub_strs = strs[core_mask & (count == nocc)]
addrs = cistring.strs2addr(norb, neleca, sub_strs)
fcivec1 = numpy.zeros((na, na))
fcivec1[addrs[:, None], addrs] = fcivec
fcivec1[parity, :] *= -1
fcivec1[:, parity] *= -1
return fcivec1
def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True):
dovov, dovOV, dOVov, dOVOV = d2[0]
dvvvv, dvvVV, dVVvv, dVVVV = d2[1]
doooo, dooOO, dOOoo, dOOOO = d2[2]
doovv, dooVV, dOOvv, dOOVV = d2[3]
dovvo, dovVO, dOVvo, dOVVO = d2[4]
dvvov, dvvOV, dVVov, dVVOV = d2[5]
dovvv, dovVV, dOVvv, dOVVV = d2[6]
dooov, dooOV, dOOov, dOOOV = d2[7]
nocca, nvira, noccb, nvirb = dovOV.shape
nmoa = nocca + nvira
nmob = noccb + nvirb
dm2aa = numpy.empty((nmoa,nmoa,nmoa,nmoa), dtype=doovv.dtype)
dm2ab = numpy.empty((nmoa,nmoa,nmob,nmob), dtype=doovv.dtype)
dm2bb = numpy.empty((nmob,nmob,nmob,nmob), dtype=doovv.dtype)
# dm2aa
dovov = numpy.asarray(dovov)
dm2aa[:nocca,nocca:,:nocca,nocca:] = dovov
dm2aa[nocca:,:nocca,nocca:,:nocca] = dm2aa[:nocca,nocca:,:nocca,nocca:].transpose(1,0,3,2).conj()
dovov = None
#assert(abs(doovv+dovvo.transpose(0,3,2,1)).max() == 0)
dovvo = numpy.asarray(dovvo)
dm2aa[:nocca,:nocca,nocca:,nocca:] =-dovvo.transpose(0,3,2,1)
dm2aa[nocca:,nocca:,:nocca,:nocca] = dm2aa[:nocca,:nocca,nocca:,nocca:].transpose(2,3,0,1)
dm2aa[:nocca,nocca:,nocca:,:nocca] = dovvo
dm2aa[nocca:,:nocca,:nocca,nocca:] = dm2aa[:nocca,nocca:,nocca:,:nocca].transpose(1,0,3,2).conj()
dovvo = None
if len(dvvvv.shape) == 2:
dvvvv = ao2mo.restore(1, dvvvv, nvira)
dm2aa[nocca:,nocca:,nocca:,nocca:] = dvvvv
dm2aa[:nocca,:nocca,:nocca,:nocca] = doooo
dovvv = numpy.asarray(dovvv)
dm2aa[:nocca,nocca:,nocca:,nocca:] = dovvv
dm2aa[nocca:,nocca:,:nocca,nocca:] = dovvv.transpose(2,3,0,1)
dm2aa[nocca:,nocca:,nocca:,:nocca] = dovvv.transpose(3,2,1,0).conj()
dm2aa[nocca:,:nocca,nocca:,nocca:] = dovvv.transpose(1,0,3,2).conj()
dovvv = None
dooov = numpy.asarray(dooov)
dm2aa[:nocca,:nocca,:nocca,nocca:] = dooov
dm2aa[:nocca,nocca:,:nocca,:nocca] = dooov.transpose(2,3,0,1)
dm2aa[:nocca,:nocca,nocca:,:nocca] = dooov.transpose(1,0,3,2).conj()
dm2aa[nocca:,:nocca,:nocca,:nocca] = dooov.transpose(3,2,1,0).conj()
dooov = None
# dm2bb
dOVOV = numpy.asarray(dOVOV)
dm2bb[:noccb,noccb:,:noccb,noccb:] = dOVOV
dm2bb[noccb:,:noccb,noccb:,:noccb] = dm2bb[:noccb,noccb:,:noccb,noccb:].transpose(1,0,3,2).conj()
dOVOV = None
dOVVO = numpy.asarray(dOVVO)
dm2bb[:noccb,:noccb,noccb:,noccb:] =-dOVVO.transpose(0,3,2,1)
dm2bb[noccb:,noccb:,:noccb,:noccb] = dm2bb[:noccb,:noccb,noccb:,noccb:].transpose(2,3,0,1)
dm2bb[:noccb,noccb:,noccb:,:noccb] = dOVVO
dm2bb[noccb:,:noccb,:noccb,noccb:] = dm2bb[:noccb,noccb:,noccb:,:noccb].transpose(1,0,3,2).conj()
dOVVO = None
if len(dVVVV.shape) == 2:
dVVVV = ao2mo.restore(1, dVVVV, nvirb)
dm2bb[noccb:,noccb:,noccb:,noccb:] = dVVVV
dm2bb[:noccb,:noccb,:noccb,:noccb] = dOOOO
dOVVV = numpy.asarray(dOVVV)
dm2bb[:noccb,noccb:,noccb:,noccb:] = dOVVV
dm2bb[noccb:,noccb:,:noccb,noccb:] = dOVVV.transpose(2,3,0,1)
dm2bb[noccb:,noccb:,noccb:,:noccb] = dOVVV.transpose(3,2,1,0).conj()
dm2bb[noccb:,:noccb,noccb:,noccb:] = dOVVV.transpose(1,0,3,2).conj()
dOVVV = None
dOOOV = numpy.asarray(dOOOV)
dm2bb[:noccb,:noccb,:noccb,noccb:] = dOOOV
dm2bb[:noccb,noccb:,:noccb,:noccb] = dOOOV.transpose(2,3,0,1)
dm2bb[:noccb,:noccb,noccb:,:noccb] = dOOOV.transpose(1,0,3,2).conj()
dm2bb[noccb:,:noccb,:noccb,:noccb] = dOOOV.transpose(3,2,1,0).conj()
dOOOV = None
# dm2ab
dovOV = numpy.asarray(dovOV)
dm2ab[:nocca,nocca:,:noccb,noccb:] = dovOV
dm2ab[nocca:,:nocca,noccb:,:noccb] = dm2ab[:nocca,nocca:,:noccb,noccb:].transpose(1,0,3,2).conj()
dovOV = None
dovVO = numpy.asarray(dovVO)
dm2ab[:nocca,:nocca,noccb:,noccb:] = dooVV
dm2ab[nocca:,nocca:,:noccb,:noccb] = dOOvv.transpose(2,3,0,1)
dm2ab[:nocca,nocca:,noccb:,:noccb] = dovVO
dm2ab[nocca:,:nocca,:noccb,noccb:] = dovVO.transpose(1,0,3,2).conj()
dovVO = None
if len(dvvVV.shape) == 2:
idxa = numpy.tril_indices(nvira)
dvvVV1 = lib.unpack_tril(dvvVV)
dvvVV = numpy.empty((nvira,nvira,nvirb,nvirb))
dvvVV[idxa] = dvvVV1
dvvVV[idxa[1],idxa[0]] = dvvVV1
dvvVV1 = None
dm2ab[nocca:,nocca:,noccb:,noccb:] = dvvVV
dm2ab[:nocca,:nocca,:noccb,:noccb] = dooOO
dovVV = numpy.asarray(dovVV)
dm2ab[:nocca,nocca:,noccb:,noccb:] = dovVV
dm2ab[nocca:,nocca:,:noccb,noccb:] = dOVvv.transpose(2,3,0,1)
dm2ab[nocca:,nocca:,noccb:,:noccb] = dOVvv.transpose(3,2,1,0).conj()
dm2ab[nocca:,:nocca,noccb:,noccb:] = dovVV.transpose(1,0,3,2).conj()
dovVV = None
dooOV = numpy.asarray(dooOV)
dm2ab[:nocca,:nocca,:noccb,noccb:] = dooOV
dm2ab[:nocca,nocca:,:noccb,:noccb] = dOOov.transpose(2,3,0,1)
dm2ab[:nocca,:nocca,noccb:,:noccb] = dooOV.transpose(1,0,3,2).conj()
dm2ab[nocca:,:nocca,:noccb,:noccb] = dOOov.transpose(3,2,1,0).conj()
dooOV = None
if with_frozen and not (mycc.frozen is 0 or mycc.frozen is None):
nmoa0 = dm2aa.shape[0]
nmob0 = dm2bb.shape[0]
nmoa = mycc.mo_occ[0].size
nmob = mycc.mo_occ[1].size
nocca = numpy.count_nonzero(mycc.mo_occ[0] > 0)
noccb = numpy.count_nonzero(mycc.mo_occ[1] > 0)
rdm2aa = numpy.zeros((nmoa,nmoa,nmoa,nmoa), dtype=dm2aa.dtype)
rdm2ab = numpy.zeros((nmoa,nmoa,nmob,nmob), dtype=dm2ab.dtype)
rdm2bb = numpy.zeros((nmob,nmob,nmob,nmob), dtype=dm2bb.dtype)
moidxa, moidxb = mycc.get_frozen_mask()
moidxa = numpy.where(moidxa)[0]
moidxb = numpy.where(moidxb)[0]
idxa = (moidxa.reshape(-1,1) * nmoa + moidxa).ravel()
idxb = (moidxb.reshape(-1,1) * nmob + moidxb).ravel()
lib.takebak_2d(rdm2aa.reshape(nmoa**2,nmoa**2),
dm2aa.reshape(nmoa0**2,nmoa0**2), idxa, idxa)
lib.takebak_2d(rdm2bb.reshape(nmob**2,nmob**2),
dm2bb.reshape(nmob0**2,nmob0**2), idxb, idxb)
lib.takebak_2d(rdm2ab.reshape(nmoa**2,nmob**2),
dm2ab.reshape(nmoa0**2,nmob0**2), idxa, idxb)
dm2aa, dm2ab, dm2bb = rdm2aa, rdm2ab, rdm2bb
if with_dm1:
dm1a, dm1b = _make_rdm1(mycc, d1, with_frozen=True)
dm1a[numpy.diag_indices(nocca)] -= 1
dm1b[numpy.diag_indices(noccb)] -= 1
for i in range(nocca):
dm2aa[i,i,:,:] += dm1a
dm2aa[:,:,i,i] += dm1a
dm2aa[:,i,i,:] -= dm1a
dm2aa[i,:,:,i] -= dm1a.T
dm2ab[i,i,:,:] += dm1b
for i in range(noccb):
dm2bb[i,i,:,:] += dm1b
dm2bb[:,:,i,i] += dm1b
dm2bb[:,i,i,:] -= dm1b
dm2bb[i,:,:,i] -= dm1b.T
dm2ab[:,:,i,i] += dm1a
for i in range(nocca):
for j in range(nocca):
dm2aa[i,i,j,j] += 1
dm2aa[i,j,j,i] -= 1
for i in range(noccb):
for j in range(noccb):
dm2bb[i,i,j,j] += 1
dm2bb[i,j,j,i] -= 1
for i in range(nocca):
for j in range(noccb):
dm2ab[i,i,j,j] += 1
dm2aa = dm2aa.transpose(1,0,3,2)
dm2ab = dm2ab.transpose(1,0,3,2)
dm2bb = dm2bb.transpose(1,0,3,2)
return dm2aa, dm2ab, dm2bb
def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True, ao_repr=False):
r'''
dm2[p,q,r,s] = \sum_{sigma,tau} <p_sigma^\dagger r_tau^\dagger s_tau q_sigma>
Note the contraction between ERIs (in Chemist's notation) and rdm2 is
E = einsum('pqrs,pqrs', eri, rdm2)
'''
dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = d2
nocc, nvir = dovov.shape[:2]
nmo = nocc + nvir
dm2 = numpy.empty((nmo,nmo,nmo,nmo), dtype=doovv.dtype)
dovov = numpy.asarray(dovov)
dm2[:nocc,nocc:,:nocc,nocc:] = dovov
dm2[:nocc,nocc:,:nocc,nocc:]+= dovov.transpose(2,3,0,1)
dm2[nocc:,:nocc,nocc:,:nocc] = dm2[:nocc,nocc:,:nocc,nocc:].transpose(1,0,3,2).conj()
dovov = None
doovv = numpy.asarray(doovv)
dm2[:nocc,:nocc,nocc:,nocc:] = doovv
dm2[:nocc,:nocc,nocc:,nocc:]+= doovv.transpose(1,0,3,2).conj()
dm2[nocc:,nocc:,:nocc,:nocc] = dm2[:nocc,:nocc,nocc:,nocc:].transpose(2,3,0,1)
doovv = None
dovvo = numpy.asarray(dovvo)
dm2[:nocc,nocc:,nocc:,:nocc] = dovvo
dm2[:nocc,nocc:,nocc:,:nocc]+= dovvo.transpose(3,2,1,0).conj()
dm2[nocc:,:nocc,:nocc,nocc:] = dm2[:nocc,nocc:,nocc:,:nocc].transpose(1,0,3,2).conj()
dovvo = None
if len(dvvvv.shape) == 2:
# To handle the case of compressed vvvv, which is used in nuclear gradients
dvvvv = ao2mo.restore(1, dvvvv, nvir)
dm2[nocc:,nocc:,nocc:,nocc:] = dvvvv
dm2[nocc:,nocc:,nocc:,nocc:]*= 4
else:
dvvvv = numpy.asarray(dvvvv)
dm2[nocc:,nocc:,nocc:,nocc:] = dvvvv
dm2[nocc:,nocc:,nocc:,nocc:]+= dvvvv.transpose(1,0,3,2).conj()
dm2[nocc:,nocc:,nocc:,nocc:]*= 2
dvvvv = None
doooo = numpy.asarray(doooo)
dm2[:nocc,:nocc,:nocc,:nocc] = doooo
dm2[:nocc,:nocc,:nocc,:nocc]+= doooo.transpose(1,0,3,2).conj()
dm2[:nocc,:nocc,:nocc,:nocc]*= 2
doooo = None
dovvv = numpy.asarray(dovvv)
dm2[:nocc,nocc:,nocc:,nocc:] = dovvv
dm2[nocc:,nocc:,:nocc,nocc:] = dovvv.transpose(2,3,0,1)
dm2[nocc:,nocc:,nocc:,:nocc] = dovvv.transpose(3,2,1,0).conj()
dm2[nocc:,:nocc,nocc:,nocc:] = dovvv.transpose(1,0,3,2).conj()
dovvv = None
dooov = numpy.asarray(dooov)
dm2[:nocc,:nocc,:nocc,nocc:] = dooov
dm2[:nocc,nocc:,:nocc,:nocc] = dooov.transpose(2,3,0,1)
dm2[:nocc,:nocc,nocc:,:nocc] = dooov.transpose(1,0,3,2).conj()
dm2[nocc:,:nocc,:nocc,:nocc] = dooov.transpose(3,2,1,0).conj()
if with_frozen and mycc.frozen is not None:
nmo, nmo0 = mycc.mo_occ.size, nmo
nocc = numpy.count_nonzero(mycc.mo_occ > 0)
rdm2 = numpy.zeros((nmo,nmo,nmo,nmo), dtype=dm2.dtype)
moidx = numpy.where(mycc.get_frozen_mask())[0]
idx = (moidx.reshape(-1,1) * nmo + moidx).ravel()
lib.takebak_2d(rdm2.reshape(nmo**2,nmo**2),
dm2.reshape(nmo0**2,nmo0**2), idx, idx)
dm2 = rdm2
if with_dm1:
dm1 = _make_rdm1(mycc, d1, with_frozen)
dm1[numpy.diag_indices(nocc)] -= 2
for i in range(nocc):
dm2[i,i,:,:] += dm1 * 2
dm2[:,:,i,i] += dm1 * 2
dm2[:,i,i,:] -= dm1
dm2[i,:,:,i] -= dm1.T
for i in range(nocc):
for j in range(nocc):
dm2[i,i,j,j] += 4
dm2[i,j,j,i] -= 2
# dm2 was computed as dm2[p,q,r,s] = < p^\dagger r^\dagger s q > in the
# above. Transposing it so that it be contracted with ERIs (in Chemist's
# notation):
# E = einsum('pqrs,pqrs', eri, rdm2)
dm2 = dm2.transpose(1,0,3,2)
if ao_repr:
dm2 = _rdm2_mo2ao(dm2.transpose(1,0,3,2), mycc.mo_coeff)
return dm2
def contract_2e(eri, fcivec, norb, nelec, link_index=None, orbsym=None, wfnsym=0):
if orbsym is None:
return direct_spin1.contract_2e(eri, fcivec, norb, nelec, link_index)
eri = ao2mo.restore(4, eri, norb)
neleca, nelecb = _unpack_nelec(nelec)
link_indexa, link_indexb = direct_spin1._unpack(norb, nelec, link_index)
na, nlinka = link_indexa.shape[:2]
nb, nlinkb = link_indexb.shape[:2]
eri_irs, rank_eri, irrep_eri = reorder_eri(eri, norb, orbsym)
strsa = cistring.gen_strings4orblist(range(norb), neleca)
aidx, link_indexa = gen_str_irrep(strsa, orbsym, link_indexa, rank_eri, irrep_eri)
if neleca == nelecb:
bidx, link_indexb = aidx, link_indexa
else:
strsb = cistring.gen_strings4orblist(range(norb), nelecb)
bidx, link_indexb = gen_str_irrep(strsb, orbsym, link_indexb, rank_eri, irrep_eri)
Tirrep = ctypes.c_void_p*TOTIRREPS
linka_ptr = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in link_indexa])
linkb_ptr = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in link_indexb])
eri_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in eri_irs])
dimirrep = (ctypes.c_int*TOTIRREPS)(*[x.shape[0] for x in eri_irs])
fcivec_shape = fcivec.shape
fcivec = fcivec.reshape((na,nb), order='C')
ci1new = numpy.zeros_like(fcivec)
nas = (ctypes.c_int*TOTIRREPS)(*[x.size for x in aidx])
nbs = (ctypes.c_int*TOTIRREPS)(*[x.size for x in bidx])
# aa, ab
ci0 = []
ci1 = []
for ir in range(TOTIRREPS):
ma, mb = aidx[ir].size, bidx[wfnsym^ir].size
ci0.append(numpy.zeros((ma,mb)))
ci1.append(numpy.zeros((ma,mb)))
if ma > 0 and mb > 0:
lib.take_2d(fcivec, aidx[ir], bidx[wfnsym^ir], out=ci0[ir])
ci0_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in ci0])
ci1_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in ci1])
libfci.FCIcontract_2e_symm1(eri_ptrs, ci0_ptrs, ci1_ptrs,
ctypes.c_int(norb), nas, nbs,
ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
linka_ptr, linkb_ptr, dimirrep,
ctypes.c_int(wfnsym))
for ir in range(TOTIRREPS):
if ci0[ir].size > 0:
lib.takebak_2d(ci1new, ci1[ir], aidx[ir], bidx[wfnsym^ir])
# bb, ba
ci0T = []
for ir in range(TOTIRREPS):
mb, ma = bidx[ir].size, aidx[wfnsym^ir].size
ci0T.append(numpy.zeros((mb,ma)))
if ma > 0 and mb > 0:
lib.transpose(ci0[wfnsym^ir], out=ci0T[ir])
ci0, ci0T = ci0T, None
ci1 = [numpy.zeros_like(x) for x in ci0]
ci0_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in ci0])
ci1_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in ci1])
libfci.FCIcontract_2e_symm1(eri_ptrs, ci0_ptrs, ci1_ptrs,
ctypes.c_int(norb), nbs, nas,
ctypes.c_int(nlinkb), ctypes.c_int(nlinka),
linkb_ptr, linka_ptr, dimirrep,
ctypes.c_int(wfnsym))
for ir in range(TOTIRREPS):
if ci0[ir].size > 0:
lib.takebak_2d(ci1new, lib.transpose(ci1[ir]), aidx[wfnsym^ir], bidx[ir])
return ci1new.reshape(fcivec_shape)
def spatial2spin(tx, orbspin):
'''call orbspin_of_sorted_mo_energy to get orbspin'''
if len(tx) == 2: # t1
t1a, t1b = tx
nocc_a, nvir_a = t1a.shape
nocc_b, nvir_b = t1b.shape
else:
t2aa, t2ab, t2bb = tx
nocc_a, nocc_b, nvir_a, nvir_b = t2ab.shape
nocc = nocc_a + nocc_b
nvir = nvir_a + nvir_b
idxoa = numpy.where(orbspin[:nocc] == 0)[0]
idxob = numpy.where(orbspin[:nocc] == 1)[0]
idxva = numpy.where(orbspin[nocc:] == 0)[0]
idxvb = numpy.where(orbspin[nocc:] == 1)[0]
if len(tx) == 2: # t1
t1 = numpy.zeros((nocc, nvir), dtype=t1a.dtype)
lib.takebak_2d(t1, t1a, idxoa, idxva)
lib.takebak_2d(t1, t1b, idxob, idxvb)
return t1
else:
t2 = numpy.zeros((nocc**2, nvir**2), dtype=t2aa.dtype)
idxoaa = idxoa[:, None] * nocc + idxoa
idxoab = idxoa[:, None] * nocc + idxob
idxoba = idxob[:, None] * nocc + idxoa
idxobb = idxob[:, None] * nocc + idxob
idxvaa = idxva[:, None] * nvir + idxva
idxvab = idxva[:, None] * nvir + idxvb
idxvba = idxvb[:, None] * nvir + idxva
idxvbb = idxvb[:, None] * nvir + idxvb
t2aa = t2aa.reshape(nocc_a * nocc_a, nvir_a * nvir_a)
t2ab = t2ab.reshape(nocc_a * nocc_b, nvir_a * nvir_b)
t2bb = t2bb.reshape(nocc_b * nocc_b, nvir_b * nvir_b)
lib.takebak_2d(t2, t2aa, idxoaa.ravel(), idxvaa.ravel())
lib.takebak_2d(t2, t2bb, idxobb.ravel(), idxvbb.ravel())
lib.takebak_2d(t2, t2ab, idxoab.ravel(), idxvab.ravel())
lib.takebak_2d(t2, t2ab, idxoba.T.ravel(), idxvba.T.ravel())
abba = -t2ab
lib.takebak_2d(t2, abba, idxoab.ravel(), idxvba.T.ravel())
lib.takebak_2d(t2, abba, idxoba.T.ravel(), idxvab.ravel())
return t2.reshape(nocc, nocc, nvir, nvir)
def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True):
dovov, dovOV, dOVov, dOVOV = d2[0]
dvvvv, dvvVV, dVVvv, dVVVV = d2[1]
doooo, dooOO, dOOoo, dOOOO = d2[2]
doovv, dooVV, dOOvv, dOOVV = d2[3]
dovvo, dovVO, dOVvo, dOVVO = d2[4]
dvvov, dvvOV, dVVov, dVVOV = d2[5]
dovvv, dovVV, dOVvv, dOVVV = d2[6]
dooov, dooOV, dOOov, dOOOV = d2[7]
nocca, nvira, noccb, nvirb = dovOV.shape
nmoa = nocca + nvira
nmob = noccb + nvirb
dm2aa = numpy.empty((nmoa, nmoa, nmoa, nmoa), dtype=doovv.dtype)
dm2ab = numpy.empty((nmoa, nmoa, nmob, nmob), dtype=doovv.dtype)
dm2bb = numpy.empty((nmob, nmob, nmob, nmob), dtype=doovv.dtype)
# dm2aa
dovov = numpy.asarray(dovov)
dm2aa[:nocca, nocca:, :nocca, nocca:] = dovov
dm2aa[nocca:, :nocca,
nocca:, :nocca] = dm2aa[:nocca, nocca:, :nocca,
nocca:].transpose(1, 0, 3, 2).conj()
dovov = None
#assert(abs(doovv+dovvo.transpose(0,3,2,1)).max() == 0)
dovvo = numpy.asarray(dovvo)
dm2aa[:nocca, :nocca, nocca:, nocca:] = -dovvo.transpose(0, 3, 2, 1)
dm2aa[nocca:,
nocca:, :nocca, :nocca] = dm2aa[:nocca, :nocca, nocca:,
nocca:].transpose(2, 3, 0, 1)
dm2aa[:nocca, nocca:, nocca:, :nocca] = dovvo
dm2aa[nocca:, :nocca, :nocca,
nocca:] = dm2aa[:nocca, nocca:,
nocca:, :nocca].transpose(1, 0, 3, 2).conj()
dovvo = None
if len(dvvvv.shape) == 2:
dvvvv = ao2mo.restore(1, dvvvv, nvira)
dm2aa[nocca:, nocca:, nocca:, nocca:] = dvvvv
dm2aa[:nocca, :nocca, :nocca, :nocca] = doooo
dovvv = numpy.asarray(dovvv)
dm2aa[:nocca, nocca:, nocca:, nocca:] = dovvv
dm2aa[nocca:, nocca:, :nocca, nocca:] = dovvv.transpose(2, 3, 0, 1)
dm2aa[nocca:, nocca:, nocca:, :nocca] = dovvv.transpose(3, 2, 1, 0).conj()
dm2aa[nocca:, :nocca, nocca:, nocca:] = dovvv.transpose(1, 0, 3, 2).conj()
dovvv = None
dooov = numpy.asarray(dooov)
dm2aa[:nocca, :nocca, :nocca, nocca:] = dooov
dm2aa[:nocca, nocca:, :nocca, :nocca] = dooov.transpose(2, 3, 0, 1)
dm2aa[:nocca, :nocca, nocca:, :nocca] = dooov.transpose(1, 0, 3, 2).conj()
dm2aa[nocca:, :nocca, :nocca, :nocca] = dooov.transpose(3, 2, 1, 0).conj()
dooov = None
# dm2bb
dOVOV = numpy.asarray(dOVOV)
dm2bb[:noccb, noccb:, :noccb, noccb:] = dOVOV
dm2bb[noccb:, :noccb,
noccb:, :noccb] = dm2bb[:noccb, noccb:, :noccb,
noccb:].transpose(1, 0, 3, 2).conj()
dOVOV = None
dOVVO = numpy.asarray(dOVVO)
dm2bb[:noccb, :noccb, noccb:, noccb:] = -dOVVO.transpose(0, 3, 2, 1)
dm2bb[noccb:,
noccb:, :noccb, :noccb] = dm2bb[:noccb, :noccb, noccb:,
noccb:].transpose(2, 3, 0, 1)
dm2bb[:noccb, noccb:, noccb:, :noccb] = dOVVO
dm2bb[noccb:, :noccb, :noccb,
noccb:] = dm2bb[:noccb, noccb:,
noccb:, :noccb].transpose(1, 0, 3, 2).conj()
dOVVO = None
if len(dVVVV.shape) == 2:
dVVVV = ao2mo.restore(1, dVVVV, nvirb)
dm2bb[noccb:, noccb:, noccb:, noccb:] = dVVVV
dm2bb[:noccb, :noccb, :noccb, :noccb] = dOOOO
dOVVV = numpy.asarray(dOVVV)
dm2bb[:noccb, noccb:, noccb:, noccb:] = dOVVV
dm2bb[noccb:, noccb:, :noccb, noccb:] = dOVVV.transpose(2, 3, 0, 1)
dm2bb[noccb:, noccb:, noccb:, :noccb] = dOVVV.transpose(3, 2, 1, 0).conj()
dm2bb[noccb:, :noccb, noccb:, noccb:] = dOVVV.transpose(1, 0, 3, 2).conj()
dOVVV = None
dOOOV = numpy.asarray(dOOOV)
dm2bb[:noccb, :noccb, :noccb, noccb:] = dOOOV
dm2bb[:noccb, noccb:, :noccb, :noccb] = dOOOV.transpose(2, 3, 0, 1)
dm2bb[:noccb, :noccb, noccb:, :noccb] = dOOOV.transpose(1, 0, 3, 2).conj()
dm2bb[noccb:, :noccb, :noccb, :noccb] = dOOOV.transpose(3, 2, 1, 0).conj()
dOOOV = None
# dm2ab
dovOV = numpy.asarray(dovOV)
dm2ab[:nocca, nocca:, :noccb, noccb:] = dovOV
dm2ab[nocca:, :nocca,
noccb:, :noccb] = dm2ab[:nocca, nocca:, :noccb,
noccb:].transpose(1, 0, 3, 2).conj()
dovOV = None
dovVO = numpy.asarray(dovVO)
dm2ab[:nocca, :nocca, noccb:, noccb:] = dooVV
dm2ab[nocca:, nocca:, :noccb, :noccb] = dOOvv.transpose(2, 3, 0, 1)
dm2ab[:nocca, nocca:, noccb:, :noccb] = dovVO
dm2ab[nocca:, :nocca, :noccb, noccb:] = dovVO.transpose(1, 0, 3, 2).conj()
dovVO = None
if len(dvvVV.shape) == 2:
idxa = numpy.tril_indices(nvira)
dvvVV1 = lib.unpack_tril(dvvVV)
dvvVV = numpy.empty((nvira, nvira, nvirb, nvirb))
dvvVV[idxa] = dvvVV1
dvvVV[idxa[1], idxa[0]] = dvvVV1
dvvVV1 = None
dm2ab[nocca:, nocca:, noccb:, noccb:] = dvvVV
dm2ab[:nocca, :nocca, :noccb, :noccb] = dooOO
dovVV = numpy.asarray(dovVV)
dm2ab[:nocca, nocca:, noccb:, noccb:] = dovVV
dm2ab[nocca:, nocca:, :noccb, noccb:] = dOVvv.transpose(2, 3, 0, 1)
dm2ab[nocca:, nocca:, noccb:, :noccb] = dOVvv.transpose(3, 2, 1, 0).conj()
dm2ab[nocca:, :nocca, noccb:, noccb:] = dovVV.transpose(1, 0, 3, 2).conj()
dovVV = None
dooOV = numpy.asarray(dooOV)
dm2ab[:nocca, :nocca, :noccb, noccb:] = dooOV
dm2ab[:nocca, nocca:, :noccb, :noccb] = dOOov.transpose(2, 3, 0, 1)
dm2ab[:nocca, :nocca, noccb:, :noccb] = dooOV.transpose(1, 0, 3, 2).conj()
dm2ab[nocca:, :nocca, :noccb, :noccb] = dOOov.transpose(3, 2, 1, 0).conj()
dooOV = None
if with_frozen and not (mycc.frozen is 0 or mycc.frozen is None):
nmoa0 = dm2aa.shape[0]
nmob0 = dm2bb.shape[0]
nmoa = mycc.mo_occ[0].size
nmob = mycc.mo_occ[1].size
nocca = numpy.count_nonzero(mycc.mo_occ[0] > 0)
noccb = numpy.count_nonzero(mycc.mo_occ[1] > 0)
rdm2aa = numpy.zeros((nmoa, nmoa, nmoa, nmoa), dtype=dm2aa.dtype)
rdm2ab = numpy.zeros((nmoa, nmoa, nmob, nmob), dtype=dm2ab.dtype)
rdm2bb = numpy.zeros((nmob, nmob, nmob, nmob), dtype=dm2bb.dtype)
moidxa, moidxb = mycc.get_frozen_mask()
moidxa = numpy.where(moidxa)[0]
moidxb = numpy.where(moidxb)[0]
idxa = (moidxa.reshape(-1, 1) * nmoa + moidxa).ravel()
idxb = (moidxb.reshape(-1, 1) * nmob + moidxb).ravel()
lib.takebak_2d(rdm2aa.reshape(nmoa**2, nmoa**2),
dm2aa.reshape(nmoa0**2, nmoa0**2), idxa, idxa)
lib.takebak_2d(rdm2bb.reshape(nmob**2, nmob**2),
dm2bb.reshape(nmob0**2, nmob0**2), idxb, idxb)
lib.takebak_2d(rdm2ab.reshape(nmoa**2, nmob**2),
dm2ab.reshape(nmoa0**2, nmob0**2), idxa, idxb)
dm2aa, dm2ab, dm2bb = rdm2aa, rdm2ab, rdm2bb
if with_dm1:
dm1a, dm1b = _make_rdm1(mycc, d1, with_frozen=True)
dm1a[numpy.diag_indices(nocca)] -= 1
dm1b[numpy.diag_indices(noccb)] -= 1
for i in range(nocca):
dm2aa[i, i, :, :] += dm1a
dm2aa[:, :, i, i] += dm1a
dm2aa[:, i, i, :] -= dm1a
dm2aa[i, :, :, i] -= dm1a.T
dm2ab[i, i, :, :] += dm1b
for i in range(noccb):
dm2bb[i, i, :, :] += dm1b
dm2bb[:, :, i, i] += dm1b
dm2bb[:, i, i, :] -= dm1b
dm2bb[i, :, :, i] -= dm1b.T
dm2ab[:, :, i, i] += dm1a
for i in range(nocca):
for j in range(nocca):
dm2aa[i, i, j, j] += 1
dm2aa[i, j, j, i] -= 1
for i in range(noccb):
for j in range(noccb):
dm2bb[i, i, j, j] += 1
dm2bb[i, j, j, i] -= 1
for i in range(nocca):
for j in range(noccb):
dm2ab[i, i, j, j] += 1
dm2aa = dm2aa.transpose(1, 0, 3, 2)
dm2ab = dm2ab.transpose(1, 0, 3, 2)
dm2bb = dm2bb.transpose(1, 0, 3, 2)
return dm2aa, dm2ab, dm2bb
def spatial2spin_ea_doublet(r1, r2, kconserv, kshift, orbspin=None):
'''Convert R1/R2 of spatial orbital representation to R1/R2 of
spin orbital representation'''
r1a, r1b = r1
r2aaa, r2aba, r2bab, r2bbb = r2
nkpts, nocc_a, nvir_a = np.array(r2aaa.shape)[[0, 2, 3]]
nkpts, nocc_b, nvir_b = np.array(r2bbb.shape)[[0, 2, 3]]
if orbspin is None:
orbspin = np.zeros((nocc_a+nvir_a)*2, dtype=int)
orbspin[1::2] = 1
nocc = nocc_a + nocc_b
nvir = nvir_a + nvir_b
idxoa = [np.where(orbspin[k][:nocc] == 0)[0] for k in range(nkpts)]
idxob = [np.where(orbspin[k][:nocc] == 1)[0] for k in range(nkpts)]
idxva = [np.where(orbspin[k][nocc:] == 0)[0] for k in range(nkpts)]
idxvb = [np.where(orbspin[k][nocc:] == 1)[0] for k in range(nkpts)]
r1 = np.zeros((nvir), dtype=r1a.dtype)
r1[idxva[kshift]] = r1a
r1[idxvb[kshift]] = r1b
r2 = np.zeros((nkpts,nkpts,nocc,nvir**2), dtype=r2aaa.dtype)
for kj, ka in itertools.product(range(nkpts), repeat=2):
kb = kconserv[kshift, ka, kj]
idxvaa = idxva[ka][:,None] * nvir + idxva[kb]
idxvab = idxva[ka][:,None] * nvir + idxvb[kb]
idxvba = idxvb[ka][:,None] * nvir + idxva[kb]
idxvbb = idxvb[ka][:,None] * nvir + idxvb[kb]
r2aaa_tmp = r2aaa[kj,ka].reshape(nocc_a, nvir_a*nvir_a)
r2aba_tmp = r2aba[kj,ka].reshape(nocc_a, nvir_b*nvir_a)
r2bab_tmp = r2bab[kj,ka].reshape(nocc_b, nvir_a*nvir_b)
r2bbb_tmp = r2bbb[kj,ka].reshape(nocc_b, nvir_b*nvir_b)
lib.takebak_2d(r2[kj,ka], r2aaa_tmp, idxoa[kj], idxvaa.ravel())
lib.takebak_2d(r2[kj,ka], r2aba_tmp, idxoa[kj], idxvba.ravel())
lib.takebak_2d(r2[kj,ka], r2bab_tmp, idxob[kj], idxvab.ravel())
lib.takebak_2d(r2[kj,ka], r2bbb_tmp, idxob[kj], idxvbb.ravel())
r2aab_tmp = -r2aba[kj,kb].reshape(nocc_a, nvir_b*nvir_a)
r2bba_tmp = -r2bab[kj,kb].reshape(nocc_b, nvir_a*nvir_b)
lib.takebak_2d(r2[kj,ka], r2bba_tmp, idxob[kj], idxvba.T.ravel())
lib.takebak_2d(r2[kj,ka], r2aab_tmp, idxoa[kj], idxvab.T.ravel())
r2 = r2.reshape(nkpts, nkpts, nocc, nvir, nvir)
return r1, r2
def to_fcivec(cisdvec, norb, nelec, frozen=0):
from pyscf.ci.gcisd import t2strs
if isinstance(nelec, (int, numpy.number)):
nelecb = nelec // 2
neleca = nelec - nelecb
else:
neleca, nelecb = nelec
frozena_mask = numpy.zeros(norb, dtype=bool)
frozenb_mask = numpy.zeros(norb, dtype=bool)
if isinstance(frozen, (int, numpy.integer)):
nfroza = nfrozb = frozen
frozena_mask[:frozen] = True
frozenb_mask[:frozen] = True
else:
nfroza = len(frozen[0])
nfrozb = len(frozen[1])
frozena_mask[frozen[0]] = True
frozenb_mask[frozen[1]] = True
# if nfroza != nfrozb:
# raise NotImplementedError
nocca = numpy.count_nonzero(~frozena_mask[:neleca])
noccb = numpy.count_nonzero(~frozenb_mask[:nelecb])
nmo = nmoa, nmob = norb - nfroza, norb - nfrozb
nocc = nocca, noccb
nvira, nvirb = nmoa - nocca, nmob - noccb
c0, c1, c2 = cisdvec_to_amplitudes(cisdvec, nmo, nocc)
c1a, c1b = c1
c2aa, c2ab, c2bb = c2
t1addra, t1signa = cisd.t1strs(nmoa, nocca)
t1addrb, t1signb = cisd.t1strs(nmob, noccb)
na = cistring.num_strings(nmoa, nocca)
nb = cistring.num_strings(nmob, noccb)
fcivec = numpy.zeros((na, nb))
fcivec[0, 0] = c0
fcivec[t1addra, 0] = c1a[::-1].T.ravel() * t1signa
fcivec[0, t1addrb] = c1b[::-1].T.ravel() * t1signb
c2ab = c2ab[::-1, ::-1].transpose(2, 0, 3, 1).reshape(nocca * nvira, -1)
c2ab = numpy.einsum('i,j,ij->ij', t1signa, t1signb, c2ab)
lib.takebak_2d(fcivec, c2ab, t1addra, t1addrb)
if nocca > 1 and nvira > 1:
ooidx = numpy.tril_indices(nocca, -1)
vvidx = numpy.tril_indices(nvira, -1)
c2aa = c2aa[ooidx][:, vvidx[0], vvidx[1]]
t2addra, t2signa = t2strs(nmoa, nocca)
fcivec[t2addra, 0] = c2aa[::-1].T.ravel() * t2signa
if noccb > 1 and nvirb > 1:
ooidx = numpy.tril_indices(noccb, -1)
vvidx = numpy.tril_indices(nvirb, -1)
c2bb = c2bb[ooidx][:, vvidx[0], vvidx[1]]
t2addrb, t2signb = t2strs(nmob, noccb)
fcivec[0, t2addrb] = c2bb[::-1].T.ravel() * t2signb
if nfroza == nfrozb == 0:
return fcivec
assert (norb < 63)
strsa = cistring.gen_strings4orblist(range(norb), neleca)
strsb = cistring.gen_strings4orblist(range(norb), nelecb)
na = len(strsa)
nb = len(strsb)
count_a = numpy.zeros(na, dtype=int)
count_b = numpy.zeros(nb, dtype=int)
parity_a = numpy.zeros(na, dtype=bool)
parity_b = numpy.zeros(nb, dtype=bool)
core_a_mask = numpy.ones(na, dtype=bool)
core_b_mask = numpy.ones(nb, dtype=bool)
for i in range(norb):
if frozena_mask[i]:
if i < neleca:
core_a_mask &= (strsa & (1 << i)) != 0
parity_a ^= (count_a & 1) == 1
else:
core_a_mask &= (strsa & (1 << i)) == 0
else:
count_a += (strsa & (1 << i)) != 0
if frozenb_mask[i]:
if i < nelecb:
core_b_mask &= (strsb & (1 << i)) != 0
parity_b ^= (count_b & 1) == 1
else:
core_b_mask &= (strsb & (1 << i)) == 0
else:
count_b += (strsb & (1 << i)) != 0
sub_strsa = strsa[core_a_mask & (count_a == nocca)]
sub_strsb = strsb[core_b_mask & (count_b == noccb)]
addrsa = cistring.strs2addr(norb, neleca, sub_strsa)
addrsb = cistring.strs2addr(norb, nelecb, sub_strsb)
fcivec1 = numpy.zeros((na, nb))
fcivec1[addrsa[:, None], addrsb] = fcivec
fcivec1[parity_a, :] *= -1
fcivec1[:, parity_b] *= -1
return fcivec1
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
def spatial2spin(tx, orbspin=None):
'''Convert T1/T2 of spatial orbital representation to T1/T2 of
spin-orbital representation
'''
if isinstance(tx, numpy.ndarray) and tx.ndim == 2:
# RCCSD t1 amplitudes
return spatial2spin((tx,tx), orbspin)
elif isinstance(tx, numpy.ndarray) and tx.ndim == 4:
# RCCSD t2 amplitudes
t2aa = tx - tx.transpose(0,1,3,2)
return spatial2spin((t2aa,tx,t2aa), orbspin)
elif len(tx) == 2: # t1
t1a, t1b = tx
nocc_a, nvir_a = t1a.shape
nocc_b, nvir_b = t1b.shape
else:
t2aa, t2ab, t2bb = tx
nocc_a, nocc_b, nvir_a, nvir_b = t2ab.shape
if orbspin is None:
orbspin = numpy.zeros((nocc_a+nvir_a)*2, dtype=int)
orbspin[1::2] = 1
nocc = nocc_a + nocc_b
nvir = nvir_a + nvir_b
idxoa = numpy.where(orbspin[:nocc] == 0)[0]
idxob = numpy.where(orbspin[:nocc] == 1)[0]
idxva = numpy.where(orbspin[nocc:] == 0)[0]
idxvb = numpy.where(orbspin[nocc:] == 1)[0]
if len(tx) == 2: # t1
t1 = numpy.zeros((nocc,nvir), dtype=t1a.dtype)
lib.takebak_2d(t1, t1a, idxoa, idxva)
lib.takebak_2d(t1, t1b, idxob, idxvb)
t1 = lib.tag_array(t1, orbspin=orbspin)
return t1
else:
t2 = numpy.zeros((nocc**2,nvir**2), dtype=t2aa.dtype)
idxoaa = idxoa[:,None] * nocc + idxoa
idxoab = idxoa[:,None] * nocc + idxob
idxoba = idxob[:,None] * nocc + idxoa
idxobb = idxob[:,None] * nocc + idxob
idxvaa = idxva[:,None] * nvir + idxva
idxvab = idxva[:,None] * nvir + idxvb
idxvba = idxvb[:,None] * nvir + idxva
idxvbb = idxvb[:,None] * nvir + idxvb
t2aa = t2aa.reshape(nocc_a*nocc_a,nvir_a*nvir_a)
t2ab = t2ab.reshape(nocc_a*nocc_b,nvir_a*nvir_b)
t2bb = t2bb.reshape(nocc_b*nocc_b,nvir_b*nvir_b)
lib.takebak_2d(t2, t2aa, idxoaa.ravel() , idxvaa.ravel() )
lib.takebak_2d(t2, t2bb, idxobb.ravel() , idxvbb.ravel() )
lib.takebak_2d(t2, t2ab, idxoab.ravel() , idxvab.ravel() )
lib.takebak_2d(t2, t2ab, idxoba.T.ravel(), idxvba.T.ravel())
abba = -t2ab
lib.takebak_2d(t2, abba, idxoab.ravel() , idxvba.T.ravel())
lib.takebak_2d(t2, abba, idxoba.T.ravel(), idxvab.ravel() )
t2 = lib.tag_array(t2, orbspin=orbspin)
return t2.reshape(nocc,nocc,nvir,nvir)
def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True):
r'''
dm2[p,q,r,s] = <p^\dagger r^\dagger s q>
Note the contraction between ERIs (in Chemist's notation) and rdm2 is
E = einsum('pqrs,pqrs', eri, rdm2)
'''
dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = d2
nocc, nvir = dovov.shape[:2]
nmo = nocc + nvir
dm2 = numpy.empty((nmo,nmo,nmo,nmo), dtype=doooo.dtype)
dovov = numpy.asarray(dovov)
dm2[:nocc,nocc:,:nocc,nocc:] = dovov
dm2[nocc:,:nocc,nocc:,:nocc] = dm2[:nocc,nocc:,:nocc,nocc:].transpose(1,0,3,2).conj()
dovov = None
dovvo = numpy.asarray(dovvo)
dm2[:nocc,:nocc,nocc:,nocc:] =-dovvo.transpose(0,3,2,1)
dm2[nocc:,nocc:,:nocc,:nocc] =-dovvo.transpose(2,1,0,3)
dm2[:nocc,nocc:,nocc:,:nocc] = dovvo
dm2[nocc:,:nocc,:nocc,nocc:] = dovvo.transpose(1,0,3,2).conj()
dovvo = None
dm2[nocc:,nocc:,nocc:,nocc:] = dvvvv
dm2[:nocc,:nocc,:nocc,:nocc] = doooo
dovvv = numpy.asarray(dovvv)
dm2[:nocc,nocc:,nocc:,nocc:] = dovvv
dm2[nocc:,nocc:,:nocc,nocc:] = dovvv.transpose(2,3,0,1)
dm2[nocc:,nocc:,nocc:,:nocc] = dovvv.transpose(3,2,1,0).conj()
dm2[nocc:,:nocc,nocc:,nocc:] = dovvv.transpose(1,0,3,2).conj()
dovvv = None
dooov = numpy.asarray(dooov)
dm2[:nocc,:nocc,:nocc,nocc:] = dooov
dm2[:nocc,nocc:,:nocc,:nocc] = dooov.transpose(2,3,0,1)
dm2[:nocc,:nocc,nocc:,:nocc] = dooov.transpose(1,0,3,2).conj()
dm2[nocc:,:nocc,:nocc,:nocc] = dooov.transpose(3,2,1,0).conj()
if with_frozen and not (mycc.frozen is 0 or mycc.frozen is None):
nmo, nmo0 = mycc.mo_occ.size, nmo
nocc = numpy.count_nonzero(mycc.mo_occ > 0)
rdm2 = numpy.zeros((nmo,nmo,nmo,nmo), dtype=dm2.dtype)
moidx = numpy.where(mycc.get_frozen_mask())[0]
idx = (moidx.reshape(-1,1) * nmo + moidx).ravel()
lib.takebak_2d(rdm2.reshape(nmo**2,nmo**2),
dm2.reshape(nmo0**2,nmo0**2), idx, idx)
dm2 = rdm2
if with_dm1:
dm1 = _make_rdm1(mycc, d1, with_frozen)
dm1[numpy.diag_indices(nocc)] -= 1
for i in range(nocc):
dm2[i,i,:,:] += dm1
dm2[:,:,i,i] += dm1
dm2[:,i,i,:] -= dm1
dm2[i,:,:,i] -= dm1.conj()
for i in range(nocc):
for j in range(nocc):
dm2[i,i,j,j] += 1
dm2[i,j,j,i] -= 1
# dm2 was computed as dm2[p,q,r,s] = < p^\dagger r^\dagger s q > in the
# above. Transposing it so that it be contracted with ERIs (in Chemist's
# notation):
# E = einsum('pqrs,pqrs', eri, rdm2)
return dm2.transpose(1,0,3,2)
def spatial2spin(tx, orbspin, kconserv):
'''Convert T1/T2 of spatial orbital representation to T1/T2 of
spin-orbital representation
'''
if isinstance(tx, numpy.ndarray) and tx.ndim == 3:
# KRCCSD t1 amplitudes
return spatial2spin((tx,tx), orbspin, kconserv)
elif isinstance(tx, numpy.ndarray) and tx.ndim == 7:
# KRCCSD t2 amplitudes
t2aa = tx - tx.transpose(0,1,2,4,3,5,6)
return spatial2spin((t2aa,tx,t2aa), orbspin, kconserv)
elif len(tx) == 2: # KUCCSD t1
t1a, t1b = tx
nocc_a, nvir_a = t1a.shape[1:]
nocc_b, nvir_b = t1b.shape[1:]
else: # KUCCSD t2
t2aa, t2ab, t2bb = tx
nocc_a, nocc_b, nvir_a, nvir_b = t2ab.shape[3:]
nkpts = len(orbspin)
nocc = nocc_a + nocc_b
nvir = nvir_a + nvir_b
idxoa = [numpy.where(orbspin[k][:nocc] == 0)[0] for k in range(nkpts)]
idxob = [numpy.where(orbspin[k][:nocc] == 1)[0] for k in range(nkpts)]
idxva = [numpy.where(orbspin[k][nocc:] == 0)[0] for k in range(nkpts)]
idxvb = [numpy.where(orbspin[k][nocc:] == 1)[0] for k in range(nkpts)]
if len(tx) == 2: # t1
t1 = numpy.zeros((nkpts,nocc,nvir), dtype=t1a.dtype)
for k in range(nkpts):
lib.takebak_2d(t1[k], t1a[k], idxoa[k], idxva[k])
lib.takebak_2d(t1[k], t1b[k], idxob[k], idxvb[k])
t1 = lib.tag_array(t1, orbspin=orbspin)
return t1
else:
t2 = numpy.zeros((nkpts,nkpts,nkpts,nocc**2,nvir**2), dtype=t2aa.dtype)
for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
kb = kconserv[ki,ka,kj]
idxoaa = idxoa[ki][:,None] * nocc + idxoa[kj]
idxoab = idxoa[ki][:,None] * nocc + idxob[kj]
idxoba = idxob[kj][:,None] * nocc + idxoa[ki]
idxobb = idxob[ki][:,None] * nocc + idxob[kj]
idxvaa = idxva[ka][:,None] * nvir + idxva[kb]
idxvab = idxva[ka][:,None] * nvir + idxvb[kb]
idxvba = idxvb[kb][:,None] * nvir + idxva[ka]
idxvbb = idxvb[ka][:,None] * nvir + idxvb[kb]
tmp2aa = t2aa[ki,kj,ka].reshape(nocc_a*nocc_a,nvir_a*nvir_a)
tmp2bb = t2bb[ki,kj,ka].reshape(nocc_b*nocc_b,nvir_b*nvir_b)
tmp2ab = t2ab[ki,kj,ka].reshape(nocc_a*nocc_b,nvir_a*nvir_b)
lib.takebak_2d(t2[ki,kj,ka], tmp2aa, idxoaa.ravel() , idxvaa.ravel() )
lib.takebak_2d(t2[ki,kj,ka], tmp2bb, idxobb.ravel() , idxvbb.ravel() )
lib.takebak_2d(t2[ki,kj,ka], tmp2ab, idxoab.ravel() , idxvab.ravel() )
lib.takebak_2d(t2[kj,ki,kb], tmp2ab, idxoba.T.ravel(), idxvba.T.ravel())
abba = -tmp2ab
lib.takebak_2d(t2[ki,kj,kb], abba, idxoab.ravel() , idxvba.T.ravel())
lib.takebak_2d(t2[kj,ki,ka], abba, idxoba.T.ravel(), idxvab.ravel() )
t2 = t2.reshape(nkpts,nkpts,nkpts,nocc,nocc,nvir,nvir)
t2 = lib.tag_array(t2, orbspin=orbspin)
return t2
def spatial2spin_ea_doublet(r1, r2, kconserv, kshift, orbspin=None):
'''Convert R1/R2 of spatial orbital representation to R1/R2 of
spin orbital representation'''
r1a, r1b = r1
r2aaa, r2aba, r2bab, r2bbb = r2
nkpts, nocc_a, nvir_a = np.array(r2aaa.shape)[[0, 2, 3]]
nkpts, nocc_b, nvir_b = np.array(r2bbb.shape)[[0, 2, 3]]
if orbspin is None:
orbspin = np.zeros((nocc_a+nvir_a)*2, dtype=int)
orbspin[1::2] = 1
nocc = nocc_a + nocc_b
nvir = nvir_a + nvir_b
idxoa = [np.where(orbspin[k][:nocc] == 0)[0] for k in range(nkpts)]
idxob = [np.where(orbspin[k][:nocc] == 1)[0] for k in range(nkpts)]
idxva = [np.where(orbspin[k][nocc:] == 0)[0] for k in range(nkpts)]
idxvb = [np.where(orbspin[k][nocc:] == 1)[0] for k in range(nkpts)]
r1 = np.zeros((nvir), dtype=r1a.dtype)
r1[idxva[kshift]] = r1a
r1[idxvb[kshift]] = r1b
r2 = np.zeros((nkpts,nkpts,nocc,nvir**2), dtype=r2aaa.dtype)
for kj, ka in itertools.product(range(nkpts), repeat=2):
kb = kconserv[kshift, ka, kj]
idxvaa = idxva[ka][:,None] * nvir + idxva[kb]
idxvab = idxva[ka][:,None] * nvir + idxvb[kb]
idxvba = idxvb[ka][:,None] * nvir + idxva[kb]
idxvbb = idxvb[ka][:,None] * nvir + idxvb[kb]
r2aaa_tmp = r2aaa[kj,ka].reshape(nocc_a, nvir_a*nvir_a)
r2aba_tmp = r2aba[kj,ka].reshape(nocc_a, nvir_b*nvir_a)
r2bab_tmp = r2bab[kj,ka].reshape(nocc_b, nvir_a*nvir_b)
r2bbb_tmp = r2bbb[kj,ka].reshape(nocc_b, nvir_b*nvir_b)
lib.takebak_2d(r2[kj,ka], r2aaa_tmp, idxoa[kj], idxvaa.ravel())
lib.takebak_2d(r2[kj,ka], r2aba_tmp, idxoa[kj], idxvba.ravel())
lib.takebak_2d(r2[kj,ka], r2bab_tmp, idxob[kj], idxvab.ravel())
lib.takebak_2d(r2[kj,ka], r2bbb_tmp, idxob[kj], idxvbb.ravel())
r2aab_tmp = -r2aba[kj,kb].reshape(nocc_a, nvir_b*nvir_a)
r2bba_tmp = -r2bab[kj,kb].reshape(nocc_b, nvir_a*nvir_b)
lib.takebak_2d(r2[kj,ka], r2bba_tmp, idxob[kj], idxvba.T.ravel())
lib.takebak_2d(r2[kj,ka], r2aab_tmp, idxoa[kj], idxvab.T.ravel())
r2 = r2.reshape(nkpts, nkpts, nocc, nvir, nvir)
return r1, r2
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
def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True):
r'''
dm2[p,q,r,s] = <p^\dagger r^\dagger s q>
Note the contraction between ERIs (in Chemist's notation) and rdm2 is
E = einsum('pqrs,pqrs', eri, rdm2)
'''
dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = d2
nocc, nvir = dovov.shape[:2]
nmo = nocc + nvir
dm2 = numpy.empty((nmo,nmo,nmo,nmo), dtype=doooo.dtype)
dovov = numpy.asarray(dovov)
dm2[:nocc,nocc:,:nocc,nocc:] = dovov
dm2[nocc:,:nocc,nocc:,:nocc] = dm2[:nocc,nocc:,:nocc,nocc:].transpose(1,0,3,2).conj()
dovov = None
dovvo = numpy.asarray(dovvo)
dm2[:nocc,:nocc,nocc:,nocc:] =-dovvo.transpose(0,3,2,1)
dm2[nocc:,nocc:,:nocc,:nocc] =-dovvo.transpose(2,1,0,3)
dm2[:nocc,nocc:,nocc:,:nocc] = dovvo
dm2[nocc:,:nocc,:nocc,nocc:] = dovvo.transpose(1,0,3,2).conj()
dovvo = None
dm2[nocc:,nocc:,nocc:,nocc:] = dvvvv
dm2[:nocc,:nocc,:nocc,:nocc] = doooo
dovvv = numpy.asarray(dovvv)
dm2[:nocc,nocc:,nocc:,nocc:] = dovvv
dm2[nocc:,nocc:,:nocc,nocc:] = dovvv.transpose(2,3,0,1)
dm2[nocc:,nocc:,nocc:,:nocc] = dovvv.transpose(3,2,1,0).conj()
dm2[nocc:,:nocc,nocc:,nocc:] = dovvv.transpose(1,0,3,2).conj()
dovvv = None
dooov = numpy.asarray(dooov)
dm2[:nocc,:nocc,:nocc,nocc:] = dooov
dm2[:nocc,nocc:,:nocc,:nocc] = dooov.transpose(2,3,0,1)
dm2[:nocc,:nocc,nocc:,:nocc] = dooov.transpose(1,0,3,2).conj()
dm2[nocc:,:nocc,:nocc,:nocc] = dooov.transpose(3,2,1,0).conj()
if with_frozen and not (mycc.frozen is 0 or mycc.frozen is None):
nmo, nmo0 = mycc.mo_occ.size, nmo
nocc = numpy.count_nonzero(mycc.mo_occ > 0)
rdm2 = numpy.zeros((nmo,nmo,nmo,nmo), dtype=dm2.dtype)
moidx = numpy.where(mycc.get_frozen_mask())[0]
idx = (moidx.reshape(-1,1) * nmo + moidx).ravel()
lib.takebak_2d(rdm2.reshape(nmo**2,nmo**2),
dm2.reshape(nmo0**2,nmo0**2), idx, idx)
dm2 = rdm2
if with_dm1:
dm1 = _make_rdm1(mycc, d1, with_frozen)
dm1[numpy.diag_indices(nocc)] -= 1
for i in range(nocc):
# Be careful with the convention of dm1 and the transpose of dm2 at the end
dm2[i,i,:,:] += dm1
dm2[:,:,i,i] += dm1
dm2[:,i,i,:] -= dm1
dm2[i,:,:,i] -= dm1.T
for i in range(nocc):
for j in range(nocc):
dm2[i,i,j,j] += 1
dm2[i,j,j,i] -= 1
# dm2 was computed as dm2[p,q,r,s] = < p^\dagger r^\dagger s q > in the
# above. Transposing it so that it be contracted with ERIs (in Chemist's
# notation):
# E = einsum('pqrs,pqrs', eri, rdm2)
return dm2.transpose(1,0,3,2)
def contract_2e(eri,
fcivec,
norb,
nelec,
link_index=None,
orbsym=None,
wfnsym=0):
if orbsym is None:
return direct_spin1.contract_2e(eri, fcivec, norb, nelec, link_index)
eri = ao2mo.restore(4, eri, norb)
neleca, nelecb = _unpack_nelec(nelec)
link_indexa, link_indexb = direct_spin1._unpack(norb, nelec, link_index)
na, nlinka = link_indexa.shape[:2]
nb, nlinkb = link_indexb.shape[:2]
eri_irs, rank_eri, irrep_eri = reorder_eri(eri, norb, orbsym)
strsa = cistring.gen_strings4orblist(range(norb), neleca)
aidx, link_indexa = gen_str_irrep(strsa, orbsym, link_indexa, rank_eri,
irrep_eri)
if neleca == nelecb:
bidx, link_indexb = aidx, link_indexa
else:
strsb = cistring.gen_strings4orblist(range(norb), nelecb)
bidx, link_indexb = gen_str_irrep(strsb, orbsym, link_indexb, rank_eri,
irrep_eri)
Tirrep = ctypes.c_void_p * TOTIRREPS
linka_ptr = Tirrep(
*[x.ctypes.data_as(ctypes.c_void_p) for x in link_indexa])
linkb_ptr = Tirrep(
*[x.ctypes.data_as(ctypes.c_void_p) for x in link_indexb])
eri_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in eri_irs])
dimirrep = (ctypes.c_int * TOTIRREPS)(*[x.shape[0] for x in eri_irs])
fcivec_shape = fcivec.shape
fcivec = fcivec.reshape((na, nb), order='C')
ci1new = numpy.zeros_like(fcivec)
nas = (ctypes.c_int * TOTIRREPS)(*[x.size for x in aidx])
nbs = (ctypes.c_int * TOTIRREPS)(*[x.size for x in bidx])
# aa, ab
ci0 = []
ci1 = []
for ir in range(TOTIRREPS):
ma, mb = aidx[ir].size, bidx[wfnsym ^ ir].size
ci0.append(numpy.zeros((ma, mb)))
ci1.append(numpy.zeros((ma, mb)))
if ma > 0 and mb > 0:
lib.take_2d(fcivec, aidx[ir], bidx[wfnsym ^ ir], out=ci0[ir])
ci0_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in ci0])
ci1_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in ci1])
libfci.FCIcontract_2e_symm1(eri_ptrs, ci0_ptrs, ci1_ptrs,
ctypes.c_int(norb), nas, nbs,
ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
linka_ptr, linkb_ptr, dimirrep,
ctypes.c_int(wfnsym))
for ir in range(TOTIRREPS):
if ci0[ir].size > 0:
lib.takebak_2d(ci1new, ci1[ir], aidx[ir], bidx[wfnsym ^ ir])
# bb, ba
ci0T = []
for ir in range(TOTIRREPS):
mb, ma = bidx[ir].size, aidx[wfnsym ^ ir].size
ci0T.append(numpy.zeros((mb, ma)))
if ma > 0 and mb > 0:
lib.transpose(ci0[wfnsym ^ ir], out=ci0T[ir])
ci0, ci0T = ci0T, None
ci1 = [numpy.zeros_like(x) for x in ci0]
ci0_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in ci0])
ci1_ptrs = Tirrep(*[x.ctypes.data_as(ctypes.c_void_p) for x in ci1])
libfci.FCIcontract_2e_symm1(eri_ptrs, ci0_ptrs, ci1_ptrs,
ctypes.c_int(norb), nbs, nas,
ctypes.c_int(nlinkb), ctypes.c_int(nlinka),
linkb_ptr, linka_ptr, dimirrep,
ctypes.c_int(wfnsym))
for ir in range(TOTIRREPS):
if ci0[ir].size > 0:
lib.takebak_2d(ci1new, lib.transpose(ci1[ir]), aidx[wfnsym ^ ir],
bidx[ir])
return ci1new.reshape(fcivec_shape)
def restore_t2_inplace(t2T):
tmp = numpy.zeros((nvir**2, nocc**2), dtype=t2T.dtype)
lib.takebak_2d(tmp, t2T.reshape(nvir**2, nocc**2), vv_idx, oo_idx)
t2 = lib.transpose(tmp.reshape(nvir**2, nocc**2), out=t2T)
return t2.reshape(nocc, nocc, nvir, nvir)
def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True):
r'''
dm2[p,q,r,s] = \sum_{sigma,tau} <p_sigma^\dagger r_tau^\dagger s_tau q_sigma>
Note the contraction between ERIs (in Chemist's notation) and rdm2 is
E = einsum('pqrs,pqrs', eri, rdm2)
'''
dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = d2
nocc, nvir = dovov.shape[:2]
nmo = nocc + nvir
dm2 = numpy.empty((nmo,nmo,nmo,nmo), dtype=doovv.dtype)
dovov = numpy.asarray(dovov)
dm2[:nocc,nocc:,:nocc,nocc:] = dovov
dm2[:nocc,nocc:,:nocc,nocc:]+= dovov.transpose(2,3,0,1)
dm2[nocc:,:nocc,nocc:,:nocc] = dm2[:nocc,nocc:,:nocc,nocc:].transpose(1,0,3,2).conj()
dovov = None
doovv = numpy.asarray(doovv)
dm2[:nocc,:nocc,nocc:,nocc:] = doovv
dm2[:nocc,:nocc,nocc:,nocc:]+= doovv.transpose(1,0,3,2).conj()
dm2[nocc:,nocc:,:nocc,:nocc] = dm2[:nocc,:nocc,nocc:,nocc:].transpose(2,3,0,1)
doovv = None
dovvo = numpy.asarray(dovvo)
dm2[:nocc,nocc:,nocc:,:nocc] = dovvo
dm2[:nocc,nocc:,nocc:,:nocc]+= dovvo.transpose(3,2,1,0).conj()
dm2[nocc:,:nocc,:nocc,nocc:] = dm2[:nocc,nocc:,nocc:,:nocc].transpose(1,0,3,2).conj()
dovvo = None
if len(dvvvv.shape) == 2:
# To handle the case of compressed vvvv, which is used in nuclear gradients
dvvvv = ao2mo.restore(1, dvvvv, nvir)
dm2[nocc:,nocc:,nocc:,nocc:] = dvvvv
dm2[nocc:,nocc:,nocc:,nocc:]*= 4
else:
dvvvv = numpy.asarray(dvvvv)
dm2[nocc:,nocc:,nocc:,nocc:] = dvvvv
dm2[nocc:,nocc:,nocc:,nocc:]+= dvvvv.transpose(1,0,3,2).conj()
dm2[nocc:,nocc:,nocc:,nocc:]*= 2
dvvvv = None
doooo = numpy.asarray(doooo)
dm2[:nocc,:nocc,:nocc,:nocc] = doooo
dm2[:nocc,:nocc,:nocc,:nocc]+= doooo.transpose(1,0,3,2).conj()
dm2[:nocc,:nocc,:nocc,:nocc]*= 2
doooo = None
dovvv = numpy.asarray(dovvv)
dm2[:nocc,nocc:,nocc:,nocc:] = dovvv
dm2[nocc:,nocc:,:nocc,nocc:] = dovvv.transpose(2,3,0,1)
dm2[nocc:,nocc:,nocc:,:nocc] = dovvv.transpose(3,2,1,0).conj()
dm2[nocc:,:nocc,nocc:,nocc:] = dovvv.transpose(1,0,3,2).conj()
dovvv = None
dooov = numpy.asarray(dooov)
dm2[:nocc,:nocc,:nocc,nocc:] = dooov
dm2[:nocc,nocc:,:nocc,:nocc] = dooov.transpose(2,3,0,1)
dm2[:nocc,:nocc,nocc:,:nocc] = dooov.transpose(1,0,3,2).conj()
dm2[nocc:,:nocc,:nocc,:nocc] = dooov.transpose(3,2,1,0).conj()
if with_frozen and not (mycc.frozen is 0 or mycc.frozen is None):
nmo, nmo0 = mycc.mo_occ.size, nmo
nocc = numpy.count_nonzero(mycc.mo_occ > 0)
rdm2 = numpy.zeros((nmo,nmo,nmo,nmo), dtype=dm2.dtype)
moidx = numpy.where(mycc.get_frozen_mask())[0]
idx = (moidx.reshape(-1,1) * nmo + moidx).ravel()
lib.takebak_2d(rdm2.reshape(nmo**2,nmo**2),
dm2.reshape(nmo0**2,nmo0**2), idx, idx)
dm2 = rdm2
if with_dm1:
dm1 = _make_rdm1(mycc, d1, with_frozen)
dm1[numpy.diag_indices(nocc)] -= 2
for i in range(nocc):
dm2[i,i,:,:] += dm1 * 2
dm2[:,:,i,i] += dm1 * 2
dm2[:,i,i,:] -= dm1
dm2[i,:,:,i] -= dm1.T
for i in range(nocc):
for j in range(nocc):
dm2[i,i,j,j] += 4
dm2[i,j,j,i] -= 2
# dm2 was computed as dm2[p,q,r,s] = < p^\dagger r^\dagger s q > in the
# above. Transposing it so that it be contracted with ERIs (in Chemist's
# notation):
# E = einsum('pqrs,pqrs', eri, rdm2)
return dm2.transpose(1,0,3,2)
def spatial2spin(tx, orbspin=None):
'''Convert T1/T2 of spatial orbital representation to T1/T2 of
spin-orbital representation
'''
if isinstance(tx, numpy.ndarray) and tx.ndim == 2:
# RCCSD t1 amplitudes
return spatial2spin((tx, tx), orbspin)
elif isinstance(tx, numpy.ndarray) and tx.ndim == 4:
# RCCSD t2 amplitudes
t2aa = tx - tx.transpose(0, 1, 3, 2)
return spatial2spin((t2aa, tx, t2aa), orbspin)
elif len(tx) == 2: # t1
t1a, t1b = tx
nocc_a, nvir_a = t1a.shape
nocc_b, nvir_b = t1b.shape
else:
t2aa, t2ab, t2bb = tx
nocc_a, nocc_b, nvir_a, nvir_b = t2ab.shape
if orbspin is None:
orbspin = numpy.zeros((nocc_a + nvir_a) * 2, dtype=int)
orbspin[1::2] = 1
nocc = nocc_a + nocc_b
nvir = nvir_a + nvir_b
idxoa = numpy.where(orbspin[:nocc] == 0)[0]
idxob = numpy.where(orbspin[:nocc] == 1)[0]
idxva = numpy.where(orbspin[nocc:] == 0)[0]
idxvb = numpy.where(orbspin[nocc:] == 1)[0]
if len(tx) == 2: # t1
t1 = numpy.zeros((nocc, nvir), dtype=t1a.dtype)
lib.takebak_2d(t1, t1a, idxoa, idxva)
lib.takebak_2d(t1, t1b, idxob, idxvb)
t1 = lib.tag_array(t1, orbspin=orbspin)
return t1
else:
t2 = numpy.zeros((nocc**2, nvir**2), dtype=t2aa.dtype)
idxoaa = idxoa[:, None] * nocc + idxoa
idxoab = idxoa[:, None] * nocc + idxob
idxoba = idxob[:, None] * nocc + idxoa
idxobb = idxob[:, None] * nocc + idxob
idxvaa = idxva[:, None] * nvir + idxva
idxvab = idxva[:, None] * nvir + idxvb
idxvba = idxvb[:, None] * nvir + idxva
idxvbb = idxvb[:, None] * nvir + idxvb
t2aa = t2aa.reshape(nocc_a * nocc_a, nvir_a * nvir_a)
t2ab = t2ab.reshape(nocc_a * nocc_b, nvir_a * nvir_b)
t2bb = t2bb.reshape(nocc_b * nocc_b, nvir_b * nvir_b)
lib.takebak_2d(t2, t2aa, idxoaa.ravel(), idxvaa.ravel())
lib.takebak_2d(t2, t2bb, idxobb.ravel(), idxvbb.ravel())
lib.takebak_2d(t2, t2ab, idxoab.ravel(), idxvab.ravel())
lib.takebak_2d(t2, t2ab, idxoba.T.ravel(), idxvba.T.ravel())
abba = -t2ab
lib.takebak_2d(t2, abba, idxoab.ravel(), idxvba.T.ravel())
lib.takebak_2d(t2, abba, idxoba.T.ravel(), idxvab.ravel())
t2 = lib.tag_array(t2, orbspin=orbspin)
return t2.reshape(nocc, nocc, nvir, nvir)