Ejemplo n.º 1
0
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)
Ejemplo n.º 2
0
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)
Ejemplo n.º 3
0
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)
Ejemplo n.º 4
0
 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)
Ejemplo n.º 5
0
    def eig(self, h_kpts, s_kpts):
        nkpts = len(h_kpts)
        eig_kpts = []
        mo_coeff_kpts = []

        for k in range(nkpts):
            e, c = hf.RHF.eig(self, h_kpts[k], s_kpts[k])
            eig_kpts.append(e)
            mo_coeff_kpts.append(c)
        return lib.asarray(eig_kpts), lib.asarray(mo_coeff_kpts)
Ejemplo n.º 6
0
Archivo: khf.py Proyecto: eronca/pyscf
 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
Ejemplo n.º 7
0
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1, 3)), kpt_band=None):
    """Get the Coulomb (J) AO matrix at sampled k-points.

    Args:
        dm_kpts : (nkpts, nao, nao) ndarray or a list of (nkpts,nao,nao) ndarray
            Density matrix at each k-point.  If a list of k-point DMs, eg,
            UHF alpha and beta DM, the alpha and beta DMs are contracted
            separately.
        kpts : (nkpts, 3) ndarray

    Kwargs:
        kpt_band : (3,) ndarray
            An arbitrary "band" k-point at which to evalute the matrix.

    Returns:
        vj : (nkpts, nao, nao) ndarray
        or list of vj if the input dm_kpts is a list of DMs
    """
    cell = mydf.cell
    gs = mydf.gs

    dm_kpts = lib.asarray(dm_kpts, order="C")
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]

    coulG = tools.get_coulG(cell, gs=gs)
    ngs = len(coulG)

    vR = rhoR = np.zeros((nset, ngs))
    for k, aoR in mydf.aoR_loop(gs, kpts):
        for i in range(nset):
            rhoR[i] += numint.eval_rho(cell, aoR, dms[i, k])
    for i in range(nset):
        rhoR[i] *= 1.0 / nkpts
        rhoG = tools.fft(rhoR[i], gs)
        vG = coulG * rhoG
        vR[i] = tools.ifft(vG, gs).real

    if kpt_band is not None:
        for k, aoR_kband in mydf.aoR_loop(gs, kpts, kpt_band):
            pass
        vj_kpts = [cell.vol / ngs * lib.dot(aoR_kband.T.conj() * vR[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 = []
        weight = cell.vol / ngs
        for k, aoR in mydf.aoR_loop(gs, kpts):
            for i in range(nset):
                vj_kpts.append(weight * lib.dot(aoR.T.conj() * vR[i], aoR))
        vj_kpts = lib.asarray(vj_kpts).reshape(nkpts, nset, nao, nao)
        return vj_kpts.transpose(1, 0, 2, 3).reshape(dm_kpts.shape)
Ejemplo n.º 8
0
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)
Ejemplo n.º 9
0
 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()
Ejemplo n.º 10
0
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)
Ejemplo n.º 11
0
    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])
Ejemplo n.º 12
0
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)
Ejemplo n.º 13
0
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)
Ejemplo n.º 14
0
    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)])
Ejemplo n.º 15
0
 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)
Ejemplo n.º 16
0
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)
Ejemplo n.º 17
0
 def get_init_guess(self, cell=None, key='minao'):
     if cell is None: cell = self.cell
     dm = hf.RHF.get_init_guess(self, cell, key)
     if key.lower() == 'chkfile':
         dm_kpts = dm
     else:
         dm_kpts = lib.asarray([dm]*len(self.kpts))
     return dm_kpts
Ejemplo n.º 18
0
    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)
Ejemplo n.º 19
0
def make_rdm1(mo_coeff_kpts, mo_occ_kpts):
    '''One particle density matrices for all k-points.

    Returns:
        dm_kpts : (nkpts, nao, nao) ndarray
    '''
    nkpts = len(mo_occ_kpts)
    dm_kpts = [hf.make_rdm1(mo_coeff_kpts[k], mo_occ_kpts[k])
               for k in range(nkpts)]
    return lib.asarray(dm_kpts)
Ejemplo n.º 20
0
def get_veff(ks, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
             kpts=None, kpt_band=None):
    '''Coulomb + XC functional for UKS.  See pyscf/pbc/dft/uks.py
    :func:`get_veff` fore more details.
    '''
    if cell is None: cell = ks.cell
    if dm is None: dm = ks.make_rdm1()
    if kpts is None: kpts = ks.kpts
    t0 = (time.clock(), time.time())
    if ks.grids.coords is None:
        ks.grids.build()
        small_rho_cutoff = ks.small_rho_cutoff
        t0 = logger.timer(ks, 'setting up grids', *t0)
    else:
        small_rho_cutoff = 0

    dm = np.asarray(dm)
    nao = dm.shape[-1]
    # ndim = 4 : dm.shape = (alpha_beta, nkpts, nao, nao)
    ground_state = (dm.ndim == 4 and kpt_band is None)
    nkpts = len(kpts)

    if hermi == 2:  # because rho = 0
        n, ks._exc, vx = 0, 0, 0
    else:
        n, ks._exc, vx = ks._numint.nr_uks(cell, ks.grids, ks.xc, dm, 1,
                                           kpts, kpt_band)
        logger.debug(ks, 'nelec by numeric integration = %s', n)
        t0 = logger.timer(ks, 'vxc', *t0)

    hyb = ks._numint.hybrid_coeff(ks.xc, spin=(cell.spin>0)+1)
    if abs(hyb) < 1e-10:
        vj = ks.get_j(cell, dm, hermi, kpts, kpt_band)
        vhf = lib.asarray([vj[0]+vj[1]] * 2)
    else:
        vj, vk = ks.get_jk(cell, dm, hermi, kpts, kpt_band)
        vhf = pbcuhf._makevhf(vj, vk*hyb)

        if ground_state:
            ks._exc -= (np.einsum('Kij,Kji', dm[0], vk[0]) +
                        np.einsum('Kij,Kji', dm[1], vk[1])).real * .5 * hyb * (1./nkpts)

    if ground_state:
        ks._ecoul = np.einsum('Kij,Kji', dm[0]+dm[1], vj[0]+vj[1]).real * .5 * (1./nkpts)

    if small_rho_cutoff > 1e-20 and ground_state:
        # Filter grids the first time setup grids
        idx = ks._numint.large_rho_indices(cell, dm, ks.grids,
                                           small_rho_cutoff, kpts)
        logger.debug(ks, 'Drop grids %d',
                     ks.grids.weights.size - np.count_nonzero(idx))
        ks.grids.coords  = np.asarray(ks.grids.coords [idx], order='C')
        ks.grids.weights = np.asarray(ks.grids.weights[idx], order='C')
        ks._numint.non0tab = None
    return vhf + vx
Ejemplo n.º 21
0
Archivo: hf.py Proyecto: berquist/pyscf
def dot_eri_dm(eri, dm, hermi=0):
    """Compute J, K matrices in terms of the given 2-electron integrals and
    density matrix. eri or dm can be complex.

    Args:
        eri : ndarray
            complex integral array with N^4 elements (N is the number of orbitals)
        dm : ndarray or list of ndarrays
            A density matrix or a list of density matrices

    Kwargs:
        hermi : int
            Whether J, K matrix is hermitian

            | 0 : no hermitian or symmetric
            | 1 : hermitian
            | 2 : anti-hermitian

    Returns:
        Depending on the given dm, the function returns one J and one K matrix,
        or a list of J matrices and a list of K matrices, corresponding to the
        input density matrices.
    """
    dm = np.asarray(dm)
    if np.iscomplexobj(dm) or np.iscomplexobj(eri):
        nao = dm.shape[-1]
        eri = eri.reshape((nao,) * 4)

        def contract(dm):
            vj = np.einsum("ijkl,ji->kl", eri, dm)
            vk = np.einsum("ijkl,jk->il", eri, dm)
            return vj, vk

        if isinstance(dm, np.ndarray) and dm.ndim == 2:
            vj, vk = contract(dm)
        else:
            vjk = [contract(dmi) for dmi in dm]
            vj = lib.asarray([v[0] for v in vjk]).reshape(dm.shape)
            vk = lib.asarray([v[1] for v in vjk]).reshape(dm.shape)
    else:
        vj, vk = pyscf.scf.hf.dot_eri_dm(eri, dm, hermi)
    return vj, vk
Ejemplo n.º 22
0
def make_h1_fcsd(mol, mo_coeff, mo_occ, atmlst):
    '''FC + SD'''
    orboa = mo_coeff[0][:,mo_occ[0]> 0]
    orbva = mo_coeff[0][:,mo_occ[0]==0]
    orbob = mo_coeff[1][:,mo_occ[1]> 0]
    orbvb = mo_coeff[1][:,mo_occ[1]==0]
    nao = mo_coeff[0].shape[0]
    h1aa = []
    h1ab = []
    h1ba = []
    h1bb = []
    for ia in atmlst:
        h1ao = rhf_ssc._get_integrals_fcsd(mol, ia)
        # *.5 due to s = 1/2 * pauli-matrix
        h1aa.append(lib.einsum('xypq,pi,qj->xyij', h1ao, orbva.conj(), orboa) * .5)
        h1ab.append(lib.einsum('xypq,pi,qj->xyij', h1ao, orbva.conj(), orbob) * .5)
        h1ba.append(lib.einsum('xypq,pi,qj->xyij', h1ao, orbvb.conj(), orboa) * .5)
        h1bb.append(lib.einsum('xypq,pi,qj->xyij', h1ao, orbvb.conj(), orbob) *-.5)
    return (lib.asarray(h1aa), lib.asarray(h1ab),
            lib.asarray(h1ba), lib.asarray(h1bb))
Ejemplo n.º 23
0
Archivo: kuhf.py Proyecto: eronca/pyscf
 def get_init_guess(self, cell=None, key='minao'):
     if cell is None: cell = self.cell
     dm = uhf.UHF.get_init_guess(self, cell, key)
     if key.lower() == 'chkfile':
         dm_kpts = dm
     else:
         nao = dm.shape[1]
         nkpts = len(self.kpts)
         dm_kpts = lib.asarray([dm]*nkpts).reshape(nkpts,2,nao,nao)
         dm_kpts = dm_kpts.transpose(1,0,2,3)
         dm[1,:] *= .98  # To break spin symmetry
     return dm_kpts
