Exemplo n.º 1
0
def build_Lpq_pbc(mydf, auxcell, kptij_lst):
    '''Fitting coefficients for auxiliary functions'''
    kpts_ji = kptij_lst[:,1] - kptij_lst[:,0]
    uniq_kpts, uniq_index, uniq_inverse = unique(kpts_ji)
    max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]))
    if mydf.metric.upper() == 'S':
        outcore.aux_e2(mydf.cell, auxcell, mydf._cderi, 'cint3c1e_sph',
                       kptij_lst=kptij_lst, dataname='Lpq',
                       max_memory=max_memory)
        s_aux = auxcell.pbc_intor('cint1e_ovlp_sph', hermi=1, kpts=uniq_kpts)
    else:  # mydf.metric.upper() == 'T'
        outcore.aux_e2(mydf.cell, auxcell, mydf._cderi, 'cint3c1e_p2_sph',
                       kptij_lst=kptij_lst, dataname='Lpq',
                       max_memory=max_memory)
        s_aux = [x*2 for x in auxcell.pbc_intor('cint1e_kin_sph', hermi=1, kpts=uniq_kpts)]

    s_aux = [scipy.linalg.cho_factor(x) for x in s_aux]

    max_memory = mydf.max_memory - lib.current_memory()[0]
    naux = auxcell.nao_nr()
    blksize = max(int(max_memory*.5*1e6/16/naux/mydf.blockdim), 1) * mydf.blockdim
    with h5py.File(mydf._cderi) as feri:
        for k, where in enumerate(uniq_inverse):
            s_k = s_aux[where]
            key = 'Lpq/%d' % k
            Lpq = feri[key]
            nao_pair = Lpq.shape[1]
            for p0, p1 in lib.prange(0, nao_pair, blksize):
                Lpq[:,p0:p1] = scipy.linalg.cho_solve(s_k, Lpq[:,p0:p1])
Exemplo n.º 2
0
 def test_aux_e2(self):
     tmpfile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
     numpy.random.seed(1)
     kptij_lst = numpy.random.random((3,2,3))
     kptij_lst[0] = 0
     outcore.aux_e2(cell, cell, tmpfile.name, aosym='s2', comp=1,
                    kptij_lst=kptij_lst, verbose=0)
     refk = incore.aux_e2(cell, cell, aosym='s2', kptij_lst=kptij_lst)
     with h5py.File(tmpfile.name, 'r') as f:
         nao = cell.nao_nr()
         idx = numpy.tril_indices(nao)
         idx = idx[0] * nao + idx[1]
         self.assertTrue(numpy.allclose(refk[0,idx], f['eri_mo/0'].value.T))
         self.assertTrue(numpy.allclose(refk[1], f['eri_mo/1'].value.T))
         self.assertTrue(numpy.allclose(refk[2], f['eri_mo/2'].value.T))
Exemplo n.º 3
0
 def test_aux_e2(self):
     tmpfile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
     numpy.random.seed(1)
     kptij_lst = numpy.random.random((3, 2, 3))
     kptij_lst[0] = 0
     kptij_lst[1, 0] = kptij_lst[1, 1]
     outcore.aux_e2(cell,
                    cell,
                    tmpfile.name,
                    aosym='s2ij',
                    comp=1,
                    kptij_lst=kptij_lst,
                    verbose=0)
     refk0 = incore.aux_e2(cell, cell, aosym='s2ij', kpti_kptj=kptij_lst[0])
     refk1 = incore.aux_e2(cell, cell, aosym='s2ij', kpti_kptj=kptij_lst[1])
     refk2 = incore.aux_e2(cell, cell, aosym='s2ij', kpti_kptj=kptij_lst[2])
     with h5py.File(tmpfile.name, 'r') as f:
         self.assertTrue(numpy.allclose(refk0, f['eri_mo/0'].value.T))
         self.assertTrue(numpy.allclose(refk1, f['eri_mo/1'].value.T))
         self.assertTrue(numpy.allclose(refk2, f['eri_mo/2'].value.T))
