Beispiel #1
0
    def init_direct_scf(self, mol=None):
        if mol is None: mol = self.mol

        def set_vkscreen(opt, name):
            opt._this.contents.r_vkscreen = _vhf._fpointer(name)

        opt_llll = _vhf.VHFOpt(mol, 'int2e_spinor', 'CVHFrkbllll_prescreen',
                               'CVHFrkbllll_direct_scf',
                               'CVHFrkbllll_direct_scf_dm')
        opt_llll.direct_scf_tol = self.direct_scf_tol
        set_vkscreen(opt_llll, 'CVHFrkbllll_vkscreen')
        opt_ssss = _vhf.VHFOpt(mol, 'int2e_spsp1spsp2_spinor',
                               'CVHFrkbllll_prescreen',
                               'CVHFrkbssss_direct_scf',
                               'CVHFrkbssss_direct_scf_dm')
        opt_ssss.direct_scf_tol = self.direct_scf_tol
        set_vkscreen(opt_ssss, 'CVHFrkbllll_vkscreen')
        opt_ssll = _vhf.VHFOpt(mol, 'int2e_spsp1_spinor',
                               'CVHFrkbssll_prescreen',
                               'CVHFrkbssll_direct_scf',
                               'CVHFrkbssll_direct_scf_dm')
        opt_ssll.direct_scf_tol = self.direct_scf_tol
        set_vkscreen(opt_ssll, 'CVHFrkbssll_vkscreen')
        #TODO: prescreen for gaunt
        opt_gaunt = None
        return opt_llll, opt_ssll, opt_ssss, opt_gaunt
Beispiel #2
0
    def init_direct_scf(self, mol=None):
        if mol is None: mol = self.mol

        def set_vkscreen(opt, name):
            opt._this.contents.r_vkscreen = \
                ctypes.c_void_p(_ctypes.dlsym(_vhf.libcvhf._handle, name))

        opt_llll = _vhf.VHFOpt(mol, 'cint2e', 'CVHFrkbllll_prescreen',
                               'CVHFrkbllll_direct_scf',
                               'CVHFrkbllll_direct_scf_dm')
        opt_llll.direct_scf_tol = self.direct_scf_tol
        set_vkscreen(opt_llll, 'CVHFrkbllll_vkscreen')
        opt_ssss = _vhf.VHFOpt(mol, 'cint2e_spsp1spsp2',
                               'CVHFrkbllll_prescreen',
                               'CVHFrkbssss_direct_scf',
                               'CVHFrkbssss_direct_scf_dm')
        opt_ssss.direct_scf_tol = self.direct_scf_tol
        set_vkscreen(opt_ssss, 'CVHFrkbllll_vkscreen')
        opt_ssll = _vhf.VHFOpt(mol, 'cint2e_spsp1', 'CVHFrkbssll_prescreen',
                               'CVHFrkbssll_direct_scf',
                               'CVHFrkbssll_direct_scf_dm')
        opt_ssll.direct_scf_tol = self.direct_scf_tol
        set_vkscreen(opt_ssll, 'CVHFrkbssll_vkscreen')
        #TODO: prescreen for gaunt
        opt_gaunt = None
        return opt_llll, opt_ssll, opt_ssss, opt_gaunt