Ejemplo n.º 24
0
    def h_op(x1):
        p0 = 0
        x1a = []
        for k in range(nkpts):
            nocc = len(occidxa[k])
            nvir = len(viridxa[k])
            x1a.append(x1[p0:p0+nocc*nvir].reshape(nvir,nocc))
            p0 += nocc * nvir
        x1b = []
        for k in range(nkpts):
            nocc = len(occidxb[k])
            nvir = len(viridxb[k])
            x1b.append(x1[p0:p0+nocc*nvir].reshape(nvir,nocc))
            p0 += nocc * nvir

        dm1a = []
        for k in range(nkpts):
            d1 = reduce(numpy.dot, (moa[k][:,viridxa[k]], x1a[k],
                                    moa[k][:,occidxa[k]].T.conj()))
            dm1a.append(d1+d1.T.conj())
        dm1b = []
        for k in range(nkpts):
            d1 = reduce(numpy.dot, (mob[k][:,viridxa[k]], x1b[k],
                                    mob[k][:,occidxa[k]].T.conj()))
            dm1b.append(d1+d1.T.conj())
        dm1 = lib.asarray([dm1a,dm1b])
        dm1a = dm1b = None
        if hyb is None:
            v1 = mf.get_veff(cell, dm1)
        else:
            v1 = mf._numint.nr_uks_fxc(cell, mf.grids, mf.xc, dm0, dm1,
                                       0, 1, rho0, vxc, fxc)
            if abs(hyb) < 1e-10:
                vj = mf.get_j(cell, dm1)
                v1 += vj[0] + vj[1]
            else:
                vj, vk = mf.get_jk(cell, dm1)
                v1 += vj[0]+vj[1] - vk * hyb * .5

        x2a = [0] * nkpts
        x2b = [0] * nkpts
        for k in range(nkpts):
            x2a[k] = numpy.einsum('sp,sq->pq', focka[k][viridxa[k][:,None],viridxa[k]], x1a[k].conj())
            x2a[k]-= numpy.einsum('sp,rp->rs', focka[k][occidxa[k][:,None],occidxa[k]], x1a[k].conj())
            x2b[k] = numpy.einsum('sp,sq->pq', fockb[k][viridxb[k][:,None],viridxb[k]], x1b[k].conj())
            x2b[k]-= numpy.einsum('sp,rp->rs', fockb[k][occidxb[k][:,None],occidxb[k]], x1b[k].conj())

            x2a[k] += reduce(numpy.dot, (moa[k][:,occidxa[k]].T.conj(), v1[0][k],
                                         moa[k][:,viridxa[k]])).T
            x2b[k] += reduce(numpy.dot, (mob[k][:,occidxb[k]].T.conj(), v1[1][k],
                                         mob[k][:,viridxb[k]])).T

        return numpy.hstack([x.ravel() for x in (x2a+x2b)])
Ejemplo n.º 25
0
def get_hcore(mf, cell=None, kpts=None):
    '''Get the core Hamiltonian AO matrices at sampled k-points.

    Args:
        kpts : (nkpts, 3) ndarray

    Returns:
        hcore : (nkpts, nao, nao) ndarray
    '''
    if cell is None: cell = mf.cell
    if kpts is None: kpts = mf.kpts
    return lib.asarray([pbchf.get_hcore(cell, k) for k in kpts])
Ejemplo n.º 26
0
def get_ovlp(mf, cell=None, kpts=None):
    '''Get the overlap AO matrices at sampled k-points.

    Args:
        kpts : (nkpts, 3) ndarray

    Returns:
        ovlp_kpts : (nkpts, nao, nao) ndarray
    '''
    if cell is None: cell = mf.cell
    if kpts is None: kpts = mf.kpts
    return lib.asarray(cell.pbc_intor('cint1e_ovlp_sph', hermi=1, kpts=kpts))
Ejemplo n.º 27
0
def make_rdm1(mo_coeff_kpts, mo_occ_kpts, **kwargs):
    '''Alpha and beta spin one particle density matrices for all k-points.

    Returns:
        dm_kpts : (2, nkpts, nao, nao) ndarray
    '''
    nkpts = len(mo_occ_kpts[0])
    nao, nmo = mo_coeff_kpts[0][0].shape
    def make_dm(mos, occs):
        return [np.dot(mos[k]*occs[k], mos[k].T.conj()) for k in range(nkpts)]
    dm_kpts =(make_dm(mo_coeff_kpts[0], mo_occ_kpts[0]) +
              make_dm(mo_coeff_kpts[1], mo_occ_kpts[1]))
    return lib.asarray(dm_kpts).reshape(2,nkpts,nao,nao)
Ejemplo n.º 28
0
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

    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[k], f_kpts[k], level_shift_factor)
                  for k, s in enumerate(s_kpts)]
    return lib.asarray(f_kpts)
Ejemplo n.º 29
0
def get_j_for_bands(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None):
    log = logger.Logger(mydf.stdout, mydf.verbose)
    t1 = (time.clock(), time.time())

    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]

    dmsR = dms.real.reshape(nset,nkpts,nao**2)
    dmsI = dms.imag.reshape(nset,nkpts,nao**2)
    kpt_allow = numpy.zeros(3)
    mesh = mydf.mesh
    coulG = mydf.weighted_coulG(kpt_allow, False, mesh)
    ngrids = len(coulG)
    vG = numpy.zeros((nset,ngrids), dtype=numpy.complex128)
    max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8

    dmsC = dms.conj()
    for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts, max_memory=max_memory):
        #: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,p0:p1] += rho * coulG[p0:p1]
    aoao = aoaoks = p0 = p1 = None
    weight = 1./len(kpts)
    vG *= weight
    t1 = log.timer_debug1('get_j pass 1 to compute J(G)', *t1)

    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
    nband = len(kpts_band)
    vj_kpts = numpy.zeros((nset,nband,nao,nao), dtype=numpy.complex128)
    for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts_band,
                                        max_memory=max_memory):
        for k, aoao in enumerate(aoaoks):
            for i in range(nset):
                vj_kpts[i,k] += numpy.einsum('L,Lij->ij', vG[i,p0:p1],
                                             aoao.reshape(-1,nao,nao))
    aoao = aoaoks = p0 = p1 = None

    if gamma_point(kpts_band):
        vj_kpts = vj_kpts.real.copy()
    t1 = log.timer_debug1('get_j pass 2', *t1)
    return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
Ejemplo n.º 30
0
def enlarge_space(myci, civec_strs, eri, norb, nelec):
    if isinstance(civec_strs, (tuple, list)):
        nelec, (strsa, strsb) = _unpack(civec_strs[0], nelec)[1:]
        ci_coeff = lib.asarray(civec_strs)
    else:
        ci_coeff, nelec, (strsa, strsb) = _unpack(civec_strs, nelec)
    na = len(strsa)
    nb = len(strsb)
    ci0 = ci_coeff.reshape(-1,na,nb)
    civec_a_max = lib.norm(ci0, axis=2).max(axis=0)
    civec_b_max = lib.norm(ci0, axis=1).max(axis=0)
    ci_aidx = numpy.where(civec_a_max > myci.ci_coeff_cutoff)[0]
    ci_bidx = numpy.where(civec_b_max > myci.ci_coeff_cutoff)[0]
    civec_a_max = civec_a_max[ci_aidx]
    civec_b_max = civec_b_max[ci_bidx]
    strsa = strsa[ci_aidx]
    strsb = strsb[ci_bidx]

    eri = ao2mo.restore(1, eri, norb)
    eri_pq_max = abs(eri.reshape(norb**2,-1)).max(axis=1).reshape(norb,norb)

    strsa_add = select_strs(myci, eri, eri_pq_max, civec_a_max, strsa, norb, nelec[0])
    strsb_add = select_strs(myci, eri, eri_pq_max, civec_b_max, strsb, norb, nelec[1])
    strsa = numpy.append(strsa, strsa_add)
    strsb = numpy.append(strsb, strsb_add)
    aidx = numpy.argsort(strsa)
    bidx = numpy.argsort(strsb)
    ci_strs = (strsa[aidx], strsb[bidx])
    aidx = numpy.where(aidx < len(ci_aidx))[0]
    bidx = numpy.where(bidx < len(ci_bidx))[0]
    ma = len(strsa)
    mb = len(strsb)

    cs = []
    for i in range(ci0.shape[0]):
        ci1 = numpy.zeros((ma,mb))
        tmp = lib.take_2d(ci0[i], ci_aidx, ci_bidx)
        lib.takebak_2d(ci1, tmp, aidx, bidx)
        cs.append(_as_SCIvector(ci1, ci_strs))

    if not isinstance(civec_strs, (tuple, list)) and civec_strs.ndim < 3:
        cs = cs[0]
    return cs
Ejemplo n.º 31
0
Archivo: uhf.py Proyecto: pulkin/pyscf
 def eig(self, fock, s):
     e_a, c_a = self._eigh(fock[0], s)
     e_b, c_b = self._eigh(fock[1], s)
     return lib.asarray((e_a, e_b)), lib.asarray((c_a, c_b))
Ejemplo n.º 32
0
 def get_hcore(self, cell=None, kpts=None):
     hcore = khf.KSCF.get_hcore(self, cell, kpts)
     return lib.asarray([scipy.linalg.block_diag(h, h) for h in hcore])
Ejemplo n.º 33
0
 def eig(self, h_kpts, s_kpts):
     e_a, c_a = khf.KRHF.eig(self, h_kpts[0], s_kpts)
     e_b, c_b = khf.KRHF.eig(self, h_kpts[1], s_kpts)
     return lib.asarray((e_a, e_b)), lib.asarray((c_a, c_b))
Ejemplo n.º 34
0
Archivo: uks.py Proyecto: pulkin/pyscf
def get_veff(ks,
             cell=None,
             dm=None,
             dm_last=0,
             vhf_last=0,
             hermi=1,
             kpt=None,
             kpt_band=None):
    '''Coulomb + XC functional for UKS.  See pyscf/pbc/dft/uks.py
    :func:`get_veff` fore more details.
    '''
    if cell is None: cell = ks.cell
    if dm is None: dm = ks.make_rdm1()
    if kpt is None: kpt = ks.kpt
    t0 = (time.clock(), time.time())
    if ks.grids.coords is None:
        ks.grids.build(with_non0tab=True)
        small_rho_cutoff = ks.small_rho_cutoff
        t0 = logger.timer(ks, 'setting up grids', *t0)
    else:
        small_rho_cutoff = 0

    if not isinstance(dm, numpy.ndarray):
        dm = numpy.asarray(dm)
    if dm.ndim == 2:  # RHF DM
        dm = numpy.asarray((dm * .5, dm * .5))

    if hermi == 2:  # because rho = 0
        n, ks._exc, vx = (0, 0), 0, 0
    else:
        n, ks._exc, vx = ks._numint.nr_uks(cell, ks.grids, ks.xc, dm, 0, kpt,
                                           kpt_band)
        logger.debug(ks, 'nelec by numeric integration = %s', n)
        t0 = logger.timer(ks, 'vxc', *t0)

    # ndim = 3 : dm.shape = ([alpha,beta], nao, nao)
    ground_state = (dm.ndim == 3 and dm.shape[0] == 2)

    hyb = ks._numint.hybrid_coeff(ks.xc, spin=cell.spin)
    if abs(hyb) < 1e-10:
        vj = ks.get_j(cell, dm, hermi, kpt, kpt_band)
        vhf = lib.asarray([vj[0] + vj[1]] * 2)
    else:
        vj, vk = ks.get_jk(cell, dm, hermi, kpt, kpt_band)
        vhf = pbcuhf._makevhf(vj, vk * hyb)

        if ground_state:
            ks._exc -= (numpy.einsum('ij,ji', dm[0], vk[0]) +
                        numpy.einsum('ij,ji', dm[1], vk[1])).real * .5 * hyb

    if ground_state:
        ks._ecoul = numpy.einsum('ij,ji', dm[0] + dm[1],
                                 vj[0] + vj[1]).real * .5

    nelec = cell.nelec
    if (small_rho_cutoff > 1e-20 and ground_state
            and abs(n[0] - nelec[0]) < 0.01 * n[0]
            and abs(n[1] - nelec[1]) < 0.01 * n[1]):
        # Filter grids the first time setup grids
        idx = ks._numint.large_rho_indices(cell, dm, ks.grids,
                                           small_rho_cutoff, kpt)
        logger.debug(ks, 'Drop grids %d',
                     ks.grids.weights.size - numpy.count_nonzero(idx))
        ks.grids.coords = numpy.asarray(ks.grids.coords[idx], order='C')
        ks.grids.weights = numpy.asarray(ks.grids.weights[idx], order='C')
        ks.grids.non0tab = ks.grids.make_mask(cell, ks.grids.coords)
    return vhf + vx
