# Source code for pyscf.fci.cistring

#!/usr/bin/env python
#
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#
# Unless required by applicable law or agreed to in writing, software
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import sys
import ctypes
import math
import numpy
from pyscf import lib

[docs]def gen_strings4orblist(orb_list, nelec):
'''Generate string from the given orbital list.

Returns:
list of int64.  One int64 element represents one string in binary format.
The binary format takes the convention that the one bit stands for one
orbital, bit-1 means occupied and bit-0 means unoccupied.  The lowest
(right-most) bit corresponds to the lowest orbital in the orb_list.

Exampels:

>>> [bin(x) for x in gen_strings4orblist((0,1,2,3),2)]
[0b11, 0b101, 0b110, 0b1001, 0b1010, 0b1100]
>>> [bin(x) for x in gen_strings4orblist((3,1,0,2),2)]
[0b1010, 0b1001, 0b11, 0b1100, 0b110, 0b101]
'''
orb_list = list(orb_list)
if len(orb_list) > 63:
return _gen_occslst(orb_list, nelec)

assert(nelec >= 0)
if nelec == 0:
return numpy.asarray([0], dtype=numpy.int64)
elif nelec > len(orb_list):
return numpy.asarray([], dtype=numpy.int64)
def gen_str_iter(orb_list, nelec):
if nelec == 1:
res = [(1<<i) for i in orb_list]
elif nelec >= len(orb_list):
n = 0
for i in orb_list:
n = n | (1<<i)
res = [n]
else:
restorb = orb_list[:-1]
thisorb = 1 << orb_list[-1]
res = gen_str_iter(restorb, nelec)
for n in gen_str_iter(restorb, nelec-1):
res.append(n | thisorb)
return res
strings = gen_str_iter(orb_list, nelec)
assert(strings.__len__() == num_strings(len(orb_list),nelec))
return numpy.asarray(strings, dtype=numpy.int64)

def _gen_occslst(orb_list, nelec):
'''Generate occupied orbital list for each string.
'''
orb_list = list(orb_list)
assert(nelec >= 0)
if nelec == 0:
return numpy.zeros((1,nelec), dtype=numpy.int32)
elif nelec > len(orb_list):
return numpy.zeros((0,nelec), dtype=numpy.int32)
def gen_occs_iter(orb_list, nelec):
if nelec == 1:
res = [[i] for i in orb_list]
elif nelec >= len(orb_list):
res = [orb_list]
else:
restorb = orb_list[:-1]
thisorb = orb_list[-1]
res = gen_occs_iter(restorb, nelec)
for n in gen_occs_iter(restorb, nelec-1):
res.append(n + [thisorb])
return res
occslst = gen_occs_iter(orb_list, nelec)
return numpy.asarray(occslst, dtype=numpy.int32).view(OIndexList)
def _strs2occslst(strs, norb):
na = len(strs)
one_particle_strs = numpy.asarray([1<<i for i in range(norb)])
occ_masks = (strs.reshape(-1,1) & one_particle_strs) != 0
return numpy.asarray(occslst, dtype=numpy.int32).view(OIndexList)
def _occslst2strs(occslst):
na, nelec = occslst.shape
strs = numpy.zeros(na, dtype=numpy.int64)
for i in range(nelec):
strs ^= 1 << occslst[:,i]
return strs
class OIndexList(numpy.ndarray):
pass

def num_strings(n, m):
if m < 0 or m > n:
return 0
else:
return math.factorial(n) // (math.factorial(n-m)*math.factorial(m))

if strs is None:
strs = gen_strings4orblist(orb_list, nelec)
strdic = dict(zip(strs,range(strs.__len__())))
def propgate1e(str0):
occ = []
vir = []
for i in orb_list:
if str0 & (1<<i):
occ.append(i)
else:
vir.append(i)
for i in occ:
for i in occ:
for a in vir:
str1 = str0 ^ (1<<i) | (1<<a)
linktab.append((a, i, strdic[str1], cre_des_sign(a, i, str0)))

t = [propgate1e(s) for s in strs.astype(numpy.int64)]
return numpy.array(t, dtype=numpy.int32)

if nelec == 0:
return numpy.zeros((0,0,4), dtype=numpy.int32)

if strs is None:
strs = _gen_occslst(orb_list, nelec)
occslst = strs

orb_list = numpy.asarray(orb_list)
norb = len(orb_list)
assert(numpy.all(numpy.arange(norb) == orb_list))

strdic = dict((tuple(s), i) for i,s in enumerate(occslst))
nvir = norb - nelec
def propgate1e(str0):
tab = numpy.empty((nelec,4), dtype=numpy.int32)
tab[:,0] = tab[:,1] = str0
tab[:,3] = 1

