Python lib.asarray函数代码示例

def get_j_kpts(mf, cell, dm_kpts, kpts, kpt_band=None):
    coords = gen_grid.gen_uniform_grids(cell)
    nkpts = len(kpts)
    ngs = len(coords)
    dm_kpts = np.asarray(dm_kpts)
    nao = dm_kpts.shape[-1]

    ni = numint._KNumInt(kpts)
    aoR_kpts = ni.eval_ao(cell, coords, kpts)
    if kpt_band is not None:
        aoR_kband = numint.eval_ao(cell, coords, kpt_band)

    dms = dm_kpts.reshape(-1,nkpts,nao,nao)
    nset = dms.shape[0]

    vjR = [get_vjR(cell, dms[i], aoR_kpts) for i in range(nset)]
    if kpt_band is not None:
        vj_kpts = [cell.vol/ngs * lib.dot(aoR_kband.T.conj()*vjR[i], aoR_kband)
                   for i in range(nset)]
        if dm_kpts.ndim == 3:  # One set of dm_kpts for KRHF
            vj_kpts = vj_kpts[0]
        return lib.asarray(vj_kpts)
    else:
        vj_kpts = []
        for i in range(nset):
            vj = [cell.vol/ngs * lib.dot(aoR_k.T.conj()*vjR[i], aoR_k)
                  for aoR_k in aoR_kpts]
            vj_kpts.append(lib.asarray(vj))
        return lib.asarray(vj_kpts).reshape(dm_kpts.shape)

def get_jk(mf, cell, dm, hermi=1, vhfopt=None, kpt=np.zeros(3), kpt_band=None):
    dm = np.asarray(dm)
    nao = dm.shape[-1]

    coords = gen_grid.gen_uniform_grids(cell)
    if kpt_band is None:
        kpt1 = kpt2 = kpt
        aoR_k1 = aoR_k2 = numint.eval_ao(cell, coords, kpt)
    else:
        kpt1 = kpt_band
        kpt2 = kpt
        aoR_k1 = numint.eval_ao(cell, coords, kpt1)
        aoR_k2 = numint.eval_ao(cell, coords, kpt2)

    vkR_k1k2 = get_vkR(mf, cell, aoR_k1, aoR_k2, kpt1, kpt2)

    ngs, nao = aoR_k1.shape
    def contract(dm):
        vjR_k2 = get_vjR(cell, dm, aoR_k2)
        vj = (cell.vol/ngs) * np.dot(aoR_k1.T.conj(), vjR_k2.reshape(-1,1)*aoR_k1)

        #:vk = (cell.vol/ngs) * np.einsum('rs,Rp,Rqs,Rr->pq', dm, aoR_k1.conj(),
        #:                                vkR_k1k2, aoR_k2)
        aoR_dm_k2 = np.dot(aoR_k2, dm)
        tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dm_k2)
        vk = (cell.vol/ngs) * np.dot(aoR_k1.T.conj(), tmp_Rq)
        return vj, vk

    if dm.ndim == 2:
        vj, vk = contract(dm)
    else:
        jk = [contract(x) for x in dm.reshape(-1,nao,nao)]
        vj = lib.asarray([x[0] for x in jk])
        vk = lib.asarray([x[1] for x in jk])
    return vj.reshape(dm.shape), vk.reshape(dm.shape)

def get_jk_kpts(mf, cell, dm_kpts, kpts, kpt_band=None):
    coords = gen_grid.gen_uniform_grids(cell)
    nkpts = len(kpts)
    ngs = len(coords)
    dm_kpts = np.asarray(dm_kpts)
    nao = dm_kpts.shape[-1]

    dms = dm_kpts.reshape(-1,nkpts,nao,nao)
    nset = dms.shape[0]

    ni = numint._KNumInt(kpts)
    aoR_kpts = ni.eval_ao(cell, coords, kpts)
    if kpt_band is not None:
        aoR_kband = numint.eval_ao(cell, coords, kpt_band)

