Python cistring.num_strings函数代码示例

def exdiagH(h1e, g2e, norb, nelec, writefile=True):
    '''
        exactly diagonalize the hamiltonian.
    '''
    h2e = direct_spin1.absorb_h1e(h1e, g2e, norb, nelec, .5)
    if isinstance(nelec, (int, np.integer)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec

    naa = cistring.num_strings(norb, neleca)
    nbb = cistring.num_strings(norb, nelecb)
    ndim = naa*nbb
    eyebas = np.eye(ndim, ndim)
    def hop(c):
        hc = direct_spin1.contract_2e(h2e, c, norb, nelec)
        return hc.reshape(-1)
    Hmat = []
    for i in range(ndim):
        hc = hop(eyebas[i].copy())
        Hmat.append(hc)

    Hmat = np.asarray(Hmat)
#    Hmat = Hmat.T.copy()
    ew, ev = nl.eigh(Hmat.T)
    if writefile:
        np.savetxt("cards/eignE.dat", ew, fmt="%10.10f")
        np.savetxt("cards/eignV.dat", ev, fmt="%10.10f")
    return ew, ev

def from_fcivec(ci0, norb, nelec, frozen=0):
    '''Extract CISD coefficients from FCI coefficients'''
    if frozen is not 0:
        raise NotImplementedError
    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec

    norba = norbb = norb
    nocc = nocca, noccb = neleca, nelecb
    nvira = norba - nocca
    nvirb = norbb - noccb
    t1addra, t1signa = cisd.tn_addrs_signs(norba, nocca, 1)
    t1addrb, t1signb = cisd.tn_addrs_signs(norbb, noccb, 1)

    na = cistring.num_strings(norba, nocca)
    nb = cistring.num_strings(norbb, noccb)
    ci0 = ci0.reshape(na,nb)
    c0 = ci0[0,0]
    c1a = (ci0[t1addra,0] * t1signa).reshape(nocca,nvira)
    c1b = (ci0[0,t1addrb] * t1signb).reshape(noccb,nvirb)

    c2ab = numpy.einsum('i,j,ij->ij', t1signa, t1signb, ci0[t1addra[:,None],t1addrb])
    c2ab = c2ab.reshape(nocca,nvira,noccb,nvirb).transpose(0,2,1,3)
    t2addra, t2signa = cisd.tn_addrs_signs(norba, nocca, 2)
    t2addrb, t2signb = cisd.tn_addrs_signs(norbb, noccb, 2)
    c2aa = (ci0[t2addra,0] * t2signa).reshape(nocca*(nocca-1)//2, nvira*(nvira-1)//2)
    c2aa = _unpack_4fold(c2aa, nocca, nvira)
    c2bb = (ci0[0,t2addrb] * t2signb).reshape(noccb*(noccb-1)//2, nvirb*(nvirb-1)//2)
    c2bb = _unpack_4fold(c2bb, noccb, nvirb)

    return amplitudes_to_cisdvec(c0, (c1a,c1b), (c2aa,c2ab,c2bb))

def contract_2e(eri, fcivec, norb, nelec, opt=None):
    if isinstance(nelec, (int, numpy.integer)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    link_indexa = cistring.gen_linkstr_index(range(norb), neleca)
    link_indexb = cistring.gen_linkstr_index(range(norb), nelecb)
    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)
    fcivec = fcivec.reshape(na,nb)
    t1 = numpy.zeros((norb,norb,na,nb))
    for str0, tab in enumerate(link_indexa):
        for a, i, str1, sign in tab:
            t1[a,i,str1] += sign * fcivec[str0]
    for str0, tab in enumerate(link_indexb):
        for a, i, str1, sign in tab:
            t1[a,i,:,str1] += sign * fcivec[:,str0]
    t1 = numpy.dot(eri.reshape(norb*norb,-1), t1.reshape(norb*norb,-1))
    t1 = t1.reshape(norb,norb,na,nb)
    fcinew = numpy.zeros_like(fcivec)
    for str0, tab in enumerate(link_indexa):
        for a, i, str1, sign in tab:
            fcinew[str1] += sign * t1[a,i,str0]
    for str0, tab in enumerate(link_indexb):
        for a, i, str1, sign in tab:
            fcinew[:,str1] += sign * t1[a,i,:,str0]
    return fcinew

def get_init_guess(norb, nelec, nroots, hdiag):
    if isinstance(nelec, (int, numpy.integer)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)

    init_strs = []
    iroot = 0
    for addr in numpy.argsort(hdiag):
        addra = addr // nb
        addrb = addr % nb
        if (addrb,addra) not in init_strs:  # avoid initial guess linear dependency
            init_strs.append((addra,addrb))
            iroot += 1
            if iroot >= nroots:
                break
    ci0 = []
    for addra,addrb in init_strs:
        x = numpy.zeros((na,nb))
        if addra == addrb == 0:
            x[addra,addrb] = 1
        else:
            x[addra,addrb] = x[addrb,addra] = numpy.sqrt(.5)
        ci0.append(x.ravel())
    return ci0

def _make_rdm2_abba(fcivec, norb, nelec):
    if isinstance(nelec, (int, numpy.integer)):
        neleca = nelecb = nelec // 2
    else:
        neleca, nelecb = nelec
    if nelecb == norb or neleca == 0: # no intermediate determinants
        return numpy.zeros((norb,norb,norb,norb))
    acre_index = cistring.gen_cre_str_index(range(norb), neleca-1)
    bdes_index = cistring.gen_des_str_index(range(norb), nelecb+1)
    instra = cistring.num_strings(norb, neleca-1)
    nb = cistring.num_strings(norb, nelecb)
    dm1 = numpy.empty((norb,norb))
    dm2 = numpy.empty((norb,norb,norb,norb))
    fn = _ctypes.dlsym(librdm._handle, 'FCIdm2_abba_kern')
    librdm.FCIspindm12_drv(ctypes.c_void_p(fn),
                           dm1.ctypes.data_as(ctypes.c_void_p),
                           dm2.ctypes.data_as(ctypes.c_void_p),
                           fcivec.ctypes.data_as(ctypes.c_void_p),
                           fcivec.ctypes.data_as(ctypes.c_void_p),
                           ctypes.c_int(norb),
                           ctypes.c_int(instra), ctypes.c_int(nb),
                           ctypes.c_int(neleca), ctypes.c_int(nelecb),
                           acre_index.ctypes.data_as(ctypes.c_void_p),
                           bdes_index.ctypes.data_as(ctypes.c_void_p))
    return dm2

def make_rdm12e(fcivec, nsite, nelec):
    '''1-electron and 2-electron density matrices
    dm_pq = <|p^+ q|>
    dm_{pqrs} = <|p^+ r^+ q s|>  (note 2pdm is ordered in chemist notation)
    '''
    neleca, nelecb = _unpack_nelec(nelec)
    link_indexa = cistring.gen_linkstr_index(range(nsite), neleca)
    link_indexb = cistring.gen_linkstr_index(range(nsite), nelecb)
    na = cistring.num_strings(nsite, neleca)
    nb = cistring.num_strings(nsite, nelecb)

    ci0 = fcivec.reshape(na,nb,-1)
    rdm1 = numpy.zeros((nsite,nsite))
    rdm2 = numpy.zeros((nsite,nsite,nsite,nsite))
    for str0 in range(na):
        t1 = numpy.zeros((nsite,nsite,nb)+ci0.shape[2:])
        for a, i, str1, sign in link_indexa[str0]:
            t1[i,a,:] += sign * ci0[str1,:]

        for k, tab in enumerate(link_indexb):
            for a, i, str1, sign in tab:
                t1[i,a,k] += sign * ci0[str0,str1]

        rdm1 += numpy.einsum('mp,ijmp->ij', ci0[str0], t1)
        # i^+ j|0> => <0|j^+ i, so swap i and j
        #:rdm2 += numpy.einsum('ijmp,klmp->jikl', t1, t1)
        tmp = lib.dot(t1.reshape(nsite**2,-1), t1.reshape(nsite**2,-1).T)
        rdm2 += tmp.reshape((nsite,)*4).transpose(1,0,2,3)
    rdm1, rdm2 = rdm.reorder_rdm(rdm1, rdm2, True)
    return rdm1, rdm2

def ft_rdm1s(h1e, g2e, norb, nelec, T, m=50, nsamp=40, Tmin=10e-4):
    '''rdm of spin a and b at temperature T
    '''
    if T < Tmin:
       e, c = kernel(h1e, g2e, norb, nelec)
       rdma, rdmb = direct_spin1.make_rdm1s(c, norb, nelec)
       return rdma, rdmb

    h2e = direct_spin1.absorb_h1e(h1e, g2e, norb, nelec, .5)
    if isinstance(nelec, (int, numpy.integer)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)
    def vecgen(n1=na, n2=nb):
        ci0 = numpy.random.randn(n1, n2)
#        ci0[0, 0] = 1.
        return ci0.reshape(-1)
    def hop(c):
        hc = direct_spin1.contract_2e(h2e, c, norb, nelec)
        return hc.reshape(-1)
    def qud(v1, v2):
        dma, dmb = direct_spin1.trans_rdm1s(v1, v2, norb, nelec)
        return dma, dmb

#    rdma, rdmb = flan.ht_rdm1s(qud, hop, vecgen, T, norb, m, nsamp)
    rdma, rdmb = flan.ftlan_rdm1s(qud, hop, vecgen, T, norb, m, nsamp)
    return rdma, rdmb

def kernal_hubbard(t, u, norb, nelec):
    if isinstance(nelec, (int, numpy.integer)):
        neleca = nelec//2
        nelecb = nelec - neleca
    else:
        neleca, nelecb = nelec

    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)
    def mk_tridiag(t, norb):
        tmat = numpy.ones(norb-1) * (-t)
        res = numpy.diag(tmat, -1) + numpy.diag(tmat, 1)
        res[0, norb-1] = -t
        res[norb-1, 0] = -t
        return res
    f1e = mk_tridiag(t, norb)
    print f1e
    U = (0., u, 0.)
    def hop(c):
        hc = contract_2e_hubbard(U, c, norb, nelec)\
             +contract_1e(f1e, c, norb, nelec)
        return hc.reshape(-1)
    ci0 = numpy.random.random((na, nb))
    hdiag = make_hdiag_hubbard(f1e, u, norb, nelec)
    print "hdiag:,\n", hdiag
    precond = lambda x, e, *args: x/(hdiag-e+1e-4)
    e, c = pyscf.lib.davidson(hop, ci0.reshape(-1), precond)
    return e

def get_init_guess(norb, nelec, nroots, hdiag):
    '''Initial guess is the single Slater determinant
    '''
    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)

    # The "nroots" lowest determinats based on energy expectation value.
    ci0 = []
    try:
        addrs = numpy.argpartition(hdiag, nroots-1)[:nroots]
    except AttributeError:
        addrs = numpy.argsort(hdiag)[:nroots]
    for addr in addrs:
        x = numpy.zeros((na*nb))
        x[addr] = 1
        ci0.append(x.ravel())

    # Add noise
    ci0[0][0 ] += 1e-5
    ci0[0][-1] -= 1e-5
    return ci0

def contract_2e_hubbard(u, fcivec, norb, nelec, opt=None):
    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    u_aa, u_ab, u_bb = u

    strsa = numpy.asarray(cistring.gen_strings4orblist(range(norb), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(norb), nelecb))
    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)
    fcivec = fcivec.reshape(na,nb)
    fcinew = numpy.zeros_like(fcivec)

    if u_aa != 0:  # u * n_alpha^+ n_alpha
        for i in range(norb):
            maska = (strsa & (1<<i)) > 0
            fcinew[maska] += u_aa * fcivec[maska]
    if u_ab != 0:  # u * (n_alpha^+ n_beta + n_beta^+ n_alpha)
        for i in range(norb):
            maska = (strsa & (1<<i)) > 0
            maskb = (strsb & (1<<i)) > 0
            fcinew[maska[:,None]&maskb] += 2*u_ab * fcivec[maska[:,None]&maskb]
    if u_bb != 0:  # u * n_beta^+ n_beta
        for i in range(norb):
            maskb = (strsb & (1<<i)) > 0
            fcinew[:,maskb] += u_bb * fcivec[:,maskb]
    return fcinew

