Python lib.tag_array函数代码示例

 def ccsd(self, t1=None, t2=None, eris=None, **kwargs):
     if eris is None: eris = self.ao2mo(self.mo_coeff)
     e_corr, self.t1, self.t2 = ccsd.CCSD.ccsd(self, t1, t2, eris)
     if getattr(eris, 'orbspin', None) is not None:
         self.t1 = lib.tag_array(self.t1, orbspin=eris.orbspin)
         self.t2 = lib.tag_array(self.t2, orbspin=eris.orbspin)
     return e_corr, self.t1, self.t2

    def eig(self, h, s):
        mol = self.mol
        if not mol.symmetry:
            return self._eigh(h, s)

        nirrep = mol.symm_orb.__len__()
        s = symm.symmetrize_matrix(s, mol.symm_orb)
        ha = symm.symmetrize_matrix(h[0], mol.symm_orb)
        cs = []
        es = []
        orbsym = []
        for ir in range(nirrep):
            e, c = self._eigh(ha[ir], s[ir])
            cs.append(c)
            es.append(e)
            orbsym.append([mol.irrep_id[ir]] * e.size)
        ea = numpy.hstack(es)
        ca = hf_symm.so2ao_mo_coeff(mol.symm_orb, cs)
        ca = lib.tag_array(ca, orbsym=numpy.hstack(orbsym))

        hb = symm.symmetrize_matrix(h[1], mol.symm_orb)
        cs = []
        es = []
        orbsym = []
        for ir in range(nirrep):
            e, c = self._eigh(hb[ir], s[ir])
            cs.append(c)
            es.append(e)
            orbsym.append([mol.irrep_id[ir]] * e.size)
        eb = numpy.hstack(es)
        cb = hf_symm.so2ao_mo_coeff(mol.symm_orb, cs)
        cb = lib.tag_array(cb, orbsym=numpy.hstack(orbsym))
        return (ea,eb), (ca,cb)

    def ccsd(self, t1=None, t2=None, eris=None, mbpt2=False):
        '''Ground-state unrestricted (U)CCSD.

        Kwargs:
            mbpt2 : bool
                Use one-shot MBPT2 approximation to CCSD.
        '''
        if mbpt2:
            from pyscf.mp import gmp2
            pt = gmp2.GMP2(self._scf, self.frozen, self.mo_coeff, self.mo_occ)
            self.e_corr, self.t2 = pt.kernel(eris=eris)
            nocc, nvir = self.t2.shape[1:3]
            self.t1 = np.zeros((nocc,nvir))
            return self.e_corr, self.t1, self.t2

        # Initialize orbspin so that we can attach it to t1, t2
        if getattr(self.mo_coeff, 'orbspin', None) is None:
            orbspin = scf.ghf.guess_orbspin(self.mo_coeff)
            if not np.any(orbspin == -1):
                self.mo_coeff = lib.tag_array(self.mo_coeff, orbspin=orbspin)
        if eris is None: eris = self.ao2mo(self.mo_coeff)

        e_corr, self.t1, self.t2 = ccsd.CCSD.ccsd(self, t1, t2, eris)
        if getattr(eris, 'orbspin', None) is not None:
            self.t1 = lib.tag_array(self.t1, orbspin=eris.orbspin)
            self.t2 = lib.tag_array(self.t2, orbspin=eris.orbspin)
        return e_corr, self.t1, self.t2

    def _finalize(self):
        uhf.UHF._finalize(self)

        ea = numpy.hstack(self.mo_energy[0])
        eb = numpy.hstack(self.mo_energy[1])
        # Using mergesort because it is stable. We don't want to change the
        # ordering of the symmetry labels when two orbitals are degenerated.
        oa_sort = numpy.argsort(ea[self.mo_occ[0]>0 ].round(9), kind='mergesort')
        va_sort = numpy.argsort(ea[self.mo_occ[0]==0].round(9), kind='mergesort')
        ob_sort = numpy.argsort(eb[self.mo_occ[1]>0 ].round(9), kind='mergesort')
        vb_sort = numpy.argsort(eb[self.mo_occ[1]==0].round(9), kind='mergesort')
        idxa = numpy.arange(ea.size)
        idxa = numpy.hstack((idxa[self.mo_occ[0]> 0][oa_sort],
                             idxa[self.mo_occ[0]==0][va_sort]))
        idxb = numpy.arange(eb.size)
        idxb = numpy.hstack((idxb[self.mo_occ[1]> 0][ob_sort],
                             idxb[self.mo_occ[1]==0][vb_sort]))
        self.mo_energy = (ea[idxa], eb[idxb])
        orbsyma, orbsymb = get_orbsym(self.mol, self.mo_coeff)
        self.mo_coeff = (lib.tag_array(self.mo_coeff[0][:,idxa], orbsym=orbsyma[idxa]),
                         lib.tag_array(self.mo_coeff[1][:,idxb], orbsym=orbsymb[idxb]))
        self.mo_occ = (self.mo_occ[0][idxa], self.mo_occ[1][idxb])
        if self.chkfile:
            chkfile.dump_scf(self.mol, self.chkfile, self.e_tot, self.mo_energy,
                             self.mo_coeff, self.mo_occ, overwrite_mol=False)
        return self

    def _common_init_(self, mycc, mo_coeff=None):
        if mo_coeff is None:
            mo_coeff = mycc.mo_coeff
        mo_idx = ccsd.get_frozen_mask(mycc)
        if getattr(mo_coeff, 'orbspin', None) is not None:
            self.orbspin = mo_coeff.orbspin[mo_idx]
            mo_coeff = lib.tag_array(mo_coeff[:,mo_idx], orbspin=self.orbspin)
            self.mo_coeff = mo_coeff
        else:
            orbspin = scf.ghf.guess_orbspin(mo_coeff)
            self.mo_coeff = mo_coeff = mo_coeff[:,mo_idx]
            if not np.any(orbspin == -1):
                self.orbspin = orbspin[mo_idx]
                self.mo_coeff = lib.tag_array(mo_coeff, orbspin=self.orbspin)

        # Note: Recomputed fock matrix since SCF may not be fully converged.
        dm = mycc._scf.make_rdm1(mycc.mo_coeff, mycc.mo_occ)
        fockao = mycc._scf.get_fock(dm=dm)
        self.fock = reduce(np.dot, (mo_coeff.conj().T, fockao, mo_coeff))
        self.nocc = mycc.nocc

        mo_e = self.mo_energy = self.fock.diagonal().real
        gap = abs(mo_e[:self.nocc,None] - mo_e[None,self.nocc:]).min()
        if gap < 1e-5:
            logger.warn(mycc, 'HOMO-LUMO gap %s too small for GCCSD', gap)
        return self

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 __init__(self, mp, mo_coeff=None):
        self.orbspin = None

        if mo_coeff is None:
            mo_coeff = mp.mo_coeff
        mo_idx = mp.get_frozen_mask()
        if getattr(mo_coeff, 'orbspin', None) is not None:
            self.orbspin = mo_coeff.orbspin[mo_idx]
            mo_coeff = lib.tag_array(mo_coeff[:,mo_idx], orbspin=self.orbspin)
            self.mo_coeff = mo_coeff
        else:
            orbspin = scf.ghf.guess_orbspin(mo_coeff)
            self.mo_coeff = mo_coeff = mo_coeff[:,mo_idx]
            if not numpy.any(orbspin == -1):
                self.orbspin = orbspin[mo_idx]
                self.mo_coeff = lib.tag_array(mo_coeff, orbspin=self.orbspin)