str0 = numpy.asarray(str0)
# where to put vir-orb, ie how many occ-orb in the left
where_vir = numpy.sum(str0.reshape(-1,1) < vir, axis=0)
parity_occ_orb = 1  # parity for annihilating occupied orbital
for n,i in enumerate(str0):  # loop over all occupied orbitals
# o,v which index is bigger, to determine whether to annihilate occ-orb first
reorder_to_ov = vir > i
str1s = numpy.empty((nvir,nelec), dtype=int)
str1s[:] = str0
str1s[:,n] = vir
str1s.sort(axis=1)
addr = [strdic[tuple(s)] for s in str1s]
parity = (where_vir + reorder_to_ov + 1) % 2 #? +1 so that even parity has +1, odd parity = 0
parity[parity == 0] = -1
parity *= parity_occ_orb
tab = numpy.empty((nvir,4), dtype=numpy.int32)
tab[:,0] = vir
tab[:,1] = i
tab[:,3] = parity
parity_occ_orb *= -1

lidx = [propgate1e(s) for s in occslst]
lidx = numpy.asarray(lidx, dtype=numpy.int32)
if tril:
return lidx

# return [cre, des, target_address, parity]
'''Look up table, for the strings relationship in terms of a
creation-annihilating operator pair.

For given string str0, index[str0] is (nocc+nocc*nvir) x 4 array.
The first nocc rows [i(:occ),i(:occ),str0,sign] are occupied-occupied
excitations, which do not change the string. The next nocc*nvir rows
[a(:vir),i(:occ),str1,sign] are occupied-virtual exciations, starting from
str0, annihilating i, creating a, to get str1.
'''
if strs is None:
strs = gen_strings4orblist(orb_list, nocc)

if isinstance(strs, OIndexList):

strs = numpy.array(strs, dtype=numpy.int64)
assert(all(strs[:-1] < strs[1:]))
norb = len(orb_list)
nvir = norb - nocc
na = strs.shape[0]
ctypes.c_int(norb), ctypes.c_int(na),
ctypes.c_int(nocc),
strs.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(tril))

'''Compress the (a, i) pair index in linkstr_index to a lower triangular
index, to match the 4-fold symmetry of integrals.
'''
#    for j, (a, i, str1, sign) in enumerate(tab):
#        if a > i:
#            ai = a*(a+1)//2+i
#        else:
#            ai = i*(i+1)//2+a

r'''Generate linkstr_index with the assumption that :math:p^+ q|0\rangle
where :math:p > q.
So the resultant link_index has the structure [pq, *, str1, sign].
It is identical to a call to reform_linkstr_index(gen_linkstr_index(...)).
'''

# return [cre, des, target_address, parity]
[docs]def gen_cre_str_index_o0(orb_list, nelec):
'''Slow version of gen_cre_str_index function'''
cre_strs = gen_strings4orblist(orb_list, nelec+1)
if isinstance(cre_strs, OIndexList):
raise NotImplementedError('System with 64 orbitals or more')

credic = dict(zip(cre_strs,range(cre_strs.__len__())))
def progate1e(str0):
for i in orb_list:
if not str0 & (1<<i):
str1 = str0 | (1<<i)

strs = gen_strings4orblist(orb_list, nelec)
t = [progate1e(s) for s in strs.astype(numpy.int64)]
return numpy.array(t, dtype=numpy.int32)
[docs]def gen_cre_str_index_o1(orb_list, nelec):
'''C implementation of gen_cre_str_index function'''
norb = len(orb_list)
assert(nelec < norb)
strs = gen_strings4orblist(orb_list, nelec)
if isinstance(strs, OIndexList):
raise NotImplementedError('System with 64 orbitals or more')

strs = numpy.array(strs, dtype=numpy.int64)
na = strs.shape[0]
ctypes.c_int(norb), ctypes.c_int(na),
ctypes.c_int(nelec),
strs.ctypes.data_as(ctypes.c_void_p))
[docs]def gen_cre_str_index(orb_list, nelec):
'''linkstr_index to map between N electron string to N+1 electron string.
It maps the given string to the address of the string which is generated by
the creation operator.

For given string str0, index[str0] is nvir x 4 array.  Each entry
[i(cre),--,str1,sign] means starting from str0, creating i, to get str1.
'''
return gen_cre_str_index_o1(orb_list, nelec)

# return [cre, des, target_address, parity]
[docs]def gen_des_str_index_o0(orb_list, nelec):
'''Slow version of gen_des_str_index function'''
des_strs = gen_strings4orblist(orb_list, nelec-1)
if isinstance(des_strs, OIndexList):
raise NotImplementedError('System with 64 orbitals or more')

desdic = dict(zip(des_strs,range(des_strs.__len__())))
def progate1e(str0):
for i in orb_list:
if str0 & (1<<i):
str1 = str0 ^ (1<<i)

strs = gen_strings4orblist(orb_list, nelec)
t = [progate1e(s) for s in strs.astype(numpy.int64)]
return numpy.array(t, dtype=numpy.int32)
[docs]def gen_des_str_index_o1(orb_list, nelec):
'''C implementation of gen_des_str_index function'''
assert(nelec > 0)
strs = gen_strings4orblist(orb_list, nelec)
if isinstance(strs, OIndexList):
raise NotImplementedError('System with 64 orbitals or more')