# J
    vjR = [get_vjR_kpts(cell, dms[i], aoR_kpts) for i in range(nset)]
    if kpt_band is not None:
        vj_kpts = [cell.vol/ngs * lib.dot(aoR_kband.T.conj()*vjR[i], aoR_kband)
                   for i in range(nset)]
    else:
        vj_kpts = []
        for i in range(nset):
            vj = [cell.vol/ngs * lib.dot(aoR_k.T.conj()*vjR[i], aoR_k)
                  for aoR_k in aoR_kpts]
            vj_kpts.append(lib.asarray(vj))
    vj_kpts = lib.asarray(vj_kpts)
    vjR = None

# K
    weight = 1./nkpts * (cell.vol/ngs)
    vk_kpts = np.zeros_like(vj_kpts)
    if kpt_band is not None:
        for k2, kpt2 in enumerate(kpts):
            aoR_dms = [lib.dot(aoR_kpts[k2], dms[i,k2]) for i in range(nset)]
            vkR_k1k2 = get_vkR(mf, cell, aoR_kband, aoR_kpts[k2],
                               kpt_band, kpt2)
            #:vk_kpts = 1./nkpts * (cell.vol/ngs) * np.einsum('rs,Rp,Rqs,Rr->pq',
            #:            dm_kpts[k2], aoR_kband.conj(),
            #:            vkR_k1k2, aoR_kpts[k2])
            for i in range(nset):
                tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i])
                vk_kpts[i] += weight * lib.dot(aoR_kband.T.conj(), tmp_Rq)
            vkR_k1k2 = None
        if dm_kpts.ndim == 3:
            vj_kpts = vj_kpts[0]
            vk_kpts = vk_kpts[0]
        return lib.asarray(vj_kpts), lib.asarray(vk_kpts)
    else:
        for k2, kpt2 in enumerate(kpts):
            aoR_dms = [lib.dot(aoR_kpts[k2], dms[i,k2]) for i in range(nset)]
            for k1, kpt1 in enumerate(kpts):
                vkR_k1k2 = get_vkR(mf, cell, aoR_kpts[k1], aoR_kpts[k2],
                                   kpt1, kpt2)
                for i in range(nset):
                    tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dms[i])
                    vk_kpts[i,k1] += weight * lib.dot(aoR_kpts[k1].T.conj(), tmp_Rq)
            vkR_k1k2 = None
        return vj_kpts.reshape(dm_kpts.shape), vk_kpts.reshape(dm_kpts.shape)

 def get_hcore(self, cell=None, kpts=None):
     if cell is None: cell = self.cell
     if kpts is None: kpts = self.kpts
     if cell.pseudo is None:
         nuc = self.with_df.get_nuc(kpts)
     else:
         nuc = self.with_df.get_pp(kpts)
     t = cell.pbc_intor('cint1e_kin_sph', 1, 1, kpts)
     return lib.asarray(nuc) + lib.asarray(t)

 def get_hcore(self, cell=None, kpts=None):
     if cell is None: cell = self.cell
     if kpts is None: kpts = self.kpts
     if cell.pseudo:
         nuc = lib.asarray(self.with_df.get_pp(kpts))
     else:
         nuc = lib.asarray(self.with_df.get_nuc(kpts))
     if len(cell._ecpbas) > 0:
         nuc += lib.asarray(ecp.ecp_int(cell, kpts))
     t = lib.asarray(cell.pbc_intor('cint1e_kin_sph', 1, 1, kpts))
     return nuc + t

