Python cistring.gen_linkstr_index_trilidx函数代码示例

def contract_1e(f1e, fcivec, norb, nelec, link_index=None):
    assert(fcivec.flags.c_contiguous)
    if isinstance(nelec, (int, numpy.integer)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    if link_index is None:
        link_indexa = cistring.gen_linkstr_index_trilidx(range(norb), neleca)
        link_indexb = cistring.gen_linkstr_index_trilidx(range(norb), nelecb)
    else:
        link_indexa, link_indexb = link_index

    na, nlinka = link_indexa.shape[:2]
    nb, nlinkb = link_indexb.shape[:2]
    ci1 = numpy.zeros((na,nb))
    f1e_tril = pyscf.lib.pack_tril(f1e[0])
    libfci.FCIcontract_a_1e(f1e_tril.ctypes.data_as(ctypes.c_void_p),
                            fcivec.ctypes.data_as(ctypes.c_void_p),
                            ci1.ctypes.data_as(ctypes.c_void_p),
                            ctypes.c_int(norb),
                            ctypes.c_int(na), ctypes.c_int(nb),
                            ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
                            link_indexa.ctypes.data_as(ctypes.c_void_p),
                            link_indexb.ctypes.data_as(ctypes.c_void_p))
    f1e_tril = pyscf.lib.pack_tril(f1e[1])
    libfci.FCIcontract_b_1e(f1e_tril.ctypes.data_as(ctypes.c_void_p),
                            fcivec.ctypes.data_as(ctypes.c_void_p),
                            ci1.ctypes.data_as(ctypes.c_void_p),
                            ctypes.c_int(norb),
                            ctypes.c_int(na), ctypes.c_int(nb),
                            ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
                            link_indexa.ctypes.data_as(ctypes.c_void_p),
                            link_indexb.ctypes.data_as(ctypes.c_void_p))
    return ci1

def contract_2e(eri, fcivec, norb, nelec, link_index=None):
    assert(fcivec.flags.c_contiguous)
    if isinstance(nelec, (int, numpy.integer)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    g2e_aa = pyscf.ao2mo.restore(4, eri[0], norb)
    g2e_ab = pyscf.ao2mo.restore(4, eri[1], norb)
    g2e_bb = pyscf.ao2mo.restore(4, eri[2], norb)

    if link_index is None:
        link_indexa = cistring.gen_linkstr_index_trilidx(range(norb), neleca)
        link_indexb = cistring.gen_linkstr_index_trilidx(range(norb), nelecb)
    else:
        link_indexa, link_indexb = link_index

    na, nlinka = link_indexa.shape[:2]
    nb, nlinkb = link_indexb.shape[:2]
    fcivec = fcivec.reshape(na,nb)
    ci1 = numpy.empty_like(fcivec)

    libfci.FCIcontract_uhf2e(g2e_aa.ctypes.data_as(ctypes.c_void_p),
                             g2e_ab.ctypes.data_as(ctypes.c_void_p),
                             g2e_bb.ctypes.data_as(ctypes.c_void_p),
                             fcivec.ctypes.data_as(ctypes.c_void_p),
                             ci1.ctypes.data_as(ctypes.c_void_p),
                             ctypes.c_int(norb),
                             ctypes.c_int(na), ctypes.c_int(nb),
                             ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
                             link_indexa.ctypes.data_as(ctypes.c_void_p),
                             link_indexb.ctypes.data_as(ctypes.c_void_p))
    return ci1

def contract_2e(eri, fcivec, norb, nelec, link_index=None, orbsym=[]):
    if not orbsym:
        return direct_spin1.contract_2e(eri, fcivec, norb, nelec, link_index)

    eri = pyscf.ao2mo.restore(4, eri, norb)
    if link_index is None:
        if isinstance(nelec, (int, numpy.integer)):
            nelecb = nelec//2
            neleca = nelec - nelecb
        else:
            neleca, nelecb = nelec
        link_indexa = cistring.gen_linkstr_index_trilidx(range(norb), neleca)
        link_indexb = cistring.gen_linkstr_index_trilidx(range(norb), nelecb)
    else:
        link_indexa, link_indexb = link_index
    na, nlinka = link_indexa.shape[:2]
    nb, nlinkb = link_indexb.shape[:2]
    fcivec = fcivec.reshape(na,nb)
    ci1 = numpy.empty_like(fcivec)

    eri, link_indexa, dimirrep = reorder4irrep(eri, norb, link_indexa, orbsym)
    link_indexb = reorder4irrep(eri, norb, link_indexb, orbsym)[1]
    dimirrep = numpy.array(dimirrep, dtype=numpy.int32)

    libfci.FCIcontract_2e_spin1_symm(eri.ctypes.data_as(ctypes.c_void_p),
                                     fcivec.ctypes.data_as(ctypes.c_void_p),
                                     ci1.ctypes.data_as(ctypes.c_void_p),
                                     ctypes.c_int(norb),
                                     ctypes.c_int(na), ctypes.c_int(nb),
                                     ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
                                     link_indexa.ctypes.data_as(ctypes.c_void_p),
                                     link_indexb.ctypes.data_as(ctypes.c_void_p),
                                     dimirrep.ctypes.data_as(ctypes.c_void_p),
                                     ctypes.c_int(len(dimirrep)))
    return ci1

def contract_2e(eri, fcivec, norb, nelec, link_index=None):
    eri = pyscf.ao2mo.restore(4, eri, norb)
    if link_index is None:
        if isinstance(nelec, (int, numpy.integer)):
            nelecb = nelec//2
            neleca = nelec - nelecb
        else:
            neleca, nelecb = nelec
        link_indexa = cistring.gen_linkstr_index_trilidx(range(norb), neleca)
        link_indexb = cistring.gen_linkstr_index_trilidx(range(norb), nelecb)
    else:
        link_indexa, link_indexb = link_index

    na, nlinka = link_indexa.shape[:2]
    nb, nlinkb = link_indexb.shape[:2]
    fcivec = fcivec.reshape(na,nb)
    ci1 = numpy.empty_like(fcivec)

    libfci.FCIcontract_2e_spin1(eri.ctypes.data_as(ctypes.c_void_p),
                                fcivec.ctypes.data_as(ctypes.c_void_p),
                                ci1.ctypes.data_as(ctypes.c_void_p),
                                ctypes.c_int(norb),
                                ctypes.c_int(na), ctypes.c_int(nb),
                                ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
                                link_indexa.ctypes.data_as(ctypes.c_void_p),
                                link_indexb.ctypes.data_as(ctypes.c_void_p))
    return ci1

def kernel_ms1(fci, h1e, eri, norb, nelec, ci0=None, **kwargs):
    if isinstance(nelec, (int, numpy.integer)):
        nelecb = nelec//2
        neleca = nelec - nelecb
        nelec = (neleca, nelecb)
    else:
        neleca, nelecb = nelec
    link_indexa = cistring.gen_linkstr_index_trilidx(range(norb), neleca)
    link_indexb = cistring.gen_linkstr_index_trilidx(range(norb), nelecb)
    na = link_indexa.shape[0]
    nb = link_indexb.shape[0]
    hdiag = fci.make_hdiag(h1e, eri, norb, nelec)

    addr, h0 = fci.pspace(h1e, eri, norb, nelec, hdiag, fci.pspace_size)
    pw, pv = scipy.linalg.eigh(h0)
# The degenerated wfn can break symmetry.  The davidson iteration with proper
# initial guess doesn't have this issue
    if not fci.davidson_only:
        if len(pw) == na*nb:
            if na*nb == 1:
                return pw[0], pv[:,0]
            elif fci.nroots > 1:
                civec = numpy.empty((fci.nroots,na*nb))
                civec[:,addr] = pv[:,:fci.nroots].T
                return pw[:fci.nroots], civec.reshape(fci.nroots,na,nb)
            elif abs(pw[0]-pw[1]) > 1e-12:
                civec = numpy.empty((na*nb))
                civec[addr] = pv[:,0]
                return pw[0], civec.reshape(na,nb)

    precond = fci.make_precond(hdiag, pw, pv, addr)

    h2e = fci.absorb_h1e(h1e, eri, norb, nelec, .5)
    def hop(c):
        hc = fci.contract_2e(h2e, c, norb, nelec, (link_indexa,link_indexb))
        return hc.ravel()

    if ci0 is None:
        if hasattr(fci, 'get_init_guess'):
            ci0 = fci.get_init_guess(norb, nelec, fci.nroots, hdiag)
        else:
            ci0 = []
            for i in range(fci.nroots):
                x = numpy.zeros(na*nb)
                x[addr[i]] = 1
                ci0.append(x)
    else:
        if isinstance(ci0, numpy.ndarray) and ci0.size == na*nb:
            ci0 = [ci0.ravel()]
        else:
            ci0 = [x.ravel() for x in ci0]

    #e, c = pyscf.lib.davidson(hop, ci0, precond, tol=fci.conv_tol, lindep=fci.lindep)
    e, c = fci.eig(hop, ci0, precond, **kwargs)
    if fci.nroots > 1:
        return e, [ci.reshape(na,nb) for ci in c]
    else:
        return e, c.reshape(na,nb)

 def gen_linkstr(self, norb, nelec, tril=True, spin=None):
     if spin is None:
         spin = self.spin
     neleca, nelecb = _unpack_nelec(nelec, spin)
     if tril:
         link_indexa = cistring.gen_linkstr_index_trilidx(range(norb), neleca)
         link_indexb = cistring.gen_linkstr_index_trilidx(range(norb), nelecb)
     else:
         link_indexa = cistring.gen_linkstr_index(range(norb), neleca)
         link_indexb = cistring.gen_linkstr_index(range(norb), nelecb)
     return link_indexa, link_indexb

def _unpack(norb, nelec, link_index):
    if link_index is None:
        if isinstance(nelec, (int, numpy.number)):
            nelecb = nelec//2
            neleca = nelec - nelecb
        else:
            neleca, nelecb = nelec
        link_indexa = cistring.gen_linkstr_index_trilidx(range(norb), neleca)
        link_indexb = cistring.gen_linkstr_index_trilidx(range(norb), nelecb)
        return link_indexa, link_indexb
    else:
        return link_index

def contract_2e(eri, fcivec, norb, nelec, link_index=None):
    r'''Contract the 2-electron Hamiltonian with a FCI vector to get a new FCI
    vector.

    Note the input arg eri is NOT the 2e hamiltonian matrix, the 2e hamiltonian is

    .. math::

        h2e &= eri_{pq,rs} p^+ q r^+ s \\
            &= (pq|rs) p^+ r^+ s q - (pq|rs) \delta_{qr} p^+ s

    So eri is defined as

    .. math::

        eri_{pq,rs} = (pq|rs) - (1/Nelec) \sum_q (pq|qs)

    to restore the symmetry between pq and rs,

    .. math::

        eri_{pq,rs} = (pq|rs) - (.5/Nelec) [\sum_q (pq|qs) + \sum_p (pq|rp)]

    Please refer to the treatment in :func:`direct_spin1.absorb_h1e`
    '''
    assert(fcivec.flags.c_contiguous)
    eri = pyscf.ao2mo.restore(4, eri, norb)
    if link_index is None:
        if isinstance(nelec, (int, numpy.integer)):
            nelecb = nelec//2
            neleca = nelec - nelecb
        else:
            neleca, nelecb = nelec
        link_indexa = cistring.gen_linkstr_index_trilidx(range(norb), neleca)
        link_indexb = cistring.gen_linkstr_index_trilidx(range(norb), nelecb)
    else:
        link_indexa, link_indexb = link_index

    na, nlinka = link_indexa.shape[:2]
    nb, nlinkb = link_indexb.shape[:2]
    fcivec = fcivec.reshape(na,nb)
    ci1 = numpy.empty_like(fcivec)

    libfci.FCIcontract_2e_spin1(eri.ctypes.data_as(ctypes.c_void_p),
                                fcivec.ctypes.data_as(ctypes.c_void_p),
                                ci1.ctypes.data_as(ctypes.c_void_p),
                                ctypes.c_int(norb),
                                ctypes.c_int(na), ctypes.c_int(nb),
                                ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
                                link_indexa.ctypes.data_as(ctypes.c_void_p),
                                link_indexb.ctypes.data_as(ctypes.c_void_p))
    return ci1

def contract_2e(eri, fcivec, norb, nelec, link_index=None, orbsym=[]):
    assert(fcivec.flags.c_contiguous)
    if not list(orbsym):
        return direct_spin0.contract_2e(eri, fcivec, norb, nelec, link_index)

    eri = pyscf.ao2mo.restore(4, eri, norb)
    if link_index is None:
        if isinstance(nelec, (int, numpy.integer)):
            neleca = nelec//2
        else:
            neleca, nelecb = nelec
            assert(neleca == nelecb)
        link_index = cistring.gen_linkstr_index_trilidx(range(norb), neleca)
    na,nlink,_ = link_index.shape
    ci1 = numpy.empty((na,na))

    eri, link_index, dimirrep = \
            direct_spin1_symm.reorder4irrep(eri, norb, link_index, orbsym)
    dimirrep = numpy.array(dimirrep, dtype=numpy.int32)

    libfci.FCIcontract_2e_spin0_symm(eri.ctypes.data_as(ctypes.c_void_p),
                                     fcivec.ctypes.data_as(ctypes.c_void_p),
                                     ci1.ctypes.data_as(ctypes.c_void_p),
                                     ctypes.c_int(norb), ctypes.c_int(na),
                                     ctypes.c_int(nlink),
                                     link_index.ctypes.data_as(ctypes.c_void_p),
                                     dimirrep.ctypes.data_as(ctypes.c_void_p),
                                     ctypes.c_int(len(dimirrep)))
    return pyscf.lib.transpose_sum(ci1, inplace=True)

    def test_linkstr_index(self):
        idx1 = cistring.gen_linkstr_index_o0(range(4), 2)
        idx2 = cistring.gen_linkstr_index(range(4), 2)
        idx23 = numpy.array([[0, 0, 3, 1],
                             [3, 3, 3, 1],
                             [1, 0, 4, 1],
                             [2, 0, 5, 1],
                             [1, 3, 0, 1],
                             [2, 3, 1, 1],])
        self.assertTrue(numpy.all(idx1[:,:,2:] == idx2[:,:,2:]))
        self.assertTrue(numpy.all(idx23 == idx2[3]))

        idx1 = cistring.gen_linkstr_index(range(7), 3)
        idx2 = cistring.reform_linkstr_index(idx1)
        idx3 = cistring.gen_linkstr_index_trilidx(range(7), 3)
        idx3[:,:,1] = 0
        self.assertTrue(numpy.all(idx2 == idx3))

        tab1 = cistring.gen_cre_str_index_o0(range(8), 4)
        tab2 = cistring.gen_cre_str_index_o1(range(8), 4)
        self.assertAlmostEqual(abs(tab1 - tab2).max(), 0, 12)
        tab1 = cistring.gen_des_str_index_o0(range(8), 4)
        tab2 = cistring.gen_des_str_index_o1(range(8), 4)
        self.assertAlmostEqual(abs(tab1 - tab2).max(), 0, 12)

        tab1 = cistring.gen_linkstr_index_o0(range(8), 4)
        tab2 = cistring.gen_linkstr_index(range(8), 4)
        self.assertAlmostEqual(abs(tab1 - tab2).sum(), 0, 12)
        tab3 = cistring.gen_linkstr_index_o1(range(8), 4)
        self.assertAlmostEqual(abs(tab1 - tab3).sum(), 0, 12)

def kernel_ms0(fci, h1e, eri, norb, nelec, ci0=None, **kwargs):
    if isinstance(nelec, (int, numpy.integer)):
        neleca = nelec//2
    else:
        neleca, nelecb = nelec
        assert(neleca == nelecb)
    h1e = numpy.ascontiguousarray(h1e)
    eri = numpy.ascontiguousarray(eri)
    link_index = cistring.gen_linkstr_index_trilidx(range(norb), neleca)
    na = link_index.shape[0]
    hdiag = fci.make_hdiag(h1e, eri, norb, nelec)

    addr, h0 = fci.pspace(h1e, eri, norb, nelec, hdiag, fci.pspace_size)
    pw, pv = scipy.linalg.eigh(h0)
# The degenerated wfn can break symmetry.  The davidson iteration with proper
# initial guess doesn't have this issue
    if not fci.davidson_only:
        if len(addr) == 1:
            return pw, pv
        elif len(addr) == na*na:
            if fci.nroots > 1:
                ci0 = numpy.empty((fci.nroots,na*na))
                ci0[:,addr] = pv[:,:fci.nroots].T
                return pw[:fci.nroots], ci0.reshape(fci.nroots,na,na)
            elif abs(pw[0]-pw[1]) > 1e-12:
                ci0 = numpy.empty((na*na))
                ci0[addr] = pv[:,0]
                ci0 = ci0.reshape(na,na)
                ci0 = pyscf.lib.transpose_sum(ci0, inplace=True) * .5
                # direct diagonalization may lead to triplet ground state
##TODO: optimize initial guess.  Using pspace vector as initial guess may have
## spin problems.  The 'ground state' of psapce vector may have different spin
## state to the true ground state.
                if numpy.allclose(numpy.linalg.norm(ci0), 1):
                    return pw[0], ci0.reshape(na,na)

    precond = fci.make_precond(hdiag, pw, pv, addr)

    h2e = fci.absorb_h1e(h1e, eri, norb, nelec, .5)
    def hop(c):
        hc = fci.contract_2e(h2e, c, norb, nelec, link_index)
        return hc.ravel()

#TODO: check spin of initial guess
    if ci0 is None:
        # we need better initial guess
        ci0 = numpy.zeros(na*na)
        #ci0[addr] = pv[:,0]
        ci0[0] = 1
    else:
# symmetrize the initial guess, otherwise got strange numerical noise after
# couple of davidson iterations
#        ci0 = pyscf.lib.transpose_sum(ci0.reshape(na,na)).ravel()*.5
        ci0 = ci0.ravel()

    #e, c = pyscf.lib.davidson(hop, ci0, precond, tol=fci.conv_tol, lindep=fci.lindep)
    e, c = fci.eig(hop, ci0, precond, **kwargs)
    return e, c.reshape(na,na)

def _unpack(norb, nelec, link_index):
    if link_index is None:
        if isinstance(nelec, (int, numpy.number)):
            neleca = nelec//2
        else:
            neleca, nelecb = nelec
            assert(neleca == nelecb)
        return cistring.gen_linkstr_index_trilidx(range(norb), neleca)
    else:
        return link_index

 def gen_linkstr(self, norb, nelec, tril=True, spin=None):
     if isinstance(nelec, (int, numpy.number)):
         neleca = nelec//2
     else:
         neleca, nelecb = nelec
         assert(neleca == nelecb)
     if tril:
         link_index = cistring.gen_linkstr_index_trilidx(range(norb), neleca)
     else:
         link_index = cistring.gen_linkstr_index(range(norb), neleca)
     return link_index

def contract_1e(f1e, fcivec, norb, nelec, link_index=None):
    '''Contract the 1-electron Hamiltonian with a FCI vector to get a new FCI
    vector.
    '''
    if link_index is None:
        if isinstance(nelec, (int, numpy.integer)):
            nelecb = nelec//2
            neleca = nelec - nelecb
        else:
            neleca, nelecb = nelec
        link_indexa = cistring.gen_linkstr_index_trilidx(range(norb), neleca)
        link_indexb = cistring.gen_linkstr_index_trilidx(range(norb), nelecb)
    else:
        link_indexa, link_indexb = link_index

    na, nlinka = link_indexa.shape[:2]
    nb, nlinkb = link_indexb.shape[:2]
    f1e_tril = pyscf.lib.pack_tril(f1e)
    ci1 = numpy.zeros((na,nb))
    libfci.FCIcontract_a_1e(f1e_tril.ctypes.data_as(ctypes.c_void_p),
                            fcivec.ctypes.data_as(ctypes.c_void_p),
                            ci1.ctypes.data_as(ctypes.c_void_p),
                            ctypes.c_int(norb),
                            ctypes.c_int(na), ctypes.c_int(nb),
                            ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
                            link_indexa.ctypes.data_as(ctypes.c_void_p),
                            link_indexb.ctypes.data_as(ctypes.c_void_p))
    libfci.FCIcontract_b_1e(f1e_tril.ctypes.data_as(ctypes.c_void_p),
                            fcivec.ctypes.data_as(ctypes.c_void_p),
                            ci1.ctypes.data_as(ctypes.c_void_p),
                            ctypes.c_int(norb),
                            ctypes.c_int(na), ctypes.c_int(nb),
                            ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
                            link_indexa.ctypes.data_as(ctypes.c_void_p),
                            link_indexb.ctypes.data_as(ctypes.c_void_p))
    return ci1

    def test_linkstr_index(self):
        idx1 = cistring.gen_linkstr_index_o0(range(4), 2)
        idx2 = cistring.gen_linkstr_index(range(4), 2)
        idx23 = numpy.array([[0, 0, 3, 1],
                             [3, 3, 3, 1],
                             [1, 0, 4, 1],
                             [2, 0, 5, 1],
                             [1, 3, 0, 1],
                             [2, 3, 1, 1],])
        self.assertTrue(numpy.all(idx1[:,:,2:] == idx2[:,:,2:]))
        self.assertTrue(numpy.all(idx23 == idx2[3]))

        idx1 = cistring.gen_linkstr_index(range(7), 3)
        idx2 = cistring.reform_linkstr_index(idx1)
        idx3 = cistring.gen_linkstr_index_trilidx(range(7), 3)
        idx3[:,:,1] = 0
        self.assertTrue(numpy.all(idx2 == idx3))