Ejemplo n.º 35
0
def get_k_kpts(mydf,
               dm_kpts,
               hermi=1,
               kpts=numpy.zeros((1, 3)),
               kpts_band=None,
               exxdiv=None):
    cell = mydf.cell
    log = logger.Logger(mydf.stdout, mydf.verbose)
    t1 = (time.clock(), time.time())

    mesh = mydf.mesh
    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]

    swap_2e = (kpts_band is None)
    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
    nband = len(kpts_band)
    kk_table = kpts_band.reshape(-1, 1, 3) - kpts.reshape(1, -1, 3)
    kk_todo = numpy.ones(kk_table.shape[:2], dtype=bool)
    vkR = numpy.zeros((nset, nband, nao, nao))
    vkI = numpy.zeros((nset, nband, nao, nao))
    dmsR = numpy.asarray(dms.real, order='C')
    dmsI = numpy.asarray(dms.imag, order='C')

    mem_now = lib.current_memory()[0]
    max_memory = max(2000, (mydf.max_memory - mem_now)) * .8
    log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now)

    # K_pq = ( p{k1} i{k2} | i{k2} q{k1} )
    def make_kpt(kpt):  # kpt = kptj - kpti
        # search for all possible ki and kj that has ki-kj+kpt=0
        kk_match = numpy.einsum('ijx->ij', abs(kk_table + kpt)) < 1e-9
        kpti_idx, kptj_idx = numpy.where(kk_todo & kk_match)
        nkptj = len(kptj_idx)
        log.debug1('kpt = %s', kpt)
        log.debug2('kpti_idx = %s', kpti_idx)
        log.debug2('kptj_idx = %s', kptj_idx)
        kk_todo[kpti_idx, kptj_idx] = False
        if swap_2e and not is_zero(kpt):
            kk_todo[kptj_idx, kpti_idx] = False

        max_memory1 = max_memory * (nkptj + 1) / (nkptj + 5)
        #blksize = max(int(max_memory1*4e6/(nkptj+5)/16/nao**2), 16)

        #bufR = numpy.empty((blksize*nao**2))
        #bufI = numpy.empty((blksize*nao**2))
        # Use DF object to mimic KRHF/KUHF object in function get_coulG
        mydf.exxdiv = exxdiv
        vkcoulG = mydf.weighted_coulG(kpt, True, mesh)
        kptjs = kpts[kptj_idx]
        # <r|-G+k_rs|s> = conj(<s|G-k_rs|r>) = conj(<s|G+k_sr|r>)
        #buf1R = numpy.empty((blksize*nao**2))
        #buf1I = numpy.empty((blksize*nao**2))
        for aoaoks, p0, p1 in mydf.ft_loop(mesh,
                                           kpt,
                                           kptjs,
                                           max_memory=max_memory1):
            coulG = numpy.sqrt(vkcoulG[p0:p1])
            nG = p1 - p0
            bufR = numpy.empty((nG * nao**2))
            bufI = numpy.empty((nG * nao**2))
            buf1R = numpy.empty((nG * nao**2))
            buf1I = numpy.empty((nG * nao**2))

            for k, aoao in enumerate(aoaoks):
                ki = kpti_idx[k]
                kj = kptj_idx[k]

                # case 1: k_pq = (pi|iq)
                #:v4 = numpy.einsum('ijL,lkL->ijkl', pqk, pqk.conj())
                #:vk += numpy.einsum('ijkl,jk->il', v4, dm)
                pLqR = numpy.ndarray((nao, nG, nao), buffer=bufR)
                pLqI = numpy.ndarray((nao, nG, nao), buffer=bufI)
                pLqR[:] = aoao.real.reshape(nG, nao, nao).transpose(1, 0, 2)
                pLqI[:] = aoao.imag.reshape(nG, nao, nao).transpose(1, 0, 2)
                pLqR *= coulG.reshape(1, nG, 1)
                pLqI *= coulG.reshape(1, nG, 1)
                iLkR = numpy.ndarray((nao * nG, nao), buffer=buf1R)
                iLkI = numpy.ndarray((nao * nG, nao), buffer=buf1I)
                for i in range(nset):
                    iLkR, iLkI = zdotNN(pLqR.reshape(-1, nao),
                                        pLqI.reshape(-1, nao), dmsR[i, kj],
                                        dmsI[i, kj], 1, iLkR, iLkI)
                    zdotNC(iLkR.reshape(nao, -1), iLkI.reshape(nao, -1),
                           pLqR.reshape(nao, -1).T,
                           pLqI.reshape(nao, -1).T, 1, vkR[i, ki], vkI[i, ki],
                           1)

# case 2: k_pq = (iq|pi)
#:v4 = numpy.einsum('iLj,lLk->ijkl', pqk, pqk.conj())
#:vk += numpy.einsum('ijkl,li->kj', v4, dm)
                if swap_2e and not is_zero(kpt):
                    iLkR = iLkR.reshape(nao, -1)
                    iLkI = iLkI.reshape(nao, -1)
                    for i in range(nset):
                        iLkR, iLkI = zdotNN(dmsR[i, ki], dmsI[i, ki],
                                            pLqR.reshape(nao, -1),
                                            pLqI.reshape(nao, -1), 1, iLkR,
                                            iLkI)
                        zdotCN(
                            pLqR.reshape(-1, nao).T,
                            pLqI.reshape(-1, nao).T, iLkR.reshape(-1, nao),
                            iLkI.reshape(-1, nao), 1, vkR[i, kj], vkI[i, kj],
                            1)

    for ki, kpti in enumerate(kpts_band):
        for kj, kptj in enumerate(kpts):
            if kk_todo[ki, kj]:
                make_kpt(kptj - kpti)
        t1 = log.timer_debug1('get_k_kpts: make_kpt (%d,*)' % ki, *t1)

    if (gamma_point(kpts) and gamma_point(kpts_band)
            and not numpy.iscomplexobj(dm_kpts)):
        vk_kpts = vkR
    else:
        vk_kpts = vkR + vkI * 1j
    vk_kpts *= 1. / nkpts

    # G=0 was not included in the non-uniform grids
    if cell.dimension != 3 and exxdiv:
        assert (exxdiv.lower() == 'ewald')
        _ewald_exxdiv_for_G0(cell, kpts_band, dms, vk_kpts, kpts_band)

    return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