Beispiel #3
0
    def init_direct_scf(self, mol=None):
        if mol is None: mol = self.mol
        def set_vkscreen(opt, name):
            opt._this.contents.r_vkscreen = _vhf._fpointer(name)

        with mol.with_integral_screen(self.direct_scf_tol**2):
            opt_llll = _vhf.VHFOpt(mol, 'int2e_spinor', 'CVHFrkbllll_prescreen',
                                   'CVHFrkbllll_direct_scf',
                                   'CVHFrkbllll_direct_scf_dm')
            opt_llll.direct_scf_tol = self.direct_scf_tol
            set_vkscreen(opt_llll, 'CVHFrkbllll_vkscreen')
            opt_ssss = _vhf.VHFOpt(mol, 'int2e_spsp1spsp2_spinor',
                                   'CVHFrkbllll_prescreen',
                                   'CVHFrkbssss_direct_scf',
                                   'CVHFrkbssss_direct_scf_dm')
            c1 = .5 / lib.param.LIGHT_SPEED
            q_cond = opt_ssss.get_q_cond()
            q_cond *= c1**2
            opt_ssss.direct_scf_tol = self.direct_scf_tol
            set_vkscreen(opt_ssss, 'CVHFrkbllll_vkscreen')
            opt_ssll = _vhf.VHFOpt(mol, 'int2e_spsp1_spinor',
                                   'CVHFrkbssll_prescreen',
                                   'CVHFrkbssll_direct_scf',
                                   'CVHFrkbssll_direct_scf_dm')
            opt_ssll.direct_scf_tol = self.direct_scf_tol
            set_vkscreen(opt_ssll, 'CVHFrkbssll_vkscreen')
            nbas = mol.nbas
            # The second parts of q_cond corresponds to ssss integrals. They
            # need to be scaled by the factor (1/2c)^2
            q_cond = opt_ssll.get_q_cond(shape=(2, nbas, nbas))
            q_cond[1] *= c1**2

#TODO: prescreen for gaunt
            opt_gaunt = None
        return opt_llll, opt_ssll, opt_ssss, opt_gaunt
Beispiel #4
0
def get_jk(mol, dm):
    '''J = ((-nabla i) j| kl) D_lk
    K = ((-nabla i) j| kl) D_jk
    '''
    vhfopt = _vhf.VHFOpt(mol, 'int2e_ip1ip2', 'CVHFgrad_jk_prescreen',
                         'CVHFgrad_jk_direct_scf')
    dm = numpy.asarray(dm, order='C')
    if dm.ndim == 3:
        n_dm = dm.shape[0]
    else:
        n_dm = 1
    ao_loc = mol.ao_loc_nr()
    fsetdm = getattr(_vhf.libcvhf, 'CVHFgrad_jk_direct_scf_dm')
    fsetdm(vhfopt._this,
           dm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(n_dm),
           ao_loc.ctypes.data_as(ctypes.c_void_p),
           mol._atm.ctypes.data_as(ctypes.c_void_p), mol.natm,
           mol._bas.ctypes.data_as(ctypes.c_void_p), mol.nbas,
           mol._env.ctypes.data_as(ctypes.c_void_p))

    # Update the vhfopt's attributes intor.  Function direct_mapdm needs
    # vhfopt._intor and vhfopt._cintopt to compute J/K.  intor was initialized
    # as int2e_ip1ip2. It should be int2e_ip1
    vhfopt._intor = intor = mol._add_suffix('int2e_ip1')
    vhfopt._cintopt = None

    vj, vk = _vhf.direct_mapdm(intor,  # (nabla i,j|k,l)
                               's2kl', # ip1_sph has k>=l,
                               ('lk->s1ij', 'jk->s1il'),
                               dm, 3, # xyz, 3 components
                               mol._atm, mol._bas, mol._env, vhfopt=vhfopt)
    return -vj, -vk
Beispiel #5
0
 def init_direct_scf(self, mol=None):
     if mol is None: mol = self.mol
     opt = _vhf.VHFOpt(mol, 'cint2e_sph', 'CVHFnrs8_prescreen',
                       'CVHFsetnr_direct_scf',
                       'CVHFsetnr_direct_scf_dm')
     opt.direct_scf_tol = self.direct_scf_tol
     return opt
Beispiel #6
0
def _make_vhfopt(mol, dms, key, vhf_intor):
    if not hasattr(_vhf.libcvhf, vhf_intor):
        return None

    vhfopt = _vhf.VHFOpt(mol, vhf_intor, 'CVHF' + key + '_prescreen',
                         'CVHF' + key + '_direct_scf')
    dms = numpy.asarray(dms, order='C')
    if dms.ndim == 3:
        n_dm = dms.shape[0]
    else:
        n_dm = 1
    ao_loc = mol.ao_loc_nr()
    fsetdm = getattr(_vhf.libcvhf, 'CVHF' + key + '_direct_scf_dm')
    fsetdm(vhfopt._this, dms.ctypes.data_as(ctypes.c_void_p),
           ctypes.c_int(n_dm), ao_loc.ctypes.data_as(ctypes.c_void_p),
           mol._atm.ctypes.data_as(ctypes.c_void_p), mol.natm,
           mol._bas.ctypes.data_as(ctypes.c_void_p), mol.nbas,
           mol._env.ctypes.data_as(ctypes.c_void_p))

    # Update the vhfopt's attributes intor.  Function direct_mapdm needs
    # vhfopt._intor and vhfopt._cintopt to compute J/K.
    if vhf_intor != 'int2e_' + key:
        vhfopt._intor = mol._add_suffix('int2e_' + key)
        vhfopt._cintopt = None
    return vhfopt