def kernel(method, efg_nuc=None):
    log = lib.logger.Logger(method.stdout, method.verbose)
    cell = method.cell
    if efg_nuc is None:
        efg_nuc = range(cell.natm)

    dm = method.make_rdm1()
    if isinstance(method, scf.khf.KSCF):
        if isinstance(dm[0][0], numpy.ndarray) and dm[0][0].ndim == 2:
            dm = dm[0] + dm[1]  # KUHF density matrix
    elif isinstance(method, scf.hf.SCF):
        if isinstance(dm[0], numpy.ndarray) and dm[0].ndim == 2:
            dm = dm[0] + dm[1]  # UHF density matrix
    else:
        mo = method.mo_coeff
        if isinstance(dm[0][0], numpy.ndarray) and dm[0][0].ndim == 2:
            dm_a = [lib.einsum('pi,ij,qj->pq', c, dm[0][k], c.conj())
                    for k, c in enumerate(mo)]
            dm_b = [lib.einsum('pi,ij,qj->pq', c, dm[1][k], c.conj())
                    for k, c in enumerate(mo)]
            dm = lib.asarray(dm_a) + lib.asarray(dm_b)
        else:
            dm = lib.asarray([lib.einsum('pi,ij,qj->pq', c, dm[k], c.conj())
                              for k, c in enumerate(mo)])

    if isinstance(method, scf.hf.SCF):
        with_df = getattr(method, 'with_df', None)
        with_x2c = getattr(method, 'with_x2c', None)
    else:
        with_df = getattr(method._scf, 'with_df', None)
        with_x2c = getattr(method._scf, 'with_x2c', None)
    if with_x2c:
        raise NotImplementedError

    log.info('\nElectric Field Gradient Tensor Results')
    if isinstance(with_df, df.fft.FFTDF):
        efg_e = _fft_quad_integrals(with_df, dm, efg_nuc)
    else:
        efg_e = _aft_quad_integrals(with_df, dm, efg_nuc)
    efg = []
    for i, atm_id in enumerate(efg_nuc):
        efg_nuc = _get_quad_nuc(cell, atm_id)
        v = efg_nuc - efg_e[i]
        efg.append(v)

        rhf_efg._analyze(cell, atm_id, v, log)

    return numpy.asarray(efg)

 def vind(mo1):
     mo1aa, mo1ab, mo1ba, mo1bb = _split_mo1(mo1)
     dm1aa = _dm1_mo2ao(mo1aa, orbva, orboa)
     dm1ab = _dm1_mo2ao(mo1ab, orbva, orbob)
     dm1ba = _dm1_mo2ao(mo1ba, orbvb, orboa)
     dm1bb = _dm1_mo2ao(mo1bb, orbvb, orbob)
     dm1 = lib.asarray([dm1aa+dm1aa.transpose(0,2,1),
                        dm1ab+dm1ba.transpose(0,2,1),
                        dm1ba+dm1ab.transpose(0,2,1),
                        dm1bb+dm1bb.transpose(0,2,1)])
     v1 = vresp(dm1)
     v1aa = _ao2mo.nr_e2(v1[0], mo_va_oa, (0,nvira,nvira,nvira+nocca))
     v1ab = _ao2mo.nr_e2(v1[1], mo_va_ob, (0,nvira,nvira,nvira+noccb))
     v1ba = _ao2mo.nr_e2(v1[2], mo_vb_oa, (0,nvirb,nvirb,nvirb+nocca))
     v1bb = _ao2mo.nr_e2(v1[3], mo_vb_ob, (0,nvirb,nvirb,nvirb+noccb))
     v1aa = v1aa.reshape(nset,nvira,nocca)
     v1ab = v1ab.reshape(nset,nvira,noccb)
     v1ba = v1ba.reshape(nset,nvirb,nocca)
     v1bb = v1bb.reshape(nset,nvirb,noccb)
     v1aa *= eai_aa
     v1ab *= eai_ab
     v1ba *= eai_ba
     v1bb *= eai_bb
     v1mo = numpy.hstack((v1aa.reshape(nset,-1),
                          v1ab.reshape(nset,-1),
                          v1ba.reshape(nset,-1),
                          v1bb.reshape(nset,-1)))
     return v1mo.ravel()