Ejemplo n.º 36
0
def uks_j_xc(mydf,
             dm_kpts,
             xc_code,
             hermi=1,
             kpts=numpy.zeros((1, 3)),
             kpts_band=None,
             with_j=WITH_J,
             j_in_xc=J_IN_XC):
    log = lib.logger.Logger(mydf.stdout, mydf.verbose)
    cell = mydf.cell
    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]
    dms = None
    #TODO: Handle multiple sets of KUKS density matrices (2,nset,nkpts,nao,nao)
    assert (nset == 2)  # alpha and beta density matrices in KUKS

    ni = mydf._numint
    xctype = ni._xc_type(xc_code)

    if xctype == 'LDA':
        deriv = 0

        rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv=0)

        def add_j_(v, ao_l, ao_h, idx_l, idx_h, vR):
            for k in range(nkpts):
                aow = numpy.einsum('pi,p->pi', ao_l[k], vR[0])
                v[0, k, idx_l[:, None],
                  idx_h] += lib.dot(aow.conj().T, ao_h[k])
                aow = numpy.einsum('pi,p->pi', ao_l[k], vR[1])
                v[1, k, idx_l[:, None],
                  idx_h] += lib.dot(aow.conj().T, ao_h[k])

        def add_xc_(v, ao_l, ao_h, idx_l, idx_h, wv):
            add_j_(v, ao_l, ao_h, idx_l, idx_h, wv[:, 0])

    elif xctype == 'GGA':
        deriv = 1

        if RHOG_HIGH_DERIV:
            rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv)
        else:
            Gv = cell.Gv
            ngrids = Gv.shape[0]
            rhoG = numpy.empty((2, 4, ngrids), dtype=numpy.complex128)
            rhoG[:, :1] = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv=0)
            rhoG[:, 1:] = numpy.einsum('np,px->nxp', 1j * rhoG[:, 0], Gv)

        def add_j_(v, ao_l, ao_h, idx_l, idx_h, vR):
            for k in range(nkpts):
                aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[0])
                v[0, k, idx_l[:, None],
                  idx_h] += lib.dot(aow.conj().T, ao_h[k][0])
                aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[1])
                v[1, k, idx_l[:, None],
                  idx_h] += lib.dot(aow.conj().T, ao_h[k][0])

        def add_xc_(v, ao_l, ao_h, idx_l, idx_h, wv):
            wva, wvb = wv
            for k in range(nkpts):
                aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wva)
                v1 = lib.dot(aow.conj().T, ao_h[k][0])
                aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wva[1:4])
                v1 += lib.dot(ao_l[k][0].conj().T, aow)
                v[0, k, idx_l[:, None], idx_h] += v1
                aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wvb)
                v1 = lib.dot(aow.conj().T, ao_h[k][0])
                aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wvb[1:4])
                v1 += lib.dot(ao_l[k][0].conj().T, aow)
                v[1, k, idx_l[:, None], idx_h] += v1

    else:  # MGGA
        deriv = 2
        #TODO: RHOG_HIGH_DERIV:
        rhoG = _eval_rhoG(mydf, dm_kpts, hermi, kpts, deriv)

        def add_j_(v, ao_l, ao_h, idx_l, idx_h, vR):
            for k in range(nkpts):
                aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[0])
                v[0, k, idx_l[:, None],
                  idx_h] += lib.dot(aow.conj().T, ao_h[k][0])
                aow = numpy.einsum('pi,p->pi', ao_l[k][0], vR[1])
                v[1, k, idx_l[:, None],
                  idx_h] += lib.dot(aow.conj().T, ao_h[k][0])

        def add_xc_(v, ao_l, ao_h, idx_l, idx_h, wv):
            wva, wvb = wv
            for k in range(nkpts):
                aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wva[:4])
                v1 = lib.dot(aow.conj().T, ao_h[k][0])
                aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wva[1:4])
                v1 += lib.dot(ao_l[k][0].conj().T, aow)
                aow = numpy.einsum('pi,p->pi', ao_h[k][1], wva[4], out=aow)
                v1 += lib.dot(ao_l[k][1].conj().T, aow)
                aow = numpy.einsum('pi,p->pi', ao_h[k][2], wva[4], out=aow)
                v1 += lib.dot(ao_l[k][2].conj().T, aow)
                aow = numpy.einsum('pi,p->pi', ao_h[k][3], wva[4], out=aow)
                v1 += lib.dot(ao_l[k][3].conj().T, aow)
                v[0, k, idx_l[:, None], idx_h] += v1

                aow = numpy.einsum('npi,np->pi', ao_l[k][:4], wvb[:4])
                v1 = lib.dot(aow.conj().T, ao_h[k][0])
                aow = numpy.einsum('npi,np->pi', ao_h[k][1:4], wvb[1:4])
                v1 += lib.dot(ao_l[k][0].conj().T, aow)
                aow = numpy.einsum('pi,p->pi', ao_h[k][1], wvb[4], out=aow)
                v1 += lib.dot(ao_l[k][1].conj().T, aow)
                aow = numpy.einsum('pi,p->pi', ao_h[k][2], wvb[4], out=aow)
                v1 += lib.dot(ao_l[k][2].conj().T, aow)
                aow = numpy.einsum('pi,p->pi', ao_h[k][3], wvb[4], out=aow)
                v1 += lib.dot(ao_l[k][3].conj().T, aow)
                v[1, k, idx_l[:, None], idx_h] += v1

    mesh = cell.mesh
    coulG = tools.get_coulG(cell,
                            mesh=mesh,
                            low_dim_ft_type=mydf.low_dim_ft_type)
    ngrids = coulG.size
    vG = numpy.einsum('ng,g->ng', rhoG[:, 0].reshape(-1, ngrids), coulG)
    vG = vG.reshape(2, *mesh)

    weight = cell.vol / ngrids
    # *(1./weight) because rhoR is scaled by weight in _eval_rhoG.  When
    # computing rhoR with IFFT, the weight factor is not needed.
    rhoR = tools.ifft(rhoG.reshape(-1, ngrids), mesh) * (1. / weight)
    rhoR = rhoR.real.reshape(2, -1, ngrids)
    nelec = numpy.zeros(2)
    excsum = 0

    exc, vxc = ni.eval_xc(xc_code, rhoR, 1, deriv=1)[:2]
    if xctype == 'LDA':
        vrho = vxc[0]
        wva = vrho[:, 0].reshape(1, ngrids)
        wvb = vrho[:, 1].reshape(1, ngrids)
    elif xctype == 'GGA':
        vrho, vsigma = vxc[:2]
        wva = numpy.empty((4, ngrids))
        wvb = numpy.empty((4, ngrids))
        wva[0] = vrho[:, 0]
        wva[1:4] = rhoR[0, 1:4] * (vsigma[:, 0] * 2)  # sigma_uu
        wva[1:4] += rhoR[1, 1:4] * vsigma[:, 1]  # sigma_ud
        wvb[0] = vrho[:, 1]
        wvb[1:4] = rhoR[1, 1:4] * (vsigma[:, 2] * 2)  # sigma_dd
        wvb[1:4] += rhoR[0, 1:4] * vsigma[:, 1]  # sigma_ud
    else:
        vrho, vsigma, vlapl, vtau = vxc
        wva = numpy.empty((5, ngrids))
        wvb = numpy.empty((5, ngrids))
        wva[0] = vrho[:, 0]
        wva[1:4] = rhoR[0, 1:4] * (vsigma[:, 0] * 2)  # sigma_uu
        wva[1:4] += rhoR[1, 1:4] * vsigma[:, 1]  # sigma_ud
        wvb[0] = vrho[:, 1]
        wvb[1:4] = rhoR[1, 1:4] * (vsigma[:, 2] * 2)  # sigma_dd
        wvb[1:4] += rhoR[0, 1:4] * vsigma[:, 1]  # sigma_ud
        if vlapl is None:
            wvb[4] = .5 * vtau[:, 1]
            wva[4] = .5 * vtau[:, 0]
        else:
            wva[4] = (.5 * vtau[:, 0] + 2 * vlapl[:, 0])
            wvb[4] = (.5 * vtau[:, 1] + 2 * vlapl[:, 1])

    nelec[0] += rhoR[0, 0].sum() * weight
    nelec[1] += rhoR[1, 0].sum() * weight
    excsum += (rhoR[0, 0] * exc).sum() * weight
    excsum += (rhoR[1, 0] * exc).sum() * weight
    wv_freq = tools.fft(numpy.vstack((wva, wvb)), mesh) * weight
    wv_freq = wv_freq.reshape(2, -1, *mesh)
    if j_in_xc:
        wv_freq[:, 0] += vG
        vR = tools.ifft(vG.reshape(-1, ngrids), mesh)
        ecoul = numpy.einsum('ng,ng->', rhoR[:, 0].real, vR.real) * .5
        log.debug('Coulomb energy %s', ecoul)
        excsum += ecoul
    rhoR = rhoG = None

    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
    if gamma_point(kpts_band):
        veff = numpy.zeros((2, nkpts, nao, nao))
        vj = numpy.zeros((2, nkpts, nao, nao))
    else:
        veff = numpy.zeros((2, nkpts, nao, nao), dtype=numpy.complex128)
        vj = numpy.zeros((2, nkpts, nao, nao), dtype=numpy.complex128)

    for grids_high, grids_low in mydf.tasks:
        cell_high = grids_high.cell
        mesh = grids_high.mesh
        coords_idx = grids_high.coords_idx
        ngrids0 = numpy.prod(mesh)
        ngrids1 = grids_high.coords.shape[0]
        log.debug('mesh %s, ngrids %s/%s', mesh, ngrids1, ngrids0)

        gx = numpy.fft.fftfreq(mesh[0], 1. / mesh[0]).astype(int)
        gy = numpy.fft.fftfreq(mesh[1], 1. / mesh[1]).astype(int)
        gz = numpy.fft.fftfreq(mesh[2], 1. / mesh[2]).astype(int)
        sub_wvG = wv_freq[:, :, gx[:, None, None], gy[:, None],
                          gz].reshape(-1, ngrids0)
        wv = tools.ifft(sub_wvG, mesh).real.reshape(2, -1, ngrids0)
        wv = wv[:, :, coords_idx]
        if with_j:
            sub_vG = vG[:, gx[:, None, None], gy[:, None],
                        gz].reshape(-1, ngrids0)
            vR = tools.ifft(sub_vG, mesh).real.reshape(2, ngrids0)
            vR = vR[:, coords_idx]

        idx_h = grids_high.ao_idx
        if grids_low is None:
            for ao_h_etc, p0, p1 in mydf.aoR_loop(grids_high, kpts, deriv):
                ao_h = ao_h_etc[0]
                add_xc_(veff, ao_h, ao_h, idx_h, idx_h, wv[:, :, p0:p1])
                if with_j:
                    add_j_(vj, ao_h, ao_h, idx_h, idx_h, vR[:, p0:p1])
                ao_h = ao_h_etc = None
        else:
            idx_l = grids_low.ao_idx
            for ao_h_etc, ao_l_etc in zip(
                    mydf.aoR_loop(grids_high, kpts, deriv),
                    mydf.aoR_loop(grids_low, kpts, deriv)):
                p0, p1 = ao_h_etc[1:3]
                ao_h = ao_h_etc[0][0]
                ao_l = ao_l_etc[0][0]
                add_xc_(veff, ao_h, ao_h, idx_h, idx_h, wv[:, :, p0:p1])
                add_xc_(veff, ao_h, ao_l, idx_h, idx_l, wv[:, :, p0:p1])
                add_xc_(veff, ao_l, ao_h, idx_l, idx_h, wv[:, :, p0:p1])
                if with_j:
                    add_j_(vj, ao_h, ao_h, idx_h, idx_h, vR[:, p0:p1])
                    add_j_(vj, ao_h, ao_l, idx_h, idx_l, vR[:, p0:p1])
                    add_j_(vj, ao_l, ao_h, idx_l, idx_h, vR[:, p0:p1])
                ao_h = ao_l = ao_h_etc = ao_l_etc = None

    vj = _format_jks(vj, dm_kpts, input_band, kpts)
    veff = _format_jks(veff, dm_kpts, input_band, kpts)
    return nelec, excsum, veff, vj
Ejemplo n.º 37
0
 def rotate_mo(self, mo_coeff, u, log=None):
     return lib.asarray(
         [numpy.dot(mo, u[k]) for k, mo in enumerate(mo_coeff)])
Ejemplo n.º 38
0
def get_j_e1_kpts(mydf, dm_kpts, kpts=np.zeros((1, 3)), kpts_band=None):
    cell = mydf.cell
    mesh = mydf.mesh

    ni = mydf._numint
    make_rho, nset, nao = ni._gen_rho_evaluator(cell, dm_kpts, hermi=1)
    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]

    coulG = tools.get_coulG(cell, mesh=mesh)
    ngrids = len(coulG)

    if gamma_point(kpts):
        vR = rhoR = np.zeros((nset, ngrids))
        for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
            ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
            for i in range(nset):
                rhoR[i, p0:p1] += make_rho(i, ao_ks, mask, 'LDA')
            ao = ao_ks = None

        for i in range(nset):
            rhoG = tools.fft(rhoR[i], mesh)
            vG = coulG * rhoG
            vR[i] = tools.ifft(vG, mesh).real

    else:  # vR may be complex if the underlying density is complex
        vR = rhoR = np.zeros((nset, ngrids), dtype=np.complex128)
        for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
            ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
            for i in range(nset):
                for k, ao in enumerate(ao_ks):
                    ao_dm = lib.dot(ao, dms[i, k])
                    rhoR[i, p0:p1] += np.einsum('xi,xi->x', ao_dm, ao.conj())
        rhoR *= 1. / nkpts

        for i in range(nset):
            rhoG = tools.fft(rhoR[i], mesh)
            vG = coulG * rhoG
            vR[i] = tools.ifft(vG, mesh)

    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
    nband = len(kpts_band)
    weight = cell.vol / ngrids
    vR *= weight
    if gamma_point(kpts_band):
        vj_kpts = np.zeros((3, nset, nband, nao, nao))
    else:
        vj_kpts = np.zeros((3, nset, nband, nao, nao), dtype=np.complex128)
    rho = None
    for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_band, deriv=1):
        ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
        for i in range(nset):
            # ni.eval_mat can handle real vR only
            # vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i,p0:p1], mask, 'LDA')
            for k, ao in enumerate(ao_ks):
                aow = np.einsum('xi,x->xi', ao[0], vR[i, p0:p1])
                vj_kpts[:, i, k] -= lib.einsum('axi,xj->aij', ao[1:].conj(),
                                               aow)

    vj_kpts = np.asarray(
        [_format_jks(vj, dm_kpts, input_band, kpts) for vj in vj_kpts])

    return vj_kpts