Beispiel #7
0
 def init_direct_scf(self, mol=None):
     if mol is None: mol = self.mol
     def set_vkscreen(opt, name):
         opt._this.contents.r_vkscreen = _vhf._fpointer(name)
     opt = _vhf.VHFOpt(mol, 'int2e_spinor', 'CVHFrkbllll_prescreen',
                       'CVHFrkbllll_direct_scf',
                       'CVHFrkbllll_direct_scf_dm')
     opt.direct_scf_tol = self.direct_scf_tol
     set_vkscreen(opt, 'CVHFrkbllll_vkscreen')
     return opt
Beispiel #8
0
 def init_direct_scf(self, mol=None):
     if mol is None: mol = self.mol
     if mol.cart:
         intor = 'int2e_cart'
     else:
         intor = 'int2e_sph'
     opt = _vhf.VHFOpt(mol, intor, 'CVHFnrs8_prescreen',
                       'CVHFsetnr_direct_scf', 'CVHFsetnr_direct_scf_dm')
     opt.direct_scf_tol = self.direct_scf_tol
     return opt
Beispiel #9
0
def _make_opt(mol):
    '''Optimizer to genrate 3-center 2-electron integrals'''
    intor = mol._add_suffix('int3c2e')
    cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, intor)
    # intor 'int1e_ovlp' is used by the prescreen method
    # 'SGXnr_ovlp_prescreen' only. Not used again in other places.
    # It can be released early
    vhfopt = _vhf.VHFOpt(mol, 'int1e_ovlp', 'SGXnr_ovlp_prescreen',
                         'SGXsetnr_direct_scf')
    vhfopt._intor = intor
    vhfopt._cintopt = cintopt
    return vhfopt
Beispiel #10
0
 def test_direct_bindm1(self):
     numpy.random.seed(1)
     nao = mol.nao_nr(cart=True)
     dm = numpy.random.random((nao,nao))
     vhfopt = _vhf.VHFOpt(mol, 'int2e_cart', 'CVHFnrs8_prescreen',
                          'CVHFsetnr_direct_scf',
                          'CVHFsetnr_direct_scf_dm')
     vj0, vk0 = _vhf.direct(dm, mol._atm, mol._bas, mol._env,
                            vhfopt=vhfopt, hermi=0, cart=True)
     vj = _vhf.direct_bindm('int2e_cart', 's1', 'kl->s1ij', dm, 1,
                            mol._atm, mol._bas, mol._env, vhfopt)
     self.assertTrue(numpy.allclose(vj0,vj))
Beispiel #11
0
    def init_direct_scf(self, mol=None):
        if mol is None: mol = self.mol

        def set_vkscreen(opt, name):
            opt._this.contents.r_vkscreen = \
                ctypes.c_void_p(_ctypes.dlsym(_vhf.libcvhf._handle, name))

        opt = _vhf.VHFOpt(mol, 'cint2e', 'CVHFrkbllll_prescreen',
                          'CVHFrkbllll_direct_scf',
                          'CVHFrkbllll_direct_scf_dm')
        opt.direct_scf_tol = self.direct_scf_tol
        set_vkscreen(opt, 'CVHFrkbllll_vkscreen')
        return opt