def contract_2e(eri, fcivec, norb, nelec, opt=None):
    if isinstance(nelec, (int, numpy.integer)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    link_indexa = cistring.gen_linkstr_index(range(norb), neleca)
    link_indexb = cistring.gen_linkstr_index(range(norb), nelecb)
    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)
    ci0 = fcivec.reshape(na,nb)
    t1 = numpy.zeros((norb,norb,na,nb))
    for str0, tab in enumerate(link_indexa):
        for a, i, str1, sign in tab:
            t1[a,i,str1] += sign * ci0[str0]
    for str0, tab in enumerate(link_indexb):
        for a, i, str1, sign in tab:
            t1[a,i,:,str1] += sign * ci0[:,str0]

    t1 = lib.einsum('bjai,aiAB->bjAB', eri.reshape([norb]*4), t1)

    fcinew = numpy.zeros_like(ci0)
    for str0, tab in enumerate(link_indexa):
        for a, i, str1, sign in tab:
            fcinew[str1] += sign * t1[a,i,str0]
    for str0, tab in enumerate(link_indexb):
        for a, i, str1, sign in tab:
            fcinew[:,str1] += sign * t1[a,i,:,str0]
    return fcinew.reshape(fcivec.shape)

def make_rdm12e(fcivec, nsite, nelec):
    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    link_indexa = cistring.gen_linkstr_index(range(nsite), neleca)
    link_indexb = cistring.gen_linkstr_index(range(nsite), nelecb)
    na = cistring.num_strings(nsite, neleca)
    nb = cistring.num_strings(nsite, nelecb)

    ci0 = fcivec.reshape(na,nb,-1)
    rdm1 = numpy.zeros((nsite,nsite))
    rdm2 = numpy.zeros((nsite,nsite,nsite,nsite))
    for str0 in range(na):
        t1 = numpy.zeros((nsite,nsite,nb)+ci0.shape[2:])
        for a, i, str1, sign in link_indexa[str0]:
            t1[i,a,:] += sign * ci0[str1,:]

        for k, tab in enumerate(link_indexb):
            for a, i, str1, sign in tab:
                t1[i,a,k] += sign * ci0[str0,str1]

        rdm1 += numpy.einsum('mp,ijmp->ij', ci0[str0], t1)
        # i^+ j|0> => <0|j^+ i, so swap i and j
        #:rdm2 += numpy.einsum('ijmp,klmp->jikl', t1, t1)
        tmp = lib.dot(t1.reshape(nsite**2,-1), t1.reshape(nsite**2,-1).T)
        rdm2 += tmp.reshape((nsite,)*4).transpose(1,0,2,3)
    rdm1, rdm2 = pyscf.fci.rdm.reorder_rdm(rdm1, rdm2, True)
    return rdm1, rdm2

    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, 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]
            ida = aindex[maska,i,0]
            idb = bindex[maskb,i,0]
            citmp = pyscf.lib.take_2d(fcivec, maska, maskb)
            citmp = numpy.einsum('i,j,ij->ij', signa[maska], signb[maskb], citmp)
            #: t1[ida.reshape(-1,1),idb] += citmp
            pyscf.lib.takebak_2d_(t1, citmp, ida, idb)
        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]
            ida = aindex[maska,i,0]
            idb = bindex[maskb,i,0]
            citmp = pyscf.lib.take_2d(t1, ida, idb)
            citmp = numpy.einsum('i,j,ij->ij', signa[maska], signb[maskb], citmp)
            #: ci1[maska.reshape(-1,1), maskb] += citmp
            pyscf.lib.takebak_2d_(ci1, citmp, maska, maskb)

