Пример #1
0
def grids2d_int3c_jobs(cell, auxcell, kptij_lst, chunks, aosym_s2):
    ao_loc = cell.ao_loc_nr()
    if aosym_s2:
        segs = ao_loc[1:]*(ao_loc[1:]+1)//2 - ao_loc[:-1]*(ao_loc[:-1]+1)//2
        ij_ranges = balance_segs(segs, chunks[0]*chunks[1])
    else:
        segs = ao_loc[1:]-ao_loc[:-1]
        ij_ranges = balance_segs(segs, chunks[0])

    jobs = [(job_id, i0, i1) for job_id, (i0, i1, x) in enumerate(ij_ranges)]
    return jobs
Пример #2
0
def grids2d_int3c_jobs(cell, auxcell, kptij_lst, chunks, aosym_s2):
    ao_loc = cell.ao_loc_nr()
    if aosym_s2:
        segs = ao_loc[1:] * (ao_loc[1:] +
                             1) // 2 - ao_loc[:-1] * (ao_loc[:-1] + 1) // 2
        ij_ranges = balance_segs(segs, chunks[0] * chunks[1])
    else:
        segs = ao_loc[1:] - ao_loc[:-1]
        ij_ranges = balance_segs(segs, chunks[0])

    jobs = [(job_id, i0, i1) for job_id, (i0, i1, x) in enumerate(ij_ranges)]
    return jobs