Exemplo n.º 4
0
def _make_j3c(mydf, cell, auxcell, kptij_lst, cderi_file):
    t1 = (time.clock(), time.time())
    log = logger.Logger(mydf.stdout, mydf.verbose)
    max_memory = max(2000, mydf.max_memory - lib.current_memory()[0])
    fused_cell, fuse = fuse_auxcell(mydf, auxcell)
    outcore.aux_e2(cell,
                   fused_cell,
                   cderi_file,
                   'int3c2e_sph',
                   aosym='s2',
                   kptij_lst=kptij_lst,
                   dataname='j3c',
                   max_memory=max_memory)
    t1 = log.timer_debug1('3c2e', *t1)

    nao = cell.nao_nr()
    naux = auxcell.nao_nr()
    gs = mydf.gs
    Gv, Gvbase, kws = cell.get_Gv_weights(gs)
    b = cell.reciprocal_vectors()
    gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
    ngs = gxyz.shape[0]

    kptis = kptij_lst[:, 0]
    kptjs = kptij_lst[:, 1]
    kpt_ji = kptjs - kptis
    uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
    log.debug('Num uniq kpts %d', len(uniq_kpts))
    log.debug2('uniq_kpts %s', uniq_kpts)
    # j2c ~ (-kpt_ji | kpt_ji)
    j2c = fused_cell.pbc_intor('int2c2e_sph', hermi=1, kpts=uniq_kpts)
    feri = h5py.File(cderi_file)

    for k, kpt in enumerate(uniq_kpts):
        aoaux = ft_ao.ft_ao(fused_cell, Gv, None, b, gxyz, Gvbase, kpt).T
        aoaux = fuse(aoaux)
        coulG = numpy.sqrt(mydf.weighted_coulG(kpt, False, gs))
        kLR = (aoaux.real * coulG).T
        kLI = (aoaux.imag * coulG).T
        if not kLR.flags.c_contiguous: kLR = lib.transpose(kLR.T)
        if not kLI.flags.c_contiguous: kLI = lib.transpose(kLI.T)

        j2c_k = fuse(fuse(j2c[k]).T).T.copy()
        if is_zero(kpt):  # kpti == kptj
            j2c_k -= lib.dot(kLR.T, kLR)
            j2c_k -= lib.dot(kLI.T, kLI)
        else:
            # aoaux ~ kpt_ij, aoaux.conj() ~ kpt_kl
            j2cR, j2cI = zdotCN(kLR.T, kLI.T, kLR, kLI)
            j2c_k -= j2cR + j2cI * 1j
        feri['j2c/%d' % k] = j2c_k
        aoaux = kLR = kLI = j2cR = j2cI = coulG = None
    j2c = None

    def make_kpt(uniq_kptji_id):  # kpt = kptj - kpti
        kpt = uniq_kpts[uniq_kptji_id]
        log.debug1('kpt = %s', kpt)
        adapted_ji_idx = numpy.where(uniq_inverse == uniq_kptji_id)[0]
        adapted_kptjs = kptjs[adapted_ji_idx]
        nkptj = len(adapted_kptjs)
        log.debug1('adapted_ji_idx = %s', adapted_ji_idx)

        Gaux = ft_ao.ft_ao(fused_cell, Gv, None, b, gxyz, Gvbase, kpt).T
        Gaux = fuse(Gaux)
        Gaux *= mydf.weighted_coulG(kpt, False, gs)
        kLR = Gaux.T.real.copy('C')
        kLI = Gaux.T.imag.copy('C')
        j2c = numpy.asarray(feri['j2c/%d' % uniq_kptji_id])
        # Note large difference may be found in results between the CD/eig treatments.
        # In some systems, small integral errors can lead to different treatments of
        # linear dependency which can be observed in the total energy/orbital energy
        # around 4th decimal place.
        #        try:
        #            j2c = scipy.linalg.cholesky(j2c, lower=True)
        #            j2ctag = 'CD'
        #        except scipy.linalg.LinAlgError as e:
        #
        # Abandon CD treatment for better numerical stablity
        w, v = scipy.linalg.eigh(j2c)
        log.debug('MDF metric for kpt %s cond = %.4g, drop %d bfns',
                  uniq_kptji_id, w[-1] / w[0],
                  numpy.count_nonzero(w < mydf.linear_dep_threshold))
        v = v[:, w > mydf.linear_dep_threshold].T.conj()
        v /= numpy.sqrt(w[w > mydf.linear_dep_threshold]).reshape(-1, 1)
        j2c = v
        j2ctag = 'eig'
        naux0 = j2c.shape[0]

        if is_zero(kpt):  # kpti == kptj
            aosym = 's2'
            nao_pair = nao * (nao + 1) // 2

            vbar = fuse(mydf.auxbar(fused_cell))
            ovlp = cell.pbc_intor('int1e_ovlp_sph',
                                  hermi=1,
                                  kpts=adapted_kptjs)
            for k, ji in enumerate(adapted_ji_idx):
                ovlp[k] = lib.pack_tril(ovlp[k])
        else:
            aosym = 's1'
            nao_pair = nao**2

        mem_now = lib.current_memory()[0]
        log.debug2('memory = %s', mem_now)
        max_memory = max(2000, mydf.max_memory - mem_now)
        # nkptj for 3c-coulomb arrays plus 1 Lpq array
        buflen = min(
            max(int(max_memory * .6 * 1e6 / 16 / naux / (nkptj + 1)), 1),
            nao_pair)
        shranges = _guess_shell_ranges(cell, buflen, aosym)
        buflen = max([x[2] for x in shranges])
        # +1 for a pqkbuf
        if aosym == 's2':
            Gblksize = max(
                16, int(max_memory * .2 * 1e6 / 16 / buflen / (nkptj + 1)))
        else:
            Gblksize = max(
                16, int(max_memory * .4 * 1e6 / 16 / buflen / (nkptj + 1)))
        Gblksize = min(Gblksize, ngs, 16384)
        pqkRbuf = numpy.empty(buflen * Gblksize)
        pqkIbuf = numpy.empty(buflen * Gblksize)
        # buf for ft_aopair
        buf = numpy.empty((nkptj, buflen * Gblksize), dtype=numpy.complex128)

        col1 = 0
        for istep, sh_range in enumerate(shranges):
            log.debug1('int3c2e [%d/%d], AO [%d:%d], ncol = %d', \
                       istep+1, len(shranges), *sh_range)
            bstart, bend, ncol = sh_range
            col0, col1 = col1, col1 + ncol
            j3cR = []
            j3cI = []
            for k, idx in enumerate(adapted_ji_idx):
                v = fuse(numpy.asarray(feri['j3c/%d' % idx][:, col0:col1]))
                if is_zero(kpt):
                    for i, c in enumerate(vbar):
                        if c != 0:
                            v[i] -= c * ovlp[k][col0:col1]
                j3cR.append(numpy.asarray(v.real, order='C'))
                if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
                    j3cI.append(None)
                else:
                    j3cI.append(numpy.asarray(v.imag, order='C'))
                v = None

            if aosym == 's2':
                shls_slice = (bstart, bend, 0, bend)
            else:
                shls_slice = (bstart, bend, 0, cell.nbas)
            for p0, p1 in lib.prange(0, ngs, Gblksize):
                dat = ft_ao._ft_aopair_kpts(cell,
                                            Gv[p0:p1],
                                            shls_slice,
                                            aosym,
                                            b,
                                            gxyz[p0:p1],
                                            Gvbase,
                                            kpt,
                                            adapted_kptjs,
                                            out=buf)
                nG = p1 - p0
                for k, ji in enumerate(adapted_ji_idx):
                    aoao = dat[k].reshape(nG, ncol)
                    pqkR = numpy.ndarray((ncol, nG), buffer=pqkRbuf)
                    pqkI = numpy.ndarray((ncol, nG), buffer=pqkIbuf)
                    pqkR[:] = aoao.real.T
                    pqkI[:] = aoao.imag.T

                    lib.dot(kLR[p0:p1].T, pqkR.T, -1, j3cR[k], 1)
                    lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k], 1)
                    if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])):
                        lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k], 1)
                        lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k], 1)

            for k, ji in enumerate(adapted_ji_idx):
                if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
                    v = j3cR[k]
                else:
                    v = j3cR[k] + j3cI[k] * 1j
                if j2ctag == 'CD':
                    v = scipy.linalg.solve_triangular(j2c,
                                                      v,
                                                      lower=True,
                                                      overwrite_b=True)
                else:
                    v = lib.dot(j2c, v)
                feri['j3c/%d' % ji][:naux0, col0:col1] = v

        del (feri['j2c/%d' % uniq_kptji_id])
        for k, ji in enumerate(adapted_ji_idx):
            v = feri['j3c/%d' % ji][:naux0]
            del (feri['j3c/%d' % ji])
            feri['j3c/%d' % ji] = v

    for k, kpt in enumerate(uniq_kpts):
        make_kpt(k)

    feri.close()