Ejemplo n.º 39
0
def get_k_kpts(mydf,
               dm_kpts,
               hermi=1,
               kpts=np.zeros((1, 3)),
               kpts_band=None,
               exxdiv=None):
    '''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points.

    Args:
        dm_kpts : (nkpts, nao, nao) ndarray
            Density matrix at each k-point
        kpts : (nkpts, 3) ndarray

    Kwargs:
        hermi : int
            Whether K matrix is hermitian

            | 0 : not hermitian and not symmetric
            | 1 : hermitian

        kpts_band : (3,) ndarray or (*,3) ndarray
            A list of arbitrary "band" k-points at which to evalute the matrix.

    Returns:
        vj : (nkpts, nao, nao) ndarray
        vk : (nkpts, nao, nao) ndarray
        or list of vj and vk if the input dm_kpts is a list of DMs
    '''
    cell = mydf.cell
    mesh = mydf.mesh
    coords = cell.gen_uniform_grids(mesh)
    ngrids = coords.shape[0]

    if getattr(dm_kpts, 'mo_coeff', None) is not None:
        mo_coeff = dm_kpts.mo_coeff
        mo_occ = dm_kpts.mo_occ
    else:
        mo_coeff = None

    kpts = np.asarray(kpts)
    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]

    weight = 1. / nkpts * (cell.vol / ngrids)

    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
    nband = len(kpts_band)

    if gamma_point(kpts_band) and gamma_point(kpts):
        vk_kpts = np.zeros((nset, nband, nao, nao), dtype=dms.dtype)
    else:
        vk_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128)

    coords = mydf.grids.coords
    ao2_kpts = [
        np.asarray(ao.T, order='C')
        for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)
    ]
    if input_band is None:
        ao1_kpts = ao2_kpts
    else:
        ao1_kpts = [
            np.asarray(ao.T, order='C')
            for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts_band)
        ]
    if mo_coeff is not None and nset == 1:
        mo_coeff = [
            mo_coeff[k][:, occ > 0] * np.sqrt(occ[occ > 0])
            for k, occ in enumerate(mo_occ)
        ]
        ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]

    mem_now = lib.current_memory()[0]
    max_memory = mydf.max_memory - mem_now
    blksize = int(
        min(nao, max(1, (max_memory - mem_now) * 1e6 / 16 / 4 / ngrids / nao)))
    lib.logger.debug1(mydf, 'fft_jk: get_k_kpts max_memory %s  blksize %d',
                      max_memory, blksize)
    ao1_dtype = np.result_type(*ao1_kpts)
    ao2_dtype = np.result_type(*ao2_kpts)
    vR_dm = np.empty((nset, nao, ngrids), dtype=vk_kpts.dtype)

    t1 = (time.clock(), time.time())
    for k2, ao2T in enumerate(ao2_kpts):
        if ao2T.size == 0:
            continue

        kpt2 = kpts[k2]
        naoj = ao2T.shape[0]
        if mo_coeff is None or nset > 1:
            ao_dms = [lib.dot(dms[i, k2], ao2T.conj()) for i in range(nset)]
        else:
            ao_dms = [ao2T.conj()]

        for k1, ao1T in enumerate(ao1_kpts):
            kpt1 = kpts_band[k1]

            # If we have an ewald exxdiv, we add the G=0 correction near the
            # end of the function to bypass any discretization errors
            # that arise from the FFT.
            mydf.exxdiv = exxdiv
            if exxdiv == 'ewald' or exxdiv is None:
                coulG = tools.get_coulG(cell, kpt2 - kpt1, False, mydf, mesh)
            else:
                coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, mesh)
            if is_zero(kpt1 - kpt2):
                expmikr = np.array(1.)
            else:
                expmikr = np.exp(-1j * np.dot(coords, kpt2 - kpt1))

            for p0, p1 in lib.prange(0, nao, blksize):
                rho1 = np.einsum('ig,jg->ijg', ao1T[p0:p1].conj() * expmikr,
                                 ao2T)
                vG = tools.fft(rho1.reshape(-1, ngrids), mesh)
                rho1 = None
                vG *= coulG
                vR = tools.ifft(vG, mesh).reshape(p1 - p0, naoj, ngrids)
                vG = None
                if vR_dm.dtype == np.double:
                    vR = vR.real
                for i in range(nset):
                    np.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i, p0:p1])
                vR = None
            vR_dm *= expmikr.conj()

            for i in range(nset):
                vk_kpts[i, k1] += weight * lib.dot(vR_dm[i], ao1T.T)
        t1 = lib.logger.timer_debug1(mydf, 'get_k_kpts: make_kpt (%d,*)' % k2,
                                     *t1)

    # Function _ewald_exxdiv_for_G0 to add back in the G=0 component to vk_kpts
    # Note in the _ewald_exxdiv_for_G0 implementation, the G=0 treatments are
    # different for 1D/2D and 3D systems.  The special treatments for 1D and 2D
    # can only be used with AFTDF/GDF/MDF method.  In the FFTDF method, 1D, 2D
    # and 3D should use the ewald probe charge correction.
    if exxdiv == 'ewald':
        _ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band=kpts_band)

    return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
Ejemplo n.º 40
0
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1, 3)), kpts_band=None):
    '''Get the Coulomb (J) AO matrix at sampled k-points.

    Args:
        dm_kpts : (nkpts, nao, nao) ndarray or a list of (nkpts,nao,nao) ndarray
            Density matrix at each k-point.  If a list of k-point DMs, eg,
            UHF alpha and beta DM, the alpha and beta DMs are contracted
            separately.
        kpts : (nkpts, 3) ndarray

    Kwargs:
        kpts_band : (3,) ndarray or (*,3) ndarray
            A list of arbitrary "band" k-points at which to evalute the matrix.

    Returns:
        vj : (nkpts, nao, nao) ndarray
        or list of vj if the input dm_kpts is a list of DMs
    '''
    cell = mydf.cell
    mesh = mydf.mesh

    ni = mydf._numint
    make_rho, nset, nao = ni._gen_rho_evaluator(cell, dm_kpts, hermi)
    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]

    coulG = tools.get_coulG(cell, mesh=mesh)
    ngrids = len(coulG)

    if hermi == 1 or gamma_point(kpts):
        vR = rhoR = np.zeros((nset, ngrids))
        for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
            ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
            for i in range(nset):
                rhoR[i, p0:p1] += make_rho(i, ao_ks, mask, 'LDA')
            ao = ao_ks = None

        for i in range(nset):
            rhoG = tools.fft(rhoR[i], mesh)
            vG = coulG * rhoG
            vR[i] = tools.ifft(vG, mesh).real

    else:  # vR may be complex if the underlying density is complex
        vR = rhoR = np.zeros((nset, ngrids), dtype=np.complex128)
        for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
            ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
            for i in range(nset):
                for k, ao in enumerate(ao_ks):
                    ao_dm = lib.dot(ao, dms[i, k])
                    rhoR[i, p0:p1] += np.einsum('xi,xi->x', ao_dm, ao.conj())
        rhoR *= 1. / nkpts

        for i in range(nset):
            rhoG = tools.fft(rhoR[i], mesh)
            vG = coulG * rhoG
            vR[i] = tools.ifft(vG, mesh)

    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
    nband = len(kpts_band)
    weight = cell.vol / ngrids
    vR *= weight
    if gamma_point(kpts_band):
        vj_kpts = np.zeros((nset, nband, nao, nao))
    else:
        vj_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128)
    rho = None
    for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_band):
        ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
        for i in range(nset):
            # ni.eval_mat can handle real vR only
            # vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i,p0:p1], mask, 'LDA')
            for k, ao in enumerate(ao_ks):
                aow = np.einsum('xi,x->xi', ao, vR[i, p0:p1])
                vj_kpts[i, k] += lib.dot(ao.conj().T, aow)

    return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
Ejemplo n.º 41
0
def get_k_e1_kpts(mydf,
                  dm_kpts,
                  kpts=np.zeros((1, 3)),
                  kpts_band=None,
                  exxdiv=None):

    cell = mydf.cell
    mesh = mydf.mesh
    coords = cell.gen_uniform_grids(mesh)
    ngrids = coords.shape[0]

    if getattr(dm_kpts, 'mo_coeff', None) is not None:
        mo_coeff = dm_kpts.mo_coeff
        mo_occ = dm_kpts.mo_occ
    else:
        mo_coeff = None

    kpts = np.asarray(kpts)
    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]

    weight = 1. / nkpts * (cell.vol / ngrids)

    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
    nband = len(kpts_band)

    if gamma_point(kpts_band) and gamma_point(kpts):
        vk_kpts = np.zeros((3, nset, nband, nao, nao), dtype=dms.dtype)
    else:
        vk_kpts = np.zeros((3, nset, nband, nao, nao), dtype=np.complex128)

    coords = mydf.grids.coords

    if input_band is None:
        ao2_kpts = [
            np.asarray(ao.transpose(0, 2, 1), order='C')
            for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts, deriv=1)
        ]
        ao1_kpts = ao2_kpts
        ao2_kpts = [ao2_kpt[0] for ao2_kpt in ao2_kpts]
    else:
        ao2_kpts = [
            np.asarray(ao.T, order='C')
            for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)
        ]
        ao1_kpts = [
            np.asarray(ao.transpose(0, 2, 1), order='C') for ao in
            mydf._numint.eval_ao(cell, coords, kpts=kpts_band, deriv=1)
        ]

    if mo_coeff is not None and nset == 1:
        mo_coeff = [
            mo_coeff[k][:, occ > 0] * np.sqrt(occ[occ > 0])
            for k, occ in enumerate(mo_occ)
        ]
        ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]

    mem_now = lib.current_memory()[0]
    max_memory = mydf.max_memory - mem_now
    blksize = int(
        min(nao,
            max(1, (max_memory - mem_now) * 1e6 / 16 / 4 / 3 / ngrids / nao)))
    lib.logger.debug1(mydf, 'fft_jk: get_k_kpts max_memory %s  blksize %d',
                      max_memory, blksize)
    ao1_dtype = np.result_type(*ao1_kpts)
    ao2_dtype = np.result_type(*ao2_kpts)
    vR_dm = np.empty((3, nset, nao, ngrids), dtype=vk_kpts.dtype)

    t1 = (time.clock(), time.time())
    for k2, ao2T in enumerate(ao2_kpts):
        if ao2T.size == 0:
            continue

        kpt2 = kpts[k2]
        naoj = ao2T.shape[0]
        if mo_coeff is None or nset > 1:
            ao_dms = [lib.dot(dms[i, k2], ao2T.conj()) for i in range(nset)]
        else:
            ao_dms = [ao2T.conj()]

        for k1, ao1T in enumerate(ao1_kpts):
            kpt1 = kpts_band[k1]

            # If we have an ewald exxdiv, we add the G=0 correction near the
            # end of the function to bypass any discretization errors
            # that arise from the FFT.
            mydf.exxdiv = exxdiv
            if exxdiv == 'ewald' or exxdiv is None:
                coulG = tools.get_coulG(cell, kpt2 - kpt1, False, mydf, mesh)
            else:
                coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, mesh)
            if is_zero(kpt1 - kpt2):
                expmikr = np.array(1.)
            else:
                expmikr = np.exp(-1j * np.dot(coords, kpt2 - kpt1))

            for p0, p1 in lib.prange(0, nao, blksize):
                rho1 = np.einsum('aig,jg->aijg',
                                 ao1T[1:, p0:p1].conj() * expmikr, ao2T)
                vG = tools.fft(rho1.reshape(-1, ngrids), mesh)
                rho1 = None
                vG *= coulG
                vR = tools.ifft(vG, mesh).reshape(3, p1 - p0, naoj, ngrids)
                vG = None
                if vR_dm.dtype == np.double:
                    vR = vR.real
                for i in range(nset):
                    np.einsum('aijg,jg->aig',
                              vR,
                              ao_dms[i],
                              out=vR_dm[:, i, p0:p1])
                vR = None
            vR_dm *= expmikr.conj()

            for i in range(nset):
                vk_kpts[:, i, k1] -= weight * np.einsum(
                    'aig,jg->aij', vR_dm[:, i], ao1T[0])
        t1 = lib.logger.timer_debug1(mydf, 'get_k_kpts: make_kpt (%d,*)' % k2,
                                     *t1)

    # Ewald correction has no contribution to nuclear gradient unless range separted Coulomb is used
    # The gradient correction part is not added in the vk matrix
    if exxdiv == 'ewald' and cell.omega != 0:
        raise NotImplementedError("Range Separated Coulomb")
        # when cell.omega !=0: madelung constant will have a non-zero derivative
    vk_kpts = np.asarray(
        [_format_jks(vk, dm_kpts, input_band, kpts) for vk in vk_kpts])
    return vk_kpts