Пример #3
0
    def gen_int3c(auxcell, job_id, ish0, ish1):
        dataname = 'j3c-chunks/%d' % job_id
        if dataname in feri:
            del (feri[dataname])

        i0 = ao_loc[ish0]
        i1 = ao_loc[ish1]
        dii = i1 * (i1 + 1) // 2 - i0 * (i0 + 1) // 2
        dij = (i1 - i0) * nao
        if j_only:
            buflen = max(8, int(max_memory * 1e6 / 16 / (nkptij * dii + dii)))
        else:
            buflen = max(8, int(max_memory * 1e6 / 16 / (nkptij * dij + dij)))
        auxranges = balance_segs(aux_loc[1:] - aux_loc[:-1], buflen)
        buflen = max([x[2] for x in auxranges])
        buf = numpy.empty(nkptij * dij * buflen, dtype=dtype)
        buf1 = numpy.empty(dij * buflen, dtype=dtype)

        naux = aux_loc[-1]
        for kpt_id, kptij in enumerate(kptij_lst):
            key = '%s/%d' % (dataname, kpt_id)
            if aosym_s2[kpt_id]:
                shape = (naux, dii)
            else:
                shape = (naux, dij)
            if gamma_point(kptij):
                feri.create_dataset(key, shape, 'f8')
            else:
                feri.create_dataset(key, shape, 'c16')

        naux0 = 0
        for istep, auxrange in enumerate(auxranges):
            log.alldebug2("aux_e2 job_id %d step %d", job_id, istep)
            sh0, sh1, nrow = auxrange
            sub_slice = (ish0, ish1, 0, cell.nbas, sh0, sh1)
            if j_only:
                mat = numpy.ndarray((nkptij, dii, nrow),
                                    dtype=dtype,
                                    buffer=buf)
            else:
                mat = numpy.ndarray((nkptij, dij, nrow),
                                    dtype=dtype,
                                    buffer=buf)
            mat = int3c(sub_slice, mat)

            for k, kptij in enumerate(kptij_lst):
                h5dat = feri['%s/%d' % (dataname, k)]
                v = lib.transpose(mat[k], out=buf1)
                if not j_only and aosym_s2[k]:
                    idy = idxb[i0 * (i0 + 1) // 2:i1 *
                               (i1 + 1) // 2] - i0 * nao
                    out = numpy.ndarray((nrow, dii),
                                        dtype=v.dtype,
                                        buffer=mat[k])
                    v = numpy.take(v, idy, axis=1, out=out)
                if gamma_point(kptij):
                    h5dat[naux0:naux0 + nrow] = v.real
                else:
                    h5dat[naux0:naux0 + nrow] = v
            naux0 += nrow
Пример #4
0
def _guess_shell_ranges(mol, buflen, aosym):
    from pyscf.ao2mo.outcore import balance_segs
    ao_loc = mol.ao_loc_nr()
    nao = ao_loc[-1]
    if 's2' in aosym:
        segs = ao_loc[1:]*(ao_loc[1:]+1)//2 - ao_loc[:-1]*(ao_loc[:-1]+1)//2
    else:
        segs = (ao_loc[1:]-ao_loc[:-1])*nao
    return balance_segs(segs, buflen)
Пример #5
0
def _guess_shell_ranges(mol, buflen, aosym):
    from pyscf.ao2mo.outcore import balance_segs
    ao_loc = mol.ao_loc_nr()
    nao = ao_loc[-1]
    if 's2' in aosym:
        segs = ao_loc[1:] * (ao_loc[1:] +
                             1) // 2 - ao_loc[:-1] * (ao_loc[:-1] + 1) // 2
    else:
        segs = (ao_loc[1:] - ao_loc[:-1]) * nao
    return balance_segs(segs, buflen)
Пример #6
0
    def gen_int3c(job_id, ish0, ish1):
        dataname = 'j3c-chunks/%d' % job_id
        i0 = ao_loc[ish0]
        i1 = ao_loc[ish1]
        dii = i1*(i1+1)//2 - i0*(i0+1)//2
        dij = (i1 - i0) * nao
        if j_only:
            buflen = max(8, int(max_memory*1e6/16/(nkptij*dii+dii)))
        else:
            buflen = max(8, int(max_memory*1e6/16/(nkptij*dij+dij)))
        auxranges = balance_segs(aux_loc[1:]-aux_loc[:-1], buflen)
        buflen = max([x[2] for x in auxranges])
        buf = numpy.empty(nkptij*dij*buflen, dtype=dtype)
        buf1 = numpy.empty(dij*buflen, dtype=dtype)

        naux = aux_loc[-1]
        for kpt_id, kptij in enumerate(kptij_lst):
            key = '%s/%d' % (dataname, kpt_id)
            if aosym_s2[kpt_id]:
                shape = (naux, dii)
            else:
                shape = (naux, dij)
            if gamma_point(kptij):
                fswap.create_dataset(key, shape, 'f8')
            else:
                fswap.create_dataset(key, shape, 'c16')

        naux0 = 0
        for istep, auxrange in enumerate(auxranges):
            log.alldebug2("aux_e1 job_id %d step %d", job_id, istep)
            sh0, sh1, nrow = auxrange
            sub_slice = (ish0, ish1, 0, cell.nbas, sh0, sh1)
            if j_only:
                mat = numpy.ndarray((nkptij,dii,nrow), dtype=dtype, buffer=buf)
            else:
                mat = numpy.ndarray((nkptij,dij,nrow), dtype=dtype, buffer=buf)
            mat = int3c(sub_slice, mat)

            for k, kptij in enumerate(kptij_lst):
                h5dat = fswap['%s/%d'%(dataname,k)]
                v = lib.transpose(mat[k], out=buf1)
                if not j_only and aosym_s2[k]:
                    idy = idxb[i0*(i0+1)//2:i1*(i1+1)//2] - i0 * nao
                    out = numpy.ndarray((nrow,dii), dtype=v.dtype, buffer=mat[k])
                    v = numpy.take(v, idy, axis=1, out=out)
                if gamma_point(kptij):
                    h5dat[naux0:naux0+nrow] = v.real
                else:
                    h5dat[naux0:naux0+nrow] = v
            naux0 += nrow
Пример #7
0
def aux_e1(cell,
           auxcell,
           erifile,
           intor='int3c2e',
           aosym='s2ij',
           comp=None,
           kptij_lst=None,
           dataname='eri_mo',
           shls_slice=None,
           max_memory=2000,
           verbose=0):
    r'''3-center AO integrals (L|ij) with double lattice sum:
    \sum_{lm} (L[0]|i[l]j[m]), where L is the auxiliary basis.
    Three-index integral tensor (kptij_idx, naux, nao_pair) or four-index
    integral tensor (kptij_idx, comp, naux, nao_pair) are stored on disk.

    Args:
        kptij_lst : (*,2,3) array
            A list of (kpti, kptj)
    '''
    intor, comp = gto.moleintor._get_intor_and_comp(cell._add_suffix(intor),
                                                    comp)

    if isinstance(erifile, h5py.Group):
        feri = erifile
    elif h5py.is_hdf5(erifile):
        feri = h5py.File(erifile, 'a')
    else:
        feri = h5py.File(erifile, 'w')
    if dataname in feri:
        del (feri[dataname])
    if dataname + '-kptij' in feri:
        del (feri[dataname + '-kptij'])

    if kptij_lst is None:
        kptij_lst = numpy.zeros((1, 2, 3))
    feri[dataname + '-kptij'] = kptij_lst

    if shls_slice is None:
        shls_slice = (0, cell.nbas, 0, cell.nbas, 0, auxcell.nbas)

    ao_loc = cell.ao_loc_nr()
    aux_loc = auxcell.ao_loc_nr(auxcell.cart
                                or 'ssc' in intor)[:shls_slice[5] + 1]
    ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
    nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
    naux = aux_loc[shls_slice[5]] - aux_loc[shls_slice[4]]
    nkptij = len(kptij_lst)

    nii = (ao_loc[shls_slice[1]] * (ao_loc[shls_slice[1]] + 1) // 2 -
           ao_loc[shls_slice[0]] * (ao_loc[shls_slice[0]] + 1) // 2)
    nij = ni * nj

    kpti = kptij_lst[:, 0]
    kptj = kptij_lst[:, 1]
    aosym_ks2 = abs(kpti - kptj).sum(axis=1) < KPT_DIFF_TOL
    j_only = numpy.all(aosym_ks2)
    #aosym_ks2 &= (aosym[:2] == 's2' and shls_slice[:2] == shls_slice[2:4])
    aosym_ks2 &= aosym[:2] == 's2'
    for k, kptij in enumerate(kptij_lst):
        key = '%s/%d' % (dataname, k)
        if gamma_point(kptij):
            dtype = 'f8'
        else:
            dtype = 'c16'
        if aosym_ks2[k]:
            nao_pair = nii
        else:
            nao_pair = nij
        if comp == 1:
            shape = (naux, nao_pair)
        else:
            shape = (comp, naux, nao_pair)
        feri.create_dataset(key, shape, dtype)
    if naux == 0:
        feri.close()
        return erifile

    if j_only and aosym[:2] == 's2':
        assert (shls_slice[2] == 0)
        nao_pair = nii
    else:
        nao_pair = nij

    if gamma_point(kptij_lst):
        dtype = numpy.double
    else:
        dtype = numpy.complex128

    buflen = max(8, int(max_memory * 1e6 / 16 / (nkptij * ni * nj * comp)))
    auxdims = aux_loc[shls_slice[4] + 1:shls_slice[5] +
                      1] - aux_loc[shls_slice[4]:shls_slice[5]]
    auxranges = balance_segs(auxdims, buflen)
    buflen = max([x[2] for x in auxranges])
    buf = numpy.empty(nkptij * comp * ni * nj * buflen, dtype=dtype)
    buf1 = numpy.empty(ni * nj * buflen, dtype=dtype)

    int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst)

    naux0 = 0
    for istep, auxrange in enumerate(auxranges):
        sh0, sh1, nrow = auxrange
        sub_slice = (shls_slice[0], shls_slice[1], shls_slice[2],
                     shls_slice[3], shls_slice[4] + sh0, shls_slice[4] + sh1)
        mat = numpy.ndarray((nkptij, comp, nao_pair, nrow),
                            dtype=dtype,
                            buffer=buf)
        mat = int3c(sub_slice, mat)

        for k, kptij in enumerate(kptij_lst):
            h5dat = feri['%s/%d' % (dataname, k)]
            for icomp, v in enumerate(mat[k]):
                v = lib.transpose(v, out=buf1)
                if gamma_point(kptij):
                    v = v.real
                if aosym_ks2[k] and v.shape[1] == ni**2:
                    v = lib.pack_tril(v.reshape(-1, ni, ni))
                if comp == 1:
                    h5dat[naux0:naux0 + nrow] = v
                else:
                    h5dat[icomp, naux0:naux0 + nrow] = v
        naux0 += nrow

    if not isinstance(erifile, h5py.Group):
        feri.close()
    return erifile
Пример #8
0
def _aux_e2(cell,
            auxcell,
            erifile,
            intor='int3c2e',
            aosym='s2ij',
            comp=None,
            kptij_lst=None,
            dataname='eri_mo',
            shls_slice=None,
            max_memory=2000,
            verbose=0):
    r'''3-center AO integrals (ij|L) with double lattice sum:
    \sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis.
    Three-index integral tensor (kptij_idx, nao_pair, naux) or four-index
    integral tensor (kptij_idx, comp, nao_pair, naux) are stored on disk.

    **This function should be only used by df and mdf initialization function
    _make_j3c**

    Args:
        kptij_lst : (*,2,3) array
            A list of (kpti, kptj)
    '''
    intor, comp = gto.moleintor._get_intor_and_comp(cell._add_suffix(intor),
                                                    comp)

    if isinstance(erifile, h5py.Group):
        feri = erifile
    elif h5py.is_hdf5(erifile):
        feri = h5py.File(erifile, 'a')
    else:
        feri = h5py.File(erifile, 'w')
    if dataname in feri:
        del (feri[dataname])
    if dataname + '-kptij' in feri:
        del (feri[dataname + '-kptij'])

    if kptij_lst is None:
        kptij_lst = numpy.zeros((1, 2, 3))
    feri[dataname + '-kptij'] = kptij_lst

    if shls_slice is None:
        shls_slice = (0, cell.nbas, 0, cell.nbas, 0, auxcell.nbas)

    ao_loc = cell.ao_loc_nr()
    aux_loc = auxcell.ao_loc_nr(auxcell.cart
                                or 'ssc' in intor)[:shls_slice[5] + 1]
    ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
    nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
    naux = aux_loc[shls_slice[5]] - aux_loc[shls_slice[4]]
    nkptij = len(kptij_lst)

    nii = (ao_loc[shls_slice[1]] * (ao_loc[shls_slice[1]] + 1) // 2 -
           ao_loc[shls_slice[0]] * (ao_loc[shls_slice[0]] + 1) // 2)
    nij = ni * nj

    kpti = kptij_lst[:, 0]
    kptj = kptij_lst[:, 1]
    aosym_ks2 = abs(kpti - kptj).sum(axis=1) < KPT_DIFF_TOL
    j_only = numpy.all(aosym_ks2)
    #aosym_ks2 &= (aosym[:2] == 's2' and shls_slice[:2] == shls_slice[2:4])
    aosym_ks2 &= aosym[:2] == 's2'

    if j_only and aosym[:2] == 's2':
        assert (shls_slice[2] == 0)
        nao_pair = nii
    else:
        nao_pair = nij

    if gamma_point(kptij_lst):
        dtype = numpy.double
    else:
        dtype = numpy.complex128

    buflen = max(8, int(max_memory * .47e6 / 16 / (nkptij * ni * nj * comp)))
    auxdims = aux_loc[shls_slice[4] + 1:shls_slice[5] +
                      1] - aux_loc[shls_slice[4]:shls_slice[5]]
    auxranges = balance_segs(auxdims, buflen)
    buflen = max([x[2] for x in auxranges])
    buf = numpy.empty(nkptij * comp * ni * nj * buflen, dtype=dtype)
    buf1 = numpy.empty_like(buf)

    int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst)

    kptis = kptij_lst[:, 0]
    kptjs = kptij_lst[:, 1]
    kpt_ji = kptjs - kptis
    uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
    # sorted_ij_idx: Sort and group the kptij_lst according to the ordering in
    # df._make_j3c to reduce the data fragment in the hdf5 file.  When datasets
    # are written to hdf5, they are saved sequentially. If the integral data are
    # saved as the order of kptij_lst, removing the datasets in df._make_j3c will
    # lead to holes that can not be reused.
    sorted_ij_idx = numpy.hstack(
        [numpy.where(uniq_inverse == k)[0] for k, kpt in enumerate(uniq_kpts)])
    tril_idx = numpy.tril_indices(ni)
    tril_idx = tril_idx[0] * ni + tril_idx[1]

    def save(istep, mat):
        for k in sorted_ij_idx:
            v = mat[k]
            if gamma_point(kptij_lst[k]):
                v = v.real
            if aosym_ks2[k] and nao_pair == ni**2:
                v = v[:, tril_idx]
            feri['%s/%d/%d' % (dataname, k, istep)] = v

    with lib.call_in_background(save) as bsave:
        for istep, auxrange in enumerate(auxranges):
            sh0, sh1, nrow = auxrange
            sub_slice = (shls_slice[0], shls_slice[1], shls_slice[2],
                         shls_slice[3], shls_slice[4] + sh0,
                         shls_slice[4] + sh1)
            mat = numpy.ndarray((nkptij, comp, nao_pair, nrow),
                                dtype=dtype,
                                buffer=buf)
            bsave(istep, int3c(sub_slice, mat))
            buf, buf1 = buf1, buf

    if not isinstance(erifile, h5py.Group):
        feri.close()
    return erifile
Пример #9
0
def aux_e2(cell,
           auxcell,
           erifile,
           intor='cint3c2e_sph',
           aosym='s1',
           comp=1,
           kptij_lst=None,
           dataname='eri_mo',
           max_memory=2000,
           verbose=0):
    '''3-center AO integrals (ij|L) with double lattice sum:
    \sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis.
    On diks, the integrals are stored as (kptij_idx, naux, nao_pair)

    Args:
        kptij_lst : (*,2,3) array
            A list of (kpti, kptj)
    '''
    if comp > 1:
        raise NotImplementedError('comp = %d' % comp)
    if h5py.is_hdf5(erifile):
        feri = h5py.File(erifile)
        if dataname in feri:
            del (feri[dataname])
        if dataname + '-kptij' in feri:
            del (feri[dataname + '-kptij'])
    else:
        feri = h5py.File(erifile, 'w')

    if kptij_lst is None:
        kptij_lst = numpy.zeros((1, 2, 3))
    feri[dataname + '-kptij'] = kptij_lst
    nkptij = len(kptij_lst)

    # sum over largest number of images in either cell or auxcell
    nimgs = numpy.max((cell.nimgs, auxcell.nimgs), axis=0)
    Ls = cell.get_lattice_Ls(nimgs)
    logger.debug1(cell, "pbc.df.outcore.Images %s", nimgs)
    logger.debug3(cell, "Ls = %s", Ls)

    nao = cell.nao_nr()
    #naux = auxcell.nao_nr('ssc' in intor)
    naux = auxcell.nao_nr()
    aosym_s2 = numpy.zeros(nkptij, dtype=bool)
    for k, kptij in enumerate(kptij_lst):
        key = '%s/%d' % (dataname, k)
        if abs(kptij).sum() < 1e-9:  # gamma_point:
            dtype = 'f8'
        else:
            dtype = 'c16'
        aosym_s2[k] = abs(kptij[0] - kptij[1]).sum() < 1e-9
        if aosym_s2[k]:
            nao_pair = nao * (nao + 1) // 2
        else:
            nao_pair = nao * nao
        if comp == 1:
            shape = (naux, nao_pair)
        else:
            shape = (comp, naux, nao_pair)
        chunks = (min(256, naux), min(256, nao_pair))  # 512 KB
        feri.create_dataset(key, shape, dtype, chunks=chunks)
    if naux == 0:
        feri.close()
        return erifile

    aux_loc = auxcell.ao_loc_nr('ssc' in intor)
    buflen = max(8, int(max_memory * 1e6 / 16 / (nkptij * nao**2 * comp)))
    auxranges = balance_segs(aux_loc[1:] - aux_loc[:-1], buflen)
    buflen = max([x[2] for x in auxranges])
    buf = [
        numpy.zeros(nao * nao * buflen * comp, dtype=numpy.complex128)
        for k in range(nkptij)
    ]
    ints = incore._wrap_int3c(cell, auxcell, intor, comp, Ls, buf)
    atm, bas, env = ints._envs[:3]

    xyz = numpy.asarray(cell.atom_coords(), order='C')
    ptr_coordL = atm[:cell.natm, PTR_COORD]
    ptr_coordL = numpy.vstack(
        (ptr_coordL, ptr_coordL + 1, ptr_coordL + 2)).T.copy('C')
    kpti = kptij_lst[:, 0]
    kptj = kptij_lst[:, 1]

    if numpy.all(aosym_s2):

        def ccsum_or_reorder(Lpq):
            tmp = numpy.asarray(Lpq.transpose(0, 2, 1).conj(), order='C')
            tmp += Lpq
            return tmp
    else:

        def ccsum_or_reorder(Lpq):
            return numpy.asarray(Lpq, order='C')

    naux0 = 0
    for istep, auxrange in enumerate(auxranges):
        sh0, sh1, nrow = auxrange
        c_shls_slice = (ctypes.c_int * 6)(0, cell.nbas, cell.nbas,
                                          cell.nbas * 2, cell.nbas * 2 + sh0,
                                          cell.nbas * 2 + sh1)
        if numpy.all(aosym_s2):
            for l, L1 in enumerate(Ls):
                env[ptr_coordL] = xyz + L1
                e = numpy.dot(
                    Ls[:l + 1] - L1,
                    kptj.T)  # Lattice sum over half of the images {1..l}
                exp_Lk = numpy.exp(1j * numpy.asarray(e, order='C'))
                exp_Lk[l] = .5
                ints(exp_Lk, c_shls_slice)
        else:
            for l, L1 in enumerate(Ls):
                env[ptr_coordL] = xyz + L1
                e = numpy.dot(Ls, kptj.T) - numpy.dot(L1, kpti.T)
                exp_Lk = numpy.exp(1j * numpy.asarray(e, order='C'))
                ints(exp_Lk, c_shls_slice)

        for k, kptij in enumerate(kptij_lst):
            h5dat = feri['%s/%d' % (dataname, k)]
            # transpose 3201 as (comp,L,i,j)
            mat = numpy.ndarray((nao, nao, nrow, comp),
                                order='F',
                                dtype=numpy.complex128,
                                buffer=buf[k])
            for icomp, vi in enumerate(mat.transpose(3, 2, 0, 1)):
                v = ccsum_or_reorder(vi)
                if abs(kptij).sum() < 1e-9:  # gamma_point:
                    v = v.real
                if aosym_s2[k]:
                    v = lib.pack_tril(v)
                else:
                    v = v.reshape(nrow, -1)
                if comp == 1:
                    h5dat[naux0:naux0 + nrow] = v
                else:
                    h5dat[icomp, naux0:naux0 + nrow] = v
            mat[:] = 0
        naux0 += nrow

    feri.close()
    return erifile
Пример #10
0
def aux_e1(cell, auxcell, erifile, intor='int3c2e', aosym='s2ij', comp=None,
           kptij_lst=None, dataname='eri_mo', shls_slice=None, max_memory=2000,
           verbose=0):
    r'''3-center AO integrals (L|ij) with double lattice sum:
    \sum_{lm} (L[0]|i[l]j[m]), where L is the auxiliary basis.
    Three-index integral tensor (kptij_idx, naux, nao_pair) or four-index
    integral tensor (kptij_idx, comp, naux, nao_pair) are stored on disk.

    Args:
        kptij_lst : (*,2,3) array
            A list of (kpti, kptj)
    '''
    intor, comp = gto.moleintor._get_intor_and_comp(cell._add_suffix(intor), comp)

    if isinstance(erifile, h5py.Group):
        feri = erifile
    elif h5py.is_hdf5(erifile):
        feri = h5py.File(erifile)
    else:
        feri = h5py.File(erifile, 'w')
    if dataname in feri:
        del(feri[dataname])
    if dataname+'-kptij' in feri:
        del(feri[dataname+'-kptij'])

    if kptij_lst is None:
        kptij_lst = numpy.zeros((1,2,3))
    feri[dataname+'-kptij'] = kptij_lst

    if shls_slice is None:
        shls_slice = (0, cell.nbas, 0, cell.nbas, 0, auxcell.nbas)

    ao_loc = cell.ao_loc_nr()
    aux_loc = auxcell.ao_loc_nr(auxcell.cart or 'ssc' in intor)[:shls_slice[5]+1]
    ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
    nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
    naux = aux_loc[shls_slice[5]] - aux_loc[shls_slice[4]]
    nkptij = len(kptij_lst)

    nii = (ao_loc[shls_slice[1]]*(ao_loc[shls_slice[1]]+1)//2 -
           ao_loc[shls_slice[0]]*(ao_loc[shls_slice[0]]+1)//2)
    nij = ni * nj

    kpti = kptij_lst[:,0]
    kptj = kptij_lst[:,1]
    aosym_ks2 = abs(kpti-kptj).sum(axis=1) < KPT_DIFF_TOL
    j_only = numpy.all(aosym_ks2)
    #aosym_ks2 &= (aosym[:2] == 's2' and shls_slice[:2] == shls_slice[2:4])
    aosym_ks2 &= aosym[:2] == 's2'
    for k, kptij in enumerate(kptij_lst):
        key = '%s/%d' % (dataname, k)
        if gamma_point(kptij):
            dtype = 'f8'
        else:
            dtype = 'c16'
        if aosym_ks2[k]:
            nao_pair = nii
        else:
            nao_pair = nij
        if comp == 1:
            shape = (naux,nao_pair)
        else:
            shape = (comp,naux,nao_pair)
        feri.create_dataset(key, shape, dtype)
    if naux == 0:
        feri.close()
        return erifile

    if j_only and aosym[:2] == 's2':
        assert(shls_slice[2] == 0)
        nao_pair = nii
    else:
        nao_pair = nij

    if gamma_point(kptij_lst):
        dtype = numpy.double
    else:
        dtype = numpy.complex128

    buflen = max(8, int(max_memory*1e6/16/(nkptij*ni*nj*comp)))
    auxdims = aux_loc[shls_slice[4]+1:shls_slice[5]+1] - aux_loc[shls_slice[4]:shls_slice[5]]
    auxranges = balance_segs(auxdims, buflen)
    buflen = max([x[2] for x in auxranges])
    buf = numpy.empty(nkptij*comp*ni*nj*buflen, dtype=dtype)
    buf1 = numpy.empty(ni*nj*buflen, dtype=dtype)

    int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst)

    naux0 = 0
    for istep, auxrange in enumerate(auxranges):
        sh0, sh1, nrow = auxrange
        sub_slice = (shls_slice[0], shls_slice[1],
                     shls_slice[2], shls_slice[3],
                     shls_slice[4]+sh0, shls_slice[4]+sh1)
        mat = numpy.ndarray((nkptij,comp,nao_pair,nrow), dtype=dtype, buffer=buf)
        mat = int3c(sub_slice, mat)

        for k, kptij in enumerate(kptij_lst):
            h5dat = feri['%s/%d'%(dataname,k)]
            for icomp, v in enumerate(mat[k]):
                v = lib.transpose(v, out=buf1)
                if gamma_point(kptij):
                    v = v.real
                if aosym_ks2[k] and v.shape[1] == ni**2:
                    v = lib.pack_tril(v.reshape(-1,ni,ni))
                if comp == 1:
                    h5dat[naux0:naux0+nrow] = v
                else:
                    h5dat[icomp,naux0:naux0+nrow] = v
        naux0 += nrow

    if not isinstance(erifile, h5py.Group):
        feri.close()
    return erifile
Пример #11
0
def _aux_e2(cell, auxcell, erifile, intor='int3c2e', aosym='s2ij', comp=None,
            kptij_lst=None, dataname='eri_mo', shls_slice=None, max_memory=2000,
            verbose=0):
    r'''3-center AO integrals (ij|L) with double lattice sum:
    \sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis.
    Three-index integral tensor (kptij_idx, nao_pair, naux) or four-index
    integral tensor (kptij_idx, comp, nao_pair, naux) are stored on disk.

    **This function should be only used by df and mdf initialization function
    _make_j3c**

    Args:
        kptij_lst : (*,2,3) array
            A list of (kpti, kptj)
    '''
    intor, comp = gto.moleintor._get_intor_and_comp(cell._add_suffix(intor), comp)

    if isinstance(erifile, h5py.Group):
        feri = erifile
    elif h5py.is_hdf5(erifile):
        feri = h5py.File(erifile)
    else:
        feri = h5py.File(erifile, 'w')
    if dataname in feri:
        del(feri[dataname])
    if dataname+'-kptij' in feri:
        del(feri[dataname+'-kptij'])

    if kptij_lst is None:
        kptij_lst = numpy.zeros((1,2,3))
    feri[dataname+'-kptij'] = kptij_lst

    if shls_slice is None:
        shls_slice = (0, cell.nbas, 0, cell.nbas, 0, auxcell.nbas)

    ao_loc = cell.ao_loc_nr()
    aux_loc = auxcell.ao_loc_nr(auxcell.cart or 'ssc' in intor)[:shls_slice[5]+1]
    ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
    nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
    naux = aux_loc[shls_slice[5]] - aux_loc[shls_slice[4]]
    nkptij = len(kptij_lst)

    nii = (ao_loc[shls_slice[1]]*(ao_loc[shls_slice[1]]+1)//2 -
           ao_loc[shls_slice[0]]*(ao_loc[shls_slice[0]]+1)//2)
    nij = ni * nj

    kpti = kptij_lst[:,0]
    kptj = kptij_lst[:,1]
    aosym_ks2 = abs(kpti-kptj).sum(axis=1) < KPT_DIFF_TOL
    j_only = numpy.all(aosym_ks2)
    #aosym_ks2 &= (aosym[:2] == 's2' and shls_slice[:2] == shls_slice[2:4])
    aosym_ks2 &= aosym[:2] == 's2'

    if j_only and aosym[:2] == 's2':
        assert(shls_slice[2] == 0)
        nao_pair = nii
    else:
        nao_pair = nij

    if gamma_point(kptij_lst):
        dtype = numpy.double
    else:
        dtype = numpy.complex128

    buflen = max(8, int(max_memory*.47e6/16/(nkptij*ni*nj*comp)))
    auxdims = aux_loc[shls_slice[4]+1:shls_slice[5]+1] - aux_loc[shls_slice[4]:shls_slice[5]]
    auxranges = balance_segs(auxdims, buflen)
    buflen = max([x[2] for x in auxranges])
    buf = numpy.empty(nkptij*comp*ni*nj*buflen, dtype=dtype)
    buf1 = numpy.empty_like(buf)

    int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst)

    kptis = kptij_lst[:,0]
    kptjs = kptij_lst[:,1]
    kpt_ji = kptjs - kptis
    uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
# sorted_ij_idx: Sort and group the kptij_lst according to the ordering in
# df._make_j3c to reduce the data fragment in the hdf5 file.  When datasets
# are written to hdf5, they are saved sequentially. If the integral data are
# saved as the order of kptij_lst, removing the datasets in df._make_j3c will
# lead to holes that can not be reused.
    sorted_ij_idx = numpy.hstack([numpy.where(uniq_inverse == k)[0]
                                  for k, kpt in enumerate(uniq_kpts)])
    tril_idx = numpy.tril_indices(ni)
    tril_idx = tril_idx[0] * ni + tril_idx[1]
    def save(istep, mat):
        for k in sorted_ij_idx:
            v = mat[k]
            if gamma_point(kptij_lst[k]):
                v = v.real
            if aosym_ks2[k] and nao_pair == ni**2:
                v = v[:,tril_idx]
            feri['%s/%d/%d' % (dataname,k,istep)] = v

    with lib.call_in_background(save) as bsave:
        for istep, auxrange in enumerate(auxranges):
            sh0, sh1, nrow = auxrange
            sub_slice = (shls_slice[0], shls_slice[1],
                         shls_slice[2], shls_slice[3],
                         shls_slice[4]+sh0, shls_slice[4]+sh1)
            mat = numpy.ndarray((nkptij,comp,nao_pair,nrow), dtype=dtype, buffer=buf)
            bsave(istep, int3c(sub_slice, mat))
            buf, buf1 = buf1, buf

    if not isinstance(erifile, h5py.Group):
        feri.close()
    return erifile
Пример #12
0
def aux_e2(cell, auxcell, erifile, intor='cint3c2e_sph', aosym='s1', comp=1,
           kptij_lst=None, dataname='eri_mo', max_memory=2000, verbose=0):
    '''3-center AO integrals (ij|L) with double lattice sum:
    \sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis.
    On diks, the integrals are stored as (kptij_idx, naux, nao_pair)

    Args:
        kptij_lst : (*,2,3) array
            A list of (kpti, kptj)
    '''
    if comp > 1:
        raise NotImplementedError('comp = %d' % comp)
    if h5py.is_hdf5(erifile):
        feri = h5py.File(erifile)
        if dataname in feri:
            del(feri[dataname])
        if dataname+'-kptij' in feri:
            del(feri[dataname+'-kptij'])
    else:
        feri = h5py.File(erifile, 'w')

    if kptij_lst is None:
        kptij_lst = numpy.zeros((1,2,3))
    feri[dataname+'-kptij'] = kptij_lst
    nkptij = len(kptij_lst)

    # sum over largest number of images in either cell or auxcell
    nimgs = numpy.max((cell.nimgs, auxcell.nimgs), axis=0)
    Ls = cell.get_lattice_Ls(nimgs)
    logger.debug1(cell, "pbc.df.outcore.Images %s", nimgs)
    logger.debug3(cell, "Ls = %s", Ls)

    nao = cell.nao_nr()
    #naux = auxcell.nao_nr('ssc' in intor)
    naux = auxcell.nao_nr()
    aosym_s2 = numpy.zeros(nkptij, dtype=bool)
    for k, kptij in enumerate(kptij_lst):
        key = '%s/%d' % (dataname, k)
        if abs(kptij).sum() < 1e-9:  # gamma_point:
            dtype = 'f8'
        else:
            dtype = 'c16'
        aosym_s2[k] = abs(kptij[0]-kptij[1]).sum() < 1e-9
        if aosym_s2[k]:
            nao_pair = nao * (nao+1) // 2
        else:
            nao_pair = nao * nao
        if comp == 1:
            shape = (naux,nao_pair)
        else:
            shape = (comp,naux,nao_pair)
        chunks = (min(256,naux), min(256,nao_pair))  # 512 KB
        feri.create_dataset(key, shape, dtype, chunks=chunks)
    if naux == 0:
        feri.close()
        return erifile

    aux_loc = auxcell.ao_loc_nr('ssc' in intor)
    buflen = max(8, int(max_memory*1e6/16/(nkptij*nao**2*comp)))
    auxranges = balance_segs(aux_loc[1:]-aux_loc[:-1], buflen)
    buflen = max([x[2] for x in auxranges])
    buf = [numpy.zeros(nao*nao*buflen*comp, dtype=numpy.complex128)
           for k in range(nkptij)]
    ints = incore._wrap_int3c(cell, auxcell, intor, comp, Ls, buf)
    atm, bas, env = ints._envs[:3]

    xyz = numpy.asarray(cell.atom_coords(), order='C')
    ptr_coordL = atm[:cell.natm,PTR_COORD]
    ptr_coordL = numpy.vstack((ptr_coordL,ptr_coordL+1,ptr_coordL+2)).T.copy('C')
    kpti = kptij_lst[:,0]
    kptj = kptij_lst[:,1]

    if numpy.all(aosym_s2):
        def ccsum_or_reorder(Lpq):
            tmp = numpy.asarray(Lpq.transpose(0,2,1).conj(), order='C')
            tmp += Lpq
            return tmp
    else:
        def ccsum_or_reorder(Lpq):
            return numpy.asarray(Lpq, order='C')

    naux0 = 0
    for istep, auxrange in enumerate(auxranges):
        sh0, sh1, nrow = auxrange
        c_shls_slice = (ctypes.c_int*6)(0, cell.nbas, cell.nbas, cell.nbas*2,
                                        cell.nbas*2+sh0, cell.nbas*2+sh1)
        if numpy.all(aosym_s2):
            for l, L1 in enumerate(Ls):
                env[ptr_coordL] = xyz + L1
                e = numpy.dot(Ls[:l+1]-L1, kptj.T)  # Lattice sum over half of the images {1..l}
                exp_Lk = numpy.exp(1j * numpy.asarray(e, order='C'))
                exp_Lk[l] = .5
                ints(exp_Lk, c_shls_slice)
        else:
            for l, L1 in enumerate(Ls):
                env[ptr_coordL] = xyz + L1
                e = numpy.dot(Ls, kptj.T) - numpy.dot(L1, kpti.T)
                exp_Lk = numpy.exp(1j * numpy.asarray(e, order='C'))
                ints(exp_Lk, c_shls_slice)

        for k, kptij in enumerate(kptij_lst):
            h5dat = feri['%s/%d'%(dataname,k)]
            # transpose 3201 as (comp,L,i,j)
            mat = numpy.ndarray((nao,nao,nrow,comp), order='F',
                                dtype=numpy.complex128, buffer=buf[k])
            for icomp, vi in enumerate(mat.transpose(3,2,0,1)):
                v = ccsum_or_reorder(vi)
                if abs(kptij).sum() < 1e-9:  # gamma_point:
                    v = v.real
                if aosym_s2[k]:
                    v = lib.pack_tril(v)
                else:
                    v = v.reshape(nrow,-1)
                if comp == 1:
                    h5dat[naux0:naux0+nrow] = v
                else:
                    h5dat[icomp,naux0:naux0+nrow] = v
            mat[:] = 0
        naux0 += nrow

    feri.close()
    return erifile