Beispiel #12
0
def jk_ints(molA, molB, dm_ia, dm_jb):
    from pyscf.scf import jk, _vhf
    naoA = molA.nao
    naoB = molB.nao
    assert (dm_ia.shape == (naoA, naoA))
    assert (dm_jb.shape == (naoB, naoB))

    molAB = molA + molB
    vhfopt = _vhf.VHFOpt(molAB, 'int2e', 'CVHFnrs8_prescreen',
                         'CVHFsetnr_direct_scf', 'CVHFsetnr_direct_scf_dm')
    dmAB = scipy.linalg.block_diag(dm_ia, dm_jb)
    # The prescreen function CVHFnrs8_prescreen indexes q_cond and dm_cond
    # over the entire basis.  "set_dm" in function jk.get_jk/direct_bindm only
    # creates a subblock of dm_cond which is not compatible with
    # CVHFnrs8_prescreen.
    vhfopt.set_dm(dmAB, molAB._atm, molAB._bas, molAB._env)
    # Then skip the "set_dm" initialization in function jk.get_jk/direct_bindm.
    vhfopt._dmcondname = None

    with lib.temporary_env(vhfopt._this.contents,
                           fprescreen=_vhf._fpointer('CVHFnrs8_vj_prescreen')):
        shls_slice = (0, molA.nbas, 0, molA.nbas, molA.nbas, molAB.nbas,
                      molA.nbas, molAB.nbas)  # AABB
        vJ = jk.get_jk(molAB,
                       dm_jb,
                       'ijkl,lk->s2ij',
                       shls_slice=shls_slice,
                       vhfopt=vhfopt,
                       aosym='s4',
                       hermi=1)
        cJ = np.einsum('ia,ia->', vJ, dm_ia)

    with lib.temporary_env(vhfopt._this.contents,
                           fprescreen=_vhf._fpointer('CVHFnrs8_vk_prescreen')):
        shls_slice = (0, molA.nbas, molA.nbas, molAB.nbas, molA.nbas,
                      molAB.nbas, 0, molA.nbas)  # ABBA
        vK = jk.get_jk(molAB,
                       dm_jb,
                       'ijkl,jk->il',
                       shls_slice=shls_slice,
                       vhfopt=vhfopt,
                       aosym='s1',
                       hermi=0)
        cK = np.einsum('ia,ia->', vK, dm_ia)

    return cJ, cK
Beispiel #13
0
    def test_rdirect_bindm(self):
        n2c = nao * 2
        numpy.random.seed(1)
        dm = (numpy.random.random((n2c, n2c)) + numpy.random.random(
            (n2c, n2c)) * 1j)
        dm = dm + dm.conj().T

        eri0 = mol.intor('int2e_spsp1_spinor').reshape(n2c, n2c, n2c, n2c)
        vk0 = numpy.einsum('ijkl,jk->il', eri0, dm)
        vk1 = _vhf.rdirect_bindm('int2e_spsp1_spinor', 's4', 'jk->s1il', dm, 1,
                                 mol._atm, mol._bas, mol._env)
        self.assertTrue(numpy.allclose(vk0, vk1))

        opt_llll = _vhf.VHFOpt(mol, 'int2e_spinor', 'CVHFrkbllll_prescreen',
                               'CVHFrkbllll_direct_scf',
                               'CVHFrkbllll_direct_scf_dm')
        opt_llll._this.contents.r_vkscreen = _vhf._fpointer(
            'CVHFrkbllll_vkscreen')
        eri0 = mol.intor('int2e_spinor').reshape(n2c, n2c, n2c, n2c)
        vk0 = numpy.einsum('ijkl,jk->il', eri0, dm)
        vk1 = _vhf.rdirect_bindm('int2e_spinor', 's1', 'jk->s1il', dm, 1,
                                 mol._atm, mol._bas, mol._env, opt_llll)
        self.assertTrue(numpy.allclose(vk0, vk1))