Exemplo n.º 5
0
def _make_j3c(mydf, cell, auxcell, kptij_lst):
    t1 = (time.clock(), time.time())
    log = logger.Logger(mydf.stdout, mydf.verbose)
    max_memory = max(2000, mydf.max_memory - lib.current_memory()[0])
    fused_cell, fuse = fuse_auxcell(mydf, auxcell)
    outcore.aux_e2(cell,
                   fused_cell,
                   mydf._cderi,
                   'cint3c2e_sph',
                   kptij_lst=kptij_lst,
                   dataname='j3c',
                   max_memory=max_memory)
    t1 = log.timer_debug1('3c2e', *t1)

    nao = cell.nao_nr()
    naux = auxcell.nao_nr()
    gs = mydf.gs
    Gv, Gvbase, kws = cell.get_Gv_weights(gs)
    b = cell.reciprocal_vectors()
    gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
    ngs = gxyz.shape[0]

    kptis = kptij_lst[:, 0]
    kptjs = kptij_lst[:, 1]
    kpt_ji = kptjs - kptis
    uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
    # j2c ~ (-kpt_ji | kpt_ji)
    j2c = fused_cell.pbc_intor('cint2c2e_sph', hermi=1, kpts=uniq_kpts)
    kLRs = []
    kLIs = []

    # An alternative method to evalute j2c. This method might have larger numerical error?
    #    chgcell = make_modchg_basis(auxcell, mydf.eta)
    #    for k, kpt in enumerate(uniq_kpts):
    #        aoaux = ft_ao.ft_ao(chgcell, Gv, None, b, gxyz, Gvbase, kpt).T
    #        coulG = numpy.sqrt(mydf.weighted_coulG(kpt, False, gs))
    #        LkR = aoaux.real * coulG
    #        LkI = aoaux.imag * coulG
    #        j2caux = numpy.zeros_like(j2c[k])
    #        j2caux[naux:,naux:] = j2c[k][naux:,naux:]
    #        if is_zero(kpt):  # kpti == kptj
    #            j2caux[naux:,naux:] -= lib.ddot(LkR, LkR.T)
    #            j2caux[naux:,naux:] -= lib.ddot(LkI, LkI.T)
    #            j2c[k] = j2c[k][:naux,:naux] - fuse(fuse(j2caux.T).T)
    #            vbar = fuse(mydf.auxbar(fused_cell))
    #            s = (vbar != 0).astype(numpy.double)
    #            j2c[k] -= numpy.einsum('i,j->ij', vbar, s)
    #            j2c[k] -= numpy.einsum('i,j->ij', s, vbar)
    #        else:
    #            j2cR, j2cI = zdotCN(LkR, LkI, LkR.T, LkI.T)
    #            j2caux[naux:,naux:] -= j2cR + j2cI * 1j
    #            j2c[k] = j2c[k][:naux,:naux] - fuse(fuse(j2caux.T).T)
    #
    #        try:
    #            j2c[k] = scipy.linalg.cholesky(j2c[k], lower=True)
    #        except scipy.linalg.LinAlgError as e:
    #            msg =('===================================\n'
    #                  'J-metric not positive definite.\n'
    #                  'It is likely that gs is not enough.\n'
    #                  '===================================')
    #            log.error(msg)
    #            raise scipy.linalg.LinAlgError('\n'.join([e.message, msg]))
    #        kLR = LkR.T
    #        kLI = LkI.T
    #        if not kLR.flags.c_contiguous: kLR = lib.transpose(LkR)
    #        if not kLI.flags.c_contiguous: kLI = lib.transpose(LkI)
    #        kLR *= coulG.reshape(-1,1)
    #        kLI *= coulG.reshape(-1,1)
    #        kLRs.append(kLR)
    #        kLIs.append(kLI)
    #        aoaux = LkR = LkI = kLR = kLI = coulG = None

    for k, kpt in enumerate(uniq_kpts):
        aoaux = ft_ao.ft_ao(fused_cell, Gv, None, b, gxyz, Gvbase, kpt).T
        coulG = numpy.sqrt(mydf.weighted_coulG(kpt, False, gs))
        LkR = aoaux.real * coulG
        LkI = aoaux.imag * coulG

        if is_zero(kpt):  # kpti == kptj
            j2c[k][naux:] -= lib.ddot(LkR[naux:], LkR.T)
            j2c[k][naux:] -= lib.ddot(LkI[naux:], LkI.T)
            j2c[k][:naux, naux:] = j2c[k][naux:, :naux].T
        else:
            j2cR, j2cI = zdotCN(LkR[naux:], LkI[naux:], LkR.T, LkI.T)
            j2c[k][naux:] -= j2cR + j2cI * 1j
            j2c[k][:naux, naux:] = j2c[k][naux:, :naux].T.conj()

        j2c[k] = fuse(fuse(j2c[k]).T).T
        try:
            j2c[k] = ('CD', scipy.linalg.cholesky(j2c[k], lower=True))
        except scipy.linalg.LinAlgError as e:
            #msg =('===================================\n'
            #      'J-metric not positive definite.\n'
            #      'It is likely that gs is not enough.\n'
            #      '===================================')
            #log.error(msg)
            #raise scipy.linalg.LinAlgError('\n'.join([e.message, msg]))
            w, v = scipy.linalg.eigh(j2c[k])
            log.debug2('metric linear dependency for kpt %s', k)
            log.debug2('cond = %.4g, drop %d bfns', w[0] / w[-1],
                       numpy.count_nonzero(w < LINEAR_DEP_THR))
            v = v[:, w > LINEAR_DEP_THR].T.conj()
            v /= numpy.sqrt(w[w > LINEAR_DEP_THR]).reshape(-1, 1)
            j2c[k] = ('eig', v)
        kLR = LkR[naux:].T
        kLI = LkI[naux:].T
        if not kLR.flags.c_contiguous: kLR = lib.transpose(LkR[naux:])
        if not kLI.flags.c_contiguous: kLI = lib.transpose(LkI[naux:])
        kLR *= coulG.reshape(-1, 1)
        kLI *= coulG.reshape(-1, 1)
        kLRs.append(kLR)
        kLIs.append(kLI)
        aoaux = LkR = LkI = kLR = kLI = coulG = None

    nauxs = [v[1].shape[0] for v in j2c]
    feri = h5py.File(mydf._cderi)

    def make_kpt(uniq_kptji_id):  # kpt = kptj - kpti
        kpt = uniq_kpts[uniq_kptji_id]
        log.debug1('kpt = %s', kpt)
        adapted_ji_idx = numpy.where(uniq_inverse == uniq_kptji_id)[0]
        adapted_kptjs = kptjs[adapted_ji_idx]
        nkptj = len(adapted_kptjs)
        log.debug1('adapted_ji_idx = %s', adapted_ji_idx)
        kLR = kLRs[uniq_kptji_id]
        kLI = kLIs[uniq_kptji_id]

        if is_zero(kpt):  # kpti == kptj
            aosym = 's2'
            nao_pair = nao * (nao + 1) // 2

            vbar = fuse(mydf.auxbar(fused_cell))
            ovlp = cell.pbc_intor('cint1e_ovlp_sph',
                                  hermi=1,
                                  kpts=adapted_kptjs)
            for k, ji in enumerate(adapted_ji_idx):
                ovlp[k] = lib.pack_tril(ovlp[k])
        else:
            aosym = 's1'
            nao_pair = nao**2

        mem_now = lib.current_memory()[0]
        log.debug2('memory = %s', mem_now)
        max_memory = max(2000, mydf.max_memory - mem_now)
        # nkptj for 3c-coulomb arrays plus 1 Lpq array
        buflen = min(
            max(int(max_memory * .6 * 1e6 / 16 / naux / (nkptj + 1)), 1),
            nao_pair)
        shranges = _guess_shell_ranges(cell, buflen, aosym)
        buflen = max([x[2] for x in shranges])
        # +1 for a pqkbuf
        if aosym == 's2':
            Gblksize = max(
                16, int(max_memory * .2 * 1e6 / 16 / buflen / (nkptj + 1)))
        else:
            Gblksize = max(
                16, int(max_memory * .4 * 1e6 / 16 / buflen / (nkptj + 1)))
        Gblksize = min(Gblksize, ngs, 16384)
        pqkRbuf = numpy.empty(buflen * Gblksize)
        pqkIbuf = numpy.empty(buflen * Gblksize)
        # buf for ft_aopair
        buf = numpy.zeros((nkptj, buflen * Gblksize), dtype=numpy.complex128)

        col1 = 0
        for istep, sh_range in enumerate(shranges):
            log.debug1('int3c2e [%d/%d], AO [%d:%d], ncol = %d', \
                       istep+1, len(shranges), *sh_range)
            bstart, bend, ncol = sh_range
            col0, col1 = col1, col1 + ncol
            j3cR = []
            j3cI = []
            for k, idx in enumerate(adapted_ji_idx):
                v = numpy.asarray(feri['j3c/%d' % idx][:, col0:col1])
                if is_zero(kpt):
                    for i, c in enumerate(vbar):
                        if c != 0:
                            v[i] -= c * ovlp[k][col0:col1]
                j3cR.append(numpy.asarray(v.real, order='C'))
                if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
                    j3cI.append(None)
                else:
                    j3cI.append(numpy.asarray(v.imag, order='C'))
            v = None

            if aosym == 's2':
                shls_slice = (bstart, bend, 0, bend)
                for p0, p1 in lib.prange(0, ngs, Gblksize):
                    ft_ao._ft_aopair_kpts(cell,
                                          Gv[p0:p1],
                                          shls_slice,
                                          aosym,
                                          b,
                                          gxyz[p0:p1],
                                          Gvbase,
                                          kpt,
                                          adapted_kptjs,
                                          out=buf)
                    nG = p1 - p0
                    for k, ji in enumerate(adapted_ji_idx):
                        aoao = numpy.ndarray((nG, ncol),
                                             dtype=numpy.complex128,
                                             order='F',
                                             buffer=buf[k])
                        pqkR = numpy.ndarray((ncol, nG), buffer=pqkRbuf)
                        pqkI = numpy.ndarray((ncol, nG), buffer=pqkIbuf)
                        pqkR[:] = aoao.real.T
                        pqkI[:] = aoao.imag.T
                        aoao[:] = 0
                        lib.dot(kLR[p0:p1].T, pqkR.T, -1, j3cR[k][naux:], 1)
                        lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k][naux:], 1)
                        if not (is_zero(kpt)
                                and gamma_point(adapted_kptjs[k])):
                            lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k][naux:],
                                    1)
                            lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k][naux:], 1)
            else:
                shls_slice = (bstart, bend, 0, cell.nbas)
                ni = ncol // nao
                for p0, p1 in lib.prange(0, ngs, Gblksize):
                    ft_ao._ft_aopair_kpts(cell,
                                          Gv[p0:p1],
                                          shls_slice,
                                          aosym,
                                          b,
                                          gxyz[p0:p1],
                                          Gvbase,
                                          kpt,
                                          adapted_kptjs,
                                          out=buf)
                    nG = p1 - p0
                    for k, ji in enumerate(adapted_ji_idx):
                        aoao = numpy.ndarray((nG, ni, nao),
                                             dtype=numpy.complex128,
                                             order='F',
                                             buffer=buf[k])
                        pqkR = numpy.ndarray((ni, nao, nG), buffer=pqkRbuf)
                        pqkI = numpy.ndarray((ni, nao, nG), buffer=pqkIbuf)
                        pqkR[:] = aoao.real.transpose(1, 2, 0)
                        pqkI[:] = aoao.imag.transpose(1, 2, 0)
                        aoao[:] = 0
                        pqkR = pqkR.reshape(-1, nG)
                        pqkI = pqkI.reshape(-1, nG)
                        zdotCN(kLR[p0:p1].T, kLI[p0:p1].T, pqkR.T, pqkI.T, -1,
                               j3cR[k][naux:], j3cI[k][naux:], 1)

            naux0 = nauxs[uniq_kptji_id]
            for k, ji in enumerate(adapted_ji_idx):
                if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
                    v = fuse(j3cR[k])
                else:
                    v = fuse(j3cR[k] + j3cI[k] * 1j)
                if j2c[uniq_kptji_id][0] == 'CD':
                    v = scipy.linalg.solve_triangular(j2c[uniq_kptji_id][1],
                                                      v,
                                                      lower=True,
                                                      overwrite_b=True)
                else:
                    v = lib.dot(j2c[uniq_kptji_id][1], v)
                feri['j3c/%d' % ji][:naux0, col0:col1] = v

        naux0 = nauxs[uniq_kptji_id]
        for k, ji in enumerate(adapted_ji_idx):
            v = feri['j3c/%d' % ji][:naux0]
            del (feri['j3c/%d' % ji])
            feri['j3c/%d' % ji] = v

    for k, kpt in enumerate(uniq_kpts):
        make_kpt(k)

    feri.close()
