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
def jk_part(mol, grid_coords, dms, fg, weights): atm, bas, env = mol._atm, mol._bas, mol._env ngrids = grid_coords.shape[0] env = numpy.append(env, grid_coords.ravel()) env[gto.NGRIDS] = ngrids env[gto.PTR_GRIDS] = mol._env.size if pjs: sgxopt.set_dm(fg / numpy.sqrt(numpy.abs(weights[None,:])), mol._atm, mol._bas, env) ao_loc = moleintor.make_loc(bas, sgxopt._intor) shls_slice = (0, mol.nbas, 0, mol.nbas) fg = numpy.ascontiguousarray(fg.transpose(0,2,1)) 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,nao,ngrids))[:,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*4)(*shls_slice), ao_loc.ctypes.data_as(ctypes.c_void_p), sgxopt._cintopt, sgxopt._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), ctypes.c_int(env.shape[0]), ctypes.c_int(2 if aosym == 's2' else 1)) if vk is not None: vk = vk.transpose(0,2,1) vk = numpy.ascontiguousarray(vk) return vj, vk
def dot_eri_dm(eri, dms, nao_v=None, eri_dot_dm=True): assert (eri.dtype == numpy.double) eri = numpy.asarray(eri, order='C') dms = numpy.asarray(dms, order='C') dms_shape = dms.shape nao_dm = dms_shape[-1] if nao_v is None: nao_v = nao_dm dms = dms.reshape(-1, nao_dm, nao_dm) n_dm = dms.shape[0] vj = numpy.zeros((n_dm, nao_v, nao_v)) dmsptr = [] vjkptr = [] fjkptr = [] npair_v = nao_v * (nao_v + 1) // 2 npair_dm = nao_dm * (nao_dm + 1) // 2 if eri.ndim == 2 and npair_v * npair_dm == eri.size: # 4-fold symmetry eri if eri_dot_dm: # 'ijkl,kl->ij' fdrv = getattr(_vhf.libcvhf, 'CVHFnrs4_incore_drv_diff_size_v_dm') fvj = _vhf._fpointer('CVHFics4_kl_s2ij_diff_size') else: # 'ijkl,ij->kl' fdrv = getattr(_vhf.libcvhf, 'CVHFnrs4_incore_drv_diff_size_dm_v') fvj = _vhf._fpointer('CVHFics4_ij_s2kl_diff_size') 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)) fjkptr.append(fvj) else: raise RuntimeError( 'Array shape not consistent: nao_v %s, DM %s, eri %s' % (nao_v, dms_shape, eri.shape)) n_ops = len(dmsptr) fdrv(eri.ctypes.data_as(ctypes.c_void_p), (ctypes.c_void_p * n_ops)(*dmsptr), (ctypes.c_void_p * n_ops)(*vjkptr), ctypes.c_int(n_ops), ctypes.c_int(nao_v), ctypes.c_int(nao_dm), (ctypes.c_void_p * n_ops)(*fjkptr)) for i in range(n_dm): lib.hermi_triu(vj[i], 1, inplace=True) if n_dm == 1: vj = vj.reshape((nao_v, nao_v)) else: vj = vj.reshape((n_dm, nao_v, nao_v)) return vj
def _get_k_lr(mol, dm, omega=0, hermi=0, vhfopt=None): import sys sys.stderr.write('This function is deprecated. ' 'It is replaced by mol.get_k(mol, dm, omege=omega)') dm = numpy.asarray(dm) # Note, ks object caches the ERIs for small systems. The cached eris are # computed with regular Coulomb operator. ks.get_jk or ks.get_k do not evalute # the K matrix with the range separated Coulomb operator. Here jk.get_jk # function computes the K matrix with the modified Coulomb operator. nao = dm.shape[-1] dms = dm.reshape(-1, nao, nao) with mol.with_range_coulomb(omega): # Compute the long range part of ERIs temporarily with omega. Restore # the original omega when the block ends if vhfopt is None: contents = lambda: None # just a place_holder else: contents = vhfopt._this.contents with lib.temporary_env( contents, fprescreen=_vhf._fpointer('CVHFnrs8_vk_prescreen')): intor = mol._add_suffix('int2e') vklr = jk.get_jk(mol, dms, ['ijkl,jk->il'] * len(dms), intor=intor, vhfopt=vhfopt) return numpy.asarray(vklr).reshape(dm.shape)
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 jk_part(mol, grid_coords, dms, fg): 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, sgxopt._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), sgxopt._cintopt, sgxopt._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
def get_k(mol_or_mf, dm, hermi=1): if isinstance(mol_or_mf, gto.mole.Mole): mf = hf.SCF(mol_or_mf).view(SCF) else: mf = mol_or_mf # dm may be too big for mpi4py library to serialize. Broadcast dm here. if any(comm.allgather(isinstance(dm, str) and dm == 'SKIPPED_ARG')): dm = mpi.bcast_tagged_array(dm) mf.unpack_(comm.bcast(mf.pack())) if mf.opt is None: mf.opt = mf.init_direct_scf() with lib.temporary_env(mf.opt._this.contents, fprescreen=_vhf._fpointer('CVHFnrs8_vk_prescreen')): vk = _eval_jk(mf, dm, hermi, _vk_jobs_s8) return vk.reshape(dm.shape)
def get_k(mol_or_mf=None, dm=None, hermi=1, omega=None): if isinstance(mol_or_mf, gto.mole.Mole): mf = hf.SCF(mol_or_mf).view(SCF) else: mf = mol_or_mf # dm may be too big for mpi4py library to serialize. Broadcast dm here. if any(comm.allgather(dm is mpi.Message.SkippedArg)): dm = mpi.bcast_tagged_array(dm) mf.unpack_(comm.bcast(mf.pack())) if mf.opt is None: mf.opt = mf.init_direct_scf() with lib.temporary_env(mf.opt._this.contents, fprescreen=_vhf._fpointer('CVHFnrs8_vk_prescreen')): if omega is None: vk = _eval_jk(mf, dm, hermi, _vk_jobs_s8) else: with mf.mol.with_range_coulomb(omega): vk = _eval_jk(mf, dm, hermi, _vk_jobs_s8) return vk.reshape(dm.shape)
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 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 set_vkscreen(opt, name): opt._this.contents.r_vkscreen = _vhf._fpointer(name)
def set_vkscreen(opt, name): opt._this.contents.r_vkscreen = _vhf._fpointer(name)
def _gen_jk_direct(mol, aosym, with_j, with_k, direct_scf_tol, sgxopt=None): '''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) ''' if sgxopt is None: from pyscf.sgx import sgx sgxopt = sgx._make_opt(mol) sgxopt.direct_scf_tol = direct_scf_tol ncomp = 1 nao = mol.nao cintor = _vhf._fpointer(sgxopt._intor) fdot = _vhf._fpointer('SGXdot_nr'+aosym) drv = _vhf.libcvhf.SGXnr_direct_drv def jk_part(mol, grid_coords, dms, fg): 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, sgxopt._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), sgxopt._cintopt, sgxopt._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
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
def _gen_jk_direct(mol, aosym, with_j, with_k, direct_scf_tol, sgxopt=None, pjs=False): '''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) ''' if sgxopt is None: from pyscf.sgx import sgx sgxopt = sgx._make_opt(mol, pjs=pjs) sgxopt.direct_scf_tol = direct_scf_tol ncomp = 1 nao = mol.nao cintor = _vhf._fpointer(sgxopt._intor) fdot = _vhf._fpointer('SGXdot_nrk') drv = _vhf.libcvhf.SGXnr_direct_drv def jk_part(mol, grid_coords, dms, fg, weights): atm, bas, env = mol._atm, mol._bas, mol._env ngrids = grid_coords.shape[0] env = numpy.append(env, grid_coords.ravel()) env[gto.NGRIDS] = ngrids env[gto.PTR_GRIDS] = mol._env.size if pjs: sgxopt.set_dm(fg / numpy.sqrt(numpy.abs(weights[None,:])), mol._atm, mol._bas, env) ao_loc = moleintor.make_loc(bas, sgxopt._intor) shls_slice = (0, mol.nbas, 0, mol.nbas) fg = numpy.ascontiguousarray(fg.transpose(0,2,1)) 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,nao,ngrids))[:,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*4)(*shls_slice), ao_loc.ctypes.data_as(ctypes.c_void_p), sgxopt._cintopt, sgxopt._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), ctypes.c_int(env.shape[0]), ctypes.c_int(2 if aosym == 's2' else 1)) if vk is not None: vk = vk.transpose(0,2,1) vk = numpy.ascontiguousarray(vk) return vj, vk return jk_part