strs = numpy.array(strs, dtype=numpy.int64)
norb = len(orb_list)
na = strs.shape[0]
ctypes.c_int(norb), ctypes.c_int(na),
ctypes.c_int(nelec),
strs.ctypes.data_as(ctypes.c_void_p))
[docs]def gen_des_str_index(orb_list, nelec):
'''linkstr_index to map between N electron string to N-1 electron string.
It maps the given string to the address of the string which is generated by
the annihilation operator.

For given string str0, index[str0] is nvir x 4 array.  Each entry
[--,i(des),str1,sign] means starting from str0, annihilating i, to get str1.
'''
return gen_des_str_index_o1(orb_list, nelec)

# Determine the sign of  p^+ q |string0>
def cre_des_sign(p, q, string0):
if p == q:
return 1
else:
if (string0 & (1<<p)) or (not (string0 & (1<<q))):
return 0
elif p > q:
mask = (1 << p) - (1 << (q+1))
else:
mask = (1 << q) - (1 << (p+1))
return (-1) ** bin(string0 & mask).count('1')

# Determine the sign of  p^+ |string0>
def cre_sign(p, string0):
if (string0 & (1<<p)):
return 0
else:
return (-1) ** bin(string0>>(p+1)).count('1')

# Determine the sign of  p |string0>
def des_sign(p, string0):
if (not (string0 & (1<<p))):
return 0
else:
return (-1) ** bin(string0>>(p+1)).count('1')

# Determine the sign of  string1 = p^+ q |string0>
def parity(string0, string1):
sys.stderr.write('Function cistring.parity is deprecated\n')
ss = string1 - string0
def count_bit1(n):
# see Hamming weight problem and K&R C program
return bin(n).count('1')
if ss > 0:
# string1&ss gives the number of 1s between two strings
return (-1) ** (count_bit1(string1&ss))
elif ss == 0:
return 1
else:
return (-1) ** (count_bit1(string0&(-ss)))

if addr == 0 or nelec == norb or nelec == 0:
return (1<<nelec) - 1   # ..0011..11
else:
for i in reversed(range(norb)):
if addr == 0 or nelec == norb or nelec == 0:
return (1<<nelec) - 1   # ..0011..11
str1 = 0
nelec_left = nelec
for norb_left in reversed(range(norb)):
if nelec_left == 0:
break
str1 |= (1<<nelec_left) - 1
break
str1 |= 1<<norb_left
nelec_left -= 1
return str1
'''Convert CI determinant address to string'''
'''Convert a list of CI determinant address to string'''
strs = numpy.empty(count, dtype=numpy.int64)
ctypes.c_int(count),
ctypes.c_int(norb), ctypes.c_int(nelec))
return strs

#    if norb <= nelec or nelec == 0:
#        return 0
#    elif (1<<(norb-1)) & string:  # remove the first bit
#        return num_strings(norb-1, nelec) \
#    else:
#    #TODO: assert norb > first-bit-in-string, nelec == num-1-in-string
#    nelec_left = nelec
#    for norb_left in reversed(range(norb)):
#        if nelec_left == 0 or norb_left < nelec_left:
#            break
#        elif (1<<norb_left) & string:
#            nelec_left -= 1
'''Convert string to CI determinant address'''
if isinstance(string, str):
assert(string.count('1') == nelec)
string = int(string, 2)
else:
assert(bin(string).count('1') == nelec)
ctypes.c_ulonglong(string))
'''Convert a list of string to CI determinant address'''
strings = numpy.asarray(strings, dtype=numpy.int64)
count = strings.size
strings.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(count),
ctypes.c_int(norb), ctypes.c_int(nelec))

'''The addresses of the determinants which include the specified orbital
indices. The size of the returned addresses is equal to the number of
determinants of (norb, nelec) system.
'''
assert(norb < 63)
if sub_nelec == 0:
strs = gen_strings4orblist(orbital_indices, nelec)
else:
strs = gen_strings4orblist(range(norb), nelec)
counts = numpy.zeros(len(strs), dtype=int)
for i in orbital_indices:
counts += (strs & (1<<i)) != 0
sub_strs = strs[counts == sub_nelec]

[docs]def tn_strs(norb, nelec, n):
'''Generate strings for Tn amplitudes.  Eg n=1 (T1) has nvir*nocc strings,
n=2 (T2) has nvir*(nvir-1)/2 * nocc*(nocc-1)/2 strings.
'''
if nelec < n or norb-nelec < n:
return numpy.zeros(0, dtype=int)
occs_allow = numpy.asarray(gen_strings4orblist(range(nelec), n)[::-1])
virs_allow = numpy.asarray(gen_strings4orblist(range(nelec,norb), n))
hf_str = int('1'*nelec, 2)
tns = (hf_str | virs_allow.reshape(-1,1)) ^ occs_allow
return tns.ravel()

if __name__ == '__main__':
#print([bin(i) for i in gen_strings4orblist(range(2,5), 2)])
#print(gen_strings4orblist(range(4), 2))
#    idx16 = index[:16]
#    print(idx16[:,:,2])
strs = gen_strings4orblist(range(8), 4)
occlst = _gen_occslst(range(8), 4)
print(abs(occlst - _strs2occslst(strs, 8)).sum())
print(abs(strs - _occslst2strs(occlst)).sum())
print(abs(tab1 - tab2).sum())
print(abs(tab1 - tab3).sum())