Exemplo n.º 6
0
Arquivo: mdf.py Projeto: lichen5/pyscf
def _make_j3c(mydf, cell, auxcell, kptij_lst):
    t1 = (time.clock(), time.time())
    log = logger.Logger(mydf.stdout, mydf.verbose)
    max_memory = max(2000, mydf.max_memory - lib.current_memory()[0])
    fused_cell, fuse = fuse_auxcell_(mydf, mydf.auxcell)
    outcore.aux_e2(cell,
                   fused_cell,
                   mydf._cderi,
                   'cint3c2e_sph',
                   kptij_lst=kptij_lst,
                   dataname='j3c',
                   max_memory=max_memory)
    t1 = log.timer_debug1('3c2e', *t1)

    nao = cell.nao_nr()
    naux = auxcell.nao_nr()
    gs = mydf.gs
    gxyz = lib.cartesian_prod(
        (numpy.append(range(gs[0] + 1), range(-gs[0], 0)),
         numpy.append(range(gs[1] + 1), range(-gs[1], 0)),
         numpy.append(range(gs[2] + 1), range(-gs[2], 0))))
    invh = numpy.linalg.inv(cell._h)
    Gv = 2 * numpy.pi * numpy.dot(gxyz, invh)
    ngs = gxyz.shape[0]

    kptis = kptij_lst[:, 0]
    kptjs = kptij_lst[:, 1]
    kpt_ji = kptjs - kptis
    uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
    # j2c ~ (-kpt_ji | kpt_ji)
    j2c = fused_cell.pbc_intor('cint2c2e_sph', hermi=1, kpts=uniq_kpts)
    kLRs = []
    kLIs = []
    for k, kpt in enumerate(uniq_kpts):
        aoaux = ft_ao.ft_ao(fused_cell, Gv, None, invh, gxyz, gs, kpt).T
        aoaux = fuse(aoaux)
        coulG = numpy.sqrt(tools.get_coulG(cell, kpt, gs=gs) / cell.vol)
        kLR = (aoaux.real * coulG).T
        kLI = (aoaux.imag * coulG).T
        if not kLR.flags.c_contiguous: kLR = lib.transpose(kLR.T)
        if not kLI.flags.c_contiguous: kLI = lib.transpose(kLI.T)

        j2c[k] = fuse(fuse(j2c[k]).T).T.copy()
        if is_zero(kpt):  # kpti == kptj
            j2c[k] -= lib.dot(kLR.T, kLR)
            j2c[k] -= lib.dot(kLI.T, kLI)
        else:
            # aoaux ~ kpt_ij, aoaux.conj() ~ kpt_kl
            j2cR, j2cI = zdotCN(kLR.T, kLI.T, kLR, kLI)
            j2c[k] -= j2cR + j2cI * 1j

        kLR *= coulG.reshape(-1, 1)
        kLI *= coulG.reshape(-1, 1)
        kLRs.append(kLR)
        kLIs.append(kLI)
        aoaux = kLR = kLI = j2cR = j2cI = coulG = None

    feri = h5py.File(mydf._cderi)

    # Expand approx Lpq for aosym='s1'.  The approx Lpq are all in aosym='s2' mode
    if mydf.approx_sr_level > 0 and len(kptij_lst) > 1:
        Lpq_fake = _fake_Lpq_kpts(mydf, feri, naux, nao)

    def save(label, dat, col0, col1):
        nrow = dat.shape[0]
        feri[label][:nrow, col0:col1] = dat

    def make_kpt(uniq_kptji_id):  # kpt = kptj - kpti
        kpt = uniq_kpts[uniq_kptji_id]
        log.debug1('kpt = %s', kpt)
        adapted_ji_idx = numpy.where(uniq_inverse == uniq_kptji_id)[0]
        adapted_kptjs = kptjs[adapted_ji_idx]
        nkptj = len(adapted_kptjs)
        log.debug1('adapted_ji_idx = %s', adapted_ji_idx)
        kLR = kLRs[uniq_kptji_id]
        kLI = kLIs[uniq_kptji_id]

        if is_zero(kpt):  # kpti == kptj
            aosym = 's2'
            nao_pair = nao * (nao + 1) // 2

            vbar = fuse(mydf.auxbar(fused_cell))
            ovlp = cell.pbc_intor('cint1e_ovlp_sph',
                                  hermi=1,
                                  kpts=adapted_kptjs)
            for k, ji in enumerate(adapted_ji_idx):
                ovlp[k] = lib.pack_tril(ovlp[k])
        else:
            aosym = 's1'
            nao_pair = nao**2

        max_memory = max(2000, mydf.max_memory - lib.current_memory()[0])
        # nkptj for 3c-coulomb arrays plus 1 Lpq array
        buflen = min(
            max(int(max_memory * .6 * 1e6 / 16 / naux / (nkptj + 1)), 1),
            nao_pair)
        shranges = pyscf.df.outcore._guess_shell_ranges(cell, buflen, aosym)
        buflen = max([x[2] for x in shranges])
        # +1 for a pqkbuf
        if aosym == 's2':
            Gblksize = max(
                16, int(max_memory * .2 * 1e6 / 16 / buflen / (nkptj + 1)))
        else:
            Gblksize = max(
                16, int(max_memory * .4 * 1e6 / 16 / buflen / (nkptj + 1)))
        Gblksize = min(Gblksize, ngs)
        pqkRbuf = numpy.empty(buflen * Gblksize)
        pqkIbuf = numpy.empty(buflen * Gblksize)
        # buf for ft_aopair
        buf = numpy.zeros((nkptj, buflen * Gblksize), dtype=numpy.complex128)

        col1 = 0
        for istep, sh_range in enumerate(shranges):
            log.debug1('int3c2e [%d/%d], AO [%d:%d], ncol = %d', \
                       istep+1, len(shranges), *sh_range)
            bstart, bend, ncol = sh_range
            col0, col1 = col1, col1 + ncol
            j3cR = []
            j3cI = []
            for k, idx in enumerate(adapted_ji_idx):
                v = fuse(numpy.asarray(feri['j3c/%d' % idx][:, col0:col1]))

                if mydf.approx_sr_level == 0:
                    Lpq = numpy.asarray(feri['Lpq/%d' % idx][:, col0:col1])
                elif aosym == 's2':
                    Lpq = numpy.asarray(feri['Lpq/0'][:, col0:col1])
                else:
                    Lpq = numpy.asarray(Lpq_fake[:, col0:col1])
                lib.dot(j2c[uniq_kptji_id], Lpq, -.5, v, 1)
                if is_zero(kpt):
                    for i, c in enumerate(vbar):
                        if c != 0:
                            v[i] -= c * ovlp[k][col0:col1]

                j3cR.append(numpy.asarray(v.real, order='C'))
                if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
                    j3cI.append(None)
                else:
                    j3cI.append(numpy.asarray(v.imag, order='C'))
            v = Lpq = None

            if aosym == 's2':
                shls_slice = (bstart, bend, 0, bend)
                for p0, p1 in lib.prange(0, ngs, Gblksize):
                    ft_ao._ft_aopair_kpts(cell,
                                          Gv[p0:p1],
                                          shls_slice,
                                          aosym,
                                          invh,
                                          gxyz[p0:p1],
                                          gs,
                                          kpt,
                                          adapted_kptjs,
                                          out=buf)
                    nG = p1 - p0
                    for k, ji in enumerate(adapted_ji_idx):
                        aoao = numpy.ndarray((nG, ncol),
                                             dtype=numpy.complex128,
                                             order='F',
                                             buffer=buf[k])
                        pqkR = numpy.ndarray((ncol, nG), buffer=pqkRbuf)
                        pqkI = numpy.ndarray((ncol, nG), buffer=pqkIbuf)
                        pqkR[:] = aoao.real.T
                        pqkI[:] = aoao.imag.T
                        aoao[:] = 0
                        lib.dot(kLR[p0:p1].T, pqkR.T, -1, j3cR[k], 1)
                        lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k], 1)
                        if not (is_zero(kpt)
                                and gamma_point(adapted_kptjs[k])):
                            lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k], 1)
                            lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k], 1)
            else:
                shls_slice = (bstart, bend, 0, cell.nbas)
                ni = ncol // nao
                for p0, p1 in lib.prange(0, ngs, Gblksize):
                    ft_ao._ft_aopair_kpts(cell,
                                          Gv[p0:p1],
                                          shls_slice,
                                          aosym,
                                          invh,
                                          gxyz[p0:p1],
                                          gs,
                                          kpt,
                                          adapted_kptjs,
                                          out=buf)
                    nG = p1 - p0
                    for k, ji in enumerate(adapted_ji_idx):
                        aoao = numpy.ndarray((nG, ni, nao),
                                             dtype=numpy.complex128,
                                             order='F',
                                             buffer=buf[k])
                        pqkR = numpy.ndarray((ni, nao, nG), buffer=pqkRbuf)
                        pqkI = numpy.ndarray((ni, nao, nG), buffer=pqkIbuf)
                        pqkR[:] = aoao.real.transpose(1, 2, 0)
                        pqkI[:] = aoao.imag.transpose(1, 2, 0)
                        aoao[:] = 0
                        pqkR = pqkR.reshape(-1, nG)
                        pqkI = pqkI.reshape(-1, nG)
                        zdotCN(kLR[p0:p1].T, kLI[p0:p1].T, pqkR.T, pqkI.T, -1,
                               j3cR[k], j3cI[k], 1)

            for k, ji in enumerate(adapted_ji_idx):
                if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
                    save('j3c/%d' % ji, j3cR[k], col0, col1)
                else:
                    save('j3c/%d' % ji, j3cR[k] + j3cI[k] * 1j, col0, col1)

    for k, kpt in enumerate(uniq_kpts):
        make_kpt(k)

    feri.close()