def bcast_tagged_array(arr):
    '''Broadcast big nparray or tagged array.'''
    if comm.bcast(not isinstance(arr, numpy.ndarray)):
        return comm.bcast(arr)

    new_arr = bcast(arr)

    if comm.bcast(isinstance(arr, lib.NPArrayWithTag)):
        new_arr = lib.tag_array(new_arr)
        if rank == 0:
            kv = []
            for k, v in arr.__dict__.items():
                if isinstance(v, numpy.ndarray) and v.nbytes > 1e5:
                    kv.append((k, 'NPARRAY_TO_BCAST'))
                else:
                    kv.append((k, v))
            comm.bcast(kv)
        else:
            kv = comm.bcast(None)
            new_arr.__dict__.update(kv)

        for k, v in kv:
            if v is 'NPARRAY_TO_BCAST':
                new_arr.k = bcast(v)

    if rank != 0:
        arr = new_arr
    return arr

def get_roothaan_fock(focka_fockb, dma_dmb, s):
    '''Roothaan's effective fock.
    Ref. http://www-theor.ch.cam.ac.uk/people/ross/thesis/node15.html

    ======== ======== ====== =========
    space     closed   open   virtual
    ======== ======== ====== =========
    closed      Fc      Fb     Fc
    open        Fb      Fc     Fa
    virtual     Fc      Fa     Fc
    ======== ======== ====== =========

    where Fc = (Fa + Fb) / 2

    Returns:
        Roothaan effective Fock matrix
    '''
    nao = s.shape[0]
    focka, fockb = focka_fockb
    dma, dmb = dma_dmb
    fc = (focka + fockb) * .5