def get_fock(
    mf,
    h1e_kpts,
    s_kpts,
    vhf_kpts,
    dm_kpts,
    cycle=-1,
    adiis=None,
    diis_start_cycle=None,
    level_shift_factor=None,
    damp_factor=None,
):
    if diis_start_cycle is None:
        diis_start_cycle = mf.diis_start_cycle
    if level_shift_factor is None:
        level_shift_factor = mf.level_shift
    if damp_factor is None:
        damp_factor = mf.damp

    if isinstance(level_shift_factor, (tuple, list, np.ndarray)):
        shifta, shiftb = level_shift_factor
    else:
        shifta = shiftb = level_shift_factor

    f_kpts = h1e_kpts + vhf_kpts
    if adiis and cycle >= diis_start_cycle:
        f_kpts = adiis.update(s_kpts, dm_kpts, f_kpts)
    if abs(level_shift_factor) > 1e-4:
        f_kpts = [hf.level_shift(s, dm_kpts[0, k], f_kpts[0, k], shifta) for k, s in enumerate(s_kpts)] + [
            hf.level_shift(s, dm_kpts[1, k], f_kpts[1, k], shiftb) for k, s in enumerate(s_kpts)
        ]
    return lib.asarray(f_kpts)

 def eval_mat(self, cell, ao_kpts, weight, rho, vxc,
              non0tab=None, xctype='LDA', spin=0, verbose=None):
     nkpts = len(ao_kpts)
     mat = [eval_mat(cell, ao_kpts[k], weight, rho, vxc,
                     non0tab, xctype, spin, verbose)
            for k in range(nkpts)]
     return lib.asarray(mat)

def get_j(cell, dm, hermi=1, vhfopt=None, kpt=np.zeros(3), kpt_band=None):
    dm = np.asarray(dm)
    nao = dm.shape[-1]

    coords = gen_grid.gen_uniform_grids(cell)
    if kpt_band is None:
        kpt1 = kpt2 = kpt
        aoR_k1 = aoR_k2 = numint.eval_ao(cell, coords, kpt)
    else:
        kpt1 = kpt_band
        kpt2 = kpt
        aoR_k1 = numint.eval_ao(cell, coords, kpt1)
        aoR_k2 = numint.eval_ao(cell, coords, kpt2)
    ngs, nao = aoR_k1.shape

    def contract(dm):
        vjR_k2 = get_vjR(cell, dm, aoR_k2)
        vj = (cell.vol/ngs) * np.dot(aoR_k1.T.conj(), vjR_k2.reshape(-1,1)*aoR_k1)
        return vj

    if dm.ndim == 2:
        vj = contract(dm)
    else:
        vj = lib.asarray([contract(x) for x in dm.reshape(-1,nao,nao)])
    return vj.reshape(dm.shape)

