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