def kernel(h1e, eri, norb, nelec, ecore=0, verbose=logger.NOTE):
    if isinstance(nelec, (int, numpy.integer)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    h2e = direct_spin1.absorb_h1e(h1e, eri, norb, nelec, .5)
    namax = cistring.num_strings(norb, neleca)
    nbmax = cistring.num_strings(norb, nelecb)

    myci = SelectedCI()

    strsa = [int('1'*neleca, 2)]
    strsb = [int('1'*nelecb, 2)]
    ci_strs = (strsa, strsb)
    ci0 = numpy.ones((1,1))
    ci0, ci_strs = enlarge_space(myci, (ci0, ci_strs), h2e, norb, nelec)

    def all_linkstr_index(ci_strs):
        cd_indexa = cre_des_linkstr(ci_strs[0], norb, neleca)
        dd_indexa = des_des_linkstr(ci_strs[0], norb, neleca)
        cd_indexb = cre_des_linkstr(ci_strs[1], norb, nelecb)
        dd_indexb = des_des_linkstr(ci_strs[1], norb, nelecb)
        return cd_indexa, dd_indexa, cd_indexb, dd_indexb

    def hop(c):
        hc = contract_2e(h2e, (c, ci_strs), norb, nelec, link_index)
        return hc.reshape(-1)
    precond = lambda x, e, *args: x/(hdiag-e+1e-4)

    e_last = 0
    tol = 1e-2
    conv = False
    for icycle in range(norb):
        tol = max(tol*1e-2, myci.float_tol)
        link_index = all_linkstr_index(ci_strs)
        hdiag = make_hdiag(h1e, eri, ci_strs, norb, nelec)
        e, ci0 = lib.davidson(hop, ci0.reshape(-1), precond, tol=tol,
                              verbose=verbose)
        print('icycle %d  ci.shape %s  E = %.15g' %
              (icycle, (len(ci_strs[0]), len(ci_strs[1])), e))
        if ci0.shape == (namax,nbmax) or abs(e-e_last) < myci.float_tol*10:
            conv = True
            break
        ci1, ci_strs = enlarge_space(myci, (ci0, ci_strs), h2e, norb, nelec)
        if ci1.size < ci0.size*1.02:
            conv = True
            break
        e_last = e
        ci0 = ci1

    link_index = all_linkstr_index(ci_strs)
    hdiag = make_hdiag(h1e, eri, ci_strs, norb, nelec)
    e, ci0 = lib.davidson(hop, ci0.reshape(-1), precond, tol=myci.conv_tol,
                          verbose=verbose)

    na = len(ci_strs[0])
    nb = len(ci_strs[1])
    return e+ecore, (ci0.reshape(na,nb), ci_strs)