Exemplo n.º 7
0
Arquivo: mdf.py Projeto: armunoz/pyscf
def _make_j3c(mydf, cell, auxcell, kptij_lst):
    t1 = (time.clock(), time.time())
    log = logger.Logger(mydf.stdout, mydf.verbose)
    max_memory = max(2000, mydf.max_memory-lib.current_memory()[0])
    fused_cell, fuse = fuse_auxcell(mydf, mydf.auxcell)

    nao = cell.nao_nr()
    naux = auxcell.nao_nr()
    gs = mydf.gs
    Gv, Gvbase, kws = cell.get_Gv_weights(gs)
    b = cell.reciprocal_vectors()
    gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
    ngs = gxyz.shape[0]

    kptis = kptij_lst[:,0]
    kptjs = kptij_lst[:,1]
    kpt_ji = kptjs - kptis
    uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
    # j2c ~ (-kpt_ji | kpt_ji)
    j2c = fused_cell.pbc_intor('cint2c2e_sph', hermi=1, kpts=uniq_kpts)
    kLRs = []
    kLIs = []
    for k, kpt in enumerate(uniq_kpts):
        aoaux = ft_ao.ft_ao(fused_cell, Gv, None, b, gxyz, Gvbase, kpt).T
        aoaux = fuse(aoaux)
        coulG = numpy.sqrt(mydf.weighted_coulG(kpt, False, gs))
        kLR = (aoaux.real * coulG).T
        kLI = (aoaux.imag * coulG).T
        if not kLR.flags.c_contiguous: kLR = lib.transpose(kLR.T)
        if not kLI.flags.c_contiguous: kLI = lib.transpose(kLI.T)

        j2c[k] = fuse(fuse(j2c[k]).T).T.copy()
        if is_zero(kpt):  # kpti == kptj
            j2c[k] -= lib.dot(kLR.T, kLR)
            j2c[k] -= lib.dot(kLI.T, kLI)
        else:
             # aoaux ~ kpt_ij, aoaux.conj() ~ kpt_kl
            j2cR, j2cI = zdotCN(kLR.T, kLI.T, kLR, kLI)
            j2c[k] -= j2cR + j2cI * 1j

        w, v = scipy.linalg.eigh(j2c[k])
        log.debug('MDF metric for kpt %s cond = %.4g, drop %d bfns',
                  k, w[0]/w[-1], numpy.count_nonzero(w<df.LINEAR_DEP_THR))
        v = v[:,w>df.LINEAR_DEP_THR].T.conj()
        v /= numpy.sqrt(w[w>df.LINEAR_DEP_THR]).reshape(-1,1)
        j2c[k] = ('eig', v)

        kLR *= coulG.reshape(-1,1)
        kLI *= coulG.reshape(-1,1)
        kLRs.append(kLR)
        kLIs.append(kLI)
        aoaux = kLR = kLI = j2cR = j2cI = coulG = None

    outcore.aux_e2(cell, fused_cell, mydf._cderi, 'cint3c2e_sph',
                   kptij_lst=kptij_lst, dataname='j3c', max_memory=max_memory)
    t1 = log.timer_debug1('3c2e', *t1)
    nauxs = [v[1].shape[0] for v in j2c]
    feri = h5py.File(mydf._cderi)

    def make_kpt(uniq_kptji_id):  # kpt = kptj - kpti
        kpt = uniq_kpts[uniq_kptji_id]
        log.debug1('kpt = %s', kpt)
        adapted_ji_idx = numpy.where(uniq_inverse == uniq_kptji_id)[0]
        adapted_kptjs = kptjs[adapted_ji_idx]
        nkptj = len(adapted_kptjs)
        log.debug1('adapted_ji_idx = %s', adapted_ji_idx)
        kLR = kLRs[uniq_kptji_id]
        kLI = kLIs[uniq_kptji_id]

        if is_zero(kpt):  # kpti == kptj
            aosym = 's2'
            nao_pair = nao*(nao+1)//2

            vbar = fuse(mydf.auxbar(fused_cell))
            ovlp = cell.pbc_intor('cint1e_ovlp_sph', hermi=1, kpts=adapted_kptjs)
            for k, ji in enumerate(adapted_ji_idx):
                ovlp[k] = lib.pack_tril(ovlp[k])
        else:
            aosym = 's1'
            nao_pair = nao**2

        mem_now = lib.current_memory()[0]
        log.debug2('memory = %s', mem_now)
        max_memory = max(2000, mydf.max_memory-mem_now)
        # nkptj for 3c-coulomb arrays plus 1 Lpq array
        buflen = min(max(int(max_memory*.6*1e6/16/naux/(nkptj+1)), 1), nao_pair)
        shranges = _guess_shell_ranges(cell, buflen, aosym)
        buflen = max([x[2] for x in shranges])
        # +1 for a pqkbuf
        if aosym == 's2':
            Gblksize = max(16, int(max_memory*.2*1e6/16/buflen/(nkptj+1)))
        else:
            Gblksize = max(16, int(max_memory*.4*1e6/16/buflen/(nkptj+1)))
        Gblksize = min(Gblksize, ngs, 16384)
        pqkRbuf = numpy.empty(buflen*Gblksize)
        pqkIbuf = numpy.empty(buflen*Gblksize)
        # buf for ft_aopair
        buf = numpy.zeros((nkptj,buflen*Gblksize), dtype=numpy.complex128)

        col1 = 0
        for istep, sh_range in enumerate(shranges):
            log.debug1('int3c2e [%d/%d], AO [%d:%d], ncol = %d', \
                       istep+1, len(shranges), *sh_range)
            bstart, bend, ncol = sh_range
            col0, col1 = col1, col1+ncol
            j3cR = []
            j3cI = []
            for k, idx in enumerate(adapted_ji_idx):
                v = fuse(numpy.asarray(feri['j3c/%d'%idx][:,col0:col1]))
                if is_zero(kpt):
                    for i, c in enumerate(vbar):
                        if c != 0:
                            v[i] -= c * ovlp[k][col0:col1]
                j3cR.append(numpy.asarray(v.real, order='C'))
                if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
                    j3cI.append(None)
                else:
                    j3cI.append(numpy.asarray(v.imag, order='C'))
                v = None

            if aosym == 's2':
                shls_slice = (bstart, bend, 0, bend)
                for p0, p1 in lib.prange(0, ngs, Gblksize):
                    ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym,
                                          b, gxyz[p0:p1], Gvbase, kpt,
                                          adapted_kptjs, out=buf)
                    nG = p1 - p0
                    for k, ji in enumerate(adapted_ji_idx):
                        aoao = numpy.ndarray((nG,ncol), dtype=numpy.complex128,
                                             order='F', buffer=buf[k])
                        pqkR = numpy.ndarray((ncol,nG), buffer=pqkRbuf)
                        pqkI = numpy.ndarray((ncol,nG), buffer=pqkIbuf)
                        pqkR[:] = aoao.real.T
                        pqkI[:] = aoao.imag.T
                        aoao[:] = 0
                        lib.dot(kLR[p0:p1].T, pqkR.T, -1, j3cR[k], 1)
                        lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k], 1)
                        if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])):
                            lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k], 1)
                            lib.dot(kLI[p0:p1].T, pqkR.T,  1, j3cI[k], 1)
            else:
                shls_slice = (bstart, bend, 0, cell.nbas)
                ni = ncol // nao
                for p0, p1 in lib.prange(0, ngs, Gblksize):
                    ft_ao._ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym,
                                          b, gxyz[p0:p1], Gvbase, kpt,
                                          adapted_kptjs, out=buf)
                    nG = p1 - p0
                    for k, ji in enumerate(adapted_ji_idx):
                        aoao = numpy.ndarray((nG,ni,nao), dtype=numpy.complex128,
                                             order='F', buffer=buf[k])
                        pqkR = numpy.ndarray((ni,nao,nG), buffer=pqkRbuf)
                        pqkI = numpy.ndarray((ni,nao,nG), buffer=pqkIbuf)
                        pqkR[:] = aoao.real.transpose(1,2,0)
                        pqkI[:] = aoao.imag.transpose(1,2,0)
                        aoao[:] = 0
                        pqkR = pqkR.reshape(-1,nG)
                        pqkI = pqkI.reshape(-1,nG)
                        zdotCN(kLR[p0:p1].T, kLI[p0:p1].T, pqkR.T, pqkI.T,
                               -1, j3cR[k], j3cI[k], 1)

            naux0 = nauxs[uniq_kptji_id]
            for k, ji in enumerate(adapted_ji_idx):
                if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
                    v = j3cR[k]
                else:
                    v = j3cR[k] + j3cI[k] * 1j
                if j2c[uniq_kptji_id][0] == 'CD':
                    v = scipy.linalg.solve_triangular(j2c[uniq_kptji_id][1], v,
                                                      lower=True, overwrite_b=True)
                else:
                    v = lib.dot(j2c[uniq_kptji_id][1], v)
                feri['j3c/%d'%ji][:naux0,col0:col1] = v

        naux0 = nauxs[uniq_kptji_id]
        for k, ji in enumerate(adapted_ji_idx):
            v = feri['j3c/%d'%ji][:naux0]
            del(feri['j3c/%d'%ji])
            feri['j3c/%d'%ji] = v

    for k, kpt in enumerate(uniq_kpts):
        make_kpt(k)

    feri.close()