def get_fock(mf, h1e=None, s1e=None, vhf=None, dm=None, cycle=-1, diis=None,
             diis_start_cycle=None, level_shift_factor=None, damp_factor=None):
    h1e_kpts, s_kpts, vhf_kpts, dm_kpts = h1e, s1e, vhf, dm
    if h1e_kpts is None: h1e_kpts = mf.get_hcore()
    if vhf_kpts is None: vhf_kpts = mf.get_veff(mf.cell, dm_kpts)
    f_kpts = h1e_kpts + vhf_kpts
    if cycle < 0 and diis is None:  # Not inside the SCF iteration
        return f_kpts

    if diis_start_cycle is None:
        diis_start_cycle = mf.diis_start_cycle
    if level_shift_factor is None:
        level_shift_factor = mf.level_shift
    if damp_factor is None:
        damp_factor = mf.damp
    if s_kpts is None: s_kpts = mf.get_ovlp()
    if dm_kpts is None: dm_kpts = mf.make_rdm1()

    if isinstance(level_shift_factor, (tuple, list, np.ndarray)):
        shifta, shiftb = level_shift_factor
    else:
        shifta = shiftb = level_shift_factor

    if diis and cycle >= diis_start_cycle:
        f_kpts = diis.update(s_kpts, dm_kpts, f_kpts, mf, h1e_kpts, vhf_kpts)
    if abs(level_shift_factor) > 1e-4:
        f_kpts =([mol_hf.level_shift(s, dm_kpts[0,k], f_kpts[0,k], shifta)
                  for k, s in enumerate(s_kpts)],
                 [mol_hf.level_shift(s, dm_kpts[1,k], f_kpts[1,k], shiftb)
                  for k, s in enumerate(s_kpts)])
    return lib.asarray(f_kpts)

    def h_op(x1):
        xtmp = []
        p0 = 0
        for k in range(nkpts):
            nocc = len(occidx[k])
            nvir = len(viridx[k])
            xtmp.append(x1[p0:p0+nocc*nvir].reshape(nvir,nocc))
            p0 += nocc * nvir
        x1 = xtmp

        dm1 = []
        for k in range(nkpts):
            d1 = reduce(numpy.dot, (mo_coeff[k][:,viridx[k]], x1[k],
                                    mo_coeff[k][:,occidx[k]].T.conj()))
            dm1.append(d1+d1.T.conj())
        dm1 = lib.asarray(dm1)
        if hyb is None:
            v1 = mf.get_veff(cell, dm1)
        else:
            v1 = mf._numint.nr_rks_fxc(cell, mf.grids, mf.xc, dm0, dm1,
                                       0, 1, rho0, vxc, fxc)
            if abs(hyb) < 1e-10:
                v1 += mf.get_j(cell, dm1)
            else:
                vj, vk = mf.get_jk(cell, dm1)
                v1 += vj - vk * hyb * .5

        x2 = [0] * nkpts
        for k in range(nkpts):
            x2[k] = numpy.einsum('sp,sq->pq', fvv[k], x1[k].conj()) * 2
            x2[k]-= numpy.einsum('sp,rp->rs', foo[k], x1[k].conj()) * 2

            x2[k] += reduce(numpy.dot, (mo_coeff[k][:,occidx[k]].T.conj(), v1[k],
                                        mo_coeff[k][:,viridx[k]])).T * 4
        return numpy.hstack([x.ravel() for x in x2])