Ejemplo n.º 42
0
def get_jk_kpts(mf, cell, dm_kpts, kpts, kpts_band=None):
    coords = gen_grid.gen_uniform_grids(cell)
    nkpts = len(kpts)
    ngrids = 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 kpts_band is not None:
        aoR_kband = numint.eval_ao(cell, coords, kpts_band)


# J
    vjR = [get_vjR_kpts(cell, dms[i], aoR_kpts) for i in range(nset)]
    if kpts_band is not None:
        vj_kpts = [
            cell.vol / ngrids * 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 / ngrids * 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 / ngrids)
    vk_kpts = np.zeros_like(vj_kpts)
    if kpts_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], kpts_band,
                               kpt2)
            #:vk_kpts = 1./nkpts * (cell.vol/ngrids) * 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)
Ejemplo n.º 43
0
def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)),
               kpts_band=None, exxdiv=None):
    mydf = _sync_mydf(mydf)
    cell = mydf.cell
    mesh = mydf.mesh
    coords = cell.gen_uniform_grids(mesh)
    ngrids = coords.shape[0]

    if hasattr(dm_kpts, 'mo_coeff'):
        if dm_kpts.ndim == 3:  # KRHF
            mo_coeff = [dm_kpts.mo_coeff]
            mo_occ   = [dm_kpts.mo_occ  ]
        else:  # KUHF
            mo_coeff = dm_kpts.mo_coeff
            mo_occ   = dm_kpts.mo_occ
    elif hasattr(dm_kpts[0], 'mo_coeff'):
        mo_coeff = [dm.mo_coeff for dm in dm_kpts]
        mo_occ   = [dm.mo_occ   for dm in dm_kpts]
    else:
        mo_coeff = None

    kpts = numpy.asarray(kpts)
    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]

    weight = 1./nkpts * (cell.vol/ngrids)

    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
    nband = len(kpts_band)

    if gamma_point(kpts_band) and gamma_point(kpts):
        vk_kpts = numpy.zeros((nset,nband,nao,nao), dtype=dms.dtype)
    else:
        vk_kpts = numpy.zeros((nset,nband,nao,nao), dtype=numpy.complex128)

    coords = mydf.grids.coords
    ao2_kpts = [numpy.asarray(ao.T, order='C')
                for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)]
    if input_band is None:
        ao1_kpts = ao2_kpts
    else:
        ao1_kpts = [numpy.asarray(ao.T, order='C')
                    for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts_band)]

    mem_now = lib.current_memory()[0]
    max_memory = mydf.max_memory - mem_now
    blksize = int(min(nao, max(1, (max_memory-mem_now)*1e6/16/4/ngrids/nao)))
    lib.logger.debug1(mydf, 'max_memory %s  blksize %d', max_memory, blksize)
    ao1_dtype = numpy.result_type(*ao1_kpts)
    ao2_dtype = numpy.result_type(*ao2_kpts)
    vR_dm = numpy.empty((nset,nao,ngrids), dtype=vk_kpts.dtype)

    ao_dms_buf = [None] * nkpts
    tasks = [(k1,k2) for k2 in range(nkpts) for k1 in range(nband)]
    for k1, k2 in mpi.static_partition(tasks):
        ao1T = ao1_kpts[k1]
        ao2T = ao2_kpts[k2]
        kpt1 = kpts_band[k1]
        kpt2 = kpts[k2]
        if ao2T.size == 0 or ao1T.size == 0:
            continue

        # If we have an ewald exxdiv, we add the G=0 correction near the
        # end of the function to bypass any discretization errors
        # that arise from the FFT.
        if exxdiv == 'ewald' or exxdiv is None:
            coulG = tools.get_coulG(cell, kpt2-kpt1, False, mydf, mesh)
        else:
            coulG = tools.get_coulG(cell, kpt2-kpt1, exxdiv, mydf, mesh)
        if is_zero(kpt1-kpt2):
            expmikr = numpy.array(1.)
        else:
            expmikr = numpy.exp(-1j * numpy.dot(coords, kpt2-kpt1))

        if ao_dms_buf[k2] is None:
            if mo_coeff is None:
                ao_dms = [lib.dot(dm[k2], ao2T.conj()) for dm in dms]
            else:
                ao_dms = []
                for i, dm in enumerate(dms):
                    occ = mo_occ[i][k2]
                    mo_scaled = mo_coeff[i][k2][:,occ>0] * numpy.sqrt(occ[occ>0])
                    ao_dms.append(lib.dot(mo_scaled.T, ao2T).conj())
            ao_dms_buf[k2] = ao_dms
        else:
            ao_dms = ao_dms_buf[k2]

        if mo_coeff is None:
            for p0, p1 in lib.prange(0, nao, blksize):
                rho1 = numpy.einsum('ig,jg->ijg', ao1T[p0:p1].conj()*expmikr, ao2T)
                vG = tools.fft(rho1.reshape(-1,ngrids), mesh)
                rho1 = None
                vG *= coulG
                vR = tools.ifft(vG, mesh).reshape(p1-p0,nao,ngrids)
                vG = None
                if vR_dm.dtype == numpy.double:
                    vR = vR.real
                for i in range(nset):
                    numpy.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i,p0:p1])
                vR = None
        else:
            for p0, p1 in lib.prange(0, nao, blksize):
                for i in range(nset):
                    rho1 = numpy.einsum('ig,jg->ijg',
                                        ao1T[p0:p1].conj()*expmikr,
                                        ao_dms[i].conj())
                    vG = tools.fft(rho1.reshape(-1,ngrids), mesh)
                    rho1 = None
                    vG *= coulG
                    vR = tools.ifft(vG, mesh).reshape(p1-p0,-1,ngrids)
                    vG = None
                    if vR_dm.dtype == numpy.double:
                        vR = vR.real
                    numpy.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i,p0:p1])
                    vR = None
        vR_dm *= expmikr.conj()

        for i in range(nset):
            vk_kpts[i,k1] += weight * lib.dot(vR_dm[i], ao1T.T)

    vk_kpts = mpi.reduce(lib.asarray(vk_kpts))
    if gamma_point(kpts_band) and gamma_point(kpts):
        vk_kpts = vk_kpts.real

    if rank == 0:
        if exxdiv == 'ewald':
            _ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band=kpts_band)
        return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
Ejemplo n.º 44
0
def get_k_kpts(mydf,
               dm_kpts,
               hermi=1,
               kpts=np.zeros((1, 3)),
               kpts_band=None,
               exxdiv=None):
    '''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points.

    Args:
        dm_kpts : (nkpts, nao, nao) ndarray
            Density matrix at each k-point
        kpts : (nkpts, 3) ndarray

    Kwargs:
        kpts_band : (3,) ndarray or (*,3) ndarray
            A list of arbitrary "band" k-points at which to evalute the matrix.

    Returns:
        vj : (nkpts, nao, nao) ndarray
        vk : (nkpts, nao, nao) ndarray
        or list of vj and vk if the input dm_kpts is a list of DMs
    '''
    cell = mydf.cell
    gs = mydf.gs
    coords = cell.gen_uniform_grids(gs)
    ngs = coords.shape[0]

    if hasattr(dm_kpts, 'mo_coeff'):
        mo_coeff = dm_kpts.mo_coeff
        mo_occ = dm_kpts.mo_occ
    else:
        mo_coeff = None

    kpts = np.asarray(kpts)
    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]

    weight = 1. / nkpts * (cell.vol / ngs)

    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
    nband = len(kpts_band)

    if gamma_point(kpts_band) and gamma_point(kpts):
        vk_kpts = np.zeros((nset, nband, nao, nao), dtype=dms.dtype)
    else:
        vk_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128)

    ao2_kpts = mydf._numint.eval_ao(cell, coords, kpts, non0tab=mydf.non0tab)
    ao2_kpts = [np.asarray(ao.T, order='C') for ao in ao2_kpts]
    if input_band is None:
        ao1_kpts = ao2_kpts
    else:
        ao1_kpts = mydf._numint.eval_ao(cell,
                                        coords,
                                        kpts_band,
                                        non0tab=mydf.non0tab)
        ao1_kpts = [np.asarray(ao.T, order='C') for ao in ao1_kpts]
    if mo_coeff is not None and nset == 1:
        mo_coeff = [
            mo_coeff[k][:, occ > 0] * np.sqrt(occ[occ > 0])
            for k, occ in enumerate(mo_occ)
        ]
        ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]

    max_memory = mydf.max_memory - lib.current_memory()[0]
    blksize = int(min(nao, max_memory * 1e6 / 16 / 2 / ngs / nao + 1))
    ao1_dtype = np.result_type(*ao1_kpts)
    ao2_dtype = np.result_type(*ao2_kpts)
    vR_dm = np.empty((nset, nao, ngs), dtype=vk_kpts.dtype)

    for k2, ao2T in enumerate(ao2_kpts):
        if ao2T.size == 0:
            continue

        kpt2 = kpts[k2]
        if mo_coeff is None or nset > 1:
            ao_dms = [lib.dot(dms[i, k2], ao2T.conj()) for i in range(nset)]
        else:
            ao_dms = [ao2T.conj()]

        for k1, ao1T in enumerate(ao1_kpts):
            kpt1 = kpts_band[k1]
            mydf.exxdiv = exxdiv
            coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, gs)
            if is_zero(kpt1 - kpt2):
                expmikr = np.array(1.)
            else:
                expmikr = np.exp(-1j * np.dot(coords, kpt2 - kpt1))

            for p0, p1 in lib.prange(0, nao, blksize):
                rho1 = np.einsum('ig,jg->ijg', ao1T[p0:p1].conj() * expmikr,
                                 ao2T)
                vG = tools.fft(rho1.reshape(-1, ngs), gs)
                vG *= coulG
                vR = tools.ifft(vG, gs).reshape(p1 - p0, -1, ngs)
                vG = None
                if vR_dm.dtype == np.double:
                    vR = vR.real
                for i in range(nset):
                    np.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i, p0:p1])
                vR = None
            vR_dm *= expmikr.conj()

            for i in range(nset):
                vk_kpts[i, k1] += weight * lib.dot(vR_dm[i], ao1T.T)

    return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