Exemplo n.º 8
0
def build_Lpq_pbc(mydf, auxcell, kptij_lst):
    '''Fitting coefficients for auxiliary functions'''
    kptis = kptij_lst[:,0]
    kptjs = kptij_lst[:,1]
    kpts_ji = kptjs - kptis
    uniq_kpts, uniq_index, uniq_inverse = unique(kpts_ji)
    max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]))
    if mydf.metric.upper() == 'S':
        outcore.aux_e2(mydf.cell, auxcell, mydf._cderi, 'cint3c1e_sph',
                       kptij_lst=kptij_lst, dataname='Lpq',
                       max_memory=max_memory)
        s_aux = auxcell.pbc_intor('cint1e_ovlp_sph', hermi=1, kpts=uniq_kpts)
    elif mydf.metric.upper() == 'T':
        outcore.aux_e2(mydf.cell, auxcell, mydf._cderi, 'cint3c1e_p2_sph',
                       kptij_lst=kptij_lst, dataname='Lpq', max_memory=max_memory)
        s_aux = [x*2 for x in auxcell.pbc_intor('cint1e_kin_sph', hermi=1, kpts=uniq_kpts)]
    elif mydf.metric.upper() == 'J':
        fused_cell, fuse = fuse_auxcell(mydf, auxcell)
        outcore.aux_e2(mydf.cell, fused_cell, mydf._cderi, 'cint3c2e_sph',
                       kptij_lst=kptij_lst, dataname='j3c', max_memory=max_memory)
        vbar = fuse(mydf.auxbar(fused_cell))
        with h5py.File(mydf._cderi) as f:
            f['Lpq-kptij'] = kptij_lst
            for k_uniq, kpt_uniq in enumerate(uniq_kpts):
                adapted_ji_idx = numpy.where(uniq_inverse == k_uniq)[0]
                adapted_kptjs = kptjs[adapted_ji_idx]
                if is_zero(kpt_uniq):
                    ovlp = mydf.cell.pbc_intor('cint1e_ovlp_sph', hermi=1, kpts=adapted_kptjs)
                    for k, ji in enumerate(adapted_ji_idx):
                        ovlp[k] = lib.pack_tril(ovlp[k])
                for k, idx in enumerate(adapted_ji_idx):
                    v = fuse(numpy.asarray(f['j3c/%d'%idx]))
                    if is_zero(kpt_uniq):
                        for i, c in enumerate(vbar):
                            if c != 0:
                                v[i] -= c * ovlp[k]
                    f['Lpq/%d'%idx] = v
        v = ovlp = vbar = None

        j2c = fused_cell.pbc_intor('cint2c2e_sph', hermi=1, kpts=uniq_kpts)
        for k, kpt in enumerate(uniq_kpts):
            j2c[k] = fuse(fuse(j2c[k]).T).T.copy()
        s_aux = j2c

