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
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
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
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
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
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
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
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
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
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))
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
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
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))
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)
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
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