Ejemplo n.º 45
0
def _eval_rhoG(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1, 3)), deriv=0):
    log = lib.logger.Logger(mydf.stdout, mydf.verbose)
    cell = mydf.cell

    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]

    tasks = getattr(mydf, 'tasks', None)
    if tasks is None:
        mydf.tasks = tasks = multi_grids_tasks(cell, log)
        log.debug('Multigrid ntasks %s', len(tasks))

    assert (deriv <= 2)
    if abs(dms - dms.transpose(0, 1, 3, 2).conj()).max() < 1e-9:

        def dot_bra(bra, aodm):
            rho = numpy.einsum('pi,pi->p', bra.real, aodm.real)
            if aodm.dtype == numpy.complex:
                rho += numpy.einsum('pi,pi->p', bra.imag, aodm.imag)
            return rho

        if deriv == 0:
            xctype = 'LDA'
            rhodim = 1

            def make_rho(ao_l, ao_h, dm_lh, dm_hl):
                c0 = lib.dot(ao_l, dm_lh)
                rho = dot_bra(ao_h, c0)
                return rho * 2

        elif deriv == 1:
            xctype = 'GGA'
            rhodim = 4

            def make_rho(ao_l, ao_h, dm_lh, dm_hl):
                ngrids = ao_l[0].shape[0]
                rho = numpy.empty((4, ngrids))
                c0 = lib.dot(ao_l[0], dm_lh)
                rho[0] = dot_bra(ao_h[0], c0)
                for i in range(1, 4):
                    rho[i] = dot_bra(ao_h[i], c0)
                c0 = lib.dot(ao_h[0], dm_hl)
                for i in range(1, 4):
                    rho[i] += dot_bra(ao_l[i], c0)
                return rho * 2  # *2 for dm_lh+dm_hl.T

        elif deriv == 2:
            xctype = 'MGGA'
            rhodim = 6

            def make_rho(ao_l, ao_h, dm_lh, dm_hl):
                ngrids = ao_l[0].shape[0]
                rho = numpy.empty((6, ngrids))
                c = [lib.dot(ao_l[i], dm_lh) for i in range(4)]
                rho[0] = dot_bra(ao_h[0], c[0])
                rho[5] = 0
                for i in range(1, 4):
                    rho[i] = dot_bra(ao_h[i], c[0])
                    rho[i] += dot_bra(ao_h[0], c[i])
                    rho[5] += dot_bra(ao_h[i], c[i]) * 2
                XX, YY, ZZ = 4, 7, 9
                ao2 = ao_h[XX] + ao_h[YY] + ao_h[ZZ]
                rho[4] = dot_bra(ao2, c[0])
                ao2 = lib.dot(ao_l[XX] + ao_l[YY] + ao_l[ZZ], dm_lh)
                rho[4] += dot_bra(ao2, ao_h[0])
                rho[4] += rho[5] * 2
                rho[5] *= .5
                return rho * 2  # *2 for dm_lh+dm_hl.T
    else:
        raise NotImplementedError('Non-hermitian density matrices')

    ni = mydf._numint
    nx, ny, nz = cell.mesh
    rhoG = numpy.zeros((nset * rhodim, nx, ny, nz), dtype=numpy.complex)
    for grids_high, grids_low in tasks:
        cell_high = grids_high.cell
        mesh = grids_high.mesh
        coords_idx = grids_high.coords_idx
        ngrids0 = numpy.prod(mesh)
        ngrids1 = grids_high.coords.shape[0]
        log.debug('mesh %s, ngrids %s/%s', mesh, ngrids1, ngrids0)

        idx_h = grids_high.ao_idx
        dms_hh = numpy.asarray(dms[:, :, idx_h[:, None], idx_h], order='C')
        if grids_low is not None:
            idx_l = grids_low.ao_idx
            dms_hl = numpy.asarray(dms[:, :, idx_h[:, None], idx_l], order='C')
            dms_lh = numpy.asarray(dms[:, :, idx_l[:, None], idx_h], order='C')

        rho = numpy.zeros((nset, rhodim, ngrids1))
        if grids_low is None:
            for ao_h_etc, p0, p1 in mydf.aoR_loop(grids_high, kpts, deriv):
                ao_h, mask = ao_h_etc[0], ao_h_etc[2]
                for k in range(nkpts):
                    for i in range(nset):
                        rho_sub = numint.eval_rho(cell_high, ao_h[k],
                                                  dms_hh[i, k], mask, xctype,
                                                  hermi)
                        rho[i, :, p0:p1] += rho_sub.real
                ao_h = ao_h_etc = None
        else:
            for ao_h_etc, ao_l_etc in zip(
                    mydf.aoR_loop(grids_high, kpts, deriv),
                    mydf.aoR_loop(grids_low, kpts, deriv)):
                p0, p1 = ao_h_etc[1:3]
                ao_h, mask = ao_h_etc[0][0], ao_h_etc[0][2]
                ao_l = ao_l_etc[0][0]
                for k in range(nkpts):
                    for i in range(nset):
                        rho_sub = numint.eval_rho(cell_high, ao_h[k],
                                                  dms_hh[i, k], mask, xctype,
                                                  hermi)
                        rho[i, :, p0:p1] += rho_sub.real
                        rho_sub = make_rho(ao_l[k], ao_h[k], dms_lh[i, k],
                                           dms_hl[i, k])
                        rho[i, :, p0:p1] += rho_sub.real
                ao_h = ao_l = ao_h_etc = ao_l_etc = None

        rho *= 1. / nkpts
        rhoR = numpy.zeros((nset * rhodim, ngrids0))
        rhoR[:, coords_idx] = rho.reshape(nset * rhodim, ngrids1)
        gx = numpy.fft.fftfreq(mesh[0], 1. / mesh[0]).astype(int)
        gy = numpy.fft.fftfreq(mesh[1], 1. / mesh[1]).astype(int)
        gz = numpy.fft.fftfreq(mesh[2], 1. / mesh[2]).astype(int)
        rho_freq = tools.fft(rhoR, mesh) * cell.vol / ngrids0
        for i in range(nset * rhodim):
            rhoG[i, gx[:, None, None], gy[:, None],
                 gz] += rho_freq[i].reshape(mesh)

    return rhoG.reshape(nset, rhodim, ngrids0)
Ejemplo n.º 46
0
def get_k_kpts(mydf,
               dm_kpts,
               hermi=1,
               kpts=numpy.zeros((1, 3)),
               kpts_band=None,
               exxdiv=None):
    cell = mydf.cell
    log = logger.Logger(mydf.stdout, mydf.verbose)
    t1 = (time.clock(), time.time())

    mesh = mydf.mesh
    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]

    swap_2e = (kpts_band is None)
    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
    nband = len(kpts_band)
    kk_table = kpts_band.reshape(-1, 1, 3) - kpts.reshape(1, -1, 3)
    kk_todo = numpy.ones(kk_table.shape[:2], dtype=bool)
    vkR = numpy.zeros((nset, nband, nao, nao))
    vkI = numpy.zeros((nset, nband, nao, nao))
    dmsR = numpy.asarray(dms.real, order='C')
    dmsI = numpy.asarray(dms.imag, order='C')

    mem_now = lib.current_memory()[0]
    max_memory = max(2000, (mydf.max_memory - mem_now)) * .8
    log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now)

    # K_pq = ( p{k1} i{k2} | i{k2} q{k1} )
    def make_kpt(kpt):  # kpt = kptj - kpti
        # search for all possible ki and kj that has ki-kj+kpt=0
        kk_match = numpy.einsum('ijx->ij', abs(kk_table + kpt)) < 1e-9
        kpti_idx, kptj_idx = numpy.where(kk_todo & kk_match)
        nkptj = len(kptj_idx)
        log.debug1('kpt = %s', kpt)
        log.debug2('kpti_idx = %s', kpti_idx)
        log.debug2('kptj_idx = %s', kptj_idx)
        kk_todo[kpti_idx, kptj_idx] = False
        if swap_2e and not is_zero(kpt):
            kk_todo[kptj_idx, kpti_idx] = False

        max_memory1 = max_memory * (nkptj + 1) / (nkptj + 5)
        #blksize = max(int(max_memory1*4e6/(nkptj+5)/16/nao**2), 16)

        #bufR = numpy.empty((blksize*nao**2))
        #bufI = numpy.empty((blksize*nao**2))
        # Use DF object to mimic KRHF/KUHF object in function get_coulG
        vkcoulG = mydf.weighted_coulG(kpt, exxdiv, mesh)
        kptjs = kpts[kptj_idx]
        weight = 1. / len(kpts)
        perm_sym = swap_2e and not is_zero(kpt)
        for aoaoks, p0, p1 in mydf.ft_loop(mesh,
                                           kpt,
                                           kptjs,
                                           max_memory=max_memory1):
            _update_vk_((vkR, vkI), aoaoks, (dmsR, dmsI), vkcoulG[p0:p1],
                        weight, kpti_idx, kptj_idx, perm_sym)

    for ki, kpti in enumerate(kpts_band):
        for kj, kptj in enumerate(kpts):
            if kk_todo[ki, kj]:
                make_kpt(kptj - kpti)
        t1 = log.timer_debug1('get_k_kpts: make_kpt (%d,*)' % ki, *t1)

    if (gamma_point(kpts) and gamma_point(kpts_band)
            and not numpy.iscomplexobj(dm_kpts)):
        vk_kpts = vkR
    else:
        vk_kpts = vkR + vkI * 1j

    # Add ewald_exxdiv contribution because G=0 was not included in the
    # non-uniform grids
    if (exxdiv == 'ewald' and
        (cell.dimension < 2 or  # 0D and 1D are computed with inf_vacuum
         (cell.dimension == 2 and cell.low_dim_ft_type == 'inf_vacuum'))):
        _ewald_exxdiv_for_G0(cell, kpts_band, dms, vk_kpts, kpts_band)

    return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
Ejemplo n.º 47
0
def get_pp(mydf, kpts=None):
    mydf = _sync_mydf(mydf)
    cell = mydf.cell
    if kpts is None:
        kpts_lst = numpy.zeros((1,3))
    else:
        kpts_lst = numpy.reshape(kpts, (-1,3))
    if abs(kpts_lst).sum < 1e-9:
        dtype = numpy.float64
    else:
        dtype = numpy.complex128

    gs = mydf.gs
    SI = cell.get_SI()
    Gv = cell.get_Gv(gs)
    vpplocG = pseudo.get_vlocG(cell, Gv)
    vpplocG = -numpy.einsum('ij,ij->j', SI, vpplocG)
    vpplocG[0] = numpy.sum(pseudo.get_alphas(cell)) # from get_jvloc_G0 function
    ngs = len(vpplocG)
    nao = cell.nao_nr()

    # vpploc evaluated in real-space
    vpplocR = tools.ifft(vpplocG, cell.gs).real
    vpp = [lib.dot(aoR.T.conj()*vpplocR, aoR)
           for k, aoR in mydf.mpi_aoR_loop(gs, kpts_lst)]
    vpp = mpi.gather(lib.asarray(vpp, dtype=dtype))

    # vppnonloc evaluated in reciprocal space
    fakemol = gto.Mole()
    fakemol._atm = numpy.zeros((1,gto.ATM_SLOTS), dtype=numpy.int32)
    fakemol._bas = numpy.zeros((1,gto.BAS_SLOTS), dtype=numpy.int32)
    ptr = gto.PTR_ENV_START
    fakemol._env = numpy.zeros(ptr+10)
    fakemol._bas[0,gto.NPRIM_OF ] = 1
    fakemol._bas[0,gto.NCTR_OF  ] = 1
    fakemol._bas[0,gto.PTR_EXP  ] = ptr+3
    fakemol._bas[0,gto.PTR_COEFF] = ptr+4

    # buf for SPG_lmi upto l=0..3 and nl=3
    buf = numpy.empty((48,ngs), dtype=numpy.complex128)
    def vppnl_by_k(kpt):
        Gk = Gv + kpt
        G_rad = lib.norm(Gk, axis=1)
        aokG = ft_ao.ft_ao(cell, Gv, kpt=kpt) * (ngs/cell.vol)
        vppnl = 0
        for ia in range(cell.natm):
            symb = cell.atom_symbol(ia)
            if symb not in cell._pseudo:
                continue
            pp = cell._pseudo[symb]
            p1 = 0
            for l, proj in enumerate(pp[5:]):
                rl, nl, hl = proj
                if nl > 0:
                    fakemol._bas[0,gto.ANG_OF] = l
                    fakemol._env[ptr+3] = .5*rl**2
                    fakemol._env[ptr+4] = rl**(l+1.5)*numpy.pi**1.25
                    pYlm_part = dft.numint.eval_ao(fakemol, Gk, deriv=0)

                    p0, p1 = p1, p1+nl*(l*2+1)
                    # pYlm is real, SI[ia] is complex
                    pYlm = numpy.ndarray((nl,l*2+1,ngs), dtype=numpy.complex128, buffer=buf[p0:p1])
                    for k in range(nl):
                        qkl = pseudo.pp._qli(G_rad*rl, l, k)
                        pYlm[k] = pYlm_part.T * qkl
                    #:SPG_lmi = numpy.einsum('g,nmg->nmg', SI[ia].conj(), pYlm)
                    #:SPG_lm_aoG = numpy.einsum('nmg,gp->nmp', SPG_lmi, aokG)
                    #:tmp = numpy.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
                    #:vppnl += numpy.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp)
            SPG_lmi = buf[:p1]
            SPG_lmi *= SI[ia].conj()
            SPG_lm_aoGs = lib.zdot(SPG_lmi, aokG)
            p1 = 0
            for l, proj in enumerate(pp[5:]):
                rl, nl, hl = proj
                if nl > 0:
                    p0, p1 = p1, p1+nl*(l*2+1)
                    hl = numpy.asarray(hl)
                    SPG_lm_aoG = SPG_lm_aoGs[p0:p1].reshape(nl,l*2+1,-1)
                    tmp = numpy.einsum('ij,jmp->imp', hl, SPG_lm_aoG)
                    vppnl += numpy.einsum('imp,imq->pq', SPG_lm_aoG.conj(), tmp)
        return vppnl * (1./ngs**2)

    vppnl = []
    for kpt in mpi.static_partition(kpts_lst):
        vppnl.append(vppnl_by_k(kpt))
    vppnl = mpi.gather(lib.asarray(vppnl, dtype=dtype))

    if rank == 0:
        vpp += vppnl
        if kpts is None or numpy.shape(kpts) == (3,):
            vpp = vpp[0]
        return vpp