Beispiel #14
0
def get_j(dfobj, dm, hermi=1, direct_scf_tol=1e-13):
    from pyscf.scf import _vhf
    from pyscf.scf import jk
    from pyscf.df import addons
    t0 = t1 = (time.clock(), time.time())

    mol = dfobj.mol
    if dfobj._vjopt is None:
        dfobj.auxmol = auxmol = addons.make_auxmol(mol, dfobj.auxbasis)
        opt = _vhf.VHFOpt(mol, 'int3c2e', 'CVHFnr3c2e_schwarz_cond')
        opt.direct_scf_tol = direct_scf_tol

        # q_cond part 1: the regular int2e (ij|ij) for mol's basis
        opt.init_cvhf_direct(mol, 'int2e', 'CVHFsetnr_direct_scf')
        mol_q_cond = lib.frompointer(opt._this.contents.q_cond, mol.nbas**2)

        # Update q_cond to include the 2e-integrals (auxmol|auxmol)
        j2c = auxmol.intor('int2c2e', hermi=1)
        j2c_diag = numpy.sqrt(abs(j2c.diagonal()))
        aux_loc = auxmol.ao_loc
        aux_q_cond = [
            j2c_diag[i0:i1].max() for i0, i1 in zip(aux_loc[:-1], aux_loc[1:])
        ]
        q_cond = numpy.hstack((mol_q_cond, aux_q_cond))
        fsetqcond = _vhf.libcvhf.CVHFset_q_cond
        fsetqcond(opt._this, q_cond.ctypes.data_as(ctypes.c_void_p),
                  ctypes.c_int(q_cond.size))

        try:
            opt.j2c = j2c = scipy.linalg.cho_factor(j2c, lower=True)
            opt.j2c_type = 'cd'
        except scipy.linalg.LinAlgError:
            opt.j2c = j2c
            opt.j2c_type = 'regular'

        # jk.get_jk function supports 4-index integrals. Use bas_placeholder
        # (l=0, nctr=1, 1 function) to hold the last index.
        bas_placeholder = numpy.array([0, 0, 1, 1, 0, 0, 0, 0],
                                      dtype=numpy.int32)
        fakemol = mol + auxmol
        fakemol._bas = numpy.vstack((fakemol._bas, bas_placeholder))
        opt.fakemol = fakemol
        dfobj._vjopt = opt
        t1 = logger.timer_debug1(dfobj, 'df-vj init_direct_scf', *t1)

    opt = dfobj._vjopt
    fakemol = opt.fakemol
    dm = numpy.asarray(dm, order='C')
    dm_shape = dm.shape
    nao = dm_shape[-1]
    dm = dm.reshape(-1, nao, nao)
    n_dm = dm.shape[0]

    # First compute the density in auxiliary basis
    # j3c = fauxe2(mol, auxmol)
    # jaux = numpy.einsum('ijk,ji->k', j3c, dm)
    # rho = numpy.linalg.solve(auxmol.intor('int2c2e'), jaux)
    nbas = mol.nbas
    nbas1 = mol.nbas + dfobj.auxmol.nbas
    shls_slice = (0, nbas, 0, nbas, nbas, nbas1, nbas1, nbas1 + 1)
    with lib.temporary_env(opt,
                           prescreen='CVHFnr3c2e_vj_pass1_prescreen',
                           _dmcondname='CVHFsetnr_direct_scf_dm'):
        jaux = jk.get_jk(fakemol,
                         dm, ['ijkl,ji->kl'] * n_dm,
                         'int3c2e',
                         aosym='s2ij',
                         hermi=0,
                         shls_slice=shls_slice,
                         vhfopt=opt)
    # remove the index corresponding to bas_placeholder
    jaux = numpy.array(jaux)[:, :, 0]
    t1 = logger.timer_debug1(dfobj, 'df-vj pass 1', *t1)

    if opt.j2c_type == 'cd':
        rho = scipy.linalg.cho_solve(opt.j2c, jaux.T)
    else:
        rho = scipy.linalg.solve(opt.j2c, jaux.T)
    # transform rho to shape (:,1,naux), to adapt to 3c2e integrals (ij|k)
    rho = rho.T[:, numpy.newaxis, :]
    t1 = logger.timer_debug1(dfobj, 'df-vj solve ', *t1)

    # Next compute the Coulomb matrix
    # j3c = fauxe2(mol, auxmol)
    # vj = numpy.einsum('ijk,k->ij', j3c, rho)
    with lib.temporary_env(opt,
                           prescreen='CVHFnr3c2e_vj_pass2_prescreen',
                           _dmcondname=None):
        # CVHFnr3c2e_vj_pass2_prescreen requires custom dm_cond
        aux_loc = dfobj.auxmol.ao_loc
        dm_cond = [
            abs(rho[:, :, i0:i1]).max()
            for i0, i1 in zip(aux_loc[:-1], aux_loc[1:])
        ]
        dm_cond = numpy.array(dm_cond)
        fsetcond = _vhf.libcvhf.CVHFset_dm_cond
        fsetcond(opt._this, dm_cond.ctypes.data_as(ctypes.c_void_p),
                 ctypes.c_int(dm_cond.size))

        vj = jk.get_jk(fakemol,
                       rho, ['ijkl,lk->ij'] * n_dm,
                       'int3c2e',
                       aosym='s2ij',
                       hermi=1,
                       shls_slice=shls_slice,
                       vhfopt=opt)

    t1 = logger.timer_debug1(dfobj, 'df-vj pass 2', *t1)
    logger.timer(dfobj, 'df-vj', *t0)
    return numpy.asarray(vj).reshape(dm_shape)