# Projector for core, open-shell, and virtual
    pc = numpy.dot(dmb, s)
    po = numpy.dot(dma-dmb, s)
    pv = numpy.eye(nao) - numpy.dot(dma, s)
    fock  = reduce(numpy.dot, (pc.conj().T, fc, pc)) * .5
    fock += reduce(numpy.dot, (po.conj().T, fc, po)) * .5
    fock += reduce(numpy.dot, (pv.conj().T, fc, pv)) * .5
    fock += reduce(numpy.dot, (po.conj().T, fockb, pc))
    fock += reduce(numpy.dot, (po.conj().T, focka, pv))
    fock += reduce(numpy.dot, (pv.conj().T, fc, pc))
    fock = fock + fock.conj().T
    fock = lib.tag_array(fock, focka=focka, fockb=fockb)
    return fock

 def eig(self, fock, s):
     e, c = self._eigh(fock, s)
     if getattr(fock, 'focka', None) is not None:
         mo_ea = numpy.einsum('pi,pi->i', c.conj(), fock.focka.dot(c)).real
         mo_eb = numpy.einsum('pi,pi->i', c.conj(), fock.fockb.dot(c)).real
         e = lib.tag_array(e, mo_ea=mo_ea, mo_eb=mo_eb)
     return e, c

def _make_eris_incore(mp, mo_coeff=None, ao2mofn=None, verbose=None):
    eris = _PhysicistsERIs(mp, mo_coeff)

    nocc = mp.nocc
    nao, nmo = eris.mo_coeff.shape
    nvir = nmo - nocc
    orbspin = eris.orbspin

    if callable(ao2mofn):
        orbo = eris.mo_coeff[:,:nocc]
        orbv = eris.mo_coeff[:,nocc:]
        orbo = lib.tag_array(orbo, orbspin=orbspin)
        eri = ao2mofn((orbo,orbv,orbo,orbv)).reshape(nocc,nvir,nocc,nvir)
    else:
        orboa = eris.mo_coeff[:nao//2,:nocc]
        orbob = eris.mo_coeff[nao//2:,:nocc]
        orbva = eris.mo_coeff[:nao//2,nocc:]
        orbvb = eris.mo_coeff[nao//2:,nocc:]
        if orbspin is None:
            eri  = ao2mo.kernel(mp._scf._eri, (orboa,orbva,orboa,orbva))
            eri += ao2mo.kernel(mp._scf._eri, (orbob,orbvb,orbob,orbvb))
            eri1 = ao2mo.kernel(mp._scf._eri, (orboa,orbva,orbob,orbvb))
            eri += eri1
            eri += eri1.T
            eri = eri.reshape(nocc,nvir,nocc,nvir)
        else:
            co = orboa + orbob
            cv = orbva + orbvb
            eri = ao2mo.kernel(mp._scf._eri, (co,cv,co,cv)).reshape(nocc,nvir,nocc,nvir)
            sym_forbid = (orbspin[:nocc,None] != orbspin[nocc:])
            eri[sym_forbid,:,:] = 0
            eri[:,:,sym_forbid] = 0

    eris.oovv = eri.transpose(0,2,1,3) - eri.transpose(0,2,3,1)
    return eris

 def test_uhf_symm_get_occ(self):
     pmol = n2sym.copy()
     pmol.spin = 2
     mf = scf.UHF(pmol).set(verbose = 0)
     orbsym = numpy.array([[0 , 5 , 0 , 5 , 6 , 7 , 0 , 2 , 3 , 5 , 0 , 6 , 7 , 0 , 2 , 3 , 5 , 10, 11, 5],
                           [5 , 0 , 6 , 7 , 5 , 10, 11, 0 , 5 , 0 , 5 , 5 , 6 , 7 , 0 , 2 , 3 , 0 , 2 , 3]])
     energy = numpy.array([[34, 2 , 54, 43, 42, 33, 20, 61, 29, 26, 62, 52, 13, 51, 18, 78, 85, 49, 84, 7],
                           [29, 26, 13, 54, 18, 78, 85, 49, 84, 62, 42, 74, 20, 61, 51, 34, 2 , 33, 52, 3]])
     mf.irrep_nelec = {'A1g':7, 'A1u':3, 'E1ux':2, 'E1uy':2}
     mo_coeff = lib.tag_array([numpy.eye(energy.size)]*2, orbsym=orbsym)
     self.assertTrue(numpy.allclose(mf.get_occ(energy, mo_coeff),
             [[1, 1, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1],
              [0, 1, 1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0]]))
     mf.irrep_nelec = {'A1g':(5,2), 'A1u':(1,2)}
     self.assertTrue(numpy.allclose(mf.get_occ(energy, mo_coeff),
             [[1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0],
              [1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1]]))
     mf.irrep_nelec = {'E1ux':2, 'E1uy':2}
     self.assertTrue(numpy.allclose(mf.get_occ(energy, mo_coeff),
             [[0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1],
              [0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1]]))
     mf.irrep_nelec = {}
     self.assertTrue(numpy.allclose(mf.get_occ(energy, mo_coeff),
             [[0, 1, 0, 0, 0, 1, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1],
              [0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1]]))

 def rotate_mo(self, mo_coeff, u, log=None):
     mo = numpy.asarray((numpy.dot(mo_coeff[0], u[0]),
                         numpy.dot(mo_coeff[1], u[1])))
     if self._scf.mol.symmetry:
         orbsym = uhf_symm.get_orbsym(self._scf.mol, mo_coeff)
         mo = lib.tag_array(mo, orbsym=orbsym)
     return mo

    def test_canonicalize1(self):
        numpy.random.seed(1)
        f1 = numpy.random.random(mcr.mo_coeff.shape)
        u1 = numpy.linalg.svd(f1)[0]
        mo1 = numpy.dot(mcr.mo_coeff, u1)
        mo1 = lib.tag_array(mo1, orbsym=mcr.mo_coeff.orbsym)
        mo, ci, mo_e = mcr.canonicalize(mo1)
        e1 = numpy.einsum('ji,jk,ki', mo, f1, mo)
        self.assertAlmostEqual(e1, 44.2658681077, 7)
        self.assertAlmostEqual(lib.finger(mo_e), 5.1364166175063097, 7)

        mo, ci, mo_e = mcr.canonicalize(mo1, eris=mcr.ao2mo(mcr.mo_coeff))
        e1 = numpy.einsum('ji,jk,ki', mo, f1, mo)
        self.assertAlmostEqual(e1, 44.2658681077, 7)
        self.assertAlmostEqual(lib.finger(mo_e), 4.1206025804989173, 7)

        mcr1 = copy.copy(mcr)
        mcr1.frozen = 2
        mo, ci, mo_e = mcr1.canonicalize(mo1)
        self.assertAlmostEqual(lib.finger(mo_e), 6.6030999409178577, 7)

        mcr1.frozen = [0,1]
        mo, ci, mo_e = mcr1.canonicalize(mo1)
        self.assertAlmostEqual(lib.finger(mo_e), 6.6030999409178577, 7)

        mcr1.frozen = [1,12]
        mo, ci, mo_e = mcr1.canonicalize(mo1)
        self.assertAlmostEqual(lib.finger(mo_e), 5.2182584355788162, 7)

def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
    if getattr(dm, 'mo_coeff', None) is not None:
        mo_coeff = dm.mo_coeff
        mo_occ_a = (dm.mo_occ > 0).astype(numpy.double)
        mo_occ_b = (dm.mo_occ ==2).astype(numpy.double)
        dm = lib.tag_array(dm, mo_coeff=(mo_coeff,mo_coeff),
                           mo_occ=(mo_occ_a,mo_occ_b))
    return uks.get_veff(ks, mol, dm, dm_last, vhf_last, hermi)