Ejemplo n.º 48
0
 def eig(self, fock, s):
     e_a, c_a = hf.SCF.eig(self, fock[0], s)
     e_b, c_b = hf.SCF.eig(self, fock[1], s)
     return lib.asarray((e_a, e_b)), lib.asarray((c_a, c_b))
Ejemplo n.º 49
0
 def get_ovlp(self, cell=None, kpts=None):
     s = khf.KSCF.get_ovlp(self, cell, kpts)
     return lib.asarray([scipy.linalg.block_diag(x, x) for x in s])
Ejemplo n.º 50
0
def get_k_kpts(mydf,
               dm_kpts,
               hermi=1,
               kpts=numpy.zeros((1, 3)),
               kpts_band=None,
               exxdiv=None):
    cell = mydf.cell
    log = logger.Logger(mydf.stdout, mydf.verbose)

    if exxdiv is not None and exxdiv != 'ewald':
        log.warn(
            'GDF does not support exxdiv %s. '
            'exxdiv needs to be "ewald" or None', exxdiv)
        raise RuntimeError('GDF does not support exxdiv %s' % exxdiv)

    t1 = (logger.process_clock(), logger.perf_counter())
    if mydf._cderi is None or not mydf.has_kpts(kpts_band):
        if mydf._cderi is not None:
            log.warn(
                'DF integrals for band k-points were not found %s. '
                'DF integrals will be rebuilt to include band k-points.',
                mydf._cderi)
        mydf.build(kpts_band=kpts_band)
        t1 = log.timer_debug1('Init get_k_kpts', *t1)

    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]

    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
    nband = len(kpts_band)
    vkR = numpy.zeros((nset, nband, nao, nao))
    vkI = numpy.zeros((nset, nband, nao, nao))
    dmsR = numpy.asarray(dms.real, order='C')
    dmsI = numpy.asarray(dms.imag, order='C')

    # K_pq = ( p{k1} i{k2} | i{k2} q{k1} )
    bufR = numpy.empty((mydf.blockdim * nao**2))
    bufI = numpy.empty((mydf.blockdim * nao**2))
    max_memory = max(2000, mydf.max_memory - lib.current_memory()[0])

    def make_kpt(ki, kj, swap_2e):
        kpti = kpts[ki]
        kptj = kpts_band[kj]

        for LpqR, LpqI, sign in mydf.sr_loop((kpti, kptj), max_memory, False):
            nrow = LpqR.shape[0]
            pLqR = numpy.ndarray((nao, nrow, nao), buffer=bufR)
            pLqI = numpy.ndarray((nao, nrow, nao), buffer=bufI)
            tmpR = numpy.ndarray((nao, nrow * nao), buffer=LpqR)
            tmpI = numpy.ndarray((nao, nrow * nao), buffer=LpqI)
            pLqR[:] = LpqR.reshape(-1, nao, nao).transpose(1, 0, 2)
            pLqI[:] = LpqI.reshape(-1, nao, nao).transpose(1, 0, 2)

            for i in range(nset):
                zdotNN(dmsR[i, ki], dmsI[i, ki], pLqR.reshape(nao, -1),
                       pLqI.reshape(nao, -1), 1, tmpR, tmpI)
                zdotCN(
                    pLqR.reshape(-1, nao).T,
                    pLqI.reshape(-1, nao).T, tmpR.reshape(-1, nao),
                    tmpI.reshape(-1, nao), sign, vkR[i, kj], vkI[i, kj], 1)

            if swap_2e:
                tmpR = tmpR.reshape(nao * nrow, nao)
                tmpI = tmpI.reshape(nao * nrow, nao)
                for i in range(nset):
                    zdotNN(pLqR.reshape(-1, nao), pLqI.reshape(-1, nao),
                           dmsR[i, kj], dmsI[i, kj], 1, tmpR, tmpI)
                    zdotNC(tmpR.reshape(nao, -1), tmpI.reshape(nao, -1),
                           pLqR.reshape(nao, -1).T,
                           pLqI.reshape(nao, -1).T, sign, vkR[i, ki],
                           vkI[i, ki], 1)

    if kpts_band is kpts:  # normal k-points HF/DFT
        for ki in range(nkpts):
            for kj in range(ki):
                make_kpt(ki, kj, True)
            make_kpt(ki, ki, False)
            t1 = log.timer_debug1('get_k_kpts: make_kpt ki>=kj (%d,*)' % ki,
                                  *t1)
    else:
        for ki in range(nkpts):
            for kj in range(nband):
                make_kpt(ki, kj, False)
            t1 = log.timer_debug1('get_k_kpts: make_kpt (%d,*)' % ki, *t1)

    if (gamma_point(kpts) and gamma_point(kpts_band)
            and not numpy.iscomplexobj(dm_kpts)):
        vk_kpts = vkR
    else:
        vk_kpts = vkR + vkI * 1j
    vk_kpts *= 1. / nkpts

    if exxdiv == 'ewald':
        _ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band)

    return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
Ejemplo n.º 51
0
def get_j_kpts(mydf,
               dm_kpts,
               hermi=1,
               kpts=numpy.zeros((1, 3)),
               kpts_band=None):
    log = logger.Logger(mydf.stdout, mydf.verbose)
    t1 = (logger.process_clock(), logger.perf_counter())
    if mydf._cderi is None or not mydf.has_kpts(kpts_band):
        if mydf._cderi is not None:
            log.warn(
                'DF integrals for band k-points were not found %s. '
                'DF integrals will be rebuilt to include band k-points.',
                mydf._cderi)
        mydf.build(kpts_band=kpts_band)
        t1 = log.timer_debug1('Init get_j_kpts', *t1)

    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]
    if mydf.auxcell is None:
        # If mydf._cderi is the file that generated from another calculation,
        # guess naux based on the contents of the integral file.
        naux = mydf.get_naoaux()
    else:
        naux = mydf.auxcell.nao_nr()
    nao_pair = nao * (nao + 1) // 2

    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
    nband = len(kpts_band)
    j_real = gamma_point(kpts_band) and not numpy.iscomplexobj(dms)

    dmsR = dms.real.transpose(0, 1, 3, 2).reshape(nset, nkpts, nao**2)
    dmsI = dms.imag.transpose(0, 1, 3, 2).reshape(nset, nkpts, nao**2)
    rhoR = numpy.zeros((nset, naux))
    rhoI = numpy.zeros((nset, naux))
    max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]))
    for k, kpt in enumerate(kpts):
        kptii = numpy.asarray((kpt, kpt))
        p1 = 0
        for LpqR, LpqI, sign in mydf.sr_loop(kptii, max_memory, False):
            p0, p1 = p1, p1 + LpqR.shape[0]
            #:Lpq = (LpqR + LpqI*1j).reshape(-1,nao,nao)
            #:rhoR[:,p0:p1] += numpy.einsum('Lpq,xqp->xL', Lpq, dms[:,k]).real
            #:rhoI[:,p0:p1] += numpy.einsum('Lpq,xqp->xL', Lpq, dms[:,k]).imag
            rhoR[:,
                 p0:p1] += sign * numpy.einsum('Lp,xp->xL', LpqR, dmsR[:, k])
            rhoI[:,
                 p0:p1] += sign * numpy.einsum('Lp,xp->xL', LpqR, dmsI[:, k])
            if LpqI is not None:
                rhoR[:, p0:p1] -= sign * numpy.einsum('Lp,xp->xL', LpqI,
                                                      dmsI[:, k])
                rhoI[:, p0:p1] += sign * numpy.einsum('Lp,xp->xL', LpqI,
                                                      dmsR[:, k])
            LpqR = LpqI = None
    t1 = log.timer_debug1('get_j pass 1', *t1)

    weight = 1. / nkpts
    rhoR *= weight
    rhoI *= weight
    vjR = numpy.zeros((nset, nband, nao_pair))
    vjI = numpy.zeros((nset, nband, nao_pair))
    for k, kpt in enumerate(kpts_band):
        kptii = numpy.asarray((kpt, kpt))
        p1 = 0
        for LpqR, LpqI, sign in mydf.sr_loop(kptii, max_memory, True):
            p0, p1 = p1, p1 + LpqR.shape[0]
            #:Lpq = (LpqR + LpqI*1j)#.reshape(-1,nao,nao)
            #:vjR[:,k] += numpy.dot(rho[:,p0:p1], Lpq).real
            #:vjI[:,k] += numpy.dot(rho[:,p0:p1], Lpq).imag
            vjR[:, k] += numpy.dot(rhoR[:, p0:p1], LpqR)
            if not j_real:
                vjI[:, k] += numpy.dot(rhoI[:, p0:p1], LpqR)
                if LpqI is not None:
                    vjR[:, k] -= numpy.dot(rhoI[:, p0:p1], LpqI)
                    vjI[:, k] += numpy.dot(rhoR[:, p0:p1], LpqI)
            LpqR = LpqI = None
    t1 = log.timer_debug1('get_j pass 2', *t1)

    if j_real:
        vj_kpts = vjR
    else:
        vj_kpts = vjR + vjI * 1j
    vj_kpts = lib.unpack_tril(vj_kpts.reshape(-1, nao_pair))
    vj_kpts = vj_kpts.reshape(nset, nband, nao, nao)

    return _format_jks(vj_kpts, dm_kpts, input_band, kpts)