def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None):
    if kpts_band is not None:
        return get_j_for_bands(mydf, dm_kpts, hermi, kpts, kpts_band)

    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]
    vj_kpts = numpy.zeros((nset,nkpts,nao,nao), dtype=numpy.complex128)
    kpt_allow = numpy.zeros(3)
    mesh = mydf.mesh
    coulG = mydf.weighted_coulG(kpt_allow, False, mesh)
    max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8
    weight = 1./len(kpts)
    dmsC = dms.conj()
    for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts, max_memory=max_memory):
        vG = [0] * nset
        #:rho = numpy.einsum('lkL,lk->L', pqk.conj(), dm)
        for k, aoao in enumerate(aoaoks):
            for i in range(nset):
                rho = numpy.einsum('ij,Lij->L', dmsC[i,k],
                                   aoao.reshape(-1,nao,nao)).conj()
                vG[i] += rho * coulG[p0:p1]
        for i in range(nset):
            vG[i] *= weight
        for k, aoao in enumerate(aoaoks):
            for i in range(nset):
                vj_kpts[i,k] += numpy.einsum('L,Lij->ij', vG[i],
                                             aoao.reshape(-1,nao,nao))
    aoao = aoaoks = p0 = p1 = None

    if gamma_point(kpts):
        vj_kpts = vj_kpts.real.copy()
    return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts)

    def hop_uhf2ghf(x1):
        x1ab = []
        x1ba = []
        ip = 0
        for k in range(nkpts):
            nv = nvira[k]
            no = noccb[k]
            x1ab.append(x1[ip:ip+nv*no].reshape(nv,no))
            ip += nv * no
        for k in range(nkpts):
            nv = nvirb[k]
            no = nocca[k]
            x1ba.append(x1[ip:ip+nv*no].reshape(nv,no))
            ip += nv * no

        dm1ab = []
        dm1ba = []
        for k in range(nkpts):
            d1ab = reduce(numpy.dot, (orbva[k], x1ab[k], orbob[k].T.conj()))
            d1ba = reduce(numpy.dot, (orbvb[k], x1ba[k], orboa[k].T.conj()))
            dm1ab.append(d1ab+d1ba.T.conj())
            dm1ba.append(d1ba+d1ab.T.conj())

        v1ao = vresp1(lib.asarray([dm1ab,dm1ba]))
        x2ab = [0] * nkpts
        x2ba = [0] * nkpts
        for k in range(nkpts):
            x2ab[k] = numpy.einsum('pr,rq->pq', fvva[k], x1ab[k])
            x2ab[k]-= numpy.einsum('sq,ps->pq', foob[k], x1ab[k])
            x2ba[k] = numpy.einsum('pr,rq->pq', fvvb[k], x1ba[k])
            x2ba[k]-= numpy.einsum('qs,ps->pq', fooa[k], x1ba[k])
            x2ab[k] += reduce(numpy.dot, (orbva[k].T.conj(), v1ao[0][k], orbob[k]))
            x2ba[k] += reduce(numpy.dot, (orbvb[k].T.conj(), v1ao[1][k], orboa[k]))
        return numpy.hstack([x.real.ravel() for x in (x2ab+x2ba)])

    def get_hcore(self, cell=None, kpts=None):
        if cell is None: cell = self.cell
        if kpts is None:
            kpts_lst = numpy.zeros((1,3))
        else:
            kpts_lst = numpy.reshape(kpts, (-1,3))

        xcell, contr_coeff = self.get_xmol(cell)
        with_df = aft.AFTDF(xcell)
        c = lib.param.LIGHT_SPEED
        assert('1E' in self.approx.upper())
        if 'ATOM' in self.approx.upper():
            atom_slices = xcell.offset_nr_by_atom()
            nao = xcell.nao_nr()
            x = numpy.zeros((nao,nao))
            vloc = numpy.zeros((nao,nao))
            wloc = numpy.zeros((nao,nao))
            for ia in range(xcell.natm):
                ish0, ish1, p0, p1 = atom_slices[ia]
                shls_slice = (ish0, ish1, ish0, ish1)
                t1 = xcell.intor('int1e_kin', shls_slice=shls_slice)
                s1 = xcell.intor('int1e_ovlp', shls_slice=shls_slice)
                with xcell.with_rinv_as_nucleus(ia):
                    z = -xcell.atom_charge(ia)
                    v1 = z * xcell.intor('int1e_rinv', shls_slice=shls_slice)
                    w1 = z * xcell.intor('int1e_prinvp', shls_slice=shls_slice)
                vloc[p0:p1,p0:p1] = v1
                wloc[p0:p1,p0:p1] = w1
                x[p0:p1,p0:p1] = x2c._x2c1e_xmatrix(t1, v1, w1, s1, c)
        else:
            raise NotImplementedError

        t = xcell.pbc_intor('int1e_kin', 1, lib.HERMITIAN, kpts_lst)
        s = xcell.pbc_intor('int1e_ovlp', 1, lib.HERMITIAN, kpts_lst)
        v = with_df.get_nuc(kpts_lst)
        #w = get_pnucp(with_df, kpts_lst)
        if self.basis is not None:
            s22 = s
            s21 = pbcgto.intor_cross('int1e_ovlp', xcell, cell, kpts=kpts_lst)

        h1_kpts = []
        for k in range(len(kpts_lst)):
# The treatment of pnucp local part has huge effects to hcore
            #h1 = x2c._get_hcore_fw(t[k], vloc, wloc, s[k], x, c) - vloc + v[k]
            #h1 = x2c._get_hcore_fw(t[k], v[k], w[k], s[k], x, c)
            h1 = x2c._get_hcore_fw(t[k], v[k], wloc, s[k], x, c)
            if self.basis is not None:
                c = lib.cho_solve(s22[k], s21[k])
                h1 = reduce(numpy.dot, (c.T, h1, c))
            if self.xuncontract and contr_coeff is not None:
                h1 = reduce(numpy.dot, (contr_coeff.T, h1, contr_coeff))
            h1_kpts.append(h1)

        if kpts is None or numpy.shape(kpts) == (3,):
            h1_kpts = h1_kpts[0]
        return lib.asarray(h1_kpts)