Пример #1
0
    def sr_loop(self,
                kpti_kptj=numpy.zeros((2, 3)),
                max_memory=2000,
                compact=True,
                blksize=None):
        '''Short range part'''
        if self._cderi is None:
            self.build()
        cell = self.cell
        kpti, kptj = kpti_kptj
        unpack = is_zero(kpti - kptj) and not compact
        is_real = is_zero(kpti_kptj)
        nao = cell.nao_nr()
        if blksize is None:
            if is_real:
                blksize = max_memory * 1e6 / 8 / (nao**2 * 2)
            else:
                blksize = max_memory * 1e6 / 16 / (nao**2 * 2)
            blksize /= 2  # For prefetch
            blksize = max(16, min(int(blksize), self.blockdim))
            logger.debug3(self, 'max_memory %d MB, blksize %d', max_memory,
                          blksize)

        def load(aux_slice):
            b0, b1 = aux_slice
            if is_real:
                LpqR = numpy.asarray(j3c[b0:b1])
                if unpack:
                    LpqR = lib.unpack_tril(LpqR).reshape(-1, nao**2)
                LpqI = numpy.zeros_like(LpqR)
            else:
                Lpq = numpy.asarray(j3c[b0:b1])
                LpqR = numpy.asarray(Lpq.real, order='C')
                LpqI = numpy.asarray(Lpq.imag, order='C')
                Lpq = None
                if unpack:
                    LpqR = lib.unpack_tril(LpqR).reshape(-1, nao**2)
                    LpqI = lib.unpack_tril(LpqI,
                                           lib.ANTIHERMI).reshape(-1, nao**2)
            return LpqR, LpqI

        with _load3c(self._cderi, 'j3c', kpti_kptj, 'j3c-kptij') as j3c:
            slices = lib.prange(0, j3c.shape[0], blksize)
            for LpqR, LpqI in lib.map_with_prefetch(load, slices):
                yield LpqR, LpqI, 1
                LpqR = LpqI = None

        if cell.dimension == 2 and cell.low_dim_ft_type != 'inf_vacuum':
            # Truncated Coulomb operator is not postive definite. Load the
            # CDERI tensor of negative part.
            with _load3c(self._cderi,
                         'j3c-',
                         kpti_kptj,
                         'j3c-kptij',
                         ignore_key_error=True) as j3c:
                slices = lib.prange(0, j3c.shape[0], blksize)
                for LpqR, LpqI in lib.map_with_prefetch(load, slices):
                    yield LpqR, LpqI, -1
                    LpqR = LpqI = None