#    else: # T+S
#        outcore.aux_e2(mydf.cell, auxcell, mydf._cderi, 'cint3c1e_sph',
#                       kptij_lst=kptij_lst, dataname='Lpq_s',
#                       max_memory=max_memory)
#        outcore.aux_e2(mydf.cell, auxcell, mydf._cderi, 'cint3c1e_p2_sph',
#                       kptij_lst=kptij_lst, dataname='Lpq',
#                       max_memory=max_memory)
#        with h5py.File(mydf._cderi) as f:
#            for k in range(len(kptij_lst)):
#                f['Lpq/%d'%k][:] = f['Lpq/%d'%k].value + f['Lpq_s/%d'%k].value
#                del(f['Lpq_s/%d'%k])
#        s_aux = auxcell.pbc_intor('cint1e_ovlp_sph', hermi=1, kpts=uniq_kpts)
#        s_aux = [x+y*2 for x,y in zip(s_aux, auxcell.pbc_intor('cint1e_kin_sph', hermi=1, kpts=uniq_kpts))]

    try:
        s_aux = [scipy.linalg.cho_factor(x) for x in s_aux]
    except scipy.linalg.LinAlgError:
        eigs = [scipy.linalg.eigh(x)[0] for x in s_aux]
        conds = [x[-1]/max(1e-16, x[0]) for x in eigs]
        n = eigs[0].size
        shift = [0] * len(s_aux)
        for i, x in enumerate(s_aux):
            if conds[i] > 1e15:
                shift[i] = max(abs(eigs[i][0])*2, eigs[i][-1]*1e-18)
                x += numpy.eye(n) * shift[i]
        logger.warn(mydf, 'Ill condition number %s found in metric %s.\n'
                    'Level shift %s is applied.',
                    conds, mydf.metric, shift)
        s_aux = [scipy.linalg.cho_factor(x) for x in s_aux]

    max_memory = mydf.max_memory - lib.current_memory()[0]
    naux = auxcell.nao_nr()
    blksize = max(int(max_memory*.5*1e6/16/naux/mydf.blockdim), 1) * mydf.blockdim
    with h5py.File(mydf._cderi) as feri:
        for k, where in enumerate(uniq_inverse):
            s_k = s_aux[where]
            key = 'Lpq/%d' % k
            Lpq = feri[key]
            nao_pair = Lpq.shape[1]
            for p0, p1 in lib.prange(0, nao_pair, blksize):
                Lpq[:,p0:p1] = scipy.linalg.cho_solve(s_k, Lpq[:,p0:p1])