Beispiel #15
0
    def build(self, omega=None, direct_scf_tol=None):
        cpu0 = (time.clock(), time.time())
        cell = self.cell
        kpts = self.kpts

        k_scaled = cell.get_scaled_kpts(kpts).sum(axis=0)
        k_mod_to_half = k_scaled * 2 - (k_scaled * 2).round(0)
        if abs(k_mod_to_half).sum() > 1e-5:
            raise NotImplementedError('k-points must be symmetryic')

        if omega is not None:
            self.omega = omega

        if self.omega is None:
            # Search a proper range-separation parameter omega that can balance the
            # computational cost between the real space integrals and moment space
            # integrals
            self.omega, self.mesh, self.ke_cutoff = _guess_omega(
                cell, kpts, self.mesh)
        else:
            self.ke_cutoff = aft.estimate_ke_cutoff_for_omega(cell, self.omega)
            self.mesh = pbctools.cutoff_to_mesh(cell.lattice_vectors(),
                                                self.ke_cutoff)

        logger.info(self, 'omega = %.15g  ke_cutoff = %s  mesh = %s',
                    self.omega, self.ke_cutoff, self.mesh)

        if direct_scf_tol is None:
            direct_scf_tol = cell.precision**1.5
            logger.debug(self, 'Set direct_scf_tol %g', direct_scf_tol)

        self.cell_rs = cell_rs = _re_contract_cell(cell, self.ke_cutoff)
        self.bvk_kmesh = kmesh = k2gamma.kpts_to_kmesh(cell_rs, kpts)
        bvkcell, phase = k2gamma.get_phase(cell_rs, kpts, kmesh)
        self.bvkmesh_Ls = Ks = k2gamma.translation_vectors_for_kmesh(
            cell_rs, kmesh)
        self.bvkcell = bvkcell
        self.phase = phase

        # Given ke_cutoff, eta corresponds to the most steep Gaussian basis
        # of which the Coulomb integrals can be accurately computed in moment
        # space.
        eta = aft.estimate_eta_for_ke_cutoff(cell,
                                             self.ke_cutoff,
                                             precision=cell.precision)
        # * Assuming the most steep function in smooth basis has exponent eta,
        # with attenuation parameter omega, rcut_sr is the distance of which
        # the value of attenuated Coulomb integrals of four shells |eta> is
        # smaller than the required precision.
        # * The attenuated coulomb integrals between four s-type Gaussians
        # (2*a/pi)^{3/4}exp(-a*r^2) is
        #   (erfc(omega*a^0.5/(omega^2+a)^0.5*R) - erfc(a^0.5*R)) / R
        # if two Gaussians on one center and the other two on another center
        # and the distance between the two centers are R.
        # * The attenuated coulomb integrals between two spherical charge
        # distributions is
        #   ~(pi/eta)^3/2 (erfc(tau*(eta/2)^0.5*R) - erfc((eta/2)^0.5*R)) / R
        #       tau = omega/sqrt(omega^2 + eta/2)
        # if the spherical charge distribution is the product of above s-type
        # Gaussian with exponent eta and a very smooth function.
        # When R is large, the attenuated Coulomb integral is
        #   ~= (pi/eta)^3/2 erfc(tau*(eta/2)^0.5*R) / R
        #   ~= pi/(tau*eta^2*R^2) exp(-tau^2*eta*R^2/2)
        tau = self.omega / (self.omega**2 + eta / 2)**.5
        rcut_sr = 10  # initial guess
        rcut_sr = (-np.log(direct_scf_tol * tau * (eta * rcut_sr)**2 / np.pi) /
                   (tau**2 * eta / 2))**.5
        logger.debug(self, 'eta = %g  rcut_sr = %g', eta, rcut_sr)

        # Ls is the translation vectors to mimic periodicity of a cell
        Ls = bvkcell.get_lattice_Ls(rcut=cell.rcut + rcut_sr)
        self.supmol_Ls = Ls = Ls[np.linalg.norm(Ls, axis=1).argsort()]

        supmol = _make_extended_mole(cell_rs, Ls, Ks, self.omega,
                                     direct_scf_tol)
        self.supmol = supmol

        nkpts = len(self.bvkmesh_Ls)
        nbas = cell_rs.nbas
        n_steep, n_local, n_diffused = cell_rs._nbas_each_set
        n_compact = n_steep + n_local
        bas_mask = supmol._bas_mask

        self.bvk_bas_mask = bvk_bas_mask = bas_mask.any(axis=2)
        # Some basis in bvk-cell are not presented in the supmol. They can be
        # skipped when computing SR integrals
        self.bvkcell._bas = bvkcell._bas[bvk_bas_mask.ravel()]

        # Record the mapping between the dense bvkcell basis and the
        # original sparse bvkcell basis
        bvk_cell_idx = np.repeat(np.arange(nkpts)[:, None], nbas, axis=1)
        self.bvk_cell_id = bvk_cell_idx[bvk_bas_mask].astype(np.int32)
        cell0_shl_idx = np.repeat(np.arange(nbas)[None, :], nkpts, axis=0)
        self.cell0_shl_id = cell0_shl_idx[bvk_bas_mask].astype(np.int32)

        logger.timer_debug1(self, 'initializing supmol', *cpu0)
        logger.info(self, 'sup-mol nbas = %d cGTO = %d pGTO = %d', supmol.nbas,
                    supmol.nao, supmol.npgto_nr())

        supmol.omega = -self.omega  # Set short range coulomb
        with supmol.with_integral_screen(direct_scf_tol**2):
            vhfopt = _vhf.VHFOpt(supmol,
                                 'int2e_sph',
                                 qcondname=libpbc.PBCVHFsetnr_direct_scf)
        vhfopt.direct_scf_tol = direct_scf_tol
        self.vhfopt = vhfopt
        logger.timer(self, 'initializing vhfopt', *cpu0)

        q_cond = vhfopt.get_q_cond((supmol.nbas, supmol.nbas))
        idx = supmol._images_loc
        bvk_q_cond = lib.condense('NP_absmax', q_cond, idx, idx)
        ovlp_mask = bvk_q_cond > direct_scf_tol
        # Remove diffused-diffused block
        if n_diffused > 0:
            diffused_mask = np.zeros_like(bvk_bas_mask)
            diffused_mask[:, n_compact:] = True
            diffused_mask = diffused_mask[bvk_bas_mask]
            ovlp_mask[diffused_mask[:, None] & diffused_mask] = False
        self.ovlp_mask = ovlp_mask.astype(np.int8)

        # mute rcut_threshold, divide basis into two sets only
        cell_lr_aft = _re_contract_cell(cell, self.ke_cutoff, -1, verbose=0)
        self.lr_aft = lr_aft = _LongRangeAFT(cell_lr_aft, kpts, self.omega,
                                             self.bvk_kmesh)
        lr_aft.ke_cutoff = self.ke_cutoff
        lr_aft.mesh = self.mesh
        lr_aft.eta = eta
        return self