Пример #2
0
    def sr_loop(self,
                kpti_kptj=numpy.zeros((2, 3)),
                max_memory=2000,
                compact=True,
                blksize=None):
        '''Short range part'''
        if self._cderi is None:
            self.build()
        kpti, kptj = kpti_kptj
        unpack = is_zero(kpti - kptj) and not compact
        is_real = is_zero(kpti_kptj)
        nao = self.cell.nao_nr()
        if blksize is None:
            if is_real:
                if unpack:
                    blksize = max_memory * 1e6 / 8 / (nao *
                                                      (nao + 1) // 2 + nao**2)
                else:
                    blksize = max_memory * 1e6 / 8 / (nao * (nao + 1))
            else:
                blksize = max_memory * 1e6 / 16 / (nao**2 * 2)
            blksize = max(16, min(int(blksize), self.blockdim))
            logger.debug3(self, 'max_memory %d MB, blksize %d', max_memory,
                          blksize)

        if unpack:
            buf = numpy.empty((blksize, nao * (nao + 1) // 2))

        def load(Lpq, b0, b1, bufR, bufI):
            Lpq = numpy.asarray(Lpq[b0:b1])
            if is_real:
                if unpack:
                    LpqR = lib.unpack_tril(Lpq, out=bufR).reshape(-1, nao**2)
                else:
                    LpqR = Lpq
                LpqI = numpy.zeros_like(LpqR)
            else:
                shape = Lpq.shape
                if unpack:
                    tmp = numpy.ndarray(shape, buffer=buf)
                    tmp[:] = Lpq.real
                    LpqR = lib.unpack_tril(tmp, out=bufR).reshape(-1, nao**2)
                    tmp[:] = Lpq.imag
                    LpqI = lib.unpack_tril(tmp, lib.ANTIHERMI,
                                           out=bufI).reshape(-1, nao**2)
                else:
                    LpqR = numpy.ndarray(shape, buffer=bufR)
                    LpqR[:] = Lpq.real
                    LpqI = numpy.ndarray(shape, buffer=bufI)
                    LpqI[:] = Lpq.imag
            return LpqR, LpqI

        LpqR = LpqI = None
        with _load3c(self._cderi, 'j3c', kpti_kptj) as j3c:
            naux = j3c.shape[0]
            for b0, b1 in lib.prange(0, naux, blksize):
                LpqR, LpqI = load(j3c, b0, b1, LpqR, LpqI)
                yield LpqR, LpqI
Пример #3
0
    def __init__(self, cell, kpts=numpy.zeros((1, 3))):
        self.cell = cell
        self.stdout = cell.stdout
        self.verbose = cell.verbose
        self.max_memory = cell.max_memory

        self.kpts = kpts  # default is gamma point
        self.kpts_band = None
        self._auxbasis = None

        # Search for optimized eta and mesh here.
        if cell.dimension == 0:
            self.eta = 0.2
            self.mesh = cell.mesh
        else:
            ke_cutoff = tools.mesh_to_cutoff(cell.lattice_vectors(), cell.mesh)
            ke_cutoff = ke_cutoff[:cell.dimension].min()
            eta_cell = aft.estimate_eta_for_ke_cutoff(cell, ke_cutoff,
                                                      cell.precision)
            eta_guess = estimate_eta(cell, cell.precision)
            logger.debug3(self, 'eta_guess = %g', eta_guess)
            if eta_cell < eta_guess:
                self.eta = eta_cell
                self.mesh = cell.mesh
            else:
                self.eta = eta_guess
                ke_cutoff = aft.estimate_ke_cutoff_for_eta(
                    cell, self.eta, cell.precision)
                self.mesh = tools.cutoff_to_mesh(cell.lattice_vectors(),
                                                 ke_cutoff)
                if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum':
                    self.mesh[cell.dimension:] = cell.mesh[cell.dimension:]
        self.mesh = _round_off_to_odd_mesh(self.mesh)

        # exp_to_discard to remove diffused fitting functions. The diffused
        # fitting functions may cause linear dependency in DF metric. Removing
        # the fitting functions whose exponents are smaller than exp_to_discard
        # can improve the linear dependency issue. However, this parameter
        # affects the quality of the auxiliary basis. The default value of
        # this parameter was set to 0.2 in v1.5.1 or older and was changed to
        # 0 since v1.5.2.
        self.exp_to_discard = cell.exp_to_discard

        # The following attributes are not input options.
        self.exxdiv = None  # to mimic KRHF/KUHF object in function get_coulG
        self.auxcell = None
        self.blockdim = getattr(__config__, 'pbc_df_df_DF_blockdim', 240)
        self.linear_dep_threshold = LINEAR_DEP_THR
        self._j_only = False
        # If _cderi_to_save is specified, the 3C-integral tensor will be saved in this file.
        self._cderi_to_save = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
        # If _cderi is specified, the 3C-integral tensor will be read from this file
        self._cderi = None
        self._rsh_df = {}  # Range separated Coulomb DF objects
        self._keys = set(self.__dict__.keys())
Пример #4
0
Файл: df.py Проект: eronca/pyscf
    def sr_loop(self, kpti_kptj=numpy.zeros((2,3)), max_memory=2000,
                compact=True, blksize=None):
        '''Short range part'''
        kpti, kptj = kpti_kptj
        unpack = is_zero(kpti-kptj) and not compact
        is_real = is_zero(kpti_kptj)
        nao = self.cell.nao_nr()
        if blksize is None:
            if is_real:
                if unpack:
                    blksize = max_memory*1e6/8/(nao*(nao+1)//2+nao**2)
                else:
                    blksize = max_memory*1e6/8/(nao*(nao+1))
            else:
                blksize = max_memory*1e6/16/(nao**2*2)
            blksize = max(16, min(int(blksize), self.blockdim))
            logger.debug3(self, 'max_memory %d MB, blksize %d', max_memory, blksize)

        if unpack:
            buf = numpy.empty((blksize,nao*(nao+1)//2))
        def load(Lpq, b0, b1, bufR, bufI):
            Lpq = numpy.asarray(Lpq[b0:b1])
            if is_real:
                if unpack:
                    LpqR = lib.unpack_tril(Lpq, out=bufR).reshape(-1,nao**2)
                else:
                    LpqR = Lpq
                LpqI = numpy.zeros_like(LpqR)
            else:
                shape = Lpq.shape
                if unpack:
                    tmp = numpy.ndarray(shape, buffer=buf)
                    tmp[:] = Lpq.real
                    LpqR = lib.unpack_tril(tmp, out=bufR).reshape(-1,nao**2)
                    tmp[:] = Lpq.imag
                    LpqI = lib.unpack_tril(tmp, lib.ANTIHERMI, out=bufI).reshape(-1,nao**2)
                else:
                    LpqR = numpy.ndarray(shape, buffer=bufR)
                    LpqR[:] = Lpq.real
                    LpqI = numpy.ndarray(shape, buffer=bufI)
                    LpqI[:] = Lpq.imag
            return LpqR, LpqI

        LpqR = LpqI = None
        with _load3c(self._cderi, 'j3c', kpti_kptj) as j3c:
            naux = j3c.shape[0]
            for b0, b1 in lib.prange(0, naux, blksize):
                LpqR, LpqI = load(j3c, b0, b1, LpqR, LpqI)
                yield LpqR, LpqI
Пример #5
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
Пример #6
0
    def sr_loop(self,
                kpti_kptj=numpy.zeros((2, 3)),
                max_memory=2000,
                compact=True,
                blksize=None):
        '''Short range part'''
        if self._cderi is None:
            self.build()
        cell = self.cell
        kpti, kptj = kpti_kptj
        unpack = is_zero(kpti - kptj) and not compact
        is_real = is_zero(kpti_kptj)
        nao = cell.nao_nr()
        if blksize is None:
            if is_real:
                if unpack:
                    blksize = max_memory * 1e6 / 8 / (nao *
                                                      (nao + 1) // 2 + nao**2)
                else:
                    blksize = max_memory * 1e6 / 8 / (nao * (nao + 1))
            else:
                blksize = max_memory * 1e6 / 16 / (nao**2 * 2)
            blksize = max(16, min(int(blksize), self.blockdim))
            logger.debug3(self, 'max_memory %d MB, blksize %d', max_memory,
                          blksize)

        def load(Lpq, b0, b1, bufR, bufI):
            Lpq = numpy.asarray(Lpq[b0:b1])
            if is_real:
                if unpack:
                    LpqR = lib.unpack_tril(Lpq, out=bufR).reshape(-1, nao**2)
                else:
                    LpqR = Lpq
                LpqI = numpy.zeros_like(LpqR)
            else:
                shape = Lpq.shape
                if unpack:
                    tmp = numpy.ndarray(shape, buffer=buf)
                    tmp[:] = Lpq.real
                    LpqR = lib.unpack_tril(tmp, out=bufR).reshape(-1, nao**2)
                    tmp[:] = Lpq.imag
                    LpqI = lib.unpack_tril(tmp, lib.ANTIHERMI,
                                           out=bufI).reshape(-1, nao**2)
                else:
                    LpqR = numpy.ndarray(shape, buffer=bufR)
                    LpqR[:] = Lpq.real
                    LpqI = numpy.ndarray(shape, buffer=bufI)
                    LpqI[:] = Lpq.imag
            return LpqR, LpqI

        LpqR = LpqI = None
        with _load3c(self._cderi, 'j3c', kpti_kptj, 'j3c-kptij') as j3c:
            naux = j3c.shape[0]
            if unpack:
                buf = numpy.empty((min(blksize, naux), nao * (nao + 1) // 2))
            for b0, b1 in lib.prange(0, naux, blksize):
                LpqR, LpqI = load(j3c, b0, b1, LpqR, LpqI)
                yield LpqR, LpqI, 1

        if cell.dimension == 2 and cell.low_dim_ft_type != 'inf_vacuum':
            # Truncated Coulomb operator is not postive definite. Load the
            # CDERI tensor of negative part.
            LpqR = LpqI = None
            with _load3c(self._cderi,
                         'j3c-',
                         kpti_kptj,
                         'j3c-kptij',
                         ignore_key_error=True) as j3c:
                naux = j3c.shape[0]
                if unpack:
                    buf = numpy.empty((min(blksize,
                                           naux), nao * (nao + 1) // 2))
                for b0, b1 in lib.prange(0, naux, blksize):
                    LpqR, LpqI = load(j3c, b0, b1, LpqR, LpqI)
                    yield LpqR, LpqI, -1
Пример #7
0
def aux_e2(cell, auxcell, intor='cint3c2e_sph', aosym='s1', comp=1,
           kpti_kptj=numpy.zeros((2,3))):
    '''3-center AO integrals (ij|L) with double lattice sum:
    \sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis.

    Returns:
        (nao_pair, naux) array
    '''
    assert(comp == 1)
    # sum over largest number of images in either cell or auxcell
    rcut = max(cell.rcut, auxcell.rcut)
    Ls = cell.get_lattice_Ls(rcut=rcut)
    logger.debug1(cell, "rcut %s", rcut)
    logger.debug3(cell, "Ls = %s", Ls)

    kpti, kptj = kpti_kptj
    expkL = numpy.exp(1j*numpy.asarray(numpy.dot(Ls, numpy.reshape(kpti, (1,3)).T), order='C'))
    expkR = numpy.exp(1j*numpy.asarray(numpy.dot(Ls, numpy.reshape(kptj, (1,3)).T), order='C'))
    gamma_point = abs(kpti).sum() < 1e-9 and abs(kptj).sum() < 1e-9

    nao = cell.nao_nr()
    #naux = auxcell.nao_nr('ssc' in intor)
    naux = auxcell.nao_nr()
    if naux == 0:
        if aosym == 's1' or abs(kpti-kptj).sum() > 1e-9:
            nao_pair = nao * (nao+1) // 2
        else:
            nao_pair = nao * nao
        mat = numpy.zeros((comp,nao_pair,naux))
        return mat

    buf = [numpy.zeros((nao,nao,naux,comp), order='F', dtype=numpy.complex128)]
    ints = _wrap_int3c(cell, auxcell, intor, comp, Ls, buf)
    atm, bas, env = ints._envs[:3]
    shls_slice = (0, cell.nbas, cell.nbas, cell.nbas*2,
                  cell.nbas*2, cell.nbas*2+auxcell.nbas)
    c_shls_slice = (ctypes.c_int*6)(*(shls_slice[:6]))

    xyz = cell.atom_coords().copy('C')
    ptr_coordL = atm[:cell.natm,pyscf.gto.PTR_COORD]
    ptr_coordL = numpy.vstack((ptr_coordL,ptr_coordL+1,ptr_coordL+2)).T.copy('C')

    if aosym == 's1' or abs(kpti-kptj).sum() > 1e-9:
        for l, L1 in enumerate(Ls):
            env[ptr_coordL] = xyz + L1
            exp_Lk = expkL[l].conj() * expkR
            ints(exp_Lk, c_shls_slice)
        mat, buf = buf[0], None
    else:
        for l, L1 in enumerate(Ls):
            env[ptr_coordL] = xyz + L1
            exp_Lk = expkL[l].conj() * expkR[:l+1]
            exp_Lk[l] = .5
            ints(exp_Lk, c_shls_slice)
        mat, buf = buf[0], None
        if gamma_point:
            mat = mat.real + mat.real.swapaxes(0,1)
        else:
            mat = mat + mat.swapaxes(0,1).conj()
        mat = mat[numpy.tril_indices(nao)]
    if comp == 1:
        mat = mat.reshape(-1,naux)
    else:
        mat = numpy.rollaxis(mat, -1, 0).reshape(comp,-1,naux)
    return mat
Пример #8
0
def aux_e2(cell,
           auxcell,
           intor='cint3c2e_sph',
           aosym='s1',
           comp=1,
           kpti_kptj=numpy.zeros((2, 3))):
    '''3-center AO integrals (ij|L) with double lattice sum:
    \sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis.

    Returns:
        (nao_pair, naux) array
    '''
    assert (comp == 1)
    # 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, "Images %s", nimgs)
    logger.debug3(cell, "Ls = %s", Ls)

    kpti, kptj = kpti_kptj
    expkL = numpy.exp(
        1j *
        numpy.asarray(numpy.dot(Ls,
                                numpy.reshape(kpti, (1, 3)).T), order='C'))
    expkR = numpy.exp(
        1j *
        numpy.asarray(numpy.dot(Ls,
                                numpy.reshape(kptj, (1, 3)).T), order='C'))
    gamma_point = abs(kpti).sum() < 1e-9 and abs(kptj).sum() < 1e-9

    nao = cell.nao_nr()
    #naux = auxcell.nao_nr('ssc' in intor)
    naux = auxcell.nao_nr()
    if naux == 0:
        if aosym == 's1' or abs(kpti - kptj).sum() > 1e-9:
            nao_pair = nao * (nao + 1) // 2
        else:
            nao_pair = nao * nao
        mat = numpy.zeros((comp, nao_pair, naux))
        return mat

    buf = [
        numpy.zeros((nao, nao, naux, comp), order='F', dtype=numpy.complex128)
    ]
    ints = _wrap_int3c(cell, auxcell, intor, comp, Ls, buf)
    atm, bas, env = ints._envs[:3]
    shls_slice = (0, cell.nbas, cell.nbas, cell.nbas * 2, cell.nbas * 2,
                  cell.nbas * 2 + auxcell.nbas)
    c_shls_slice = (ctypes.c_int * 6)(*(shls_slice[:6]))

    xyz = cell.atom_coords().copy('C')
    ptr_coordL = atm[:cell.natm, pyscf.gto.PTR_COORD]
    ptr_coordL = numpy.vstack(
        (ptr_coordL, ptr_coordL + 1, ptr_coordL + 2)).T.copy('C')

    if aosym == 's1' or abs(kpti - kptj).sum() > 1e-9:
        for l, L1 in enumerate(Ls):
            env[ptr_coordL] = xyz + L1
            exp_Lk = expkL[l].conj() * expkR
            ints(exp_Lk, c_shls_slice)
        mat, buf = buf[0], None
    else:
        for l, L1 in enumerate(Ls):
            env[ptr_coordL] = xyz + L1
            exp_Lk = expkL[l].conj() * expkR[:l + 1]
            exp_Lk[l] = .5
            ints(exp_Lk, c_shls_slice)
        mat, buf = buf[0], None
        if gamma_point:
            mat = mat.real + mat.real.swapaxes(0, 1)
        else:
            mat = mat + mat.swapaxes(0, 1).conj()
        mat = mat[numpy.tril_indices(nao)]
    if comp == 1:
        mat = mat.reshape(-1, naux)
    else:
        mat = numpy.rollaxis(mat, -1, 0).reshape(comp, -1, naux)
    return mat
Пример #9
0
    def sr_loop(self, kpti_kptj=numpy.zeros((2,3)), max_memory=2000,
                compact=True, blksize=None):
        '''Short range part'''
        if self._cderi is None:
            self.build()
        cell = self.cell
        kpti, kptj = kpti_kptj
        unpack = is_zero(kpti-kptj) and not compact
        is_real = is_zero(kpti_kptj)
        nao = cell.nao_nr()
        if blksize is None:
            if is_real:
                if unpack:
                    blksize = max_memory*1e6/8/(nao*(nao+1)//2+nao**2)
                else:
                    blksize = max_memory*1e6/8/(nao*(nao+1))
            else:
                blksize = max_memory*1e6/16/(nao**2*2)
            blksize = max(16, min(int(blksize), self.blockdim))
            logger.debug3(self, 'max_memory %d MB, blksize %d', max_memory, blksize)

        if unpack:
            buf = numpy.empty((blksize,nao*(nao+1)//2))
        def load(Lpq, b0, b1, bufR, bufI):
            Lpq = numpy.asarray(Lpq[b0:b1])
            if is_real:
                if unpack:
                    LpqR = lib.unpack_tril(Lpq, out=bufR).reshape(-1,nao**2)
                else:
                    LpqR = Lpq
                LpqI = numpy.zeros_like(LpqR)
            else:
                shape = Lpq.shape
                if unpack:
                    tmp = numpy.ndarray(shape, buffer=buf)
                    tmp[:] = Lpq.real
                    LpqR = lib.unpack_tril(tmp, out=bufR).reshape(-1,nao**2)
                    tmp[:] = Lpq.imag
                    LpqI = lib.unpack_tril(tmp, lib.ANTIHERMI, out=bufI).reshape(-1,nao**2)
                else:
                    LpqR = numpy.ndarray(shape, buffer=bufR)
                    LpqR[:] = Lpq.real
                    LpqI = numpy.ndarray(shape, buffer=bufI)
                    LpqI[:] = Lpq.imag
            return LpqR, LpqI

        LpqR = LpqI = None
        with _load3c(self._cderi, 'j3c', kpti_kptj, 'j3c-kptij') as j3c:
            naux = j3c.shape[0]
            for b0, b1 in lib.prange(0, naux, blksize):
                LpqR, LpqI = load(j3c, b0, b1, LpqR, LpqI)
                yield LpqR, LpqI, 1

        if cell.dimension == 2 and cell.low_dim_ft_type != 'inf_vacuum':
            # Truncated Coulomb operator is not postive definite. Load the
            # CDERI tensor of negative part.
            LpqR = LpqI = None
            with _load3c(self._cderi, 'j3c-', kpti_kptj, 'j3c-kptij',
                         ignore_key_error=True) as j3c:
                naux = j3c.shape[0]
                for b0, b1 in lib.prange(0, naux, blksize):
                    LpqR, LpqI = load(j3c, b0, b1, LpqR, LpqI)
                    yield LpqR, LpqI, -1
Пример #10
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