Exemplo n.º 9
0
def _make_j3c(mydf, cell, auxcell, kptij_lst, cderi_file):
    t1 = (time.clock(), time.time())
    log = logger.Logger(mydf.stdout, mydf.verbose)
    max_memory = max(2000, mydf.max_memory - lib.current_memory()[0])
    fused_cell, fuse = fuse_auxcell(mydf, auxcell)
    outcore.aux_e2(cell,
                   fused_cell,
                   cderi_file,
                   'int3c2e',
                   aosym='s2',
                   kptij_lst=kptij_lst,
                   dataname='j3c',
                   max_memory=max_memory)
    t1 = log.timer_debug1('3c2e', *t1)

    nao = cell.nao_nr()
    naux = auxcell.nao_nr()
    mesh = mydf.mesh
    Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
    b = cell.reciprocal_vectors()
    gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])
    ngrids = gxyz.shape[0]

    kptis = kptij_lst[:, 0]
    kptjs = kptij_lst[:, 1]
    kpt_ji = kptjs - kptis
    uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
    log.debug('Num uniq kpts %d', len(uniq_kpts))
    log.debug2('uniq_kpts %s', uniq_kpts)
    # j2c ~ (-kpt_ji | kpt_ji)
    j2c = fused_cell.pbc_intor('int2c2e', hermi=1, kpts=uniq_kpts)
    feri = h5py.File(cderi_file)

    # An alternative method to evalute j2c. This method might have larger numerical error?
    #    chgcell = make_modchg_basis(auxcell, mydf.eta)
    #    for k, kpt in enumerate(uniq_kpts):
    #        aoaux = ft_ao.ft_ao(chgcell, Gv, None, b, gxyz, Gvbase, kpt).T
    #        coulG = numpy.sqrt(mydf.weighted_coulG(kpt, False, mesh))
    #        LkR = aoaux.real * coulG
    #        LkI = aoaux.imag * coulG
    #        j2caux = numpy.zeros_like(j2c[k])
    #        j2caux[naux:,naux:] = j2c[k][naux:,naux:]
    #        if is_zero(kpt):  # kpti == kptj
    #            j2caux[naux:,naux:] -= lib.ddot(LkR, LkR.T)
    #            j2caux[naux:,naux:] -= lib.ddot(LkI, LkI.T)
    #            j2c[k] = j2c[k][:naux,:naux] - fuse(fuse(j2caux.T).T)
    #            vbar = fuse(mydf.auxbar(fused_cell))
    #            s = (vbar != 0).astype(numpy.double)
    #            j2c[k] -= numpy.einsum('i,j->ij', vbar, s)
    #            j2c[k] -= numpy.einsum('i,j->ij', s, vbar)
    #        else:
    #            j2cR, j2cI = zdotCN(LkR, LkI, LkR.T, LkI.T)
    #            j2caux[naux:,naux:] -= j2cR + j2cI * 1j
    #            j2c[k] = j2c[k][:naux,:naux] - fuse(fuse(j2caux.T).T)
    #        feri['j2c/%d'%k] = fuse(fuse(j2c[k]).T).T
    #        aoaux = LkR = LkI = coulG = None

    if cell.dimension == 1 or cell.dimension == 2:
        plain_ints = _gaussian_int(fused_cell)

    max_memory = max(2000, mydf.max_memory - lib.current_memory()[0])
    blksize = max(2048, int(max_memory * .5e6 / 16 / fused_cell.nao_nr()))
    log.debug2('max_memory %s (MB)  blocksize %s', max_memory, blksize)
    for k, kpt in enumerate(uniq_kpts):
        coulG = numpy.sqrt(mydf.weighted_coulG(kpt, False, mesh))
        for p0, p1 in lib.prange(0, ngrids, blksize):
            aoaux = ft_ao.ft_ao(fused_cell, Gv[p0:p1], None, b, gxyz[p0:p1],
                                Gvbase, kpt)
            if (cell.dimension == 1 or cell.dimension == 2) and is_zero(kpt):
                G0idx, SI_on_z = pbcgto.cell._SI_for_uniform_model_charge(
                    cell, Gv[p0:p1])
                aoaux[G0idx] -= numpy.einsum('g,i->gi', SI_on_z, plain_ints)

            aoaux = aoaux.T
            LkR = aoaux.real * coulG[p0:p1]
            LkI = aoaux.imag * coulG[p0:p1]
            aoaux = None

            if is_zero(kpt):  # kpti == kptj
                j2c[k][naux:] -= lib.ddot(LkR[naux:], LkR.T)
                j2c[k][naux:] -= lib.ddot(LkI[naux:], LkI.T)
                j2c[k][:naux, naux:] = j2c[k][naux:, :naux].T
            else:
                j2cR, j2cI = zdotCN(LkR[naux:], LkI[naux:], LkR.T, LkI.T)
                j2c[k][naux:] -= j2cR + j2cI * 1j
                j2c[k][:naux, naux:] = j2c[k][naux:, :naux].T.conj()
            LkR = LkI = None
        feri['j2c/%d' % k] = fuse(fuse(j2c[k]).T).T
    j2c = coulG = None

    def make_kpt(uniq_kptji_id):  # kpt = kptj - kpti
        kpt = uniq_kpts[uniq_kptji_id]
        log.debug1('kpt = %s', kpt)
        adapted_ji_idx = numpy.where(uniq_inverse == uniq_kptji_id)[0]
        adapted_kptjs = kptjs[adapted_ji_idx]
        nkptj = len(adapted_kptjs)
        log.debug1('adapted_ji_idx = %s', adapted_ji_idx)

        shls_slice = (auxcell.nbas, fused_cell.nbas)
        Gaux = ft_ao.ft_ao(fused_cell, Gv, shls_slice, b, gxyz, Gvbase, kpt)
        if (cell.dimension == 1 or cell.dimension == 2) and is_zero(kpt):
            G0idx, SI_on_z = pbcgto.cell._SI_for_uniform_model_charge(cell, Gv)
            s = plain_ints[-Gaux.shape[1]:]  # Only compensated Gaussians
            Gaux[G0idx] -= numpy.einsum('g,i->gi', SI_on_z, s)

        wcoulG = mydf.weighted_coulG(kpt, False, mesh)
        Gaux *= wcoulG.reshape(-1, 1)
        kLR = Gaux.real.copy('C')
        kLI = Gaux.imag.copy('C')
        Gaux = None
        j2c = numpy.asarray(feri['j2c/%d' % uniq_kptji_id])
        try:
            j2c = scipy.linalg.cholesky(j2c, lower=True)
            j2ctag = 'CD'
        except scipy.linalg.LinAlgError as e:
            #msg =('===================================\n'
            #      'J-metric not positive definite.\n'
            #      'It is likely that mesh is not enough.\n'
            #      '===================================')
            #log.error(msg)
            #raise scipy.linalg.LinAlgError('\n'.join([e.message, msg]))
            w, v = scipy.linalg.eigh(j2c)
            log.debug('DF metric linear dependency for kpt %s', uniq_kptji_id)
            log.debug('cond = %.4g, drop %d bfns', w[-1] / w[0],
                      numpy.count_nonzero(w < mydf.linear_dep_threshold))
            v = v[:, w > mydf.linear_dep_threshold].T.conj()
            v /= numpy.sqrt(w[w > mydf.linear_dep_threshold]).reshape(-1, 1)
            j2c = v
            j2ctag = 'eig'
        naux0 = j2c.shape[0]

        if is_zero(kpt):  # kpti == kptj
            aosym = 's2'
            nao_pair = nao * (nao + 1) // 2

            vbar = mydf.auxbar(fused_cell)
            ovlp = cell.pbc_intor('int1e_ovlp', hermi=1, kpts=adapted_kptjs)
            ovlp = [lib.pack_tril(s) for s in ovlp]
        else:
            aosym = 's1'
            nao_pair = nao**2

        mem_now = lib.current_memory()[0]
        log.debug2('memory = %s', mem_now)
        max_memory = max(2000, mydf.max_memory - mem_now)
        # nkptj for 3c-coulomb arrays plus 1 Lpq array
        buflen = min(
            max(int(max_memory * .6 * 1e6 / 16 / naux / (nkptj + 1)), 1),
            nao_pair)
        shranges = _guess_shell_ranges(cell, buflen, aosym)
        buflen = max([x[2] for x in shranges])
        # +1 for a pqkbuf
        if aosym == 's2':
            Gblksize = max(
                16, int(max_memory * .2 * 1e6 / 16 / buflen / (nkptj + 1)))
        else:
            Gblksize = max(
                16, int(max_memory * .4 * 1e6 / 16 / buflen / (nkptj + 1)))
        Gblksize = min(Gblksize, ngrids, 16384)
        pqkRbuf = numpy.empty(buflen * Gblksize)
        pqkIbuf = numpy.empty(buflen * Gblksize)
        # buf for ft_aopair
        buf = numpy.empty(nkptj * buflen * Gblksize, dtype=numpy.complex128)

        col1 = 0
        for istep, sh_range in enumerate(shranges):
            log.debug1('int3c2e [%d/%d], AO [%d:%d], ncol = %d', \
                       istep+1, len(shranges), *sh_range)
            bstart, bend, ncol = sh_range
            col0, col1 = col1, col1 + ncol
            j3cR = []
            j3cI = []
            for k, idx in enumerate(adapted_ji_idx):
                v = numpy.asarray(feri['j3c/%d' % idx][:, col0:col1])
                if is_zero(kpt) and cell.dimension == 3:
                    for i in numpy.where(vbar != 0)[0]:
                        v[i] -= vbar[i] * ovlp[k][col0:col1]
                j3cR.append(numpy.asarray(v.real, order='C'))
                if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
                    j3cI.append(None)
                else:
                    j3cI.append(numpy.asarray(v.imag, order='C'))
            v = None

            if aosym == 's2':
                shls_slice = (bstart, bend, 0, bend)
            else:
                shls_slice = (bstart, bend, 0, cell.nbas)
            for p0, p1 in lib.prange(0, ngrids, Gblksize):
                dat = ft_ao._ft_aopair_kpts(cell,
                                            Gv[p0:p1],
                                            shls_slice,
                                            aosym,
                                            b,
                                            gxyz[p0:p1],
                                            Gvbase,
                                            kpt,
                                            adapted_kptjs,
                                            out=buf)

                if (cell.dimension == 1
                        or cell.dimension == 2) and is_zero(kpt):
                    G0idx, SI_on_z = pbcgto.cell._SI_for_uniform_model_charge(
                        cell, Gv[p0:p1])
                    if SI_on_z.size > 0:
                        for k, aoao in enumerate(dat):
                            aoao[G0idx] -= numpy.einsum(
                                'g,i->gi', SI_on_z, ovlp[k])
                            aux = fuse(
                                ft_ao.ft_ao(fused_cell, Gv[p0:p1][G0idx]).T)
                            vG_mod = numpy.einsum('ig,g,g->i', aux.conj(),
                                                  wcoulG[p0:p1][G0idx],
                                                  SI_on_z)
                            if gamma_point(adapted_kptjs[k]):
                                j3cR[k][:naux] -= vG_mod[:,
                                                         None].real * ovlp[k]
                            else:
                                tmp = vG_mod[:, None] * ovlp[k]
                                j3cR[k][:naux] -= tmp.real
                                j3cI[k][:naux] -= tmp.imag
                            tmp = aux = vG_mod

                nG = p1 - p0
                for k, ji in enumerate(adapted_ji_idx):
                    aoao = dat[k].reshape(nG, ncol)
                    pqkR = numpy.ndarray((ncol, nG), buffer=pqkRbuf)
                    pqkI = numpy.ndarray((ncol, nG), buffer=pqkIbuf)
                    pqkR[:] = aoao.real.T
                    pqkI[:] = aoao.imag.T

                    lib.dot(kLR[p0:p1].T, pqkR.T, -1, j3cR[k][naux:], 1)
                    lib.dot(kLI[p0:p1].T, pqkI.T, -1, j3cR[k][naux:], 1)
                    if not (is_zero(kpt) and gamma_point(adapted_kptjs[k])):
                        lib.dot(kLR[p0:p1].T, pqkI.T, -1, j3cI[k][naux:], 1)
                        lib.dot(kLI[p0:p1].T, pqkR.T, 1, j3cI[k][naux:], 1)

            for k, ji in enumerate(adapted_ji_idx):
                if is_zero(kpt) and gamma_point(adapted_kptjs[k]):
                    v = fuse(j3cR[k])
                else:
                    v = fuse(j3cR[k] + j3cI[k] * 1j)
                if j2ctag == 'CD':
                    v = scipy.linalg.solve_triangular(j2c,
                                                      v,
                                                      lower=True,
                                                      overwrite_b=True)
                else:
                    v = lib.dot(j2c, v)
                feri['j3c/%d' % ji][:naux0, col0:col1] = v

        del (feri['j2c/%d' % uniq_kptji_id])
        for k, ji in enumerate(adapted_ji_idx):
            v = feri['j3c/%d' % ji][:naux0]
            del (feri['j3c/%d' % ji])
            feri['j3c/%d' % ji] = v

    for k, kpt in enumerate(uniq_kpts):
        make_kpt(k)

    feri.close()