Beispiel #16
0
def _gen_jk_direct(mol, aosym, with_j, with_k, direct_scf_tol):
    '''Contraction between sgX Coulomb integrals and density matrices
    J: einsum('guv,xg->xuv', gbn, dms) if dms == rho at grid
       einsum('gij,xij->xg', gbn, dms) if dms are density matrices
    K: einsum('gtv,xgt->xgv', gbn, fg)
    '''
    intor = mol._add_suffix('int3c2e')
    cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, intor)
    ncomp = 1
    nao = mol.nao

    vhfopt = _vhf.VHFOpt(mol, 'int1e_ovlp', 'SGXnr_ovlp_prescreen',
                         'SGXsetnr_direct_scf')
    vhfopt.direct_scf_tol = direct_scf_tol
    cintor = _vhf._fpointer(intor)
    fdot = _vhf._fpointer('SGXdot_nr' + aosym)
    drv = _vhf.libcvhf.SGXnr_direct_drv

    # for linsgx, from _vhf.VHFOpt
    libcvhf = lib.load_library('libcvhf')
    intor = mol._add_suffix('int1e_ovlp')
    c_atm = numpy.asarray(mol._atm, dtype=numpy.int32, order='C')
    c_bas = numpy.asarray(mol._bas, dtype=numpy.int32, order='C')
    c_env = numpy.asarray(mol._env, dtype=numpy.double, order='C')
    natm = ctypes.c_int(c_atm.shape[0])
    nbas = ctypes.c_int(c_bas.shape[0])

    @profile
    def jk_part(mol, grid_coords, dms, fg):

        # transfer bvv to SGXsetnr_direct_scf_blk. from _vhf.VHFOpt
        # need add mol._bvv in scf.mole.py
        c_bvv = numpy.asarray(mol._bvv, dtype=numpy.int32, order='C')
        nbvv = ctypes.c_int(c_bvv.shape[0])
        ao_loc = make_loc(c_bas, intor)
        fsetqcond = getattr(libcvhf, 'SGXsetnr_direct_scf_blk')
        fsetqcond(vhfopt._this, getattr(libcvhf, intor), lib.c_null_ptr(),
                  ao_loc.ctypes.data_as(ctypes.c_void_p),
                  c_atm.ctypes.data_as(ctypes.c_void_p), natm,
                  c_bas.ctypes.data_as(ctypes.c_void_p), nbas,
                  c_env.ctypes.data_as(ctypes.c_void_p),
                  c_bvv.ctypes.data_as(ctypes.c_void_p), nbvv)

        fakemol = gto.fakemol_for_charges(grid_coords)
        atm, bas, env = gto.mole.conc_env(mol._atm, mol._bas, mol._env,
                                          fakemol._atm, fakemol._bas,
                                          fakemol._env)

        ao_loc = moleintor.make_loc(bas, intor)
        shls_slice = (0, mol.nbas, 0, mol.nbas, mol.nbas, len(bas))
        ngrids = grid_coords.shape[0]

        vj = vk = None
        fjk = []
        dmsptr = []
        vjkptr = []
        if with_j:
            if dms[0].ndim == 1:  # the value of density at each grid
                vj = numpy.zeros((len(dms), ncomp, nao, nao))[:, 0]
                for i, dm in enumerate(dms):
                    dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
                    vjkptr.append(vj[i].ctypes.data_as(ctypes.c_void_p))
                    fjk.append(_vhf._fpointer('SGXnr' + aosym + '_ijg_g_ij'))
            else:
                vj = numpy.zeros((len(dms), ncomp, ngrids))[:, 0]
                for i, dm in enumerate(dms):
                    dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
                    vjkptr.append(vj[i].ctypes.data_as(ctypes.c_void_p))
                    fjk.append(_vhf._fpointer('SGXnr' + aosym + '_ijg_ji_g'))
        if with_k:
            vk = numpy.zeros((len(fg), ncomp, ngrids, nao))[:, 0]
            for i, dm in enumerate(fg):
                dmsptr.append(dm.ctypes.data_as(ctypes.c_void_p))
                vjkptr.append(vk[i].ctypes.data_as(ctypes.c_void_p))
                fjk.append(_vhf._fpointer('SGXnr' + aosym + '_ijg_gj_gi'))

        n_dm = len(fjk)
        fjk = (ctypes.c_void_p * (n_dm))(*fjk)
        dmsptr = (ctypes.c_void_p * (n_dm))(*dmsptr)
        vjkptr = (ctypes.c_void_p * (n_dm))(*vjkptr)

        drv(cintor, fdot, fjk, dmsptr, vjkptr, n_dm, ncomp,
            (ctypes.c_int * 6)(*shls_slice),
            ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt, vhfopt._this,
            atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
            bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
            env.ctypes.data_as(ctypes.c_void_p))
        return vj, vk

    return jk_part